af-agent-runtime 0.5.0

use std::sync::Arc;
use std::time::{Duration, Instant};

use af_agent::testing::ScriptedModel;
use af_agent::ChatModel;
use af_agent_runtime::{CancellationToken, Compactor, ModelCompactor, RuntimeError};
use af_llm::{ChatMessage, CompletionRequest, CompletionResponse, LlmError};
use async_trait::async_trait;

struct PendingModel;

#[async_trait]
impl ChatModel for PendingModel {
    async fn complete_streaming(
        &self,
        _: &CompletionRequest,
        _: tokio::sync::mpsc::UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError> {
        std::future::pending().await
    }
}

fn messages() -> Vec<ChatMessage> {
    vec![
        ChatMessage::user("what happened so far?"),
        ChatMessage::assistant("a lot of things happened here"),
    ]
}

fn far_deadline() -> Instant {
    Instant::now() + Duration::from_secs(30)
}

#[tokio::test]
async fn summary_usage_falls_back_to_character_estimates_without_provider_usage() {
    let mut response = ScriptedModel::response(ChatMessage::assistant("  summary text  "));
    response.usage = None;
    let model = Arc::new(ScriptedModel::new([Ok(response)]));
    let compactor = ModelCompactor::new(model.clone());
    let result = compactor
        .summarize(
            "m",
            &messages(),
            "compaction:1:attempt:1",
            CancellationToken::default(),
            far_deadline(),
        )
        .await
        .unwrap();
    assert_eq!(result.summary, "summary text");
    // Every prompt message, including the compaction instructions the
    // compactor prepends, is estimated at four characters per token.
    let expected_prompt: u64 = model
        .requests()
        .pop()
        .unwrap()
        .messages
        .iter()
        .filter_map(|message| message.content.as_deref())
        .map(|content| content.chars().count().div_ceil(4) as u64)
        .sum();
    assert_eq!(result.prompt_tokens, expected_prompt);
    assert_eq!(result.completion_tokens, 3);
}

#[tokio::test]
async fn summary_reports_provider_usage_and_pins_the_operation_id() {
    let model = Arc::new(ScriptedModel::replies([ChatMessage::assistant("gist")]));
    let compactor = ModelCompactor::new(model.clone());
    let result = compactor
        .summarize(
            "m",
            &messages(),
            "compaction:2:attempt:1",
            CancellationToken::default(),
            far_deadline(),
        )
        .await
        .unwrap();
    assert_eq!((result.prompt_tokens, result.completion_tokens), (1, 1));
    let request = model.requests().pop().unwrap();
    assert_eq!(
        request.provider_attempt_id.as_deref(),
        Some("compaction:2:attempt:1")
    );
    assert_eq!(request.messages.len(), 3);
}

#[tokio::test]
async fn blank_summaries_model_errors_cancellation_and_deadlines_fail_closed() {
    let blank = ModelCompactor::new(Arc::new(ScriptedModel::replies([ChatMessage::assistant(
        "   ",
    )])));
    assert!(matches!(
        blank
            .summarize("m", &messages(), "op", CancellationToken::default(), far_deadline())
            .await,
        Err(RuntimeError::Model(message)) if message.contains("no summary")
    ));

    let failing = ModelCompactor::new(Arc::new(ScriptedModel::new([Err(
        LlmError::StreamProtocol("upstream down".into()),
    )])));
    assert!(matches!(
        failing
            .summarize("m", &messages(), "op", CancellationToken::default(), far_deadline())
            .await,
        Err(RuntimeError::Model(message)) if message.contains("upstream down")
    ));

    let pending = ModelCompactor::new(Arc::new(PendingModel));
    let cancellation = CancellationToken::default();
    cancellation.cancel();
    assert!(matches!(
        pending
            .summarize("m", &messages(), "op", cancellation, far_deadline())
            .await,
        Err(RuntimeError::Cancelled)
    ));
    assert!(matches!(
        pending
            .summarize(
                "m",
                &messages(),
                "op",
                CancellationToken::default(),
                Instant::now() + Duration::from_millis(20),
            )
            .await,
        Err(RuntimeError::Model(message)) if message.contains("deadline")
    ));
}