openhuman/tests/raw_coverage/agent_harness_raw_coverage_e2e.rs

use anyhow::Result;
use async_trait::async_trait;
use openhuman_core::openhuman::agent::dispatcher::NativeToolDispatcher;
use openhuman_core::openhuman::agent::harness::definition::AgentDefinitionRegistry;
use openhuman_core::openhuman::agent::harness::session::Agent;
use openhuman_core::openhuman::agent::harness::{
    run_subagent, with_parent_context, AgentDefinition, ParentExecutionContext, PromptSource,
    SandboxMode, SubagentRunOptions, ToolScope,
};
use openhuman_core::openhuman::agent::progress::AgentProgress;
use openhuman_core::openhuman::config::AgentConfig;
use openhuman_core::openhuman::context::prompt::ToolCallFormat;
use openhuman_core::openhuman::memory::{Memory, MemoryCategory, MemoryEntry, NamespaceSummary};
use openhuman_core::openhuman::tokenjuice::AgentTokenjuiceCompression;
use openhuman_core::openhuman::tools::SpawnSubagentTool;
use openhuman_core::openhuman::tools::{Tool, ToolResult};
use parking_lot::Mutex;
use serde_json::json;
use std::path::{Path, PathBuf};
use std::sync::Arc;
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 ScriptedModel {
    responses: Mutex<Vec<ModelResponse>>,
    requests: Mutex<Vec<Vec<Message>>>,
}

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

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

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

    async fn invoke(
        &self,
        _state: &(),
        request: ModelRequest,
    ) -> tinyagents::Result<ModelResponse> {
        self.requests.lock().push(request.messages);
        let mut responses = self.responses.lock();
        Ok(if responses.is_empty() {
            ModelResponse::assistant("fallback final").with_usage(usage(7, 3))
        } else {
            responses.remove(0)
        })
    }
}

struct StubMemory;

#[async_trait]
impl Memory for StubMemory {
    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: openhuman_core::openhuman::memory::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
    }

    fn name(&self) -> &str {
        "stub-memory"
    }
}

struct EchoTool;

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

    fn description(&self) -> &str {
        "Echo a deterministic payload for harness coverage"
    }

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

    async fn execute(&self, args: serde_json::Value) -> Result<ToolResult> {
        let message = args
            .get("message")
            .and_then(|value| value.as_str())
            .unwrap_or("(missing)");
        Ok(ToolResult::success(format!("echoed:{message}")))
    }
}

fn usage(input_tokens: u64, output_tokens: u64) -> Usage {
    usage_with_cached(input_tokens, output_tokens, input_tokens / 2)
}

fn usage_with_cached(input_tokens: u64, output_tokens: u64, cached_input_tokens: u64) -> Usage {
    let mut usage = Usage::new(input_tokens, output_tokens);
    usage.cache_read_tokens = cached_input_tokens;
    usage
}

fn tool_call(id: &str, name: &str, arguments: serde_json::Value) -> ToolCall {
    ToolCall::new(id, name, arguments)
}

fn response(
    text: Option<&str>,
    tool_calls: Vec<ToolCall>,
    input: u64,
    output: u64,
) -> ModelResponse {
    let usage = usage(input, output);
    ModelResponse {
        message: AssistantMessage {
            id: None,
            content: text
                .map(|text| vec![ContentBlock::Text(text.to_string())])
                .unwrap_or_default(),
            tool_calls,
            usage: Some(usage),
        },
        usage: Some(usage),
        finish_reason: None,
        raw: None,
        resolved_model: None,
        continue_turn: None,
    }
}

fn response_with_cached(
    text: Option<&str>,
    tool_calls: Vec<ToolCall>,
    input: u64,
    output: u64,
    cached: u64,
) -> ModelResponse {
    let usage = usage_with_cached(input, output, cached);
    ModelResponse {
        message: AssistantMessage {
            id: None,
            content: text
                .map(|text| vec![ContentBlock::Text(text.to_string())])
                .unwrap_or_default(),
            tool_calls,
            usage: Some(usage),
        },
        usage: Some(usage),
        finish_reason: None,
        raw: None,
        resolved_model: None,
        continue_turn: None,
    }
}

