roboticus-api 0.11.3

//! Decomposition gate, specialist proposal types, and delegation planning.

use std::collections::HashSet;

use serde_json::json;

use roboticus_agent::orchestration::{OrchestrationPattern, Orchestrator};

use super::AppState;

#[derive(Debug, Clone)]
pub(super) struct SpecialistProposal {
    pub name: String,
    pub display_name: String,
    pub description: String,
    pub skills: Vec<String>,
    pub model: String,
}

#[derive(Debug, Clone)]
pub(super) struct DelegationPlan {
    pub subtasks: Vec<String>,
    pub rationale: String,
    pub expected_utility_margin: f64,
}

#[derive(Debug, Clone)]
pub(super) enum DecompositionDecision {
    Centralized {
        rationale: String,
        expected_utility_margin: f64,
    },
    Delegated(DelegationPlan),
    RequiresSpecialistCreation {
        proposal: SpecialistProposal,
        rationale: String,
    },
}

#[derive(Debug, Clone, Default)]
pub(super) struct DelegationProvenance {
    pub subagent_task_started: bool,
    pub subagent_task_completed: bool,
    pub subagent_result_attached: bool,
}

pub(super) fn split_subtasks(input: &str) -> Vec<String> {
    let mut seen = std::collections::HashSet::new();
    let mut out = Vec::new();
    for part in input
        .split(&['\n', ';'][..])
        .flat_map(|p| p.split(" then "))
        .flat_map(|p| p.split(" and "))
    {
        let trimmed = part.trim();
        if !trimmed.is_empty() && seen.insert(trimmed.to_string()) {
            out.push(trimmed.to_string());
        }
    }
    out
}

pub(super) fn capability_tokens(text: &str) -> Vec<String> {
    text.to_ascii_lowercase()
        .split(|c: char| !c.is_ascii_alphanumeric())
        .filter(|t| t.len() >= 4)
        .map(|s| s.to_string())
        .collect()
}

pub(super) fn utility_margin_for_delegation(
    complexity_score: f64,
    subtask_count: usize,
    capability_fit_ratio: f64,
) -> f64 {
    let complexity_gain = complexity_score * 0.5;
    let parallel_gain = ((subtask_count.saturating_sub(1)) as f64) * 0.12;
    let fit_gain = capability_fit_ratio * 0.45;
    let orchestration_cost = 0.25 + ((subtask_count as f64) * 0.04);
    complexity_gain + parallel_gain + fit_gain - orchestration_cost
}

pub(super) fn proposal_to_json(
    proposal: &SpecialistProposal,
    rationale: &str,
) -> serde_json::Value {
    json!({
        "name": proposal.name,
        "display_name": proposal.display_name,
        "description": proposal.description,
        "skills": proposal.skills,
        "model": proposal.model,
        "rationale": rationale,
    })
}

