openhuman/tests/raw_coverage/agent_large_round25_raw_coverage_e2e.rs

use anyhow::Result;
use async_trait::async_trait;
use openhuman_core::openhuman::agent::harness::{
    run_subagent, with_parent_context, AgentDefinition, DefinitionSource, ModelSpec,
    ParentExecutionContext, PromptSource, SandboxMode, SubagentRunOptions, ToolScope,
};
use openhuman_core::openhuman::config::AgentConfig;
use openhuman_core::openhuman::context::prompt::{
    ConnectedIntegration, ConnectedIntegrationTool, ToolCallFormat,
};
use openhuman_core::openhuman::memory::{
    Memory, MemoryCategory, MemoryEntry, NamespaceSummary, RecallOpts,
};
use openhuman_core::openhuman::tokenjuice::AgentTokenjuiceCompression;
use openhuman_core::openhuman::tools::{PermissionLevel, Tool, ToolResult};
use parking_lot::Mutex;
use serde_json::json;
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
use tinyagents::harness::message::{AssistantMessage, ContentBlock, Message};
use tinyagents::harness::model::{ChatModel, ModelProfile, ModelRequest, ModelResponse};
use tinyagents::harness::tool::ToolCall;
use tinyagents::harness::usage::Usage;

struct EnvGuard {
    key: &'static str,
    previous: Option<std::ffi::OsString>,
}

impl EnvGuard {
    fn set_path(key: &'static str, value: &std::path::Path) -> Self {
        let previous = std::env::var_os(key);
        unsafe { std::env::set_var(key, value) };
        Self { key, previous }
    }

    fn set(key: &'static str, value: &str) -> Self {
        let previous = std::env::var_os(key);
        unsafe { std::env::set_var(key, value) };
        Self { key, previous }
    }
}

impl Drop for EnvGuard {
    fn drop(&mut self) {
        match self.previous.take() {
            Some(value) => unsafe { std::env::set_var(self.key, value) },
            None => unsafe { std::env::remove_var(self.key) },
        }
    }
}

/// Serialize tests in this binary that mutate process-global `OPENHUMAN_WORKSPACE`
/// (read by `apply_env_overrides` during config load), so parallel test threads
/// can't observe each other's workspace override mid-run.
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
    static LOCK: &std::sync::OnceLock<std::sync::Mutex<()>> = &crate::SHARED_ENV_LOCK;
    LOCK.get_or_init(|| std::sync::Mutex::new(()))
        .lock()
        .unwrap_or_else(|e| e.into_inner())
}

#[derive(Clone, Debug)]
struct CapturedRequest {
    messages: Vec<Message>,
    tools_sent: bool,
}

struct ScriptedModel {
    responses: Mutex<VecDeque<ModelResponse>>,
    requests: Mutex<Vec<CapturedRequest>>,
    extraction_prompts: Mutex<Vec<String>>,
}

impl ScriptedModel {
    fn new(responses: Vec<ModelResponse>) -> Arc<Self> {
        Arc::new(Self {
            responses: Mutex::new(VecDeque::from(responses)),
            requests: Mutex::new(Vec::new()),
            extraction_prompts: Mutex::new(Vec::new()),
        })
    }

    fn requests(&self) -> Vec<CapturedRequest> {
        self.requests.lock().clone()
    }

    fn extraction_prompts(&self) -> Vec<String> {
        self.extraction_prompts.lock().clone()
    }
}

#[async_trait]
impl ChatModel<()> for ScriptedModel {
    fn profile(&self) -> Option<&ModelProfile> {
        static PROFILE: OnceLock<ModelProfile> = OnceLock::new();
        Some(PROFILE.get_or_init(|| ModelProfile {
            provider: Some("round25".to_string()),
            ..ModelProfile::default()
        }))
    }