fn agent_config() -> AgentConfig {
    AgentConfig {
        max_tool_iterations: 4,
        max_history_messages: 12,
        ..AgentConfig::default()
    }
}

fn build_agent(workspace: &Path, provider: Arc<ScriptedModel>, agent_name: &str) -> Result<Agent> {
    build_agent_with_tools(workspace, provider, agent_name, vec![Box::new(EchoTool)])
}

fn build_agent_with_tools(
    workspace: &Path,
    provider: Arc<ScriptedModel>,
    agent_name: &str,
    tools: Vec<Box<dyn Tool>>,
) -> Result<Agent> {
    let mut agent = Agent::builder()
        .chat_model(provider)
        .tools(tools)
        .memory(Arc::new(StubMemory))
        .tool_dispatcher(Box::new(NativeToolDispatcher))
        .config(agent_config())
        .model_name("coverage-model".to_string())
        .temperature(0.0)
        .workspace_dir(workspace.to_path_buf())
        .workflows(Vec::new())
        .auto_save(false)
        .event_context("coverage-session", "coverage-channel")
        .agent_definition_name(agent_name)
        .omit_profile(true)
        .omit_memory_md(true)
        .build()?;
    agent.set_connected_integrations(Vec::new());
    Ok(agent)
}

fn parent_context(workspace: PathBuf, provider: Arc<ScriptedModel>) -> ParentExecutionContext {
    let tools: Vec<Box<dyn Tool>> = vec![Box::new(EchoTool)];
    let tool_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(
            provider,
        ),
        all_tools: Arc::new(tools),
        all_tool_specs: Arc::new(tool_specs),
        visible_tool_names: std::collections::HashSet::new(),
        subagent_tool_ceiling_names: std::collections::HashSet::new(),
        model_name: "coverage-model".to_string(),
        temperature: 0.0,
        workspace_dir: workspace,
        workspace_descriptor: None,
        memory: Arc::new(StubMemory),
        agent_config: agent_config(),
        workflows: Arc::new(Vec::new()),
        memory_context: Arc::new(Some("parent memory context".to_string())),
        session_id: "parent-session".to_string(),
        channel: "coverage-channel".to_string(),
        connected_integrations: Vec::new(),
        tool_call_format: ToolCallFormat::Native,
        session_key: "1700000000_parent".to_string(),
        session_parent_prefix: Some("root-chain".to_string()),
        on_progress: None,
        run_queue: None,
    }
}

fn coverage_definition() -> AgentDefinition {
    AgentDefinition {
        id: "coverage_worker".to_string(),
        when_to_use: "Used by raw integration coverage tests".to_string(),
        display_name: Some("Coverage Worker".to_string()),
        system_prompt: PromptSource::Inline("Answer only from deterministic test tools.".into()),
        omit_identity: true,
        omit_memory_context: false,
        omit_safety_preamble: true,
        omit_skills_catalog: true,
        omit_profile: true,
        omit_memory_md: true,
        model: Default::default(),
        temperature: 0.0,
        tools: ToolScope::Named(vec!["echo".to_string()]),
        disallowed_tools: Vec::new(),
        skill_filter: None,
        extra_tools: Vec::new(),
        max_iterations: 3,
        iteration_policy: Default::default(),
        max_result_chars: Some(18),
        max_turn_output_tokens: None,
        timeout_secs: None,
        sandbox_mode: SandboxMode::ReadOnly,
        background: false,
        trigger_memory_agent: Default::default(),
        tokenjuice_compression: AgentTokenjuiceCompression::Auto,
        subagents: Vec::new(),
        delegate_name: None,
        agent_tier: Default::default(),
        source: Default::default(),
        graph: Default::default(),
    }
}

fn transcript_jsonl_files(workspace: &Path) -> Vec<PathBuf> {
    let session_raw = workspace.join("session_raw");
    let mut files = std::fs::read_dir(session_raw)
        .ok()
        .into_iter()
        .flat_map(|entries| entries.filter_map(Result::ok))
        .map(|entry| entry.path())
        .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("jsonl"))
        .collect::<Vec<_>>();
    files.sort();
    files
}

