a3s-code-core 9.0.0

use super::*;

#[tokio::test]
async fn test_history_empty_on_new_session() {
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = agent
        .session_async("/tmp/test-workspace", None)
        .await
        .unwrap();
    assert!(session.history().is_empty());
}

#[tokio::test]
async fn test_stream_updates_history_and_auto_saves() {
    let store = Arc::new(crate::store::MemorySessionStore::new());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_store(store.clone())
        .with_session_id("stream-history-test")
        .with_auto_save(true);
    crate::fact_control::reset_session_fact_log("/tmp/test-stream-history");
    let session = agent
        .build_session(
            "/tmp/test-stream-history".into(),
            Arc::new(StaticStreamingClient::new("streamed answer")),
            &opts,
        )
        .unwrap();

    let (mut rx, handle) = session.stream("hello", None).await.unwrap();
    let mut saw_end = false;
    while let Some(event) = rx.recv().await {
        if matches!(event, AgentEvent::End { .. }) {
            saw_end = true;
            break;
        }
    }
    handle.await.unwrap();

    assert!(saw_end);
    let history = session.history();
    assert_eq!(history.len(), 2);
    assert_eq!(history[0].text(), "hello");
    assert_eq!(history[1].text(), "streamed answer");

    let saved = store
        .load("stream-history-test")
        .await
        .unwrap()
        .expect("saved session");
    assert_eq!(saved.messages.len(), 2);
    assert_eq!(saved.messages[1].text(), "streamed answer");

    let run_records = store
        .load_run_records("stream-history-test")
        .await
        .unwrap()
        .expect("saved run records");
    assert_eq!(run_records.len(), 1);
    assert_eq!(
        run_records[0].snapshot.status,
        crate::run::RunStatus::Completed
    );
    assert!(run_records[0]
        .events
        .iter()
        .any(|record| matches!(record.event, AgentEvent::End { .. })));
}

/// After a completed tool round, `auto_save` must flush `sessions` JSON at the
/// checkpoint boundary — not only on stream `End`. Otherwise a process kill
/// after `write` leaves the workspace file without a durable session marker.
#[tokio::test(flavor = "multi_thread")]
async fn test_stream_tool_round_checkpoint_flushes_session_before_end() {
    let store = Arc::new(crate::store::MemorySessionStore::new());
    let workspace = std::env::temp_dir().join(format!(
        "a3s-ckpt-flush-{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis()
    ));
    std::fs::create_dir_all(&workspace).unwrap();
    let marker = "marker-ckpt-flush-xyz";
    let client = Arc::new(ScriptedStreamingClient::new(vec![
        scripted_tool_call_response(
            "call-write",
            "write",
            serde_json::json!({
                "path": "a3s-ckpt-flush.md",
                "content": "7"
            }),
        ),
        // Deliberately never consumed: we abort after ToolEnd to prove the
        // session snapshot does not depend on End autosave.
        scripted_text_response(marker),
    ]));
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_store(store.clone())
        .with_session_id("stream-ckpt-flush-test")
        .with_auto_save(true)
        .with_confirmation_policy(crate::hitl::ConfirmationPolicy::default())
        .with_planning_mode(crate::prompts::PlanningMode::Disabled);
    let session = agent
        .build_session(workspace.display().to_string(), client, &opts)
        .unwrap();

    let (mut rx, handle) = session
        .stream("write a3s-ckpt-flush.md then stop", None)
        .await
        .unwrap();
    let mut flushed_before_end = false;
    while let Some(event) = rx.recv().await {
        match &event {
            AgentEvent::ToolEnd { id, .. } if id == "call-write" => {
                // Checkpoint commit acks after the event sink drains.
                for _ in 0..50 {
                    tokio::time::sleep(std::time::Duration::from_millis(40)).await;
                    if let Ok(Some(saved)) = store.load("stream-ckpt-flush-test").await {
                        let hay = serde_json::to_string(&saved.messages).unwrap_or_default();
                        if hay.contains("a3s-ckpt-flush.md") || hay.contains("call-write") {
                            flushed_before_end = true;
                            break;
                        }
                    }
                }
                handle.abort();
                break;
            }
            AgentEvent::End { .. } => break,
            _ => {}
        }
    }
    let _ = handle.await;

    assert!(
        flushed_before_end,
        "session JSON must flush at the tool-round checkpoint before stream End"
    );
    let _ = std::fs::remove_dir_all(&workspace);
}