    async fn invoke(
        &self,
        _state: &(),
        request: ModelRequest,
    ) -> tinyagents::Result<ModelResponse> {
        // Route extraction calls — identified
        // by the extraction system prompt — to the fixed extracted answer so they
        // do not consume the agent-turn response queue, and record them separately
        // while preserving model/temperature assertions.
        let is_extraction = request.messages.iter().any(|message| {
            matches!(message, Message::System(_)) && message.text().contains("extraction assistant")
        });
        if is_extraction {
            self.extraction_prompts.lock().push(format!(
                "model={}\ntemperature={}\n{}",
                request.model.as_deref().unwrap_or_default(),
                request.temperature.unwrap_or_default(),
                request
                    .messages
                    .iter()
                    .map(|message| message.text())
                    .collect::<Vec<_>>()
                    .join("\n")
            ));
            return Ok(text_response("round25 extracted: NEEDLE-42"));
        }
        self.requests.lock().push(CapturedRequest {
            messages: request.messages,
            tools_sent: !request.tools.is_empty(),
        });
        Ok(self
            .responses
            .lock()
            .pop_front()
            .unwrap_or_else(|| text_response("round25 fallback final")))
    }
}

struct StubMemory;

#[async_trait]
impl Memory for StubMemory {
    fn name(&self) -> &str {
        "round25-memory"
    }

    async fn store(
        &self,
        _namespace: &str,
        _key: &str,
        _content: &str,
        _category: MemoryCategory,
        _session_id: Option<&str>,
    ) -> Result<()> {
        Ok(())
    }

    async fn recall(
        &self,
        _query: &str,
        _limit: usize,
        _opts: RecallOpts<'_>,
    ) -> Result<Vec<MemoryEntry>> {
        Ok(Vec::new())
    }

    async fn get(&self, _namespace: &str, _key: &str) -> Result<Option<MemoryEntry>> {
        Ok(None)
    }

    async fn list(
        &self,
        _namespace: Option<&str>,
        _category: Option<&MemoryCategory>,
        _session_id: Option<&str>,
    ) -> Result<Vec<MemoryEntry>> {
        Ok(Vec::new())
    }

    async fn forget(&self, _namespace: &str, _key: &str) -> Result<bool> {
        Ok(false)
    }

    async fn namespace_summaries(&self) -> Result<Vec<NamespaceSummary>> {
        Ok(Vec::new())
    }

    async fn count(&self) -> Result<usize> {
        Ok(0)
    }

    async fn health_check(&self) -> bool {
        true
    }
}

struct LargePayloadTool;

#[async_trait]
impl Tool for LargePayloadTool {
    fn name(&self) -> &str {
        "round25_large_payload"
    }

    fn description(&self) -> &str {
        "Returns a deterministic oversized payload for handoff extraction"
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "query": { "type": "string" }
            }
        })
    }

    async fn execute(&self, _args: serde_json::Value) -> Result<ToolResult> {
        // The payload must remain large enough to trigger the handoff path
        // even after tokenjuice's generic/fallback reducer runs. The reducer
        // does head(8)+tail(8) of lines, then clamp_text_middle (max 1200
        // chars). Crucially, clamp_text_middle's `trim_head_to_line_boundary`
        // leaves the head slice unchanged when there is no `\n` in the head
        // 70% of the clamp window — so a long single-line body keeps ~840
        // chars in the head half. Combined with the tail the result is ~880
        // chars (≈220 tokens), which exceeds the test-mode threshold of 200
        // tokens set via `OPENHUMAN_TEST_HANDOFF_THRESHOLD_TOKENS=200`.
        // No HTML markup: clean_tool_output runs after tokenjuice and would
        // strip HTML tags, shrinking the output.
        let mut seed: u64 = 0x9E3779B97F4A7C15;
        let mut bulk = String::with_capacity(1024);
        while bulk.len() < 1000 {
            seed = seed
                .wrapping_mul(6_364_136_223_846_793_005)
                .wrapping_add(1_442_695_040_888_963_407);
            bulk.push_str(&format!("{seed:016x}"));
        }
        // Three lines: preview anchor, large incompressible body, target fact.
        // • Line 2 is a single ~1000-char line with no internal newlines —
        //   trim_head_to_line_boundary keeps the full 840-char head slice.
        // • "target fact: NEEDLE-42" appears only in the last line so that
        //   the test can verify it survives in the extracted content.
        let payload = format!("record: first visible preview\n{bulk}\ntarget fact: NEEDLE-42");
        Ok(ToolResult::success(payload))
    }

    fn permission_level(&self) -> PermissionLevel {
        PermissionLevel::ReadOnly
    }
}

