Kigi-CLI/crates/codegen/kigi-shell/src/extensions/task.rs

use agent_client_protocol as acp;
use kigi_tools::types::{KillOutcome, TaskSnapshot};
use serde::{Deserialize, Serialize};

use kigi_tools::implementations::grok_build::task::types::{
    SubagentCancelOutcome, SubagentSnapshot, SubagentSnapshotStatus,
};

use crate::agent::MvpAgent;
use crate::agent::subagent::{ResolvedRunningSubagent, is_running, resolve_running_list};
use crate::session::ExtMethodResult;

type ExtResult = Result<acp::ExtResponse, acp::Error>;

/// Wire DTO for the `x.ai/task/kill` ext request.
///
/// `pub` (with both serde directions) so ACP clients (kigi-tui) build
/// the request from the same type the agent parses — keeping the wire
/// contract typed end-to-end instead of duplicated `json!` literals.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct KillTaskRequest {
    pub session_id: String,
    pub task_id: String,
}

/// Wire DTO for the `x.ai/task/kill` ext response payload (nested under
/// `result` in the `ExtMethodResult` envelope).
///
/// `pub` (with both serde directions) so ACP clients deserialize the typed
/// outcome instead of probing raw JSON.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct KillTaskResponse {
    pub task_id: String,
    pub outcome: KillOutcome,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListTasksRequest {
    session_id: String,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct ListTasksResponse {
    tasks: Vec<TaskSnapshot>,
}

/// Wire DTO for the `x.ai/subagent/cancel` ext request.
///
/// `pub` (with both serde directions) so ACP clients (kigi-tui) build
/// the request from the same type the agent parses.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelSubagentRequest {
    pub subagent_id: String,
}

/// Wire mirror of the coordinator's [`SubagentCancelOutcome`], `kind`-tagged so
/// a client can branch and read the already-finished `status`. Sent alongside
/// the legacy `cancelled` bool: a new pager prefers this, an old one ignores it.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SubagentCancelOutcomeDto {
    /// A live subagent was cancelled — a real `SubagentFinished` is coming.
    Cancelled,
    /// Already finished — no finish coming; `status` is the real terminal status.
    AlreadyFinished { status: String },
    /// The id is unknown (never existed / evicted) — no finish coming.
    NotFound,
    /// Unknown future `kind` (`#[serde(other)]`): lets an old client still parse
    /// and fall back to the legacy bool. Never produced by `From`.
    #[serde(other)]
    Unknown,
}

impl SubagentCancelOutcomeDto {
    /// Legacy bool for older pagers: true only when a live subagent was stopped.
    /// Already-finished / not-found → false so an old pager finalizes the row.
    fn cancelled_bool(&self) -> bool {
        matches!(self, Self::Cancelled)
    }
}

impl From<SubagentCancelOutcome> for SubagentCancelOutcomeDto {
    fn from(outcome: SubagentCancelOutcome) -> Self {
        match outcome {
            SubagentCancelOutcome::Cancelled => Self::Cancelled,
            SubagentCancelOutcome::AlreadyFinished { status } => Self::AlreadyFinished { status },
            SubagentCancelOutcome::NotFound => Self::NotFound,
        }
    }
}

/// Wire DTO for the `x.ai/subagent/cancel` response payload (under `result` in
/// the `ExtMethodResult` envelope). `pub` + both serde dirs so clients read it typed.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelSubagentResponse {
    pub subagent_id: String,
    /// Legacy wire-compat flag for older pagers; new clients prefer `outcome`.
    pub cancelled: bool,
    /// Typed outcome; `None` only from an older shell. This shell always sets it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub outcome: Option<SubagentCancelOutcomeDto>,
}