#[tokio::test(flavor = "multi_thread")]
async fn test_stream_bridges_subagent_lifecycle_events() {
    use crate::prompts::PlanningMode;
    use crate::subagent_task_tracker::SubagentStatus;

    let client = Arc::new(ScriptedStreamingClient::new(vec![
        scripted_tool_call_response(
            "call-parallel",
            "task",
            serde_json::json!({
                "tasks": [
                    {
                        "agent": "explore",
                        "description": "Find auth code",
                        "prompt": "Find the auth code."
                    },
                    {
                        "agent": "explore",
                        "description": "Find docs",
                        "prompt": "Find the docs."
                    }
                ]
            }),
        ),
        scripted_text_response("auth child result"),
        scripted_text_response("docs child result"),
        scripted_text_response("final answer"),
    ]));
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_id("stream-subagents-test")
        .with_confirmation_policy(crate::hitl::ConfirmationPolicy::default())
        .with_planning_mode(PlanningMode::Disabled);
    crate::fact_control::reset_session_fact_log("/tmp/test-stream-subagents");
    let session = agent
        .build_session("/tmp/test-stream-subagents".into(), client, &opts)
        .unwrap();

    let (mut rx, handle) = session.stream("fan out this work", None).await.unwrap();
    let mut subagent_starts = 0;
    let mut subagent_ends = 0;
    let mut event_index = 0usize;
    let mut last_subagent_end = None;
    let mut parent_tool_end = None;
    while let Some(event) = rx.recv().await {
        match event {
            AgentEvent::SubagentStart { .. } => subagent_starts += 1,
            AgentEvent::SubagentEnd { .. } => {
                subagent_ends += 1;
                last_subagent_end = Some(event_index);
            }
            AgentEvent::ToolEnd { id, .. } if id == "call-parallel" => {
                parent_tool_end = Some(event_index);
            }
            AgentEvent::End { .. } => break,
            _ => {}
        }
        event_index += 1;
    }
    handle.await.unwrap();

    assert_eq!(subagent_starts, 2);
    assert_eq!(subagent_ends, 2);
    assert!(
        last_subagent_end.expect("foreground tasks must emit SubagentEnd")
            < parent_tool_end.expect("task fan-out must emit ToolEnd"),
        "all foreground SubagentEnd events must precede the parent ToolEnd"
    );

    let tasks = session.subagent_tasks().await;
    assert_eq!(tasks.len(), 2);
    assert!(tasks
        .iter()
        .all(|task| task.status == SubagentStatus::Completed));
}

#[tokio::test]
async fn test_stream_with_custom_history_does_not_update_session_history() {
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = agent
        .build_session(
            "/tmp/test-stream-custom-history".into(),
            Arc::new(StaticStreamingClient::new("custom history answer")),
            &SessionOptions::new(),
        )
        .unwrap();
    let custom_history = vec![Message::user("custom prompt")];

    let (mut rx, handle) = session
        .stream("ignored", Some(&custom_history))
        .await
        .unwrap();
    while let Some(event) = rx.recv().await {
        if matches!(event, AgentEvent::End { .. }) {
            break;
        }
    }
    handle.await.unwrap();

    assert!(session.history().is_empty());
}

#[tokio::test]
async fn test_stream_error_does_not_update_history_or_auto_save() {
    let store = Arc::new(crate::store::MemorySessionStore::new());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_store(store.clone())
        .with_session_id("stream-error-test")
        .with_auto_save(true);
    let session = agent
        .build_session(
            "/tmp/test-stream-error".into(),
            Arc::new(FailingStreamingClient),
            &opts,
        )
        .unwrap();

    let (mut rx, handle) = session.stream("hello", None).await.unwrap();
    let mut saw_error = false;
    while let Some(event) = rx.recv().await {
        if matches!(event, AgentEvent::Error { .. }) {
            saw_error = true;
            break;
        }
    }
    handle.await.unwrap();

    assert!(saw_error);
    assert!(session.history().is_empty());
    assert!(store.load("stream-error-test").await.unwrap().is_none());
}

