tiny-agent 0.1.0

use super::*;
use crate::{
    checkpoint::CheckpointStorage,
    error::StorageError,
    middleware::{Middleware, MiddlewareChain, MiddlewareCtx, ToolDecision},
    providers::{
        FinishReason, LLMRequest, LlmProvider, ProviderError, ProviderStreamChunk, TokenUsage,
    },
    sandbox::{ExecOutput, HostSandbox, Sandbox, SandboxError},
    shared::{ContentBlock, Message, Role, ToolCall, UserInteraction},
    state::AgentState,
    tools::{ToolOutcome, ToolRegistry, register_default_tools, register_spawn_agent},
    trajectory::{NoopSink, TrajectoryEvent, TrajectoryEventKind, TrajectorySink},
    transcript::{Transcript, TranscriptStorage},
};
use async_trait::async_trait;
use futures_util::{StreamExt, stream};
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::json;
use std::{
    collections::HashMap,
    future::pending,
    path::PathBuf,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    time::Duration,
};
use tokio_util::sync::CancellationToken;

#[derive(Default)]
struct InMemoryStorage {
    transcripts: Mutex<HashMap<String, Transcript>>,
    checkpoints: Mutex<HashMap<String, String>>,
    fail_transcript_append: AtomicBool,
}

#[async_trait]
impl TranscriptStorage for InMemoryStorage {
    async fn save_transcript(
        &self,
        session_id: &str,
        transcript: &Transcript,
    ) -> Result<(), StorageError> {
        self.transcripts
            .lock()
            .unwrap()
            .insert(session_id.to_string(), transcript.clone());
        Ok(())
    }

    async fn get_transcript(&self, session_id: &str) -> Result<Option<Transcript>, StorageError> {
        Ok(self.transcripts.lock().unwrap().get(session_id).cloned())
    }

    async fn append_messages(
        &self,
        session_id: &str,
        messages: Vec<Message>,
    ) -> Result<(), StorageError> {
        if self.fail_transcript_append.load(Ordering::SeqCst) {
            return Err(StorageError::unavailable("transcript append failed"));
        }
        let mut transcripts = self.transcripts.lock().unwrap();
        let transcript = transcripts
            .entry(session_id.to_string())
            .or_insert_with(|| Transcript::new(session_id));
        transcript.messages.extend(messages);
        Ok(())
    }

    async fn delete_transcript(&self, session_id: &str) -> Result<(), StorageError> {
        self.transcripts.lock().unwrap().remove(session_id);
        Ok(())
    }
}

#[async_trait]
impl CheckpointStorage for InMemoryStorage {
    async fn save_checkpoint(&self, session_id: &str, agent: &Agent) -> Result<(), StorageError> {
        let checkpoint =
            serde_json::to_string(agent).map_err(|e| StorageError::serde(e.to_string()))?;
        self.checkpoints
            .lock()
            .unwrap()
            .insert(session_id.to_string(), checkpoint);
        Ok(())
    }

    async fn get_checkpoint(&self, session_id: &str) -> Result<Option<Agent>, StorageError> {
        self.checkpoints
            .lock()
            .unwrap()
            .get(session_id)
            .map(|checkpoint| {
                serde_json::from_str(checkpoint).map_err(|e| StorageError::serde(e.to_string()))
            })
            .transpose()
    }

    async fn delete_checkpoint(&self, session_id: &str) -> Result<(), StorageError> {
        self.checkpoints.lock().unwrap().remove(session_id);
        Ok(())
    }
}

struct EmptyProvider;

#[async_trait]
impl LlmProvider for EmptyProvider {
    fn name(&self) -> &str {
        "empty"
    }

    async fn chat_completion(
        &self,
        _req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        Ok(stream::iter(vec![Ok(ProviderStreamChunk::Done {
            finish_reason: FinishReason::Stop,
            usage: Some(TokenUsage {
                input_tokens: None,
                output_tokens: None,
                total_tokens: None,
            }),
        })])
        .boxed())
    }
}

struct FollowupProvider;

#[async_trait]
impl LlmProvider for FollowupProvider {
    fn name(&self) -> &str {
        "followup"
    }

    async fn chat_completion(
        &self,
        req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        let has_tool_result = req.messages.iter().any(|msg| {
            msg.blocks
                .iter()
                .any(|block| matches!(block, ContentBlock::ToolResult { .. }))
        });
        let user_answer = req.messages.iter().rev().find_map(|msg| {
            if msg.role != Role::User {
                return None;
            }
            msg.blocks.iter().find_map(|block| match block {
                ContentBlock::Text { text } => Some(text.clone()),
                _ => None,
            })
        });

        let chunks = if has_tool_result && user_answer.as_deref() == Some("/work/note.txt") {
            vec![
                Ok(ProviderStreamChunk::ContentDelta(
                    "ok, inspecting /work/note.txt".to_string(),
                )),
                Ok(ProviderStreamChunk::Done {
                    finish_reason: FinishReason::Stop,
                    usage: None,
                }),
            ]
        } else {
            vec![
                Ok(ProviderStreamChunk::ToolCallDelta {
                    index: 0,
                    id: Some("ask-followup".to_string()),
                    name: Some("askFollowupQuestion".to_string()),
                    arguments_delta: json!({
                        "question": "Which path should I inspect?",
                        "options": ["/work/note.txt", "Cargo.toml"]
                    })
                    .to_string(),
                }),
                Ok(ProviderStreamChunk::Done {
                    finish_reason: FinishReason::ToolCalls,
                    usage: None,
                }),
            ]
        };

        Ok(stream::iter(chunks).boxed())
    }
}

struct ToolThenStopProvider;

#[async_trait]
impl LlmProvider for ToolThenStopProvider {
    fn name(&self) -> &str {
        "tool_then_stop"
    }

    async fn chat_completion(
        &self,
        req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        let has_tool_result = req.messages.iter().any(|msg| {
            msg.blocks
                .iter()
                .any(|block| matches!(block, ContentBlock::ToolResult { .. }))
        });
        let chunks = if has_tool_result {
            vec![
                Ok(ProviderStreamChunk::ContentDelta("child done".to_string())),
                Ok(ProviderStreamChunk::Done {
                    finish_reason: FinishReason::Stop,
                    usage: None,
                }),
            ]
        } else {
            vec![
                Ok(ProviderStreamChunk::ToolCallDelta {
                    index: 0,
                    id: Some("mark-call".to_string()),
                    name: Some("mark_sandbox".to_string()),
                    arguments_delta: json!({ "command": "child-write" }).to_string(),
                }),
                Ok(ProviderStreamChunk::Done {
                    finish_reason: FinishReason::ToolCalls,
                    usage: None,
                }),
            ]
        };
        Ok(stream::iter(chunks).boxed())
    }
}