pub(super) async fn evaluate_decomposition_gate(
    state: &AppState,
    user_content: &str,
    complexity_score: f64,
) -> DecompositionDecision {
    let cfg = state.config.read().await;
    if !cfg.agent.delegation_enabled {
        return DecompositionDecision::Centralized {
            rationale: "delegation disabled by configuration".to_string(),
            expected_utility_margin: -1.0,
        };
    }
    let min_complexity = cfg.agent.delegation_min_complexity;
    let min_margin = cfg.agent.delegation_min_utility_margin;
    drop(cfg);

    let subtasks = split_subtasks(user_content);
    if subtasks.len() <= 1 || complexity_score < min_complexity {
        return DecompositionDecision::Centralized {
            rationale: "task is single-step or below decomposition complexity threshold"
                .to_string(),
            expected_utility_margin: -0.1,
        };
    }

    let subagents = roboticus_db::agents::list_sub_agents(&state.db)
        .inspect_err(|e| tracing::error!(error = %e, "failed to list sub-agents for decomposition"))
        .unwrap_or_default();
    let taskable: Vec<_> = subagents
        .into_iter()
        .filter(|a| !super::is_model_proxy_role(&a.role) && a.enabled)
        .collect();
    let required = capability_tokens(user_content);
    if taskable.is_empty() {
        let proposal = SpecialistProposal {
            name: "proposed-specialist".to_string(),
            display_name: "Proposed Specialist".to_string(),
            description: "Auto-proposed specialist for uncovered capability gap".to_string(),
            skills: required.iter().take(8).cloned().collect(),
            model: "auto".to_string(),
        };
        return DecompositionDecision::RequiresSpecialistCreation {
            proposal,
            rationale:
                "no enabled taskable specialists available; compose a specialist before delegating"
                    .to_string(),
        };
    }

    let mut fit_hits = 0usize;
    for agent in &taskable {
        let skills = super::parse_skills_json(agent.skills_json.as_deref());
        let skill_tokens: HashSet<String> = skills
            .iter()
            .flat_map(|s| capability_tokens(s))
            .collect::<Vec<_>>()
            .into_iter()
            .collect();
        if required.iter().any(|t| skill_tokens.contains(t)) {
            fit_hits += 1;
        }
    }
    let capability_fit_ratio = if taskable.is_empty() {
        0.0
    } else {
        fit_hits as f64 / taskable.len() as f64
    };
    let margin =
        utility_margin_for_delegation(complexity_score, subtasks.len(), capability_fit_ratio);
    if capability_fit_ratio < 0.2 {
        let proposal = SpecialistProposal {
            name: "proposed-specialist".to_string(),
            display_name: "Proposed Specialist".to_string(),
            description: "Auto-proposed specialist for uncovered capability gap".to_string(),
            skills: required.into_iter().take(8).collect(),
            model: "auto".to_string(),
        };
        return DecompositionDecision::RequiresSpecialistCreation {
            proposal,
            rationale:
                "existing specialists do not satisfy required capability fit; proposal required"
                    .to_string(),
        };
    }

    if margin < min_margin {
        return DecompositionDecision::Centralized {
            rationale: format!(
                "delegation utility margin {:.2} below threshold {:.2}",
                margin, min_margin
            ),
            expected_utility_margin: margin,
        };
    }

    DecompositionDecision::Delegated(DelegationPlan {
        subtasks,
        rationale: format!(
            "decomposed into subtasks with estimated delegation margin {:.2}",
            margin
        ),
        expected_utility_margin: margin,
    })
}

/// Outcome of applying a decomposition gate decision.
pub(super) enum DecompositionOutcome {
    /// Task handled centrally -- no delegation.
    Centralized,
    /// A subagent creation proposal was stored; caller must relay `prompt` to the user
    /// and perform its own early-return (the return type differs between API and channel paths).
    SpecialistProposalPending { prompt: String },
    /// Task delegated via orchestrator; caller should thread `workflow_note` into context.
    Delegated { workflow_note: String },
}