// ── Subagent list_running DTOs ────────────────────────────────────────────

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListRunningSubagentsRequest {
    session_id: String,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct ListRunningSubagentsResponse {
    subagents: Vec<SubagentLiveSnapshotDto>,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct SubagentLiveSnapshotDto {
    subagent_id: String,
    parent_session_id: String,
    child_session_id: String,
    subagent_type: String,
    description: String,
    started_at_epoch_ms: u64,
    duration_ms: u64,
    turn_count: u32,
    tool_call_count: u32,
    tokens_used: u64,
    context_window_tokens: u64,
    context_usage_pct: u8,
    tools_used: Vec<String>,
    error_count: u32,
}

impl From<ResolvedRunningSubagent> for SubagentLiveSnapshotDto {
    fn from(r: ResolvedRunningSubagent) -> Self {
        Self {
            subagent_id: r.subagent_id,
            parent_session_id: r.parent_session_id,
            child_session_id: r.child_session_id,
            subagent_type: r.subagent_type,
            description: r.description,
            started_at_epoch_ms: r.started_at_epoch_ms,
            duration_ms: r.duration_ms,
            turn_count: r.turn_count,
            tool_call_count: r.tool_call_count,
            tokens_used: r.tokens_used,
            context_window_tokens: r.context_window_tokens,
            context_usage_pct: r.context_usage_pct,
            tools_used: r.tools_used,
            error_count: r.error_count,
        }
    }
}

// ── Subagent get DTOs ────────────────────────────────────────────────────

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct GetSubagentRequest {
    subagent_id: String,
    #[serde(default)]
    block: Option<bool>,
    #[serde(default)]
    timeout_ms: Option<u64>,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct GetSubagentResponse {
    snapshot: Option<SubagentSnapshotDto>,
}

/// ACP DTO for a single subagent snapshot (any status).
///
/// Extends the identity fields from `SubagentLiveSnapshotDto` with
/// status-dependent fields for completed/failed/cancelled states.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct SubagentSnapshotDto {
    subagent_id: String,
    parent_session_id: String,
    child_session_id: String,
    subagent_type: String,
    description: String,
    started_at_epoch_ms: u64,
    duration_ms: u64,
    status: String,
    // ── Running fields (present only when status == "running") ────
    #[serde(skip_serializing_if = "Option::is_none")]
    turn_count: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_call_count: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tokens_used: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    context_window_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    context_usage_pct: Option<u8>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools_used: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error_count: Option<u32>,
    // ── Completed fields ─────────────────────────────────────────
    #[serde(skip_serializing_if = "Option::is_none")]
    output: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_calls: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    turns: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    worktree_path: Option<String>,
    // ── Failed / Cancelled fields ────────────────────────────────
    #[serde(skip_serializing_if = "Option::is_none")]
    failure_error: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    cancel_reason: Option<String>,
    // ── Fork/resume provenance ─────────────────────────────────
    #[serde(skip_serializing_if = "Option::is_none")]
    fork_context_source: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    fork_parent_prompt_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    resumed_from: Option<String>,
}

impl SubagentSnapshotDto {
    /// Build a DTO from a resolved snapshot and its session identity.
    fn from_snapshot(
        snap: SubagentSnapshot,
        parent_session_id: String,
        child_session_id: String,
        provenance: crate::agent::subagent::SubagentProvenance,
    ) -> Self {
        let mut dto = SubagentSnapshotDto {
            subagent_id: snap.subagent_id,
            parent_session_id,
            child_session_id,
            subagent_type: snap.subagent_type,
            description: snap.description,
            started_at_epoch_ms: snap.started_at_epoch_ms,
            duration_ms: snap.duration_ms,
            status: String::new(),
            turn_count: None,
            tool_call_count: None,
            tokens_used: None,
            context_window_tokens: None,
            context_usage_pct: None,
            tools_used: None,
            error_count: None,
            output: None,
            tool_calls: None,
            turns: None,
            worktree_path: None,
            failure_error: None,
            cancel_reason: None,
            fork_context_source: None,
            fork_parent_prompt_id: provenance.fork_parent_prompt_id,
            resumed_from: provenance.resumed_from,
        };
        match snap.status {
            SubagentSnapshotStatus::Initializing => {
                dto.status = "initializing".into();
            }
            SubagentSnapshotStatus::Running {
                turn_count,
                tool_call_count,
                tokens_used,
                context_window_tokens,
                context_usage_pct,
                tools_used,
                error_count,
            } => {
                dto.status = "running".into();
                dto.turn_count = Some(turn_count);
                dto.tool_call_count = Some(tool_call_count);
                dto.tokens_used = Some(tokens_used);
                dto.context_window_tokens = Some(context_window_tokens);
                dto.context_usage_pct = Some(context_usage_pct);
                dto.tools_used = Some(tools_used);
                dto.error_count = Some(error_count);
            }
            SubagentSnapshotStatus::Completed {
                output,
                tool_calls,
                turns,
                worktree_path,
            } => {
                dto.status = "completed".into();
                dto.output = Some(output);
                dto.tool_calls = Some(tool_calls);
                dto.turns = Some(turns);
                dto.worktree_path = worktree_path;
            }
            SubagentSnapshotStatus::Failed { error } => {
                dto.status = "failed".into();
                dto.failure_error = Some(error);
            }
            SubagentSnapshotStatus::Cancelled { reason } => {
                dto.status = "cancelled".into();
                dto.cancel_reason = reason;
            }
        }
        dto
    }
}

// ── Helpers ──────────────────────────────────────────────────────────────

fn parse<T: serde::de::DeserializeOwned>(args: &acp::ExtRequest) -> Result<T, acp::Error> {
    serde_json::from_str(args.params.get())
        .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))
}

