Genarrative/apps/ai-game-creator-shell/src-tauri/src/provider_handoff.rs

use std::path::{Path, PathBuf};

use platform_llm::{LlmProvider, LlmRunResponse, LlmTokenUsage};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::agent::{
    agent_runtime_json_sidecar_backup_path, read_agent_runtime_json_sidecar_with_max_bytes,
    redact_agent_runtime_project_paths, redact_secret_tokens,
    remove_agent_runtime_json_sidecar_backup, strip_llm_thinking_blocks,
    write_agent_runtime_json_sidecar_with_max_bytes,
};
use crate::provider_retry::{self, validate_identity, AgentRuntimeProviderRetryIdentity};
use crate::repository_context::redact_absolute_path_tokens;

pub(crate) const PROVIDER_HANDOFF_SCHEMA_VERSION: &str = "game-creator-provider-handoff.v1";

const PROVIDER_HANDOFF_RELATIVE_DIRECTORY: &str = ".agent/runtime/provider-handoffs";
const PROVIDER_HANDOFF_SIDECAR_MAX_BYTES: usize = 512 * 1024;
const PROVIDER_HANDOFF_RESPONSE_MAX_CHARS: usize = 256 * 1024;
const PROVIDER_HANDOFF_REQUEST_ID_MAX_CHARS: usize = 256;
const PROVIDER_HANDOFF_LABEL: &str = "Agent Runtime Provider 成功响应交接记录";

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct AgentRuntimeProviderHandoffResponse {
    provider: LlmProvider,
    model: String,
    text: String,
    finish_reason: Option<String>,
    response_id: Option<String>,
    usage: Option<LlmTokenUsage>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub(crate) struct AgentRuntimeProviderHandoffRecord {
    pub(crate) schema_version: String,
    pub(crate) identity: AgentRuntimeProviderRetryIdentity,
    pub(crate) provider_request_id: String,
    pub(crate) request_slot: String,
    pub(crate) attempt: u32,
    response: AgentRuntimeProviderHandoffResponse,
    pub(crate) response_fingerprint: String,
    pub(crate) created_at_ms: u64,
}

impl AgentRuntimeProviderHandoffRecord {
    pub(crate) fn to_llm_response(&self) -> LlmRunResponse {
        LlmRunResponse {
            provider: self.response.provider,
            model: self.response.model.clone(),
            text: self.response.text.clone(),
            finish_reason: self.response.finish_reason.clone(),
            response_id: self.response.response_id.clone(),
            usage: self.response.usage.clone(),
            tool_calls: Vec::new(),
        }
    }
}

pub(crate) fn supports_request_kind(request_kind: &str) -> bool {
    matches!(
        request_kind,
        "context-compaction" | "final-reply-context-compaction" | "final-reply"
    )
}

pub(crate) fn read_for_run_at(
    root: &Path,
    agent_id: &str,
    run_id: &str,
) -> Result<Option<AgentRuntimeProviderHandoffRecord>, String> {
    validate_path_identity(agent_id, run_id)?;
    let relative_path = provider_handoff_relative_path(agent_id, run_id);
    let Some(record) = read_agent_runtime_json_sidecar_with_max_bytes(
        root,
        &relative_path,
        PROVIDER_HANDOFF_LABEL,
        PROVIDER_HANDOFF_SIDECAR_MAX_BYTES,
    )?
    else {
        return Ok(None);
    };
    validate_record(&record)?;
    if record.identity.agent_id != agent_id || record.identity.run_id != run_id {
        return Err("Provider 成功响应交接记录与路径 Agent/run 身份冲突".to_string());
    }
    Ok(Some(record))
}

pub(crate) fn write_at(
    root: &Path,
    identity: &AgentRuntimeProviderRetryIdentity,
    request_slot: &str,
    attempt: u32,
    provider_request_id: &str,
    response: &LlmRunResponse,
) -> Result<AgentRuntimeProviderHandoffRecord, String> {
    validate_identity(identity)?;
    validate_provider_request_id(provider_request_id)?;
    if !supports_request_kind(&identity.request_kind) {
        return Err("当前 Provider requestKind 不允许持久交接成功响应".to_string());
    }
    if !response.tool_calls.is_empty() {
        return Err("Provider 成功响应交接禁止保存 tool calls".to_string());
    }
    let response_text = strip_llm_thinking_blocks(&response.text);
    let response_text = redact_secret_tokens(&response_text);
    let response_text = redact_agent_runtime_project_paths(
        root,
        &response_text,
        PROVIDER_HANDOFF_RESPONSE_MAX_CHARS,
    );
    let response = AgentRuntimeProviderHandoffResponse {
        provider: response.provider,
        model: response.model.clone(),
        text: redact_absolute_path_tokens(&response_text),
        finish_reason: response.finish_reason.clone(),
        response_id: response.response_id.clone(),
        usage: response.usage.clone(),
    };
    let response_fingerprint = response_fingerprint(&response)?;
    if let Some(existing) = read_for_run_at(root, &identity.agent_id, &identity.run_id)? {
        if existing.identity == *identity
            && existing.provider_request_id == provider_request_id
            && existing.request_slot == request_slot
            && existing.attempt == attempt
            && existing.response == response
            && existing.response_fingerprint == response_fingerprint
        {
            return Ok(existing);
        }
        return Err("Provider 成功响应交接记录内容冲突".to_string());
    }
    let record = AgentRuntimeProviderHandoffRecord {
        schema_version: PROVIDER_HANDOFF_SCHEMA_VERSION.to_string(),
        identity: identity.clone(),
        provider_request_id: provider_request_id.to_string(),
        request_slot: request_slot.to_string(),
        attempt,
        response,
        response_fingerprint,
        created_at_ms: provider_retry::now_ms(),
    };
    validate_record(&record)?;
    let relative_path = provider_handoff_relative_path(&identity.agent_id, &identity.run_id);
    write_agent_runtime_json_sidecar_with_max_bytes(
        root,
        &relative_path,
        PROVIDER_HANDOFF_LABEL,
        &record,
        PROVIDER_HANDOFF_SIDECAR_MAX_BYTES,
    )?;
    let persisted = read_for_run_at(root, &identity.agent_id, &identity.run_id)?
        .ok_or_else(|| "Provider 成功响应交接记录写入后不存在".to_string())?;
    if persisted != record {
        return Err("Provider 成功响应交接记录写入后内容冲突".to_string());
    }
    Ok(persisted)
}

pub(crate) fn remove_matching_at(
    root: &Path,
    identity: &AgentRuntimeProviderRetryIdentity,
) -> Result<(), String> {
    validate_identity(identity)?;
    let Some(record) = read_for_run_at(root, &identity.agent_id, &identity.run_id)? else {
        return Ok(());
    };
    if record.identity != *identity {
        return Err("Provider 成功响应交接记录身份冲突".to_string());
    }
    remove_at(root, &identity.agent_id, &identity.run_id)
}

pub(crate) fn remove_consumed_context_at(
    root: &Path,
    agent_id: &str,
    run_id: &str,
    request_kind: &str,
    base_request_slot: &str,
) -> Result<(), String> {
    let Some(record) = read_for_run_at(root, agent_id, run_id)? else {
        return Ok(());
    };
    if record.identity.request_kind != request_kind
        || record.identity.base_request_slot != base_request_slot
    {
        return Err("已消费的上下文压缩与 Provider 成功响应交接身份冲突".to_string());
    }
    remove_at(root, agent_id, run_id)
}

pub(crate) fn remove_at(root: &Path, agent_id: &str, run_id: &str) -> Result<(), String> {
    validate_path_identity(agent_id, run_id)?;
    let path = provider_handoff_path(root, agent_id, run_id);
    let backup_path = agent_runtime_json_sidecar_backup_path(&path);
    remove_agent_runtime_json_sidecar_backup(&backup_path, PROVIDER_HANDOFF_LABEL)?;
    remove_agent_runtime_json_sidecar_backup(&path, PROVIDER_HANDOFF_LABEL)
}

fn validate_record(record: &AgentRuntimeProviderHandoffRecord) -> Result<(), String> {
    if record.schema_version != PROVIDER_HANDOFF_SCHEMA_VERSION {
        return Err(format!(
            "不支持的 Provider 成功响应交接版本：{}",
            record.schema_version
        ));
    }
    validate_identity(&record.identity)?;
    if !supports_request_kind(&record.identity.request_kind) {
        return Err("Provider 成功响应交接 requestKind 无效".to_string());
    }
    validate_provider_request_id(&record.provider_request_id)?;
    let expected_slot = request_slot_for_attempt(&record.identity, record.attempt);
    if record.request_slot != expected_slot {
        return Err("Provider 成功响应交接 requestSlot/attempt 无效".to_string());
    }
    validate_short_text("model", &record.response.model, 256, false)?;
    if record.response.text.chars().count() > PROVIDER_HANDOFF_RESPONSE_MAX_CHARS {
        return Err(format!(
            "Provider 成功响应交接正文超过 {PROVIDER_HANDOFF_RESPONSE_MAX_CHARS} 字符上限"
        ));
    }
    if let Some(finish_reason) = record.response.finish_reason.as_deref() {
        validate_short_text("finishReason", finish_reason, 80, true)?;
    }
    if let Some(response_id) = record.response.response_id.as_deref() {
        validate_short_text("responseId", response_id, 256, true)?;
    }
    if record.response_fingerprint != response_fingerprint(&record.response)? {
        return Err("Provider 成功响应交接 responseFingerprint 不匹配".to_string());
    }
    if record.created_at_ms == 0 {
        return Err("Provider 成功响应交接 createdAtMs 无效".to_string());
    }
    Ok(())
}

fn request_slot_for_attempt(identity: &AgentRuntimeProviderRetryIdentity, attempt: u32) -> String {
    if attempt == 0 {
        identity.base_request_slot.clone()
    } else {
        format!("{}-transient-{attempt}", identity.base_request_slot)
    }
}

fn validate_short_text(
    label: &str,
    value: &str,
    max_chars: usize,
    allow_empty: bool,
) -> Result<(), String> {
    if (!allow_empty && value.trim().is_empty())
        || value.chars().count() > max_chars
        || value.chars().any(char::is_control)
    {
        return Err(format!("Provider 成功响应交接 {label} 无效"));
    }
    Ok(())
}

fn validate_provider_request_id(value: &str) -> Result<(), String> {
    validate_short_text(
        "Provider lifecycle requestId",
        value,
        PROVIDER_HANDOFF_REQUEST_ID_MAX_CHARS,
        false,
    )?;
    let fingerprint = value
        .strip_prefix("provider-request-")
        .ok_or_else(|| "Provider 成功响应交接 Provider lifecycle requestId 无效".to_string())?;
    if fingerprint.len() != 64 || !fingerprint.bytes().all(|byte| byte.is_ascii_hexdigit()) {
        return Err("Provider 成功响应交接 Provider lifecycle requestId 无效".to_string());
    }
    Ok(())
}

fn response_fingerprint(response: &AgentRuntimeProviderHandoffResponse) -> Result<String, String> {
    let bytes = serde_json::to_vec(response)
        .map_err(|error| format!("序列化 Provider 成功响应指纹失败：{error}"))?;
    Ok(format!("{:x}", Sha256::digest(bytes)))
}

fn validate_path_identity(agent_id: &str, run_id: &str) -> Result<(), String> {
    if agent_id.trim().is_empty() || run_id.trim().is_empty() {
        return Err("Provider 成功响应交接路径的 Agent/run 身份不能为空".to_string());
    }
    Ok(())
}

fn provider_handoff_relative_path(agent_id: &str, run_id: &str) -> String {
    format!(
        "{PROVIDER_HANDOFF_RELATIVE_DIRECTORY}/{}/{}.json",
        path_key(agent_id),
        path_key(run_id)
    )
}

fn provider_handoff_path(root: &Path, agent_id: &str, run_id: &str) -> PathBuf {
    root.join(provider_handoff_relative_path(agent_id, run_id))
}

fn path_key(value: &str) -> String {
    format!("{:x}", Sha256::digest(value.as_bytes()))
}

#[cfg(test)]
mod tests {
    use std::fs;

    use tempfile::tempdir;

    use super::*;

    fn identity(request_kind: &str) -> AgentRuntimeProviderRetryIdentity {
        AgentRuntimeProviderRetryIdentity {
            project_id: "project-provider-handoff".to_string(),
            agent_id: "design-director".to_string(),
            task_id: "design-director".to_string(),
            session_id: "agent-session-design-director".to_string(),
            run_id: "provider-handoff-run".to_string(),
            source: "agent-background-task".to_string(),
            goal_id: None,
            goal_revision: 0,
            goal_snapshot_fingerprint: String::new(),
            applied_steer_cursor: 0,
            request_kind: request_kind.to_string(),
            base_request_slot: "final-reply-loop-1-revision-0".to_string(),
            request_fingerprint: "1".repeat(64),
            provider_config_fingerprint: "2".repeat(64),
            web_search_enabled: false,
            allow_idle_context_compaction: false,
        }
    }

    fn response(text: &str) -> LlmRunResponse {
        LlmRunResponse {
            provider: LlmProvider::OpenAiCompatible,
            model: "handoff-model".to_string(),
            text: text.to_string(),
            finish_reason: Some("stop".to_string()),
            response_id: Some("response-handoff".to_string()),
            usage: Some(LlmTokenUsage {
                prompt_tokens: 11,
                completion_tokens: 7,
                total_tokens: 18,
            }),
            tool_calls: Vec::new(),
        }
    }

    fn provider_request_id(marker: char) -> String {
        format!("provider-request-{}", marker.to_string().repeat(64))
    }

    #[test]
    fn provider_handoff_round_trips_idempotently_and_removes_both_copies() {
        let project = tempdir().expect("provider handoff project");
        let identity = identity("final-reply");
        let actual_request_id = provider_request_id('a');
        let expected = response("最终回复已持久交接。");
        let first = write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &actual_request_id,
            &expected,
        )
        .expect("write provider handoff");
        let second = write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &actual_request_id,
            &expected,
        )
        .expect("rewrite same provider handoff");
        assert_eq!(first, second);
        assert_eq!(second.provider_request_id, actual_request_id);
        assert_eq!(second.request_slot, identity.base_request_slot);
        assert_eq!(second.to_llm_response(), expected);
        assert_eq!(
            read_for_run_at(project.path(), &identity.agent_id, &identity.run_id)
                .expect("read provider handoff"),
            Some(second)
        );

        remove_matching_at(project.path(), &identity).expect("remove provider handoff");
        assert!(
            read_for_run_at(project.path(), &identity.agent_id, &identity.run_id)
                .expect("read removed provider handoff")
                .is_none()
        );
    }

    #[test]
    fn provider_handoff_rejects_tool_calls_and_content_conflicts() {
        let project = tempdir().expect("provider handoff project");
        let handoff_identity = identity("final-reply");
        let actual_request_id = provider_request_id('a');
        let expected = response("原始回复");
        write_at(
            project.path(),
            &handoff_identity,
            &handoff_identity.base_request_slot,
            0,
            &actual_request_id,
            &expected,
        )
        .expect("write original provider handoff");
        let error = write_at(
            project.path(),
            &handoff_identity,
            &handoff_identity.base_request_slot,
            0,
            &actual_request_id,
            &response("冲突回复"),
        )
        .expect_err("conflicting provider handoff must fail");
        assert!(error.contains("内容冲突"));

        let mut with_tool_call = response("工具回复");
        with_tool_call.tool_calls.push(platform_llm::LlmToolCall {
            id: "call-1".to_string(),
            name: "unsafe".to_string(),
            arguments: "{}".to_string(),
        });
        let other = identity("final-reply-context-compaction");
        let error = write_at(
            project.path(),
            &other,
            &other.base_request_slot,
            0,
            &provider_request_id('b'),
            &with_tool_call,
        )
        .expect_err("tool calls must not enter provider handoff");
        assert!(error.contains("tool calls"));
    }

    #[test]
    fn provider_handoff_rejects_request_id_conflicts() {
        let project = tempdir().expect("provider handoff project");
        let identity = identity("final-reply");
        let expected = response("同一回复");
        write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &provider_request_id('a'),
            &expected,
        )
        .expect("write original provider handoff");

        let error = write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &provider_request_id('b'),
            &expected,
        )
        .expect_err("conflicting requestId must fail");
        assert!(error.contains("内容冲突"));
    }

    #[test]
    fn provider_handoff_rejects_invalid_request_ids() {
        let project = tempdir().expect("provider handoff project");
        let identity = identity("final-reply");
        let invalid_request_ids = [
            String::new(),
            "provider-request-invalid\nidentity".to_string(),
            "provider-request-not-a-sha256".to_string(),
            "x".repeat(PROVIDER_HANDOFF_REQUEST_ID_MAX_CHARS + 1),
        ];

        for request_id in invalid_request_ids {
            let error = write_at(
                project.path(),
                &identity,
                &identity.base_request_slot,
                0,
                &request_id,
                &response("不会落盘"),
            )
            .expect_err("invalid requestId must fail");
            assert!(error.contains("requestId"));
        }
        assert!(
            read_for_run_at(project.path(), &identity.agent_id, &identity.run_id)
                .expect("read absent provider handoff")
                .is_none()
        );
    }

    #[test]
    fn provider_handoff_strips_thinking_before_persisting_response() {
        let project = tempdir().expect("provider handoff project");
        let identity = identity("final-reply");
        let private_thinking = "不得进入 handoff 的内部推理";
        let public_response = "这是公开回复。";
        let expected = response(&format!(
            "<think>{private_thinking}</think>\n{public_response}"
        ));

        let record = write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &provider_request_id('c'),
            &expected,
        )
        .expect("write provider handoff without thinking");
        assert_eq!(record.to_llm_response().text, public_response);

        let persisted = fs::read_to_string(provider_handoff_path(
            project.path(),
            &identity.agent_id,
            &identity.run_id,
        ))
        .expect("read persisted provider handoff");
        assert!(!persisted.contains(private_thinking));
        assert!(!persisted.to_ascii_lowercase().contains("<think>"));
        assert!(persisted.contains(public_response));
    }

    #[test]
    fn provider_handoff_rejects_unknown_request_kind() {
        let project = tempdir().expect("provider handoff project");
        let identity = identity("tool-plan");
        let error = write_at(
            project.path(),
            &identity,
            &identity.base_request_slot,
            0,
            &provider_request_id('d'),
            &response("tool plan"),
        )
        .expect_err("tool plan handoff must fail");
        assert!(error.contains("不允许"));
    }
}