/// Apply a decomposition gate decision: store proposals, orchestrate workflows, or log
/// centralized execution.  Returns an outcome the caller dispatches pathway-specifically.
pub(super) async fn apply_decomposition_decision(
    state: &AppState,
    gate_decision: &DecompositionDecision,
    session_id: &str,
    pathway_label: &str,
) -> DecompositionOutcome {
    match gate_decision {
        DecompositionDecision::RequiresSpecialistCreation {
            proposal,
            rationale,
        } => {
            let payload = proposal_to_json(proposal, rationale);
            {
                let mut pending = state.pending_specialist_proposals.write().await;
                pending.insert(session_id.to_string(), payload);
            }
            let prompt = format!(
                "I identified a capability gap and can create a new subagent with your approval.\n\n\
                 Proposed: `{}`\nRationale: {}\n\n\
                 Reply with:\n\
                 - `review subagent config` to inspect full config\n\
                 - `approve subagent creation` to create it\n\
                 - `deny subagent creation` to continue with main-agent execution",
                proposal.name, rationale
            );
            DecompositionOutcome::SpecialistProposalPending { prompt }
        }
        DecompositionDecision::Centralized {
            rationale,
            expected_utility_margin,
        } => {
            tracing::info!(
                decision = "centralized",
                pathway = %pathway_label,
                rationale = %rationale,
                expected_utility_margin = *expected_utility_margin,
                "decomposition gate decision"
            );
            DecompositionOutcome::Centralized
        }
        DecompositionDecision::Delegated(plan) => {
            let mut orch = Orchestrator::new();
            let wf_input = plan
                .subtasks
                .iter()
                .map(|s| (s.clone(), capability_tokens(s)))
                .collect::<Vec<_>>();
            let wf_id =
                orch.create_workflow(pathway_label, OrchestrationPattern::Parallel, wf_input);
            let agent_rows = roboticus_db::agents::list_sub_agents(&state.db)
                .inspect_err(
                    |e| tracing::error!(error = %e, "failed to list sub-agents for workflow"),
                )
                .unwrap_or_default();
            // Build a name→display_label map so that workflow notes use
            // human-readable role descriptions rather than internal identifiers.
            let agent_display: std::collections::HashMap<String, String> = agent_rows
                .iter()
                .map(|a| {
                    let label = a
                        .display_name
                        .as_deref()
                        .filter(|s| !s.is_empty())
                        .or(a.description.as_deref().filter(|s| !s.is_empty()))
                        .map(|s| s.to_string())
                        .unwrap_or_else(|| a.name.clone());
                    (a.name.clone(), label)
                })
                .collect();
            let available_agents = agent_rows
                .into_iter()
                .filter(|a| !super::is_model_proxy_role(&a.role) && a.enabled)
                .map(|a| (a.name, super::parse_skills_json(a.skills_json.as_deref())))
                .collect::<Vec<_>>();
            let matches = orch
                .match_capabilities(&wf_id, &available_agents)
                .unwrap_or_default();
            for (task_id, agent_id) in &matches {
                if let Err(e) = orch.assign_agent(&wf_id, task_id, agent_id) {
                    tracing::error!(
                        workflow = %wf_id,
                        task = %task_id,
                        agent = %agent_id,
                        error = %e,
                        "failed to assign agent to workflow task"
                    );
                }
            }
            // Format assignments using display labels, not internal names.
            let assignments = matches
                .iter()
                .map(|(task, agent)| {
                    let label = agent_display
                        .get(agent)
                        .cloned()
                        .unwrap_or_else(|| agent.clone());
                    format!("{task}->{label}")
                })
                .collect::<Vec<_>>()
                .join(", ");
            let workflow_note = format!(
                "workflow_id={wf_id}; assignments={}",
                if assignments.is_empty() {
                    "none".to_string()
                } else {
                    assignments
                }
            );
            tracing::info!(
                decision = "delegated",
                pathway = %pathway_label,
                rationale = %plan.rationale,
                subtask_count = plan.subtasks.len(),
                expected_utility_margin = plan.expected_utility_margin,
                "decomposition gate decision"
            );
            DecompositionOutcome::Delegated { workflow_note }
        }
    }
}

