PicoBot/src/command/handlers/load_task_messages.rs

use crate::command::context::CommandContext;
use crate::command::handler::{CommandHandler, CommandMetadata};
use crate::command::response::{CommandError, CommandResponse};
use crate::command::Command;
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::info!(
        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::info!(
                task_id = %task.id,
                session_id = %task.session_id,
                state = ?task.state,
                "LoadTaskMessages: task found in memory"
            );
            Some(task)
        }
        Ok(None) => {
            tracing::info!(
                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();

    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);
    }

    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 sessions = store
        .find_sessions_by_id_suffix(&format!(":{}", task_id))
        .map_err(|e| CommandError::new("DB_ERROR", e.to_string()))?;

    if sessions.is_empty() {
        return Ok(None);
    }

    let record = &sessions[0];
    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;

    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::Completed,
        created_at: record.created_at,
        updated_at: now,
        summary: None,
        error: 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 [") {
        if 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)
}