af-agent-runtime 0.9.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 request(operation: &str) -> CompletionRequest {
    let mut request = CompletionRequest::new("m", messages());
    request.provider_attempt_id = Some(operation.into());
    request
}

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(
            &request("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(
            &request("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(&request("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(&request("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(&request("op"), cancellation, far_deadline())
            .await,
        Err(RuntimeError::Cancelled)
    ));
    assert!(matches!(
        pending
            .summarize(
                &request("op"),
                CancellationToken::default(),
                Instant::now() + Duration::from_millis(20),
            )
            .await,
        Err(RuntimeError::Model(message)) if message.contains("deadline")
    ));
}

#[tokio::test]
async fn a5_compactor_preserves_caller_for_image_authorization() {
    let model = Arc::new(ScriptedModel::replies([ChatMessage::assistant(
        "image summary",
    )]));
    let mut input = request("image-compaction");
    input.context = Some(af_context::RequestContext {
        tenant_id: "tenant".parse().unwrap(),
        subject_id: "subject".parse().unwrap(),
        request_id: "request".parse().unwrap(),
        roles: Default::default(),
        entitlements: Default::default(),
        locale: "en".into(),
    });
    input.messages.last_mut().unwrap().role = af_llm::Role::User;
    input
        .messages
        .last_mut()
        .unwrap()
        .images
        .push(af_llm::InputImage {
            asset_id: "asset".parse().unwrap(),
            media_type: "image/png".into(),
        });
    ModelCompactor::new(model.clone())
        .summarize(&input, CancellationToken::default(), far_deadline())
        .await
        .unwrap();
    let sent = model.requests().pop().unwrap();
    assert_eq!(sent.context, input.context);
    assert_eq!(
        sent.messages.last().unwrap().images,
        input.messages.last().unwrap().images
    );
    assert!(!serde_json::to_string(&sent).unwrap().contains("signature"));
}