#[tokio::test]
async fn test_non_retryable_stream_error_skips_fallback_and_circuit_retries() {
    let client = Arc::new(NonRetryableStreamingClient::default());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = agent
        .build_session(
            "/tmp/test-non-retryable-stream-error".into(),
            client.clone(),
            &SessionOptions::new().with_planning_mode(PlanningMode::Disabled),
        )
        .unwrap();

    let (mut rx, handle) = session.stream("hello", None).await.unwrap();
    let mut error_message = None;
    while let Some(event) = rx.recv().await {
        if let AgentEvent::Error { message } = event {
            error_message = Some(message);
            break;
        }
    }
    handle.await.unwrap();

    assert_eq!(
        error_message.as_deref(),
        Some("Codex Pro usage limit reached. Quota resets in about 2h 45m.")
    );
    assert_eq!(
        client
            .streaming_calls
            .load(std::sync::atomic::Ordering::SeqCst),
        1,
        "a non-retryable provider error must make one streaming call"
    );
    assert_eq!(
        client
            .complete_calls
            .load(std::sync::atomic::Ordering::SeqCst),
        0,
        "a non-retryable provider error must not use non-streaming fallback"
    );
}

#[tokio::test]
async fn test_non_retryable_pre_analysis_stops_before_main_turn() {
    let client = Arc::new(NonRetryableStreamingClient::default());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = agent
        .build_session(
            "/tmp/test-non-retryable-pre-analysis".into(),
            client.clone(),
            &SessionOptions::new().with_planning_mode(PlanningMode::Auto),
        )
        .unwrap();

    let (mut rx, handle) = session
        .stream("review this repository", None)
        .await
        .unwrap();
    let mut error_message = None;
    while let Some(event) = rx.recv().await {
        if let AgentEvent::Error { message } = event {
            error_message = Some(message);
            break;
        }
    }
    handle.await.unwrap();

    assert_eq!(
        error_message.as_deref(),
        Some("Codex Pro usage limit reached. Quota resets in about 2h 45m.")
    );
    assert_eq!(
        client
            .complete_calls
            .load(std::sync::atomic::Ordering::SeqCst),
        1,
        "pre-analysis should make exactly one provider request"
    );
    assert_eq!(
        client
            .streaming_calls
            .load(std::sync::atomic::Ordering::SeqCst),
        0,
        "a terminal pre-analysis error must stop before the main turn"
    );
}

#[tokio::test]
async fn test_stream_cancel_records_interrupted_history_and_auto_saves() {
    let store = Arc::new(crate::store::MemorySessionStore::new());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_store(store.clone())
        .with_session_id("stream-cancel-test")
        .with_auto_save(true);
    crate::fact_control::reset_session_fact_log("/tmp/test-stream-cancel");
    let session = agent
        .build_session(
            "/tmp/test-stream-cancel".into(),
            Arc::new(CancellableStreamingClient::new("partial answer")),
            &opts,
        )
        .unwrap();

    let (mut rx, handle) = session.stream("hello", None).await.unwrap();
    let mut saw_delta = false;
    for _ in 0..16 {
        let event = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
            .await
            .expect("stream event before timeout")
            .expect("stream should stay open until cancelled");
        if matches!(event, AgentEvent::TextDelta { ref text } if text == "partial answer") {
            saw_delta = true;
            break;
        }
    }
    assert!(saw_delta);
    assert!(session.cancel().await);

    while rx.recv().await.is_some() {}
    handle.await.unwrap();

    let history = session.history();
    assert_eq!(history.len(), 2);
    assert_eq!(history[0].role, "user");
    assert_eq!(history[0].text(), "hello");
    assert_eq!(history[1].role, "assistant");
    assert!(history[1].text().contains("interrupted"));

    let saved = store
        .load("stream-cancel-test")
        .await
        .unwrap()
        .expect("interrupted stream should auto-save");
    assert_eq!(saved.messages.len(), 2);
    assert_eq!(saved.messages[0].text(), "hello");
    assert!(saved.messages[1].text().contains("interrupted"));
    assert!(!session.cancel().await);
}