struct PermanentProvider;

#[async_trait]
impl LlmProvider for PermanentProvider {
    fn name(&self) -> &str {
        "permanent"
    }

    async fn chat_completion(
        &self,
        _req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        Err(ProviderError::permanent("bad api key"))
    }
}

struct RecordingProvider {
    request: Arc<Mutex<Option<LLMRequest>>>,
}

#[async_trait]
impl LlmProvider for RecordingProvider {
    fn name(&self) -> &str {
        "recording"
    }

    async fn chat_completion(
        &self,
        req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        *self.request.lock().unwrap() = Some(req);
        Ok(stream::iter(vec![Ok(ProviderStreamChunk::Done {
            finish_reason: FinishReason::Stop,
            usage: None,
        })])
        .boxed())
    }
}

struct DeltaProvider;

#[async_trait]
impl LlmProvider for DeltaProvider {
    fn name(&self) -> &str {
        "delta"
    }

    async fn chat_completion(
        &self,
        _req: LLMRequest,
    ) -> Result<
        futures_util::stream::BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>,
        ProviderError,
    > {
        Ok(stream::iter(vec![
            Ok(ProviderStreamChunk::ThinkingDelta("thinking".to_string())),
            Ok(ProviderStreamChunk::ContentDelta("hello".to_string())),
            Ok(ProviderStreamChunk::ContentDelta(" world".to_string())),
            Ok(ProviderStreamChunk::Done {
                finish_reason: FinishReason::Stop,
                usage: None,
            }),
        ])
        .boxed())
    }
}

fn runtime(storage: Arc<InMemoryStorage>, tools: ToolRegistry) -> AgentRunTime {
    runtime_with_fail_limit(storage, tools, 1)
}

fn runtime_with_fail_limit(
    storage: Arc<InMemoryStorage>,
    tools: ToolRegistry,
    max_consecutive_fail_count: u32,
) -> AgentRunTime {
    AgentRunTime {
        provider: Arc::new(EmptyProvider),
        max_iter: 10,
        max_consecutive_fail_count,
        model: "test-model".to_string(),
        system_prompt: None,
        temperature: None,
        max_tokens: None,
        checkpoint_storage: storage.clone(),
        transcript_storage: storage,
        sandbox: Arc::new(HostSandbox::new()),
        tools,
        skills: None,
        prompt_composer: super::prompt::PromptComposer,
        middleware: MiddlewareChain::new(),
        trajectory: Arc::new(NoopSink),
        retry_backoff: Duration::ZERO,
        tool_timeout: Duration::from_secs(120),
        session_locks: tokio::sync::Mutex::new(HashMap::new()),
    }
}

fn test_skill_root() -> PathBuf {
    let root = std::env::temp_dir().join(format!("tiny_agent_skills_{}", uuid::Uuid::new_v4()));
    let skill_dir = root.join("docs");
    std::fs::create_dir_all(skill_dir.join("references")).unwrap();
    std::fs::write(
        skill_dir.join("SKILL.md"),
        "---\nname: Documents\ndescription: Work with document files\n---\n\n# Documents\n\nmain skill instructions",
    )
    .unwrap();
    std::fs::write(
        skill_dir.join("references").join("detail.md"),
        "detail instructions",
    )
    .unwrap();
    root
}

#[derive(Deserialize, JsonSchema)]
struct EchoArgs {
    value: String,
}

#[derive(Deserialize, JsonSchema)]
struct EmptyArgs {}

#[derive(Deserialize, JsonSchema)]
struct FailArgs {}

#[derive(Deserialize, JsonSchema)]
struct MarkSandboxArgs {
    command: String,
}

#[derive(Default)]
struct RecordingSink {
    events: Mutex<Vec<TrajectoryEvent>>,
}

#[async_trait]
impl TrajectorySink for RecordingSink {
    async fn emit(&self, event: TrajectoryEvent) {
        self.events.lock().unwrap().push(event);
    }
}

/// 护栏中间件:任何工具调用都直接拒绝,不让其真正执行。
struct DenyTools;

#[async_trait]
impl Middleware for DenyTools {
    fn name(&self) -> &str {
        "deny_tools"
    }

    async fn before_tool(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _call: &ToolCall,
    ) -> Result<ToolDecision, String> {
        Ok(ToolDecision::ShortCircuit {
            output: vec![ContentBlock::Text {
                text: "denied by guardrail".to_string(),
            }],
            is_error: true,
        })
    }
}

struct FailAfterTool;

#[async_trait]
impl Middleware for FailAfterTool {
    fn name(&self) -> &str {
        "fail_after_tool"
    }

    async fn after_tool(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _call: &ToolCall,
        _result: &mut ContentBlock,
    ) -> Result<(), String> {
        Err("redaction failed".to_string())
    }
}

#[tokio::test]
async fn skills_are_composed_into_system_without_tool_catalog() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "skill-prompt";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let request = Arc::new(Mutex::new(None));
    let mut tools = ToolRegistry::new();
    tools.register(
        "visible_tool",
        "this tool description should not be in system",
        true,
        |_s, _a: EmptyArgs| async move { Ok(json!({}).into()) },
    );

    let skill_root = test_skill_root();
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(RecordingProvider {
            request: request.clone(),
        }))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .skill_dir(skill_root)
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Success(_)));
    let req = request.lock().unwrap().clone().unwrap();
    let system = req.system.unwrap();
    assert!(system.contains("docs"));
    assert!(system.contains("readSkill"));
    assert!(!system.contains("this tool description should not be in system"));
    assert!(req.tools.iter().any(|tool| tool.name == "visible_tool"));
    assert!(req.tools.iter().any(|tool| tool.name == "readSkill"));
}

#[tokio::test]
async fn builder_model_options_flow_into_provider_request() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "model-options";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let request = Arc::new(Mutex::new(None));
    let sink = Arc::new(RecordingSink::default());

    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(RecordingProvider {
            request: request.clone(),
        }))
        .model("test-model")
        .temperature(0.7)
        .max_tokens(2048)
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .trajectory(sink.clone())
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Success(_)));
    let req = request.lock().unwrap().clone().unwrap();
    assert_eq!(req.temperature, Some(0.7));
    assert_eq!(req.max_tokens, Some(2048));

    let events = sink.events.lock().unwrap();
    let trajectory_request = events
        .iter()
        .find_map(|event| match &event.kind {
            TrajectoryEventKind::ReasoningStarted { request } => Some(request),
            _ => None,
        })
        .expect("reasoning_started should include provider request");
    assert_eq!(trajectory_request["model"], "test-model");
    assert!((trajectory_request["temperature"].as_f64().unwrap() - 0.7).abs() < 0.00001);
    assert_eq!(trajectory_request["max_tokens"], 2048);
    assert_eq!(trajectory_request["tools"], serde_json::json!([]));
    assert_eq!(trajectory_request["messages"], serde_json::json!([]));
}