#[tokio::test]
async fn agent_turn_executes_tools_persists_and_resumes_raw_transcript() -> Result<()> {
    let workspace = tempfile::tempdir()?;
    let provider = Arc::new(ScriptedModel::new(vec![
        response(
            Some("calling echo"),
            vec![tool_call("call-1", "echo", json!({"message": "alpha"}))],
            120,
            8,
        ),
        response(Some("final after echo"), Vec::new(), 132, 11),
    ]));
    let mut agent = build_agent(workspace.path(), provider.clone(), "coverage_main")?;

    let first = agent.turn("please call echo").await?;
    assert_eq!(first, "final after echo");
    assert!(agent.history().len() >= 4);

    let files = transcript_jsonl_files(workspace.path());
    assert_eq!(files.len(), 1, "expected one root transcript: {files:?}");
    let transcript = std::fs::read_to_string(&files[0])?;
    assert!(transcript.contains("\"agent\":\"coverage_main\""));
    assert!(transcript.contains("\"agent_id\":\"coverage_main\""));
    assert!(transcript.contains("\"agent_type\":\"root\""));
    assert!(transcript.contains("\"provider\":\"coverage-channel\""));
    assert!(transcript.contains("\"model\":\"coverage-model\""));
    assert!(transcript.contains("final after echo"));
    assert!(transcript.contains("\"input_tokens\":252"));
    assert!(workspace.path().join("sessions").exists());

    let resume_provider = Arc::new(ScriptedModel::new(vec![response(
        Some("resumed answer"),
        Vec::new(),
        64,
        6,
    )]));
    let mut resumed = build_agent(workspace.path(), resume_provider.clone(), "coverage_main")?;
    let second = resumed.turn("continue from transcript").await?;
    assert_eq!(second, "resumed answer");

    let requests = resume_provider.requests();
    let first_request = requests.first().expect("provider should be called");
    assert!(
        first_request
            .iter()
            .any(|message| matches!(message, Message::Assistant(_))
                && message.text() == "final after echo"),
        "resume request should include assistant message from prior transcript: {first_request:#?}"
    );

    Ok(())
}

#[tokio::test]
async fn run_subagent_filters_tools_runs_inner_loop_and_writes_child_transcript() -> Result<()> {
    let workspace = tempfile::tempdir()?;
    let provider = Arc::new(ScriptedModel::new(vec![
        response(
            Some("need a tool"),
            vec![tool_call("sub-call-1", "echo", json!({"message": "beta"}))],
            90,
            5,
        ),
        response(
            Some("subagent final response that will be capped by definition"),
            Vec::new(),
            101,
            7,
        ),
    ]));
    let parent = parent_context(workspace.path().to_path_buf(), provider.clone());
    let definition = coverage_definition();

    let outcome = with_parent_context(parent, async {
        run_subagent(
            &definition,
            "Use echo once and then summarize.",
            SubagentRunOptions {
                task_id: Some("coverage-task".to_string()),
                context: Some("caller supplied context".to_string()),
                ..SubagentRunOptions::default()
            },
        )
        .await
    })
    .await?;

    assert_eq!(outcome.agent_id, "coverage_worker");
    assert_eq!(outcome.iterations, 2);
    assert_eq!(outcome.output, "subagent final res\n[...truncated]");

    let requests = provider.requests();
    assert_eq!(requests.len(), 2);
    let first_request = requests.first().expect("subagent provider request");
    assert!(
        first_request
            .iter()
            .any(|message| matches!(message, Message::User(_))
                && message.text().contains("parent memory context")
                && message.text().contains("caller supplied context")),
        "subagent user prompt should merge parent and caller context: {first_request:#?}"
    );

    let files = transcript_jsonl_files(workspace.path());
    assert_eq!(files.len(), 1, "expected one child transcript: {files:?}");
    let stem = files[0]
        .file_stem()
        .and_then(|stem| stem.to_str())
        .unwrap_or_default();
    assert!(stem.starts_with("root-chain__1700000000_parent__"));
    assert!(stem.contains("coverage_worker"));

    let transcript = std::fs::read_to_string(&files[0])?;
    assert!(transcript.contains("\"agent\":\"coverage_worker\""));
    assert!(transcript.contains("echoed:beta"));
    assert!(transcript.contains("\"input_tokens\":191"));

    Ok(())
}