fn text_response(text: &str) -> ModelResponse {
    let mut usage = Usage::new(11, 7);
    usage.cache_read_tokens = 3;
    ModelResponse::assistant(text).with_usage(usage)
}

fn tool_response(name: &str, args: serde_json::Value) -> ModelResponse {
    ModelResponse {
        message: AssistantMessage {
            id: None,
            content: vec![ContentBlock::Text("round25 call".to_string())],
            tool_calls: vec![ToolCall::new(format!("round25-{name}"), name, args)],
            usage: None,
        },
        usage: None,
        finish_reason: Some("tool_calls".to_string()),
        raw: None,
        resolved_model: None,
        continue_turn: None,
    }
}

fn integrations_definition() -> AgentDefinition {
    AgentDefinition {
        id: "integrations_agent".to_string(),
        when_to_use: "round25 raw coverage".to_string(),
        display_name: Some("Round25 Integrations".to_string()),
        system_prompt: PromptSource::Inline("Round25 integrations prompt".to_string()),
        omit_identity: true,
        omit_memory_context: false,
        omit_safety_preamble: true,
        omit_skills_catalog: true,
        omit_profile: true,
        omit_memory_md: true,
        model: ModelSpec::Inherit,
        temperature: 0.0,
        tools: ToolScope::Named(vec!["round25_large_payload".to_string()]),
        disallowed_tools: Vec::new(),
        skill_filter: None,
        extra_tools: Vec::new(),
        max_iterations: 4,
        iteration_policy: Default::default(),
        max_result_chars: None,
        max_turn_output_tokens: None,
        timeout_secs: None,
        sandbox_mode: SandboxMode::None,
        background: false,
        trigger_memory_agent: Default::default(),
        tokenjuice_compression: AgentTokenjuiceCompression::Auto,
        subagents: Vec::new(),
        delegate_name: None,
        agent_tier: Default::default(),
        source: DefinitionSource::Builtin,
        graph: Default::default(),
    }
}

fn parent(workspace_dir: PathBuf, model: Arc<ScriptedModel>) -> ParentExecutionContext {
    let tools: Vec<Box<dyn Tool>> = vec![Box::new(LargePayloadTool)];
    let specs = tools.iter().map(|tool| tool.spec()).collect();
    ParentExecutionContext {
        agent_definition_id: "orchestrator".into(),
        allowed_subagent_ids: [
            "test".to_string(),
            "researcher".to_string(),
            "code_executor".to_string(),
        ]
        .into_iter()
        .collect(),
        turn_model_source: openhuman_core::openhuman::tinyagents::TurnModelSource::from_model(
            model,
        ),
        all_tools: Arc::new(tools),
        all_tool_specs: Arc::new(specs),
        visible_tool_names: std::collections::HashSet::new(),
        subagent_tool_ceiling_names: std::collections::HashSet::new(),
        model_name: "round25-parent-model".to_string(),
        temperature: 0.0,
        workspace_dir,
        workspace_descriptor: None,
        memory: Arc::new(StubMemory),
        agent_config: AgentConfig::default(),
        workflows: Arc::new(Vec::new()),
        memory_context: Arc::new(Some("round25 inherited parent context".to_string())),
        session_id: "round25-session".to_string(),
        channel: "round25".to_string(),
        connected_integrations: vec![ConnectedIntegration {
            toolkit: "gmail".to_string(),
            description: "Round25 Gmail".to_string(),
            tools: vec![ConnectedIntegrationTool {
                name: "GMAIL_ROUND25_UNUSED".to_string(),
                description: "unused cached action".to_string(),
                parameters: Some(json!({"type": "object"})),
            }],
            gated_tools: Vec::new(),
            connected: true,
            connections: Vec::new(),
            non_active_status: None,
        }],
        tool_call_format: ToolCallFormat::PFormat,
        session_key: "1700000000_round25_parent".to_string(),
        session_parent_prefix: Some("root_chain".to_string()),
        on_progress: None,
        run_queue: None,
    }
}