#[tokio::test]
async fn runtime_config_system_prompt_flows_into_provider_request() {
    let storage = Arc::new(InMemoryStorage::default());
    let request = Arc::new(Mutex::new(None));
    let mut tools = ToolRegistry::new();
    register_default_tools(&mut tools);
    let config = AgentRuntimeConfig::new(
        Arc::new(RecordingProvider {
            request: request.clone(),
        }),
        "config-model",
        storage.clone(),
        storage.clone(),
        Arc::new(HostSandbox::new()),
    )
    .with_tools(tools)
    .with_system_prompt("You are configured from AgentRuntimeConfig.");
    let runtime = AgentRunTime::from_config(config).unwrap();
    let session = runtime.create_session().await.unwrap();
    let ready = runtime.submit(&session, "hello").await.unwrap();

    runtime.run(ready, CancellationToken::new()).await.unwrap();

    let req = request.lock().unwrap().clone().unwrap();
    assert_eq!(
        req.system.as_deref(),
        Some("You are configured from AgentRuntimeConfig.")
    );
}

#[tokio::test]
async fn builder_system_prompt_flows_into_provider_request() {
    let storage = Arc::new(InMemoryStorage::default());
    let request = Arc::new(Mutex::new(None));
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(RecordingProvider {
            request: request.clone(),
        }))
        .model("builder-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .with_system_prompt("You are configured from AgentRunTimeBuilder.")
        .build()
        .unwrap();
    let session = runtime.create_session().await.unwrap();
    let ready = runtime.submit(&session, "hello").await.unwrap();

    runtime.run(ready, CancellationToken::new()).await.unwrap();

    let req = request.lock().unwrap().clone().unwrap();
    assert_eq!(
        req.system.as_deref(),
        Some("You are configured from AgentRunTimeBuilder.")
    );
}

#[tokio::test]
async fn trajectory_does_not_persist_streaming_deltas() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "no-delta-trajectory";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let sink = Arc::new(RecordingSink::default());
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(DeltaProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .trajectory(sink.clone())
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Success(AgentState {
            final_resp: Some(ref text),
            ..
        }) if text == "hello world"
    ));
    let events = sink.events.lock().unwrap();
    assert!(
        events
            .iter()
            .any(|event| matches!(event.kind, TrajectoryEventKind::AssistantMessage { .. }))
    );
}

#[tokio::test]
async fn read_skill_tool_reads_content_and_emits_trajectory() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "read-skill";
    let skill_root = test_skill_root();
    let mut transcript = Transcript::new(session);
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "read-skill-call".to_string(),
            name: "readSkill".to_string(),
            input: json!({ "skill_id": "docs" }),
        }],
    });
    storage.save_transcript(session, &transcript).await.unwrap();
    let sink = Arc::new(RecordingSink::default());
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .skill_dir(skill_root)
        .trajectory(sink.clone())
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .execute_pending_tools(
            &AgentState::new(Some(session.to_string())),
            ToolRunState::new(vec![ToolCall {
                call_id: "read-skill-call".to_string(),
                tool_name: "readSkill".to_string(),
                arguments: json!({ "skill_id": "docs" }),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Ready(_)));
    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert!(transcript.messages.iter().any(|message| {
        message.blocks.iter().any(|block| matches!(
            block,
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error: false,
            } if tool_use_id == "read-skill-call"
                && matches!(&content[0], ContentBlock::Text { text } if text.contains("main skill instructions"))
                && matches!(&content[0], ContentBlock::Text { text } if text.contains("references/detail.md"))
        ))
    }));
}

#[tokio::test]
async fn read_skill_tool_reads_resource_from_discovered_directory() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "read-skill-resource";
    let skill_root = test_skill_root();
    let mut transcript = Transcript::new(session);
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "read-skill-resource-call".to_string(),
            name: "readSkill".to_string(),
            input: json!({
                "skill_id": "docs",
                "path": "references/detail.md"
            }),
        }],
    });
    storage.save_transcript(session, &transcript).await.unwrap();

    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .skill_dir(skill_root)
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .execute_pending_tools(
            &AgentState::new(Some(session.to_string())),
            ToolRunState::new(vec![ToolCall {
                call_id: "read-skill-resource-call".to_string(),
                tool_name: "readSkill".to_string(),
                arguments: json!({
                    "skill_id": "docs",
                    "path": "references/detail.md"
                }),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Ready(_)));
    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert!(transcript.messages.iter().any(|message| {
        message.blocks.iter().any(|block| matches!(
            block,
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error: false,
            } if tool_use_id == "read-skill-resource-call"
                && matches!(&content[0], ContentBlock::Text { text } if text.contains("detail instructions"))
        ))
    }));
}

#[tokio::test]
async fn middleware_short_circuits_tool_and_emits_trajectory() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("guard", &Transcript::new("guard"))
        .await
        .unwrap();

    let ran = Arc::new(AtomicBool::new(false));
    let ran_in_tool = ran.clone();
    let mut tools = ToolRegistry::new();
    tools.register(
        "touch",
        "marks ran",
        false,
        move |_sandbox, _args: EmptyArgs| {
            let ran_in_tool = ran_in_tool.clone();
            async move {
                ran_in_tool.store(true, Ordering::SeqCst);
                Ok(json!({}).into())
            }
        },
    );

    let sink = Arc::new(RecordingSink::default());
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .middleware(Arc::new(DenyTools))
        .trajectory(sink.clone())
        .build()
        .unwrap();

    let state = AgentState::new(Some("guard".to_string()));
    let tool_calls = vec![ToolCall {
        call_id: "call-g".to_string(),
        tool_name: "touch".to_string(),
        arguments: json!({}),
    }];
    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Ready(_)));
    assert!(!ran.load(Ordering::SeqCst), "短路后工具本体不应执行");

    let events = sink.events.lock().unwrap();
    assert!(events.iter().any(|e| matches!(
        &e.kind,
        TrajectoryEventKind::ToolStarted { name, .. } if name == "touch"
    )));
    assert!(events.iter().any(|e| matches!(
        &e.kind,
        TrajectoryEventKind::ToolCompleted { is_error: true, .. }
    )));

    // 落库的应是中间件给的拒绝结果,而非工具真实输出。
    let transcript = storage.get_transcript("guard").await.unwrap().unwrap();
    let last = transcript.messages.last().unwrap();
    assert!(matches!(
        &last.blocks[0],
        ContentBlock::ToolResult { is_error: true, .. }
    ));
}