#[tokio::test]
async fn test_stream_with_attachments_cancel_records_interrupted_history_and_auto_saves() {
    let store = Arc::new(crate::store::MemorySessionStore::new());
    let agent = Agent::from_config(test_config()).await.unwrap();
    let opts = SessionOptions::new()
        .with_session_store(store.clone())
        .with_session_id("stream-attachments-cancel-test")
        .with_auto_save(true);
    crate::fact_control::reset_session_fact_log("/tmp/test-stream-attachments-cancel");
    let session = agent
        .build_session(
            "/tmp/test-stream-attachments-cancel".into(),
            Arc::new(CancellableStreamingClient::new("partial attachment answer")),
            &opts,
        )
        .unwrap();
    let attachments = vec![crate::llm::Attachment::png(vec![1, 2, 3])];

    let (mut rx, handle) = session
        .stream_with_attachments("hello", &attachments, None)
        .await
        .unwrap();
    let mut saw_delta = false;
    for _ in 0..16 {
        let event = tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
            .await
            .expect("stream event before timeout")
            .expect("stream should stay open until cancelled");
        if matches!(event, AgentEvent::TextDelta { .. }) {
            saw_delta = true;
            break;
        }
    }
    assert!(saw_delta);
    assert!(session.cancel().await);

    while rx.recv().await.is_some() {}
    handle.await.unwrap();

    let history = session.history();
    assert_eq!(history.len(), 2);
    assert_eq!(history[0].role, "user");
    assert_eq!(history[0].text(), "hello");
    assert_eq!(history[1].role, "assistant");
    assert!(history[1].text().contains("interrupted"));

    let saved = store
        .load("stream-attachments-cancel-test")
        .await
        .unwrap()
        .expect("interrupted attachment stream should auto-save");
    assert_eq!(saved.messages.len(), 2);
    assert_eq!(saved.messages[0].text(), "hello");
    assert!(saved.messages[1].text().contains("interrupted"));
    assert_eq!(
        session.runs().await[0].status,
        crate::run::RunStatus::Cancelled
    );
    assert!(!session.cancel().await);
}