#[tokio::test]
async fn repeated_subagent_spawns_keep_cacheable_prefix_and_record_provider_cache_hit() -> Result<()>
{
    let workspace = tempfile::tempdir()?;
    let provider = Arc::new(ScriptedModel::new(vec![
        response_with_cached(Some("first"), Vec::new(), 100, 5, 0),
        response_with_cached(Some("second"), Vec::new(), 100, 5, 88),
    ]));
    let parent = parent_context(workspace.path().to_path_buf(), provider.clone());
    let definition = coverage_definition();

    let (first, second) = with_parent_context(parent, async {
        let first = run_subagent(
            &definition,
            "Answer the stable cache probe.",
            SubagentRunOptions {
                task_id: Some("cache-probe-a".to_string()),
                ..SubagentRunOptions::default()
            },
        )
        .await?;
        let second = run_subagent(
            &definition,
            "Answer the stable cache probe.",
            SubagentRunOptions {
                task_id: Some("cache-probe-b".to_string()),
                ..SubagentRunOptions::default()
            },
        )
        .await?;
        Ok::<_, openhuman_core::openhuman::agent::harness::SubagentRunError>((first, second))
    })
    .await?;

    assert_eq!(first.output, "first");
    assert_eq!(second.output, "second");

    let requests = provider.requests();
    assert_eq!(requests.len(), 2);
    let first_system = requests[0]
        .iter()
        .find(|message| matches!(message, Message::System(_)))
        .expect("first subagent request should include a system prompt");
    let second_system = requests[1]
        .iter()
        .find(|message| matches!(message, Message::System(_)))
        .expect("second subagent request should include a system prompt");
    assert_eq!(
        first_system.text(),
        second_system.text(),
        "repeated subagent spawns of the same definition must preserve the byte-identical \
         system prefix the backend can cache"
    );

    let transcripts = transcript_jsonl_files(workspace.path());
    assert_eq!(
        transcripts.len(),
        2,
        "each subagent run should persist its own raw transcript: {transcripts:?}"
    );
    let joined = transcripts
        .iter()
        .map(std::fs::read_to_string)
        .collect::<std::io::Result<Vec<_>>>()?
        .join("\n");
    assert!(
        joined.contains("\"cached_input_tokens\":88"),
        "provider-reported cached input tokens from the second child run should be preserved \
         in subagent transcript accounting:\n{joined}"
    );
    assert!(
        joined.contains("\"agent_type\":\"subagent\"")
            && joined.contains("\"provider\":\"subagent\"")
            && joined.contains("\"model\":\"coverage-model\""),
        "subagent transcript metadata should retain agent type, provider, and model:\n{joined}"
    );

    Ok(())
}