#[tokio::test]
async fn tool_execution_appends_tool_result_and_returns_ready() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("session-1");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "call-1".to_string(),
            name: "echo_value".to_string(),
            input: json!({ "value": "hello" }),
        }],
    });
    storage
        .save_transcript("session-1", &transcript)
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register(
        "echo_value",
        "echoes a value",
        false,
        |_sandbox, args: EchoArgs| async move { Ok(json!({ "echo": args.value }).into()) },
    );
    let runtime = runtime(storage.clone(), tools);

    let state = AgentState::new(Some("session-1".to_string()));
    let tool_calls = vec![ToolCall {
        call_id: "call-1".to_string(),
        tool_name: "echo_value".to_string(),
        arguments: json!({ "value": "hello" }),
    }];
    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Ready(_)));
    let transcript = storage.get_transcript("session-1").await.unwrap().unwrap();
    let last = transcript.messages.last().unwrap();
    assert_eq!(last.role, Role::Tool);
    assert!(matches!(
        &last.blocks[0],
        ContentBlock::ToolResult {
            tool_use_id,
            is_error: false,
            ..
        } if tool_use_id == "call-1"
    ));
}

#[tokio::test]
async fn read_only_tools_in_a_batch_run_concurrently() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("session-1");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![
            ContentBlock::ToolUse {
                id: "call-1".to_string(),
                name: "r1".to_string(),
                input: json!({}),
            },
            ContentBlock::ToolUse {
                id: "call-2".to_string(),
                name: "r2".to_string(),
                input: json!({}),
            },
        ],
    });
    storage
        .save_transcript("session-1", &transcript)
        .await
        .unwrap();

    // 两个只读工具各自在同一个 Barrier(2) 上等待:只有当两者**同时在飞**时
    // barrier 才会放行。若调度退回串行,第一个工具会永远等不到第二个 —— 靠外层
    // timeout 把这种回归暴露成失败而不是挂死。
    let barrier = Arc::new(tokio::sync::Barrier::new(2));
    let mut tools = ToolRegistry::new();
    let b1 = barrier.clone();
    tools.register("r1", "readonly 1", true, move |_s, _a: EmptyArgs| {
        let b = b1.clone();
        async move {
            b.wait().await;
            Ok(json!({ "tool": "r1" }).into())
        }
    });
    let b2 = barrier.clone();
    tools.register("r2", "readonly 2", true, move |_s, _a: EmptyArgs| {
        let b = b2.clone();
        async move {
            b.wait().await;
            Ok(json!({ "tool": "r2" }).into())
        }
    });
    let runtime = runtime(storage.clone(), tools);

    let state = AgentState::new(Some("session-1".to_string()));
    let tool_calls = vec![
        ToolCall {
            call_id: "call-1".to_string(),
            tool_name: "r1".to_string(),
            arguments: json!({}),
        },
        ToolCall {
            call_id: "call-2".to_string(),
            tool_name: "r2".to_string(),
            arguments: json!({}),
        },
    ];
    let result = tokio::time::timeout(
        Duration::from_secs(5),
        runtime.execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        ),
    )
    .await
    .expect("readonly tools did not run concurrently (serialized -> barrier deadlock)")
    .unwrap();

    assert!(matches!(result, Agent::Ready(_)));
    // 段内并发,但结果仍按输入顺序落库。
    let transcript = storage.get_transcript("session-1").await.unwrap().unwrap();
    let tool_results: Vec<&str> = transcript
        .messages
        .iter()
        .filter_map(|m| match m.blocks.first() {
            Some(ContentBlock::ToolResult { tool_use_id, .. }) => Some(tool_use_id.as_str()),
            _ => None,
        })
        .collect();
    assert_eq!(tool_results, vec!["call-1", "call-2"]);
}

#[tokio::test]
async fn tool_execution_can_be_cancelled() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("session-2");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "call-2".to_string(),
            name: "hang".to_string(),
            input: json!({}),
        }],
    });
    storage
        .save_transcript("session-2", &transcript)
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register(
        "hang",
        "never returns",
        false,
        |_sandbox, _args: EmptyArgs| async move { pending::<Result<ToolOutcome, String>>().await },
    );
    let runtime = runtime(storage.clone(), tools);

    let token = CancellationToken::new();
    let child = token.child_token();
    let state = AgentState::new(Some("session-2".to_string()));
    let tool_calls = vec![ToolCall {
        call_id: "call-2".to_string(),
        tool_name: "hang".to_string(),
        arguments: json!({}),
    }];
    let task = tokio::spawn(async move {
        runtime
            .execute_pending_tools(&state, ToolRunState::new(tool_calls), child)
            .await
            .unwrap()
    });

    token.cancel();

    let result = tokio::time::timeout(std::time::Duration::from_secs(1), task)
        .await
        .expect("tool execution should stop after cancellation")
        .unwrap();

    assert!(matches!(result, Agent::Interrupted(_)));
    assert!(storage.get_checkpoint("session-2").await.unwrap().is_some());
}

#[tokio::test]
async fn tool_execution_times_out_and_records_tool_error() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("tool-timeout");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "call-timeout".to_string(),
            name: "hang".to_string(),
            input: json!({}),
        }],
    });
    storage
        .save_transcript("tool-timeout", &transcript)
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register(
        "hang",
        "never returns",
        false,
        |_sandbox, _args: EmptyArgs| async move { pending::<Result<ToolOutcome, String>>().await },
    );
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .retry_backoff(Duration::ZERO)
        .tool_timeout(Duration::from_millis(10))
        .build()
        .unwrap();

    let result = runtime
        .execute_pending_tools(
            &AgentState::new(Some("tool-timeout".to_string())),
            ToolRunState::new(vec![ToolCall {
                call_id: "call-timeout".to_string(),
                tool_name: "hang".to_string(),
                arguments: json!({}),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, .. }) if kind == FailureKind::Tool
    ));
    let transcript = storage
        .get_transcript("tool-timeout")
        .await
        .unwrap()
        .unwrap();
    assert!(transcript.messages.iter().any(|message| {
        message.blocks.iter().any(|block| matches!(
            block,
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error: true,
            } if tool_use_id == "call-timeout"
                && matches!(&content[0], ContentBlock::Text { text } if text.contains("timed out"))
        ))
    }));
}