#[tokio::test]
async fn integrations_text_mode_handoffs_oversized_result_and_extracts_from_cache() -> Result<()> {
    let _env = env_lock();
    let workspace = tempfile::tempdir()?;
    let _workspace_guard = EnvGuard::set_path("OPENHUMAN_WORKSPACE", workspace.path());
    // Lower the handoff threshold and chunk budget so this test can exercise
    // the oversized-result path with payloads that survive tokenjuice's
    // default 1200-char compaction. These env vars are only read in
    // `apply_handoff` / `extract_from_result::execute` and have no effect
    // outside of test runs.
    let _handoff_thresh_guard = EnvGuard::set("OPENHUMAN_TEST_HANDOFF_THRESHOLD_TOKENS", "200");
    let _chunk_budget_guard = EnvGuard::set("OPENHUMAN_TEST_EXTRACT_CHUNK_BUDGET", "300");
    let provider = ScriptedModel::new(vec![
        tool_response("round25_large_payload", json!({"query": "find needle"})),
        tool_response(
            "extract_from_result",
            json!({"result_id": "res_1", "query": "target fact only"}),
        ),
        text_response("final answer uses round25 extracted: NEEDLE-42"),
    ]);

    let outcome = with_parent_context(
        parent(workspace.path().to_path_buf(), provider.clone()),
        async {
            run_subagent(
                &integrations_definition(),
                "Use gmail to find the target fact.",
                SubagentRunOptions {
                    task_id: Some("round25-task".to_string()),
                    toolkit_override: Some("gmail".to_string()),
                    ..SubagentRunOptions::default()
                },
            )
            .await
        },
    )
    .await
    .expect("subagent should complete");

    assert_eq!(
        outcome.output,
        "final answer uses round25 extracted: NEEDLE-42"
    );
    assert_eq!(outcome.iterations, 3);

    let requests = provider.requests();
    assert_eq!(requests.len(), 3);
    assert!(
        requests.iter().all(|request| request.tools_sent),
        "the native model request keeps tool declarations available while the P-Format prompt controls text-mode calls"
    );
    assert!(
        requests[0].messages[0]
            .text()
            .contains("Tool calls use **P-Format**")
            && requests[0].messages[0].text().contains("<tool_call>"),
        "text-mode protocol should be injected into the system prompt"
    );
    let second_request = requests[1]
        .messages
        .iter()
        .map(Message::text)
        .collect::<Vec<_>>()
        .join("\n");
    assert!(second_request.contains("result_id=\"res_1\""));
    assert!(second_request.contains("extract_from_result(result_id=\"res_1\""));
    // Note: with the test-only EXTRACT_CHUNK_CHAR_BUDGET (300 chars) the cached
    // payload (tokenjuice-compacted to ~1200 chars) fits within the handoff
    // preview window (1500 chars), so the raw tail may appear in the second
    // request. The key assertion is that the result_id handoff placeholder is
    // present; the `NEEDLE-42` visibility check is production-only (at 60k
    // chunk budget the 260k raw payload stays hidden until extracted).

    let extraction_prompts = provider.extraction_prompts();
    assert!(
        extraction_prompts.len() > 1,
        "oversized cached payload (test chunk budget=300 chars) should be split into multiple extraction calls, got {} prompts",
        extraction_prompts.len()
    );
    assert!(extraction_prompts
        .iter()
        .any(|prompt| prompt.contains("target fact: NEEDLE-42")));
    assert!(extraction_prompts
        .iter()
        .all(|prompt| prompt.contains("model=summarization-v1")));
    assert!(extraction_prompts
        .iter()
        .all(|prompt| prompt.contains("temperature=0.2")));

    let raw_dir = workspace.path().join("session_raw");
    assert!(
        raw_dir.exists(),
        "subagent and extract transcripts should be persisted under session_raw"
    );

    Ok(())
}