owlen/crates/owlen-core/tests/compression.rs

use std::sync::Arc;

use anyhow::{Result, anyhow};
use async_trait::async_trait;
use futures::stream;
use owlen_core::config::{CompressionStrategy, Config};
use owlen_core::session::SessionController;
use owlen_core::storage::StorageManager;
use owlen_core::types::{ChatRequest, ChatResponse, Message, ModelInfo, Role};
use owlen_core::ui::NoOpUiController;
use owlen_core::{ChatStream, Provider, Result as CoreResult};
use tempfile::tempdir;

fn make_session_config(strategy: CompressionStrategy, auto: bool) -> Config {
    let mut config = Config::default();
    config.general.default_model = Some("stub-model".into());
    config.general.enable_streaming = false;
    config.chat.strategy = strategy;
    config.chat.auto_compress = auto;
    config.chat.trigger_tokens = 64;
    config.chat.retain_recent_messages = 2;
    config
}

async fn build_session(config: Config) -> Result<SessionController> {
    let temp_dir = tempdir().expect("temp dir");
    let storage = Arc::new(
        StorageManager::with_database_path(temp_dir.path().join("owlen-compression-tests.db"))
            .await
            .expect("storage"),
    );
    let provider: Arc<dyn Provider> = Arc::new(StubProvider);
    let ui = Arc::new(NoOpUiController);
    SessionController::new(provider, config, storage, ui, false, None)
        .await
        .map_err(|err| anyhow!(err))
}

struct StubProvider;

#[async_trait]
impl Provider for StubProvider {
    fn name(&self) -> &str {
        "stub-provider"
    }

    async fn list_models(&self) -> CoreResult<Vec<ModelInfo>> {
        Ok(vec![ModelInfo {
            id: "stub-model".into(),
            name: "Stub Model".into(),
            description: Some("Stub provider model".into()),
            provider: "stub-provider".into(),
            context_window: Some(8_192),
            capabilities: vec!["chat".into()],
            supports_tools: false,
        }])
    }

    async fn send_prompt(&self, _request: ChatRequest) -> CoreResult<ChatResponse> {
        Ok(ChatResponse {
            message: Message::assistant("stub completion".into()),
            usage: None,
            is_streaming: false,
            is_final: true,
        })
    }

    async fn stream_prompt(&self, _request: ChatRequest) -> CoreResult<ChatStream> {
        Ok(Box::pin(stream::empty()))
    }

    async fn health_check(&self) -> CoreResult<()> {
        Ok(())
    }

    fn as_any(&self) -> &(dyn std::any::Any + Send + Sync) {
        self
    }
}

#[tokio::test(flavor = "multi_thread")]
async fn compression_compacts_history() -> Result<()> {
    let mut session = build_session(make_session_config(CompressionStrategy::Local, true)).await?;

    for idx in 0..6 {
        session.conversation_mut().push_user_message(format!(
            "User request #{idx}: Explain the subsystem in detail."
        ));
        session.conversation_mut().push_assistant_message(format!(
            "Assistant reply #{idx}: Provided detailed explanation with follow-up tasks."
        ));
    }

    let before_len = session.conversation().messages.len();
    assert!(
        before_len > 6,
        "expected longer transcript before compression"
    );

    let report = session
        .compress_now()
        .await?
        .expect("compression should trigger");
    assert!(
        !report.automated,
        "manual compression should flag automated = false"
    );
    assert!(report.compressed_messages > 0);
    assert!(report.estimated_tokens_after < report.estimated_tokens_before);

    let after = session.conversation();
    assert!(after.messages.len() < before_len);
    let first = after
        .messages
        .first()
        .expect("summary message should exist after compression");
    assert_eq!(first.role, Role::System);
    assert!(
        first.metadata.contains_key("compression"),
        "summary message must include metadata"
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn auto_compress_respects_toggle() -> Result<()> {
    let mut session = build_session(make_session_config(CompressionStrategy::Local, false)).await?;

    for idx in 0..5 {
        session
            .conversation_mut()
            .push_user_message(format!("Message {idx} from user."));
        session
            .conversation_mut()
            .push_assistant_message(format!("Assistant reply {idx}."));
    }

    let result = session.maybe_auto_compress().await?;
    assert!(
        result.is_none(),
        "auto compression should skip when disabled"
    );

    Ok(())
}