#[tokio::test]
async fn tool_execution_failure_enters_fail_state() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("session-3", &Transcript::new("session-3"))
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register(
        "fail",
        "fails",
        false,
        |_sandbox, _args: FailArgs| async move { Err::<ToolOutcome, _>("boom".to_string()) },
    );
    let runtime = runtime(storage, tools);
    let state = AgentState::new(Some("session-3".to_string()));
    let tool_calls = vec![ToolCall {
        call_id: "call-3".to_string(),
        tool_name: "fail".to_string(),
        arguments: json!({}),
    }];

    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, message, .. })
            if kind == FailureKind::Tool && message == "boom"
    ));
}

#[tokio::test]
async fn tool_result_storage_failure_returns_storage_fail_state() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("storage-tool", &Transcript::new("storage-tool"))
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register("ok", "ok", false, |_sandbox, _args: EmptyArgs| async move {
        Ok(json!({ "ok": true }).into())
    });
    let runtime = runtime(storage.clone(), tools);
    storage.fail_transcript_append.store(true, Ordering::SeqCst);

    let result = runtime
        .execute_pending_tools(
            &AgentState::new(Some("storage-tool".to_string())),
            ToolRunState::new(vec![ToolCall {
                call_id: "call-storage".to_string(),
                tool_name: "ok".to_string(),
                arguments: json!({}),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, message, .. })
            if kind == FailureKind::Storage && message.contains("transcript append failed")
    ));
}

#[tokio::test]
async fn after_tool_middleware_error_records_executed_but_result_unavailable() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("after-tool", &Transcript::new("after-tool"))
        .await
        .unwrap();

    let ran = Arc::new(AtomicBool::new(false));
    let ran_in_tool = ran.clone();
    let mut tools = ToolRegistry::new();
    tools.register("touch", "touches", false, move |_s, _a: EmptyArgs| {
        let ran_in_tool = ran_in_tool.clone();
        async move {
            ran_in_tool.store(true, Ordering::SeqCst);
            Ok(json!({ "ok": true }).into())
        }
    });
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .middleware(Arc::new(FailAfterTool))
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .execute_pending_tools(
            &AgentState::new(Some("after-tool".to_string())),
            ToolRunState::new(vec![ToolCall {
                call_id: "call-after".to_string(),
                tool_name: "touch".to_string(),
                arguments: json!({}),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(ran.load(Ordering::SeqCst), "工具本体已经执行过");
    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, .. }) if kind == FailureKind::Middleware
    ));
    let transcript = storage.get_transcript("after-tool").await.unwrap().unwrap();
    assert!(transcript.messages.iter().any(|message| {
        message.blocks.iter().any(|block| {
            matches!(
                block,
                ContentBlock::ToolResult {
                    tool_use_id,
                    content,
                    is_error: true,
                } if tool_use_id == "call-after"
                    && matches!(
                        &content[0],
                        ContentBlock::Text { text }
                            if text.contains("tool executed but result unavailable")
                    )
            )
        })
    }));
}

#[tokio::test]
async fn handle_fail_retries_then_terminates() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("session-4", &Transcript::new("session-4"))
        .await
        .unwrap();
    let runtime = runtime_with_fail_limit(storage.clone(), ToolRegistry::new(), 2);
    let state = AgentState::new(Some("session-4".to_string()));

    // 工具失败的 tool_result 由 execute_pending_tools 负责落库,handle_fail 只计数、
    // 不再重复回应,所以 transcript 不应被它新增任何消息。
    let failure = AgentFailure::tool(
        ToolCall {
            call_id: "call-4".to_string(),
            tool_name: "fail".to_string(),
            arguments: json!({}),
        },
        "tool exploded",
    );
    let Agent::Ready(ctx) = runtime.handle_fail(state, failure).await.unwrap() else {
        panic!("first failure should retry");
    };
    assert_eq!(ctx.consecutive_fail_count, 1);
    assert!(
        storage
            .get_transcript("session-4")
            .await
            .unwrap()
            .unwrap()
            .messages
            .is_empty(),
        "工具失败不应由 handle_fail 写入 transcript"
    );

    // 非工具失败(仍在上限内):补一条 system 面包屑供模型理解并重试。
    let failure = AgentFailure::new(FailureKind::Provider, "still broken");
    let Agent::Ready(ctx) = runtime.handle_fail(ctx, failure).await.unwrap() else {
        panic!("second failure under limit should retry");
    };
    assert_eq!(ctx.consecutive_fail_count, 2);
    let transcript = storage.get_transcript("session-4").await.unwrap().unwrap();
    assert_eq!(transcript.messages.last().unwrap().role, Role::System);

    // 再失败一次越过上限 → 终态 Fail。
    let failure = AgentFailure::new(FailureKind::Provider, "still broken");
    let result = runtime.handle_fail(ctx, failure).await.unwrap();
    assert!(matches!(
        result,
        Agent::Fail(
            AgentState {
                consecutive_fail_count: 3,
                ..
            },
            AgentFailure { kind, .. },
        ) if kind == FailureKind::Provider
    ));
}

#[tokio::test]
async fn handle_fail_storage_breadcrumb_failure_is_retryable() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("storage-fail", &Transcript::new("storage-fail"))
        .await
        .unwrap();
    let runtime = runtime_with_fail_limit(storage.clone(), ToolRegistry::new(), 2);
    storage.fail_transcript_append.store(true, Ordering::SeqCst);

    let result = runtime
        .handle_fail(
            AgentState::new(Some("storage-fail".to_string())),
            AgentFailure::new(FailureKind::Provider, "provider flaked"),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Ready(AgentState {
            consecutive_fail_count: 2,
            ..
        })
    ));
}

#[tokio::test]
async fn session_lock_table_prunes_stale_entries() {
    let storage = Arc::new(InMemoryStorage::default());
    let runtime = runtime(storage, ToolRegistry::new());

    let _ = runtime.submit("lock-a", "hello").await.unwrap();
    assert_eq!(runtime.session_locks.lock().await.len(), 1);

    let _ = runtime.submit("lock-b", "hello").await.unwrap();
    assert_eq!(
        runtime.session_locks.lock().await.len(),
        1,
        "stale weak lock entries should be pruned on the next lookup"
    );
}

#[tokio::test]
async fn permanent_provider_error_does_not_retry() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "provider-permanent";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(PermanentProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(ToolRegistry::new())
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Fail(
            AgentState {
                consecutive_fail_count: 0,
                ..
            },
            AgentFailure { kind, message, .. },
        ) if kind == FailureKind::ProviderPermanent && message.contains("bad api key")
    ));
}