#[tokio::test]
async fn orchestrator_spawn_subagent_round_trip_streams_child_events_and_returns_result(
) -> Result<()> {
    let workspace = tempfile::tempdir()?;
    let agents_dir = workspace.path().join("agents");
    std::fs::create_dir_all(&agents_dir)?;
    std::fs::write(
        agents_dir.join("coverage_orchestrator.toml"),
        r#"
id = "coverage_orchestrator"
display_name = "Coverage Orchestrator"
when_to_use = "Deterministic parent agent used by harness cache tests."
temperature = 0.0
max_iterations = 3
agent_tier = "chat"
omit_identity = true
omit_memory_context = true
omit_safety_preamble = true
omit_skills_catalog = true
omit_profile = true
omit_memory_md = true

[system_prompt]
inline = "Delegate the cache probe and synthesize the result."

[subagents]
allowlist = ["cache_probe_child"]
"#,
    )?;
    std::fs::write(
        agents_dir.join("cache_probe_child.toml"),
        r#"
id = "cache_probe_child"
display_name = "Cache Probe Child"
when_to_use = "Deterministic child agent used by harness cache tests."
temperature = 0.0
max_iterations = 3
omit_identity = true
omit_memory_context = true
omit_safety_preamble = true
omit_skills_catalog = true
omit_profile = true
omit_memory_md = true

[system_prompt]
inline = "Answer the delegated cache probe directly."
"#,
    )?;
    let _ = AgentDefinitionRegistry::init_global(workspace.path());

    let child_answer = "child-cache-observation: prefix was reusable";
    let parent_final = "orchestrator final: child-cache-observation accepted";
    let provider = Arc::new(ScriptedModel::new(vec![
        response(
            Some("delegating to child"),
            vec![tool_call(
                "spawn-child-1",
                "spawn_subagent",
                json!({
                    "agent_id": "cache_probe_child",
                    "prompt": "Inspect whether the child turn can answer a cache probe.",
                    "context": "Parent observed request id cache-42.",
                    "blocking": true,
                }),
            )],
            140,
            9,
        ),
        response_with_cached(Some(child_answer), Vec::new(), 96, 7, 64),
        response(Some(parent_final), Vec::new(), 120, 10),
    ]));
    let mut agent = build_agent_with_tools(
        workspace.path(),
        provider.clone(),
        "coverage_orchestrator",
        vec![Box::new(SpawnSubagentTool::new())],
    )?;
    let (progress_tx, mut progress_rx) = tokio::sync::mpsc::channel(1024);
    agent.set_on_progress(Some(progress_tx));

    let answer = agent
        .turn("Ask a child agent for the cache observation, then respond.")
        .await?;
    assert_eq!(
        answer, parent_final,
        "orchestrator should synthesize after receiving the child result"
    );

    let mut progress = Vec::new();
    while let Ok(event) = progress_rx.try_recv() {
        progress.push(event);
    }
    assert!(
        progress.iter().any(|event| matches!(
            event,
            AgentProgress::SubagentSpawned { agent_id, task_id, .. }
                if agent_id == "cache_probe_child" && task_id.starts_with("sub-")
        )),
        "parent progress should announce the child spawn: {progress:#?}"
    );
    assert!(
        progress.iter().any(|event| matches!(
            event,
            AgentProgress::SubagentCompleted { agent_id, output_chars, .. }
                if agent_id == "cache_probe_child" && *output_chars == child_answer.chars().count()
        )),
        "parent progress should announce child completion: {progress:#?}"
    );

    let requests = provider.requests();
    assert_eq!(
        requests.len(),
        3,
        "expected parent request, child request, then parent synthesis request"
    );
    assert!(
        requests[0]
            .iter()
            .any(|message| matches!(message, Message::User(_))
                && message.text().contains("Ask a child agent")),
        "first request should be the orchestrator turn: {:#?}",
        requests[0]
    );
    assert!(
        requests[1]
            .iter()
            .any(|message| matches!(message, Message::System(_))
                && message.text().contains("Sub-agent Role Contract"))
            && requests[1]
                .iter()
                .any(|message| matches!(message, Message::User(_))
                    && message
                        .text()
                        .contains("Parent observed request id cache-42")),
        "second request should be the child subagent turn with parent-supplied context: {:#?}",
        requests[1]
    );
    assert!(
        requests[2]
            .iter()
            .any(|message| matches!(message, Message::Tool(_))
                && message.text().contains(child_answer)),
        "third request should return the child result to the orchestrator as a tool result: {:#?}",
        requests[2]
    );
    let transcripts = transcript_jsonl_files(workspace.path());
    let joined = transcripts
        .iter()
        .map(std::fs::read_to_string)
        .collect::<std::io::Result<Vec<_>>>()?
        .join("\n");
    assert!(
        joined.contains("\"agent\":\"cache_probe_child\"")
            && joined.contains("\"agent_id\":\"cache_probe_child\"")
            && joined.contains("\"agent_type\":\"subagent\"")
            && joined.contains("\"task_id\":\"sub-")
            && joined.contains(child_answer)
            && joined.contains("\"cached_input_tokens\":64"),
        "child transcript should be persisted alongside the parent turn:\n{joined}"
    );

    Ok(())
}