PicoBot/src/command/handlers/load_task_messages.rs

use crate::command::Command;
use crate::command::context::CommandContext;
use crate::command::handler::{CommandHandler, CommandMetadata};
use crate::command::handlers::list_topics::TopicTokenStats;
use crate::command::response::{CommandError, CommandResponse};
use crate::storage::SessionStore;
use crate::tools::task::repository::TaskRepository;
use crate::tools::task::types::{TaskSession, TaskSessionState};
use async_trait::async_trait;
use std::sync::Arc;

pub struct LoadTaskMessagesCommandHandler {
    task_repository: Arc<dyn TaskRepository>,
    store: Arc<SessionStore>,
}

impl LoadTaskMessagesCommandHandler {
    pub fn new(task_repository: Arc<dyn TaskRepository>, store: Arc<SessionStore>) -> Self {
        Self {
            task_repository,
            store,
        }
    }
}

#[async_trait]
impl CommandHandler for LoadTaskMessagesCommandHandler {
    fn can_handle(&self, cmd: &Command) -> bool {
        matches!(cmd, Command::LoadTaskMessages { .. })
    }

    fn metadata(&self) -> Option<CommandMetadata> {
        Some(CommandMetadata {
            name: "load_task_messages",
            description: "加载子智能体任务的消息历史",
            usage: "/load_task_messages <task_id>",
        })
    }

    async fn handle(
        &self,
        cmd: Command,
        ctx: CommandContext,
    ) -> Result<CommandResponse, CommandError> {
        match cmd {
            Command::LoadTaskMessages { task_id } => {
                handle_load_task_messages(self, task_id, ctx).await
            }
            _ => unreachable!(),
        }
    }
}

async fn handle_load_task_messages(
    handler: &LoadTaskMessagesCommandHandler,
    task_id: String,
    ctx: CommandContext,
) -> Result<CommandResponse, CommandError> {
    tracing::debug!(
        task_id = %task_id,
        request_id = %ctx.request_id,
        "LoadTaskMessages: looking up task"
    );

    // 1. Try in-memory repository first
    let task = match handler.task_repository.load_task_session(&task_id).await {
        Ok(Some(task)) => {
            tracing::debug!(
                task_id = %task.id,
                session_id = %task.session_id,
                state = ?task.state,
                "LoadTaskMessages: task found in memory"
            );
            Some(task)
        }
        Ok(None) => {
            tracing::debug!(
                task_id = %task_id,
                "LoadTaskMessages: task not in memory, searching database"
            );
            // 2. Fall back to database (survives restarts)
            reconstruct_task_from_db(&handler.store, &task_id)?
        }
        Err(e) => {
            tracing::error!(
                task_id = %task_id,
                error = %e,
                "LoadTaskMessages: repository error during lookup"
            );
            return Err(CommandError::new("LOAD_TASK_ERROR", e.to_string()));
        }
    };

    let task = task.ok_or_else(|| {
        tracing::warn!(
            task_id = %task_id,
            "LoadTaskMessages: task not found in repository or database"
        );
        CommandError::new("TASK_NOT_FOUND", format!("Task not found: {}", task_id))
    })?;

    let status = format!("{:?}", task.state).to_lowercase();

    // 查询子代理 session 的 token 统计（按 session_id 精确匹配，不过滤 sub:%）
    let token_stats = handler
        .store
        .get_session_token_stats(&task.session_id)
        .map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))?
        .map(|s| TopicTokenStats {
            prompt_tokens: s.prompt_tokens,
            completion_tokens: s.completion_tokens,
            total_tokens: s.total_tokens,
            cached_tokens: s.cached_tokens,
            last_prompt_tokens: s.last_prompt_tokens,
            context_window_tokens: s.context_window_tokens.unwrap_or(0),
        });

    let mut response = CommandResponse::success(ctx.request_id)
        .with_metadata("task_session_id", &task.session_id)
        .with_metadata("task_id", &task.id)
        .with_metadata("task_description", &task.description)
        .with_metadata("task_subagent_type", &task.subagent_type)
        .with_metadata("task_status", &status);

    if let Some(ref summary) = task.summary {
        response = response.with_metadata("task_summary", summary);
    }

    if let Some(ref stats) = token_stats {
        let stats_json = serde_json::to_string(stats)
            .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
        response = response.with_metadata("task_token_stats", &stats_json);
    }

    Ok(response)
}

/// Reconstruct a TaskSession from the database when it's not in the in-memory repository.
/// Task sessions have id format: sub:{parent_session_id}:task:{uuid}
fn reconstruct_task_from_db(
    store: &SessionStore,
    task_id: &str,
) -> Result<Option<TaskSession>, CommandError> {
    let record = store
        .find_first_session_by_id_suffix(&format!(":{}", task_id))
        .map_err(|e| CommandError::new("DB_ERROR", e.to_string()))?;

    let record = match record {
        Some(record) => record,
        None => return Ok(None),
    };

    let session_id = record.id.clone();

    // Extract parent_session_id from session_id: "sub:{parent}:task:{uuid}"
    let parent_session_id = session_id
        .strip_prefix("sub:")
        .and_then(|rest| {
            let pos = rest.find(":task:");
            pos.map(|p| rest[..p].to_string())
        })
        .unwrap_or_default();

    // Extract subagent_type and description from title
    let (subagent_type, description) = parse_subagent_title(&record.title);

    let now = record.updated_at;

    // DB 未持久化 task 状态字段，无法可靠区分 Running/Completed/Failed/Timeout。
    // 用 Unknown 表示"重启后从 DB 重建，真实状态不可知"，避免把 failed/timeout
    // 误报为 Completed 误导用户。前端会把 Unknown 显示为"未知"。
    tracing::warn!(
        task_id = %task_id,
        session_id = %session_id,
        "Reconstructing task from DB after restart; true state unknown, marking as Unknown"
    );

    Ok(Some(TaskSession {
        id: task_id.to_string(),
        session_id,
        parent_session_id,
        parent_topic_id: None,
        parent_chat_id: record.chat_id.clone(),
        parent_channel_name: record.channel_name.clone(),
        description,
        subagent_type,
        state: TaskSessionState::Unknown,
        created_at: record.created_at,
        updated_at: now,
        summary: None,
        error: None,
        tool_call_id: None,
    }))
}

/// Parse subagent title to extract type and description.
/// New format: "Subagent [type]: description"
/// Legacy format: "Subagent: description" (defaults to "general")
fn parse_subagent_title(title: &str) -> (String, String) {
    if let Some(rest) = title.strip_prefix("Subagent [")
        && let Some(bracket_pos) = rest.find("]: ")
    {
        let agent_type = rest[..bracket_pos].to_string();
        let desc = rest[bracket_pos + 3..].to_string();
        return (agent_type, desc);
    }
    let desc = title
        .strip_prefix("Subagent: ")
        .unwrap_or(title)
        .to_string();
    ("general".to_string(), desc)
}