#[tokio::test]
async fn reasoning_repairs_orphan_tool_use_before_provider_request() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "orphan-repair";
    let mut transcript = Transcript::new(session);
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![ContentBlock::ToolUse {
            id: "orphan-call".to_string(),
            name: "bash".to_string(),
            input: json!({ "command": "echo hi" }),
        }],
    });
    storage.save_transcript(session, &transcript).await.unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Success(_)));
    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert!(transcript.messages.iter().any(|message| {
        message.blocks.iter().any(|block| matches!(
            block,
            ContentBlock::ToolResult {
                tool_use_id,
                is_error: true,
                content,
            } if tool_use_id == "orphan-call"
                && matches!(&content[0], ContentBlock::Text { text } if text.contains("previous crash"))
        ))
    }));
}

/// #2 不变量:一批工具里中途失败,assistant 的每个 tool_use 都必须被回应,不留孤儿。
#[tokio::test]
async fn batch_failure_answers_every_tool_use() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("batch");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![
            ContentBlock::ToolUse {
                id: "c0".to_string(),
                name: "ok".to_string(),
                input: json!({}),
            },
            ContentBlock::ToolUse {
                id: "c1".to_string(),
                name: "boom".to_string(),
                input: json!({}),
            },
            ContentBlock::ToolUse {
                id: "c2".to_string(),
                name: "ok".to_string(),
                input: json!({}),
            },
        ],
    });
    storage.save_transcript("batch", &transcript).await.unwrap();

    let mut tools = ToolRegistry::new();
    tools.register("ok", "ok", false, |_s, _a: EmptyArgs| async move {
        Ok(json!({}).into())
    });
    tools.register("boom", "fails", false, |_s, _a: EmptyArgs| async move {
        Err::<ToolOutcome, _>("kaboom".to_string())
    });
    let runtime = runtime(storage.clone(), tools);

    let state = AgentState::new(Some("batch".to_string()));
    let tool_calls = vec![
        ToolCall {
            call_id: "c0".to_string(),
            tool_name: "ok".to_string(),
            arguments: json!({}),
        },
        ToolCall {
            call_id: "c1".to_string(),
            tool_name: "boom".to_string(),
            arguments: json!({}),
        },
        ToolCall {
            call_id: "c2".to_string(),
            tool_name: "ok".to_string(),
            arguments: json!({}),
        },
    ];
    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();
    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, .. }) if kind == FailureKind::Tool
    ));

    let transcript = storage.get_transcript("batch").await.unwrap().unwrap();
    let answered: std::collections::HashSet<String> = transcript
        .messages
        .iter()
        .flat_map(|m| m.blocks.iter())
        .filter_map(|b| match b {
            ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.clone()),
            _ => None,
        })
        .collect();
    assert_eq!(
        answered,
        ["c0", "c1", "c2"].iter().map(|s| s.to_string()).collect(),
        "三个 tool_use 都必须各有一条 tool_result"
    );
}

#[tokio::test]
async fn tool_needs_user_interaction_writes_result_and_checkpoints_waiting() {
    let storage = Arc::new(InMemoryStorage::default());
    let mut transcript = Transcript::new("ask");
    transcript.append(Message {
        role: Role::Assistant,
        blocks: vec![
            ContentBlock::ToolUse {
                id: "ask-1".to_string(),
                name: "ask_user".to_string(),
                input: json!({}),
            },
            ContentBlock::ToolUse {
                id: "later-1".to_string(),
                name: "later".to_string(),
                input: json!({}),
            },
        ],
    });
    storage.save_transcript("ask", &transcript).await.unwrap();

    let later_ran = Arc::new(AtomicBool::new(false));
    let later_ran_in_tool = later_ran.clone();
    let mut tools = ToolRegistry::new();
    tools.register(
        "ask_user",
        "asks the user",
        false,
        |_s, _a: EmptyArgs| async move {
            Ok(ToolOutcome::NeedsUserInteraction(UserInteraction {
                kind: "followup_question".to_string(),
                payload: json!({
                    "prompt": "Pick a file",
                    "source": "test",
                }),
            }))
        },
    );
    tools.register(
        "later",
        "should not run yet",
        false,
        move |_s, _a: EmptyArgs| {
            let later_ran_in_tool = later_ran_in_tool.clone();
            async move {
                later_ran_in_tool.store(true, Ordering::SeqCst);
                Ok(json!({}).into())
            }
        },
    );
    let runtime = runtime(storage.clone(), tools);

    let state = AgentState::new(Some("ask".to_string()));
    let tool_calls = vec![
        ToolCall {
            call_id: "ask-1".to_string(),
            tool_name: "ask_user".to_string(),
            arguments: json!({}),
        },
        ToolCall {
            call_id: "later-1".to_string(),
            tool_name: "later".to_string(),
            arguments: json!({}),
        },
    ];

    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        &result,
        Agent::WaitingForUser(_, interaction)
            if interaction.kind == "followup_question"
                && interaction.payload["prompt"] == "Pick a file"
    ));
    assert!(!later_ran.load(Ordering::SeqCst));
    assert!(matches!(
        storage.get_checkpoint("ask").await.unwrap(),
        Some(Agent::WaitingForUser(_, interaction))
            if interaction.kind == "followup_question"
                && interaction.payload["prompt"] == "Pick a file"
    ));

    let transcript = storage.get_transcript("ask").await.unwrap().unwrap();
    let tool_results: Vec<_> = transcript
        .messages
        .iter()
        .flat_map(|m| m.blocks.iter())
        .filter_map(|b| match b {
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error,
            } => Some((tool_use_id.as_str(), content, *is_error)),
            _ => None,
        })
        .collect();
    assert_eq!(tool_results.len(), 2);
    assert_eq!(tool_results[0].0, "ask-1");
    assert!(!tool_results[0].2);
    assert!(matches!(
        &tool_results[0].1[0],
        ContentBlock::Text { text } if text.contains("\"type\":\"needs_user_interaction\"")
    ));
    assert_eq!(tool_results[1].0, "later-1");
    assert!(tool_results[1].2, "后续未执行工具应被补 skipped result");
}