#[tokio::test]
async fn test_run_handle_cancels_send_with_attachments() {
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = Arc::new(
        agent
            .build_session(
                "/tmp/test-send-attachments-run-handle-cancel".into(),
                Arc::new(CancellableStreamingClient::new("partial answer")),
                &SessionOptions::new(),
            )
            .unwrap(),
    );
    let worker_session = Arc::clone(&session);
    let attachments = vec![crate::llm::Attachment::png(vec![1, 2, 3])];

    let worker = tokio::spawn(async move {
        worker_session
            .send_with_attachments("hello", &attachments, None)
            .await
    });

    let mut run = None;
    for _ in 0..20 {
        if let Some(current) = session.current_run().await {
            run = Some(current);
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
    let run = run.expect("current run should be visible");
    assert!(run.cancel().await);

    let result = tokio::time::timeout(std::time::Duration::from_secs(1), worker)
        .await
        .expect("send_with_attachments should stop after cancellation")
        .expect("worker should not panic");
    let result = result.expect("cancellation should preserve interrupted history");
    assert_eq!(result.messages.len(), 2);
    assert_eq!(result.messages[0].text(), "hello");
    assert!(result.messages[1].text().contains("interrupted"));
    assert_eq!(run.status().await, Some(crate::run::RunStatus::Cancelled));
    let history = session.history();
    assert_eq!(history.len(), 2);
    assert_eq!(history[0].text(), "hello");
    assert!(history[1].text().contains("interrupted"));
    assert!(!session.cancel().await);
}

#[tokio::test]
async fn test_cancel_run_only_cancels_matching_current_run() {
    let agent = Agent::from_config(test_config()).await.unwrap();
    let session = Arc::new(
        agent
            .build_session(
                "/tmp/test-cancel-run-by-id".into(),
                Arc::new(CancellableStreamingClient::new("partial answer")),
                &SessionOptions::new(),
            )
            .unwrap(),
    );
    let worker_session = Arc::clone(&session);
    let worker = tokio::spawn(async move { worker_session.send("hello", None).await });

    let mut run_id = None;
    for _ in 0..20 {
        if let Some(current) = session.current_run().await {
            run_id = Some(current.id().to_string());
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
    let run_id = run_id.expect("current run should be visible");

    assert!(!session.cancel_run("stale-run").await);
    assert!(session.cancel_run(&run_id).await);

    let result = tokio::time::timeout(std::time::Duration::from_secs(1), worker)
        .await
        .expect("send should stop after cancellation")
        .expect("worker should not panic");
    let result = result.expect("cancellation should preserve interrupted history");
    assert_eq!(result.messages.len(), 2);
    assert_eq!(result.messages[0].text(), "hello");
    assert!(result.messages[1].text().contains("interrupted"));
    assert_eq!(
        session.run_snapshot(&run_id).await.unwrap().status,
        crate::run::RunStatus::Cancelled
    );
    assert!(!session.cancel_run(&run_id).await);
}

/// A cancelled fact-log turn must close the host stream without a second
/// provider call. Memory extraction is post-turn work and follows the run
/// cancellation token.
#[tokio::test]
async fn cancelled_fact_stream_settles_without_memory_extraction() {
    let extraction_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let agent = Agent::from_config(test_config()).await.unwrap();
    let workspace = tempfile::tempdir().unwrap();
    let session = agent
        .build_session(
            workspace.path().to_string_lossy().to_string(),
            std::sync::Arc::new(SlowExtractionClient {
                extraction_calls: std::sync::Arc::clone(&extraction_calls),
            }),
            &SessionOptions::new()
                .with_session_id("cancel-skips-memory-extraction")
                .with_planning_mode(crate::prompts::PlanningMode::Disabled)
                .with_memory(std::sync::Arc::new(a3s_memory::InMemoryStore::new())),
        )
        .unwrap();

    let (mut events, worker) = session
        .stream("remember this cancelled turn", None)
        .await
        .unwrap();
    let started = tokio::time::timeout(std::time::Duration::from_secs(2), async {
        loop {
            match events.recv().await {
                Some(AgentEvent::TextDelta { .. }) => return,
                Some(AgentEvent::End { .. }) | Some(AgentEvent::Error { .. }) | None => {
                    panic!("stream ended before cancellation")
                }
                _ => {}
            }
        }
    })
    .await;
    started.expect("model delta before cancellation");
    assert!(session.cancel().await);

    tokio::time::timeout(std::time::Duration::from_secs(2), async {
        while events.recv().await.is_some() {}
    })
    .await
    .expect("cancelled stream should close without waiting for memory extraction");
    tokio::time::timeout(std::time::Duration::from_secs(2), worker)
        .await
        .expect("cancelled worker should finish without a memory extraction call")
        .expect("worker join");
    assert_eq!(
        extraction_calls.load(std::sync::atomic::Ordering::SeqCst),
        0
    );
    assert!(session
        .history()
        .iter()
        .any(|message| message.text().contains("interrupted")));
}

struct SlowExtractionClient {
    extraction_calls: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}

#[async_trait::async_trait]
impl crate::llm::LlmClient for SlowExtractionClient {
    async fn complete(
        &self,
        _messages: &[crate::llm::Message],
        _system: Option<&str>,
        _tools: &[crate::llm::ToolDefinition],
    ) -> anyhow::Result<crate::llm::LlmResponse> {
        self.extraction_calls
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        tokio::time::sleep(std::time::Duration::from_secs(30)).await;
        Ok(crate::llm::LlmResponse {
            message: crate::llm::Message::assistant("[]"),
            usage: crate::llm::TokenUsage::default(),
            stop_reason: None,
            token_logprobs: Vec::new(),
            meta: None,
        })
    }

    async fn complete_streaming(
        &self,
        _messages: &[crate::llm::Message],
        _system: Option<&str>,
        _tools: &[crate::llm::ToolDefinition],
        cancel_token: tokio_util::sync::CancellationToken,
    ) -> anyhow::Result<tokio::sync::mpsc::Receiver<crate::llm::StreamEvent>> {
        let (sender, receiver) = tokio::sync::mpsc::channel(2);
        tokio::spawn(async move {
            let _ = sender
                .send(crate::llm::StreamEvent::TextDelta("partial".into()))
                .await;
            cancel_token.cancelled().await;
        });
        Ok(receiver)
    }
}