/// Build a system note summarising the decomposition gate decision for the LLM context.
pub(super) fn build_gate_system_note(
    gate_decision: &DecompositionDecision,
    delegation_workflow_note: Option<&str>,
) -> String {
    match gate_decision {
        DecompositionDecision::Centralized {
            rationale,
            expected_utility_margin,
        } => format!(
            "Delegation decision: centralized. rationale='{}' expected_utility_margin={:.2}\n\
             For non-quick work, still follow **Subagent orchestration** in the system prompt \
             (classify turn → list-subagent-roster → compose-subagent if no fit → delegate).",
            rationale, expected_utility_margin
        ),
        DecompositionDecision::Delegated(plan) => {
            let subtask_lines = plan
                .subtasks
                .iter()
                .enumerate()
                .map(|(idx, s)| format!("{}. {}", idx + 1, s))
                .collect::<Vec<_>>()
                .join("\n");
            let mut note = format!(
                "Delegation decision: delegated.\nRationale: {}\nExpected utility margin: {:.2}\nSubtasks:\n{}",
                plan.rationale, plan.expected_utility_margin, subtask_lines
            );
            if let Some(wf_note) = delegation_workflow_note {
                note.push_str(&format!("\nWorkflow: {wf_note}"));
            }
            note.push_str(
                "\nExecution directive: follow **Subagent orchestration** — call `list-subagent-roster` \
                 this turn unless roster output is already in context; use `compose-subagent` if no \
                 specialist fits. Then emit a real tool_call for `orchestrate-subagents`, `assign-tasks`, \
                 or `delegate-subagent` with the task payload. Do not simulate delegated output.",
            );
            note
        }
        DecompositionDecision::RequiresSpecialistCreation { .. } => {
            "Delegation decision: specialist creation required with user approval.\n\
             Aligns with **Subagent orchestration** step 4: no suitable specialist in roster yet—\
             use the approval flow (or `compose-subagent` when policy allows) before delegating."
                .to_string()
        }
    }
}

/// Build tool definitions for virtual delegation and orchestration tools.
///
/// These tools are not in the ToolRegistry (they're handled by special-case dispatch
/// in `execute_tool_call`), so their schemas must be defined statically here.
/// Combined with capability/ToolRegistry-sourced definitions in `build_all_tool_definitions`.
pub(super) fn build_delegation_tool_definitions() -> Vec<roboticus_llm::format::ToolDefinition> {
    use roboticus_llm::format::ToolDefinition;

    vec![
        ToolDefinition {
            name: "orchestrate-subagents".into(),
            description: "Step 5 of **Subagent orchestration**: delegate one or more subtasks after \
                you classified a non-quick turn, reviewed `list-subagent-roster`, and created a specialist \
                if needed (`compose-subagent` or approval flow). Each subtask is matched to skills; \
                set `subagent` on an item to pin a name."
                .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "subtasks": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "task": { "type": "string", "description": "The subtask description" },
                                "subagent": { "type": "string", "description": "Optional: specific subagent name to assign to" }
                            },
                            "required": ["task"]
                        },
                        "description": "List of subtasks to delegate"
                    }
                },
                "required": ["subtasks"]
            }),
        },
        ToolDefinition {
            name: "compose-subagent".into(),
            description:
                "Step 4 of **Subagent orchestration**: create a specialist when `list-subagent-roster` \
                shows no suitable **description** + **skills** fit for a non-quick task. Requires \
                name, description (purpose), and skills keywords."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "name": { "type": "string", "description": "Unique name for the new subagent" },
                    "display_name": { "type": "string", "description": "Human-readable display name" },
                    "description": { "type": "string", "description": "What this subagent specialises in" },
                    "skills": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "List of skill/capability keywords"
                    },
                    "model": { "type": "string", "description": "Optional: preferred model for this subagent" }
                },
                "required": ["name", "description", "skills"]
            }),
        },
    ]
}