#[tokio::test]
async fn default_ask_followup_question_enters_waiting_for_user() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("default-ask", &Transcript::new("default-ask"))
        .await
        .unwrap();
    let mut tools = ToolRegistry::new();
    register_default_tools(&mut tools);
    let sink = Arc::new(RecordingSink::default());
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .trajectory(sink.clone())
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let state = AgentState::new(Some("default-ask".to_string()));
    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(vec![ToolCall {
                call_id: "ask-call".to_string(),
                tool_name: "askFollowupQuestion".to_string(),
                arguments: json!({
                    "question": "Which path should I inspect?",
                    "options": ["/work/note.txt", "Cargo.toml"]
                }),
            }]),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::WaitingForUser(_, interaction)
            if interaction.kind == "followup_question"
                && interaction.payload["question"] == "Which path should I inspect?"
                && interaction.payload["options"] == json!(["/work/note.txt", "Cargo.toml"])
    ));
    assert!(matches!(
        storage.get_checkpoint("default-ask").await.unwrap(),
        Some(Agent::WaitingForUser(_, interaction))
            if interaction.kind == "followup_question"
                && interaction.payload["question"] == "Which path should I inspect?"
                && interaction.payload["options"] == json!(["/work/note.txt", "Cargo.toml"])
    ));
    assert!(sink.events.lock().unwrap().iter().any(|event| matches!(
        &event.kind,
        TrajectoryEventKind::WaitingForUser { interaction }
            if interaction.kind == "followup_question"
                && interaction.payload["question"] == "Which path should I inspect?"
                && interaction.payload["options"] == json!(["/work/note.txt", "Cargo.toml"])
    )));
}

#[tokio::test]
async fn submit_answers_waiting_for_user_with_plain_user_message() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "waiting-submit";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    storage
        .save_checkpoint(
            session,
            &Agent::WaitingForUser(
                AgentState::new(Some(session.to_string())),
                UserInteraction {
                    kind: "approval".to_string(),
                    payload: json!({ "prompt": "Confirm?" }),
                },
            ),
        )
        .await
        .unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    let ready = runtime.submit(session, "yes").await.unwrap();

    assert!(matches!(ready, Agent::Ready(_)));
    assert!(storage.get_checkpoint(session).await.unwrap().is_none());
    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert!(matches!(
        transcript.messages.last(),
        Some(Message {
            role: Role::User,
            blocks,
        }) if matches!(&blocks[0], ContentBlock::Text { text } if text == "yes")
    ));
}

#[tokio::test]
async fn user_answer_after_followup_can_continue_to_success() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "followup-e2e";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let mut tools = ToolRegistry::new();
    register_default_tools(&mut tools);
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(FollowupProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .retry_backoff(Duration::ZERO)
        .build()
        .unwrap();

    let ready = runtime.submit(session, "please ask me").await.unwrap();
    let waiting = runtime.run(ready, CancellationToken::new()).await.unwrap();
    assert!(matches!(
        waiting,
        Agent::WaitingForUser(_, interaction)
            if interaction.kind == "followup_question"
                && interaction.payload["question"] == "Which path should I inspect?"
    ));
    assert!(matches!(
        storage.get_checkpoint(session).await.unwrap(),
        Some(Agent::WaitingForUser(_, _))
    ));

    let ready = runtime.submit(session, "/work/note.txt").await.unwrap();
    assert!(storage.get_checkpoint(session).await.unwrap().is_none());
    let result = runtime.run(ready, CancellationToken::new()).await.unwrap();

    assert!(matches!(
        result,
        Agent::Success(AgentState {
            final_resp: Some(resp),
            ..
        }) if resp == "ok, inspecting /work/note.txt"
    ));
    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert!(transcript.messages.iter().any(|msg| {
        msg.role == Role::User
            && msg.blocks.iter().any(
                |block| matches!(block, ContentBlock::Text { text } if text == "/work/note.txt"),
            )
    }));
    assert!(
        storage.get_checkpoint(session).await.unwrap().is_none(),
        "成功终态后 waiting checkpoint 应保持清理"
    );
}

/// 记录 `save()` 调用次数的 sandbox,用来验证失败路径也持久化。
struct CountingSandbox {
    saves: Arc<AtomicUsize>,
}

#[async_trait]
impl Sandbox for CountingSandbox {
    async fn open(self: Arc<Self>, _session_id: &str) -> Result<Arc<dyn Sandbox>, SandboxError> {
        Ok(self)
    }

    async fn execute(&self, _command: &str) -> Result<ExecOutput, SandboxError> {
        Ok(ExecOutput {
            stdout: String::new(),
            stderr: String::new(),
            exit_code: Some(0),
        })
    }

    async fn save(&self) -> Result<(), SandboxError> {
        self.saves.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }
}

#[derive(Default)]
struct SharedParentSandbox {
    opens: Arc<Mutex<Vec<String>>>,
    commands: Arc<Mutex<Vec<String>>>,
}

#[async_trait]
impl Sandbox for SharedParentSandbox {
    async fn open(self: Arc<Self>, session_id: &str) -> Result<Arc<dyn Sandbox>, SandboxError> {
        self.opens.lock().unwrap().push(session_id.to_string());
        Ok(self)
    }

    async fn execute(&self, command: &str) -> Result<ExecOutput, SandboxError> {
        self.commands.lock().unwrap().push(command.to_string());
        Ok(ExecOutput {
            stdout: "ok".to_string(),
            stderr: String::new(),
            exit_code: Some(0),
        })
    }
}

/// 中途失败也要持久化已成功工具的沙箱状态,否则它们的成功结果会与实际状态发散。
#[tokio::test]
async fn mid_batch_failure_persists_sandbox() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("persist", &Transcript::new("persist"))
        .await
        .unwrap();

    let mut tools = ToolRegistry::new();
    tools.register("ok", "ok", false, |_s, _a: EmptyArgs| async move {
        Ok(json!({}).into())
    });
    tools.register("boom", "fails", false, |_s, _a: EmptyArgs| async move {
        Err::<ToolOutcome, _>("kaboom".to_string())
    });

    let saves = Arc::new(AtomicUsize::new(0));
    let runtime = AgentRunTimeBuilder::new()
        .provider(Arc::new(EmptyProvider))
        .model("test-model")
        .checkpoint_storage(storage.clone())
        .transcript_storage(storage.clone())
        .sandbox(Arc::new(CountingSandbox {
            saves: saves.clone(),
        }))
        .tools(tools)
        .build()
        .unwrap();

    let state = AgentState::new(Some("persist".to_string()));
    let tool_calls = vec![
        ToolCall {
            call_id: "c0".to_string(),
            tool_name: "ok".to_string(),
            arguments: json!({}),
        },
        ToolCall {
            call_id: "c1".to_string(),
            tool_name: "boom".to_string(),
            arguments: json!({}),
        },
    ];
    let result = runtime
        .execute_pending_tools(
            &state,
            ToolRunState::new(tool_calls),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, .. }) if kind == FailureKind::Tool
    ));
    assert!(
        saves.load(Ordering::SeqCst) >= 1,
        "中途失败也应 save() 持久化已成功工具的沙箱状态"
    );
}