fn respond<T: Serialize>(result: Result<T, impl std::fmt::Display>) -> ExtResult {
    ExtMethodResult::from_result(result)
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    match args.method.as_ref() {
        "x.ai/task/kill" => {
            let req: KillTaskRequest = parse(args)?;
            let result = agent
                .kill_background_task(&req.session_id, &req.task_id)
                .await
                .map(|outcome| KillTaskResponse {
                    task_id: req.task_id,
                    outcome,
                });
            respond(result)
        }
        "x.ai/task/list" => {
            let req: ListTasksRequest = parse(args)?;
            let result = agent
                .list_tasks(&req.session_id)
                .await
                .ok_or_else(|| "session not found or no terminal backend".to_string())
                .map(|tasks| ListTasksResponse { tasks });
            respond(result)
        }
        _ => Err(acp::Error::method_not_found()),
    }
}

// ── Scheduler DTOs ────────────────────────────────────────────────────

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DeleteScheduledTaskRequest {
    session_id: String,
    task_id: String,
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct DeleteScheduledTaskResponse {
    task_id: String,
    deleted: bool,
}

/// Handle `x.ai/scheduler/*` extension methods.
pub async fn handle_scheduler(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    match args.method.as_ref() {
        "x.ai/scheduler/delete" => {
            let req: DeleteScheduledTaskRequest = parse(args)?;
            let result = agent
                .delete_scheduled_task(&req.session_id, &req.task_id)
                .await
                .map(|deleted| DeleteScheduledTaskResponse {
                    task_id: req.task_id,
                    deleted,
                });
            respond(result)
        }
        _ => Err(acp::Error::method_not_found()),
    }
}

/// Handle `x.ai/subagent/*` extension methods.
pub async fn handle_subagent(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    match args.method.as_ref() {
        "x.ai/subagent/cancel" => {
            let req: CancelSubagentRequest = parse(args)?;
            tracing::info!(subagent_id = %req.subagent_id, "Cancelling subagent via ext method");
            let outcome = SubagentCancelOutcomeDto::from(agent.cancel_subagent(&req.subagent_id));
            respond(Ok::<_, String>(CancelSubagentResponse {
                subagent_id: req.subagent_id,
                cancelled: outcome.cancelled_bool(),
                outcome: Some(outcome),
            }))
        }
        "x.ai/subagent/get" => {
            let req: GetSubagentRequest = parse(args)?;
            let block = req.block.unwrap_or(false);
            let timeout_ms = req.timeout_ms.unwrap_or(30_000);

            let ids = agent.session_ids_for_subagent(&req.subagent_id);
            let (parent_sid, child_sid) = ids.unwrap_or_default();
            let provenance = agent.provenance_for_subagent(&req.subagent_id);

            let to_dto = |snap: SubagentSnapshot| {
                SubagentSnapshotDto::from_snapshot(
                    snap,
                    parent_sid.clone(),
                    child_sid.clone(),
                    provenance.clone(),
                )
            };

            // Sync lookup, drop borrow, then resolve async.
            let lookup = agent.lookup_subagent(&req.subagent_id);
            let snapshot = crate::agent::subagent::resolve_snapshot(lookup).await;

            if block && snapshot.as_ref().is_some_and(is_running) {
                // Poll every 200ms until done or timeout.
                let deadline =
                    tokio::time::Instant::now() + tokio::time::Duration::from_millis(timeout_ms);
                loop {
                    tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
                    let lookup = agent.lookup_subagent(&req.subagent_id);
                    let snap = crate::agent::subagent::resolve_snapshot(lookup).await;
                    let still_running = snap.as_ref().is_some_and(is_running);
                    if !still_running || tokio::time::Instant::now() >= deadline {
                        return respond(Ok::<_, String>(GetSubagentResponse {
                            snapshot: snap.map(&to_dto),
                        }));
                    }
                }
            } else {
                respond(Ok::<_, String>(GetSubagentResponse {
                    snapshot: snapshot.map(to_dto),
                }))
            }
        }
        "x.ai/subagent/list_running" => {
            let req: ListRunningSubagentsRequest = parse(args)?;
            // Sync: collect seeds from coordinator, drop borrow.
            let seeds = agent.list_running_subagents(&req.session_id);
            // Async: resolve live signals concurrently.
            let resolved = resolve_running_list(seeds).await;
            let subagents = resolved
                .into_iter()
                .map(SubagentLiveSnapshotDto::from)
                .collect();
            respond(Ok::<_, String>(ListRunningSubagentsResponse { subagents }))
        }
        _ => Err(acp::Error::method_not_found()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn delete_scheduled_task_request_deserializes_camel_case() {
        let json = r#"{"sessionId":"sess-1","taskId":"task-42"}"#;
        let req: DeleteScheduledTaskRequest = serde_json::from_str(json).expect("should parse");
        assert_eq!(req.session_id, "sess-1");
        assert_eq!(req.task_id, "task-42");
    }

    #[test]
    fn delete_scheduled_task_response_serializes_camel_case() {
        let resp = DeleteScheduledTaskResponse {
            task_id: "task-42".into(),
            deleted: true,
        };
        let json = serde_json::to_value(&resp).expect("should serialize");
        assert_eq!(json["taskId"], "task-42");
        assert_eq!(json["deleted"], true);
    }

    #[test]
    fn subagent_live_snapshot_dto_serializes_camel_case() {
        let dto = SubagentLiveSnapshotDto {
            subagent_id: "sub-1".into(),
            parent_session_id: "parent-1".into(),
            child_session_id: "child-1".into(),
            subagent_type: "explore".into(),
            description: "find files".into(),
            started_at_epoch_ms: 1_700_000_000_000,
            duration_ms: 5000,
            turn_count: 2,
            tool_call_count: 7,
            tokens_used: 30_000,
            context_window_tokens: 256_000,
            context_usage_pct: 23,
            tools_used: vec!["bash".into(), "grep".into()],
            error_count: 1,
        };
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["subagentId"], "sub-1");
        assert_eq!(json["parentSessionId"], "parent-1");
        assert_eq!(json["childSessionId"], "child-1");
        assert_eq!(json["subagentType"], "explore");
        assert_eq!(json["startedAtEpochMs"], 1_700_000_000_000_u64);
        assert_eq!(json["durationMs"], 5000);
        assert_eq!(json["turnCount"], 2);
        assert_eq!(json["toolCallCount"], 7);
        assert_eq!(json["tokensUsed"], 30_000);
        assert_eq!(json["contextWindowTokens"], 256_000);
        assert_eq!(json["contextUsagePct"], 23);
        assert_eq!(json["toolsUsed"], serde_json::json!(["bash", "grep"]));
        assert_eq!(json["errorCount"], 1);
    }

    #[test]
    fn from_resolved_running_subagent_maps_all_fields() {
        let resolved = ResolvedRunningSubagent {
            subagent_id: "s".into(),
            parent_session_id: "p".into(),
            child_session_id: "c".into(),
            subagent_type: "plan".into(),
            description: "d".into(),
            started_at_epoch_ms: 100,
            duration_ms: 200,
            turn_count: 1,
            tool_call_count: 3,
            tokens_used: 500,
            context_window_tokens: 1000,
            context_usage_pct: 50,
            tools_used: vec!["read_file".into()],
            error_count: 0,
        };
        let dto = SubagentLiveSnapshotDto::from(resolved);
        assert_eq!(dto.subagent_id, "s");
        assert_eq!(dto.parent_session_id, "p");
        assert_eq!(dto.child_session_id, "c");
        assert_eq!(dto.context_usage_pct, 50);
        assert_eq!(dto.tools_used, vec!["read_file"]);
    }

    #[test]
    fn list_running_response_serializes_with_subagents_array() {
        let resp = ListRunningSubagentsResponse { subagents: vec![] };
        let json = serde_json::to_value(&resp).expect("should serialize");
        assert_eq!(json["subagents"], serde_json::json!([]));
    }

    // ── SubagentSnapshotDto serialization tests ────────────────────────

    #[test]
    fn snapshot_dto_running_serializes_with_progress_fields() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-1".into(),
            subagent_type: "explore".into(),
            description: "find files".into(),
            started_at_epoch_ms: 1000,
            duration_ms: 5000,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 3,
                tool_call_count: 12,
                tokens_used: 45_000,
                context_window_tokens: 256_000,
                context_usage_pct: 35,
                tools_used: vec!["bash".into(), "grep".into()],
                error_count: 1,
            },
        };
        let dto = SubagentSnapshotDto::from_snapshot(
            snap,
            "parent-1".into(),
            "child-1".into(),
            Default::default(),
        );
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["parentSessionId"], "parent-1");
        assert_eq!(json["childSessionId"], "child-1");
        assert_eq!(json["status"], "running");
        assert_eq!(json["turnCount"], 3);
        assert_eq!(json["toolCallCount"], 12);
        assert_eq!(json["tokensUsed"], 45_000);
        assert_eq!(json["contextUsagePct"], 35);
        assert_eq!(json["errorCount"], 1);
        // Completed-only fields should be absent
        assert!(json.get("output").is_none());
        assert!(json.get("failureError").is_none());
    }

    #[test]
    fn snapshot_dto_completed_serializes_with_output() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-2".into(),
            subagent_type: "general-purpose".into(),
            description: "refactor auth".into(),
            started_at_epoch_ms: 2000,
            duration_ms: 15_000,
            persona: None,
            status: SubagentSnapshotStatus::Completed {
                output: "Done, refactored 3 files.".into(),
                tool_calls: 8,
                turns: 2,
                worktree_path: None,
            },
        };
        let dto =
            SubagentSnapshotDto::from_snapshot(snap, "p".into(), "c".into(), Default::default());
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["status"], "completed");
        assert_eq!(json["output"], "Done, refactored 3 files.");
        assert_eq!(json["toolCalls"], 8);
        assert_eq!(json["turns"], 2);
        // Running-only fields should be absent
        assert!(json.get("turnCount").is_none());
        assert!(json.get("tokensUsed").is_none());
    }

    #[test]
    fn snapshot_dto_failed_serializes_with_error() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-3".into(),
            subagent_type: "plan".into(),
            description: "plan feature".into(),
            started_at_epoch_ms: 0,
            duration_ms: 100,
            persona: None,
            status: SubagentSnapshotStatus::Failed {
                error: "sampling error".into(),
            },
        };
        let dto =
            SubagentSnapshotDto::from_snapshot(snap, "p".into(), "c".into(), Default::default());
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["status"], "failed");
        assert_eq!(json["failureError"], "sampling error");
    }

    #[test]
    fn snapshot_dto_cancelled_serializes_with_reason() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-4".into(),
            subagent_type: "explore".into(),
            description: "search".into(),
            started_at_epoch_ms: 0,
            duration_ms: 50,
            persona: None,
            status: SubagentSnapshotStatus::Cancelled {
                reason: Some("user cancelled".into()),
            },
        };
        let dto =
            SubagentSnapshotDto::from_snapshot(snap, "p".into(), "c".into(), Default::default());
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["status"], "cancelled");
        assert_eq!(json["cancelReason"], "user cancelled");
    }

    #[test]
    fn snapshot_dto_cancelled_without_reason_omits_field() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-5".into(),
            subagent_type: "explore".into(),
            description: "search".into(),
            started_at_epoch_ms: 0,
            duration_ms: 50,
            persona: None,
            status: SubagentSnapshotStatus::Cancelled { reason: None },
        };
        let dto =
            SubagentSnapshotDto::from_snapshot(snap, "p".into(), "c".into(), Default::default());
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["status"], "cancelled");
        assert!(json.get("cancelReason").is_none());
    }

    #[test]
    fn get_subagent_response_null_snapshot() {
        let resp = GetSubagentResponse { snapshot: None };
        let json = serde_json::to_value(&resp).expect("should serialize");
        assert!(json["snapshot"].is_null());
    }

    #[test]
    fn get_subagent_response_with_running_snapshot() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-run".into(),
            subagent_type: "explore".into(),
            description: "search".into(),
            started_at_epoch_ms: 1000,
            duration_ms: 5000,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 2,
                tool_call_count: 5,
                tokens_used: 20_000,
                context_window_tokens: 256_000,
                context_usage_pct: 15,
                tools_used: vec!["bash".into()],
                error_count: 0,
            },
        };
        let resp = GetSubagentResponse {
            snapshot: Some(SubagentSnapshotDto::from_snapshot(
                snap,
                "parent-1".into(),
                "child-1".into(),
                Default::default(),
            )),
        };
        let json = serde_json::to_value(&resp).expect("should serialize");
        let s = &json["snapshot"];
        assert_eq!(s["status"], "running");
        assert_eq!(s["subagentId"], "sub-run");
        assert_eq!(s["parentSessionId"], "parent-1");
        assert_eq!(s["childSessionId"], "child-1");
        assert_eq!(s["turnCount"], 2);
        // Completed-only fields must be absent
        assert!(s.get("output").is_none());
        assert!(s.get("turns").is_none());
    }

    #[test]
    fn get_subagent_response_with_completed_snapshot() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-done".into(),
            subagent_type: "general-purpose".into(),
            description: "refactor".into(),
            started_at_epoch_ms: 0,
            duration_ms: 10_000,
            persona: None,
            status: SubagentSnapshotStatus::Completed {
                output: "Refactored 3 files.".into(),
                tool_calls: 7,
                turns: 2,
                worktree_path: None,
            },
        };
        let resp = GetSubagentResponse {
            snapshot: Some(SubagentSnapshotDto::from_snapshot(
                snap,
                "parent-2".into(),
                "child-2".into(),
                Default::default(),
            )),
        };
        let json = serde_json::to_value(&resp).expect("should serialize");
        let s = &json["snapshot"];
        assert_eq!(s["status"], "completed");
        assert_eq!(s["output"], "Refactored 3 files.");
        assert_eq!(s["toolCalls"], 7);
        assert_eq!(s["turns"], 2);
        // Running-only fields must be absent
        assert!(s.get("turnCount").is_none());
        assert!(s.get("tokensUsed").is_none());
    }

    #[test]
    fn get_subagent_request_deserializes_block_false() {
        let json = r#"{"subagentId":"sub-1","block":false}"#;
        let req: GetSubagentRequest = serde_json::from_str(json).expect("should parse");
        assert_eq!(req.subagent_id, "sub-1");
        assert_eq!(req.block, Some(false));
        assert!(req.timeout_ms.is_none());
    }

    #[test]
    fn get_subagent_request_deserializes_block_true_with_timeout() {
        let json = r#"{"subagentId":"sub-2","block":true,"timeoutMs":5000}"#;
        let req: GetSubagentRequest = serde_json::from_str(json).expect("should parse");
        assert_eq!(req.subagent_id, "sub-2");
        assert_eq!(req.block, Some(true));
        assert_eq!(req.timeout_ms, Some(5000));
    }

    #[test]
    fn get_subagent_request_defaults_block_and_timeout() {
        let json = r#"{"subagentId":"sub-3"}"#;
        let req: GetSubagentRequest = serde_json::from_str(json).expect("should parse");
        assert!(req.block.is_none());
        assert!(req.timeout_ms.is_none());
    }

    // ── Polling control-flow tests ──────────────────────────────────────

    #[test]
    fn block_true_with_completed_snapshot_returns_immediately() {
        // When block=true but the snapshot is already completed,
        // the handler should NOT enter the polling loop.
        let snap = SubagentSnapshot {
            subagent_id: "sub-done".into(),
            subagent_type: "explore".into(),
            description: "d".into(),
            started_at_epoch_ms: 0,
            duration_ms: 100,
            persona: None,
            status: SubagentSnapshotStatus::Completed {
                output: "done".into(),
                tool_calls: 1,
                turns: 1,
                worktree_path: None,
            },
        };
        // The handler's decision: `block && is_running(&snap)` → false
        let block = true;
        let should_poll = block && is_running(&snap);
        assert!(
            !should_poll,
            "completed snapshot should not trigger polling loop"
        );
    }

    #[test]
    fn block_false_with_running_snapshot_skips_polling() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-run".into(),
            subagent_type: "explore".into(),
            description: "d".into(),
            started_at_epoch_ms: 0,
            duration_ms: 100,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 1,
                tool_call_count: 2,
                tokens_used: 1000,
                context_window_tokens: 256_000,
                context_usage_pct: 1,
                tools_used: vec![],
                error_count: 0,
            },
        };
        // The handler's decision: `block && is_running(&snap)` → false
        let block = false;
        let should_poll = block && is_running(&snap);
        assert!(
            !should_poll,
            "block=false should not trigger polling loop even if running"
        );
    }

    #[test]
    fn block_true_with_running_snapshot_enters_polling() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-run".into(),
            subagent_type: "explore".into(),
            description: "d".into(),
            started_at_epoch_ms: 0,
            duration_ms: 100,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 1,
                tool_call_count: 2,
                tokens_used: 1000,
                context_window_tokens: 256_000,
                context_usage_pct: 1,
                tools_used: vec![],
                error_count: 0,
            },
        };
        // The handler's decision: `block && is_running(&snap)` → true
        let block = true;
        let should_poll = block && is_running(&snap);
        assert!(
            should_poll,
            "block=true + running should trigger polling loop"
        );
    }

    #[test]
    fn polling_loop_exits_when_snapshot_transitions_to_completed() {
        // Simulates the polling loop's exit condition when a snapshot
        // transitions from running to completed between iterations.
        let completed_snap = SubagentSnapshot {
            subagent_id: "sub-1".into(),
            subagent_type: "explore".into(),
            description: "d".into(),
            started_at_epoch_ms: 0,
            duration_ms: 500,
            persona: None,
            status: SubagentSnapshotStatus::Completed {
                output: "found it".into(),
                tool_calls: 3,
                turns: 1,
                worktree_path: None,
            },
        };
        // The polling loop checks: `!is_running(&snap) || deadline_passed`
        // When the snapshot becomes completed, `!is_running` is true → exits.
        assert!(
            !is_running(&completed_snap),
            "completed snapshot should cause polling loop exit"
        );
    }

    #[test]
    fn polling_loop_exits_on_deadline_even_if_still_running() {
        let running_snap = SubagentSnapshot {
            subagent_id: "sub-1".into(),
            subagent_type: "explore".into(),
            description: "d".into(),
            started_at_epoch_ms: 0,
            duration_ms: 100,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 1,
                tool_call_count: 1,
                tokens_used: 1000,
                context_window_tokens: 256_000,
                context_usage_pct: 1,
                tools_used: vec![],
                error_count: 0,
            },
        };
        // Simulate: deadline has passed, but snapshot is still running.
        // The polling loop checks: `!is_running(&snap) || deadline_passed`
        let deadline_passed = true;
        let should_exit = !is_running(&running_snap) || deadline_passed;
        assert!(
            should_exit,
            "deadline expiry should cause polling loop exit even if still running"
        );
    }

    #[test]
    fn snapshot_dto_resumed_provenance_serializes() {
        let snap = SubagentSnapshot {
            subagent_id: "sub-resumed".into(),
            subagent_type: "general-purpose".into(),
            description: "fix review".into(),
            started_at_epoch_ms: 2000,
            duration_ms: 3000,
            persona: None,
            status: SubagentSnapshotStatus::Running {
                turn_count: 1,
                tool_call_count: 2,
                tokens_used: 10_000,
                context_window_tokens: 256_000,
                context_usage_pct: 8,
                tools_used: vec![],
                error_count: 0,
            },
        };
        let provenance = crate::agent::subagent::SubagentProvenance {
            fork_parent_prompt_id: Some("prompt-5".into()),
            resumed_from: Some("source-agent-id".into()),
        };
        let dto = SubagentSnapshotDto::from_snapshot(
            snap,
            "parent".into(),
            "child-resumed".into(),
            provenance,
        );
        let json = serde_json::to_value(&dto).expect("should serialize");
        assert_eq!(json["resumedFrom"], "source-agent-id");
        assert_eq!(json["forkParentPromptId"], "prompt-5");
    }

    // ── x.ai/subagent/cancel outcome wire DTO ──────────────────────────

    #[test]
    fn subagent_cancel_outcome_dto_maps_from_coordinator_outcome() {
        // Cancelled → legacy bool true (a real finish is coming).
        let dto = SubagentCancelOutcomeDto::from(SubagentCancelOutcome::Cancelled);
        assert_eq!(dto, SubagentCancelOutcomeDto::Cancelled);
        assert!(dto.cancelled_bool());

        // AlreadyFinished carries the terminal status; legacy bool false.
        let dto = SubagentCancelOutcomeDto::from(SubagentCancelOutcome::AlreadyFinished {
            status: "completed".into(),
        });
        assert_eq!(
            dto,
            SubagentCancelOutcomeDto::AlreadyFinished {
                status: "completed".into()
            }
        );
        assert!(!dto.cancelled_bool());

        // NotFound → legacy bool false.
        let dto = SubagentCancelOutcomeDto::from(SubagentCancelOutcome::NotFound);
        assert_eq!(dto, SubagentCancelOutcomeDto::NotFound);
        assert!(!dto.cancelled_bool());
    }

    #[test]
    fn cancel_subagent_response_serializes_outcome_snake_case() {
        let resp = CancelSubagentResponse {
            subagent_id: "sa-1".into(),
            cancelled: false,
            outcome: Some(SubagentCancelOutcomeDto::AlreadyFinished {
                status: "failed".into(),
            }),
        };
        let json = serde_json::to_value(&resp).expect("should serialize");
        assert_eq!(json["subagentId"], "sa-1");
        assert_eq!(json["cancelled"], false);
        assert_eq!(json["outcome"]["kind"], "already_finished");
        assert_eq!(json["outcome"]["status"], "failed");
    }

    /// Wire-compat: a payload from an older shell (no `outcome`) still
    /// deserializes, leaving `outcome` as `None` so the client falls back to
    /// the legacy `cancelled` bool.
    #[test]
    fn cancel_subagent_response_deserializes_without_outcome() {
        let resp: CancelSubagentResponse =
            serde_json::from_str(r#"{"subagentId":"sa-1","cancelled":true}"#).expect("parse");
        assert_eq!(resp.subagent_id, "sa-1");
        assert!(resp.cancelled);
        assert!(resp.outcome.is_none());
    }
}