/// Build tool definitions for virtual orchestration tools (manage subagent roster).
///
/// These are the "admin" counterpart to the delegation tools above — they let
/// the orchestrator inspect and modify its subagent fleet at runtime.
fn build_orchestration_tool_definitions() -> Vec<roboticus_llm::format::ToolDefinition> {
    use roboticus_llm::format::ToolDefinition;

    vec![
        ToolDefinition {
            name: "list-subagent-roster".into(),
            description:
                "Step 3 of **Subagent orchestration**: for non-quick turns, list all subagents with \
                purpose (**description**), **skills**, models, and runtime state before delegating. \
                Compare each row to the user's task to decide fit or whether to compose a new specialist."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {},
            }),
        },
        ToolDefinition {
            name: "list-available-skills".into(),
            description:
                "List skills registered in the workspace skill catalog. \
                 Optionally filter by keyword."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "keyword": {
                        "type": "string",
                        "description": "Optional keyword to filter skills"
                    }
                },
            }),
        },
        ToolDefinition {
            name: "update-subagent-skills".into(),
            description:
                "Update the skill list for an existing subagent. \
                 Can replace or append skills."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "name": { "type": "string", "description": "Name of the subagent to update" },
                    "skills": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "New skill/capability keywords"
                    },
                    "mode": {
                        "type": "string",
                        "enum": ["replace", "append"],
                        "description": "Whether to replace all skills or append to existing (default: replace)"
                    }
                },
                "required": ["name", "skills"]
            }),
        },
        ToolDefinition {
            name: "remove-subagent".into(),
            description:
                "Remove a subagent from the roster. Stops the runtime and deletes the DB record."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "name": { "type": "string", "description": "Name of the subagent to remove" }
                },
                "required": ["name"]
            }),
        },
        ToolDefinition {
            name: "retire-unused-subagents".into(),
            description:
                "Soft-retire (set enabled=false) taskable subagents that were **never** used: no \
                 chat sessions scoped to that subagent id and no successful delegation tool output \
                 targeting them. Requires minimum age in days since row `created_at` unless \
                 min_age_days is 0. Prefer dry_run first."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "min_age_days": {
                        "type": "integer",
                        "description": "Minimum age in whole days since created_at (default 30). Use 0 to ignore age."
                    },
                    "dry_run": {
                        "type": "boolean",
                        "description": "If true, only return who would be retired (default false)"
                    },
                    "names": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "Optional filter: only evaluate these subagent names"
                    }
                }
            }),
        },
        ToolDefinition {
            name: "assign-tasks".into(),
            description:
                "Step 5 of **Subagent orchestration**: delegate one task to the best-matching **running** \
                 subagent after roster review (and creation if needed). Returns the specialist reply."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "task": { "type": "string", "description": "The task description to delegate" },
                    "subagent": { "type": "string", "description": "Optional: specific subagent name" }
                },
                "required": ["task"]
            }),
        },
        ToolDefinition {
            name: "delegate-subagent".into(),
            description:
                "Step 5 of **Subagent orchestration**: delegate after classifying a non-quick turn and \
                 reviewing the roster (use `subagent` when you already selected a name). Supports \
                 multiple `subtasks` in one call."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "task": { "type": "string", "description": "Primary task description" },
                    "subtasks": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "Optional: list of subtask strings"
                    },
                    "subagent": { "type": "string", "description": "Optional: specific subagent name" }
                },
                "required": ["task"]
            }),
        },
        ToolDefinition {
            name: "select-subagent-model".into(),
            description:
                "Select the best subagent and resolve its runtime model for a given task. \
                 Does not execute the task — returns the selection decision."
                    .into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "task": { "type": "string", "description": "Task description for matching" },
                    "subagent": { "type": "string", "description": "Optional: prefer specific subagent" }
                },
                "required": ["task"]
            }),
        },
    ]
}

/// Build tool definitions from the capability registry when populated, otherwise from `ToolRegistry`,
/// plus virtual delegation/orchestration tools.
///
/// This is the single source of truth for all tool definitions sent to the LLM provider.
pub(super) async fn build_all_tool_definitions(
    state: &AppState,
) -> Vec<roboticus_llm::format::ToolDefinition> {
    use roboticus_llm::format::ToolDefinition;

    let mut defs = build_delegation_tool_definitions();
    defs.extend(build_orchestration_tool_definitions());

    if state.capabilities.is_empty().await {
        for tool in state.tools.list() {
            defs.push(ToolDefinition {
                name: tool.name().to_string(),
                description: tool.description().to_string(),
                parameters: tool.parameters_schema(),
            });
        }
    } else {
        for summary in state.capabilities.catalog().await {
            defs.push(ToolDefinition {
                name: summary.name,
                description: summary.description,
                parameters: summary.parameters_schema,
            });
        }
    }

    defs
}