/// #1:迭代到达上限后,run 应直接进 max_iter 终态,而不是继续推理。
#[tokio::test]
async fn run_enforces_max_iter() {
    let storage = Arc::new(InMemoryStorage::default());
    storage
        .save_transcript("iter", &Transcript::new("iter"))
        .await
        .unwrap();
    let mut runtime = runtime(storage.clone(), ToolRegistry::new());
    runtime.max_iter = 1;
    let mut state = AgentState::new(Some("iter".to_string()));
    state.steps_count = 1; // 已达上限

    let result = runtime
        .run(Agent::Ready(state), CancellationToken::new())
        .await
        .unwrap();
    assert!(matches!(
        result,
        Agent::Fail(_, AgentFailure { kind, .. }) if kind == FailureKind::MaxIter
    ));
}

/// #4:submit 追加用户消息后,run 能端到端跑到 Success。
#[tokio::test]
async fn submit_then_run_reaches_success() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "chat";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    let ready = runtime.submit(session, "hello").await.unwrap();
    assert!(matches!(ready, Agent::Ready(_)));

    let result = runtime.run(ready, CancellationToken::new()).await.unwrap();
    assert!(matches!(result, Agent::Success(_)));

    let transcript = storage.get_transcript(session).await.unwrap().unwrap();
    assert_eq!(transcript.messages.first().unwrap().role, Role::User);
}

#[tokio::test]
async fn spawn_agent_tool_runs_child_runtime_and_returns_child_session() {
    let storage = Arc::new(InMemoryStorage::default());
    let child_config = AgentRuntimeConfig::new(
        Arc::new(DeltaProvider),
        "child-model",
        storage.clone(),
        storage.clone(),
        Arc::new(HostSandbox::new()),
    )
    .with_tools(ToolRegistry::new())
    .with_system_prompt("You are the child agent.")
    .with_trajectory(Arc::new(NoopSink));
    let mut tools = ToolRegistry::new();
    register_spawn_agent(&mut tools, child_config);

    let outcome = tools
        .call(
            "spawn_agent",
            json!({ "task": "say hello" }),
            Arc::new(HostSandbox::new()),
        )
        .await
        .unwrap();

    let ToolOutcome::Completed(value) = outcome else {
        panic!("spawn_agent should complete");
    };
    assert_eq!(value["status"], "success");
    assert_eq!(value["final_resp"], "hello world");
    let child_session_id = value["child_session_id"].as_str().unwrap();
    let transcript = storage
        .get_transcript(child_session_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(transcript.messages.first().unwrap().role, Role::User);
}

#[tokio::test]
async fn spawn_agent_uses_parent_sandbox() {
    let storage = Arc::new(InMemoryStorage::default());
    let root_sandbox = Arc::new(SharedParentSandbox::default());
    let parent_sandbox = Arc::new(SharedParentSandbox::default());
    let mut child_tools = ToolRegistry::new();
    child_tools.register(
        "mark_sandbox",
        "record command in sandbox",
        false,
        |sandbox, args: MarkSandboxArgs| async move {
            let output = sandbox
                .execute(&args.command)
                .await
                .map_err(|e| e.to_string())?;
            Ok(json!({ "stdout": output.stdout }).into())
        },
    );
    let child_config = AgentRuntimeConfig::new(
        Arc::new(ToolThenStopProvider),
        "child-model",
        storage.clone(),
        storage.clone(),
        root_sandbox.clone(),
    )
    .with_tools(child_tools);
    let mut tools = ToolRegistry::new();
    register_spawn_agent(&mut tools, child_config);

    let outcome = tools
        .call(
            "spawn_agent",
            json!({ "task": "mark parent sandbox" }),
            parent_sandbox.clone(),
        )
        .await
        .unwrap();

    let ToolOutcome::Completed(value) = outcome else {
        panic!("spawn_agent should complete");
    };
    assert_eq!(value["status"], "success");
    assert_eq!(value["final_resp"], "child done");
    assert_eq!(
        parent_sandbox.commands.lock().unwrap().as_slice(),
        ["child-write"]
    );
    assert!(
        root_sandbox.commands.lock().unwrap().is_empty(),
        "child config sandbox should be replaced by the parent sandbox"
    );
    assert_eq!(
        parent_sandbox.opens.lock().unwrap().len(),
        2,
        "create_session and tool execution should both open the shared parent sandbox"
    );
}

/// run 跑到终态(Success)后,应清掉该会话残留的 checkpoint,避免日后被误恢复。
#[tokio::test]
async fn run_clears_checkpoint_on_terminal() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "cleanup";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    // 预置一份上一轮中断遗留的 checkpoint。
    storage
        .save_checkpoint(
            session,
            &Agent::Interrupted(AgentState::new(Some(session.to_string()))),
        )
        .await
        .unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    let result = runtime
        .run(
            Agent::Ready(AgentState::new(Some(session.to_string()))),
            CancellationToken::new(),
        )
        .await
        .unwrap();

    assert!(matches!(result, Agent::Success(_)));
    assert!(
        storage.get_checkpoint(session).await.unwrap().is_none(),
        "终态后 checkpoint 必须被清理"
    );
}

/// resume 能从已存档的 checkpoint 续跑;没有 checkpoint 时返回 None。
#[tokio::test]
async fn resume_continues_from_checkpoint() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "resumable";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    // 没有 checkpoint:resume 返回 None。
    assert!(
        runtime
            .resume(session, CancellationToken::new())
            .await
            .unwrap()
            .is_none()
    );

    // 存一份中断 checkpoint,resume 应续上并跑到终态,且事后清理。
    storage
        .save_checkpoint(
            session,
            &Agent::Interrupted(AgentState::new(Some(session.to_string()))),
        )
        .await
        .unwrap();
    let resumed = runtime
        .resume(session, CancellationToken::new())
        .await
        .unwrap();
    assert!(matches!(resumed, Some(Agent::Success(_))));
    assert!(storage.get_checkpoint(session).await.unwrap().is_none());
}

/// submit 表示用户发新消息 → 放弃尚未恢复的中断态,删掉残留 checkpoint。
#[tokio::test]
async fn submit_discards_stale_checkpoint() {
    let storage = Arc::new(InMemoryStorage::default());
    let session = "interject";
    storage
        .save_transcript(session, &Transcript::new(session))
        .await
        .unwrap();
    storage
        .save_checkpoint(
            session,
            &Agent::Interrupted(AgentState::new(Some(session.to_string()))),
        )
        .await
        .unwrap();
    let runtime = runtime(storage.clone(), ToolRegistry::new());

    runtime.submit(session, "new message").await.unwrap();

    assert!(
        storage.get_checkpoint(session).await.unwrap().is_none(),
        "submit 后陈旧 checkpoint 必须被丢弃"
    );
}