#!/usr/bin/env python3
"""Third-party-package-free STDIO MCP server for a compiled Codex MMO mesh."""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import sys
import traceback
from typing import Any

from mmo_profiles import agent_mcp_tool_names
from mmo_runtime import (
    ADMITTING_JOB_STATUSES,
    ADMITTING_SESSION_STATUSES,
    accept_result,
    cancel_job,
    control_agent_run,
    inspect_agent_run,
    integrate_patch,
    list_agent_runs,
    list_jobs,
    load_job,
    load_session,
    read_agent_trace,
    read_agent_trace_record,
    read_result,
    reject_result,
    spawn_job,
    spawn_jobs,
    wait_for_jobs,
)
from mmo_schema import validate_instance
from mmo_snapshot import load_snapshot
from mmo_util import package_version, strict_json_loads

PROTOCOL_VERSION = "2025-06-18"
SERVER_NAME = "codex-mmo-agent-mesh"
SESSION_ID = os.environ.get("MMO_ROOT_SESSION_ID", "")
RUN_ID = os.environ.get("MMO_RUN_ID", "")
CALLER_AGENT = os.environ.get("MMO_CALLER_AGENT", "")
CALLER_JOB_ID = os.environ.get("MMO_CALLER_JOB_ID") or None
CALLER_NATIVE = os.environ.get("MMO_CALLER_NATIVE") == "1"
CALLER_TOKEN = os.environ.get("MMO_CALLER_TOKEN", "")
NATIVE_CALLER_TOKEN = os.environ.get("MMO_NATIVE_CALLER_TOKEN", "")


def _context() -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    if not SESSION_ID or not RUN_ID or not CALLER_AGENT:
        raise RuntimeError("MMO MCP server is missing session/run/caller identity")
    session = load_session(SESSION_ID)
    if session.get("status") not in ADMITTING_SESSION_STATUSES:
        raise PermissionError("root session is not active")
    if session.get("current_run_id") != RUN_ID:
        raise PermissionError("MCP caller belongs to an inactive session run")
    snapshot = load_snapshot(session["snapshot_hash"])
    resolved = snapshot["resolved"]
    if CALLER_AGENT not in resolved["agents"]:
        raise RuntimeError(f"unknown caller agent: {CALLER_AGENT}")
    if CALLER_JOB_ID:
        caller_job = load_job(CALLER_JOB_ID)
        if caller_job.get("session_id") != SESSION_ID:
            raise PermissionError("MCP caller job belongs to another session")
        if caller_job.get("run_id") != RUN_ID:
            raise PermissionError("MCP caller job belongs to another session run")
        if caller_job.get("status") not in ADMITTING_JOB_STATUSES:
            raise PermissionError("MCP caller job is not active")
        if not CALLER_NATIVE and caller_job.get("agent") != CALLER_AGENT:
            raise PermissionError("MCP caller role does not match its job")
        expected_caller_token = str(caller_job.get("mcp_caller_token_hash", ""))
    else:
        if not CALLER_NATIVE and CALLER_AGENT != session.get("root_agent"):
            raise PermissionError("only the root role may call without a job identity")
        expected_caller_token = str(session.get("root_mcp_token_hash", ""))
    observed_caller_token = hashlib.sha256(CALLER_TOKEN.encode("utf-8")).hexdigest()
    if (
        not expected_caller_token
        or not CALLER_TOKEN
        or not hmac.compare_digest(expected_caller_token, observed_caller_token)
    ):
        raise PermissionError("invalid MCP caller capability token")
    if CALLER_NATIVE:
        expected = str(session.get("native_token_hashes", {}).get(CALLER_AGENT, ""))
        observed = hashlib.sha256(NATIVE_CALLER_TOKEN.encode("utf-8")).hexdigest()
        if not expected or not NATIVE_CALLER_TOKEN or not hmac.compare_digest(expected, observed):
            raise PermissionError("invalid native-agent MCP capability token")
        if "native" not in resolved["agents"][CALLER_AGENT].get("backends", []):
            raise PermissionError("caller role is not native-enabled")
    return session, resolved, resolved["agents"][CALLER_AGENT]


def _child_description(agent_id: str, agent: dict[str, Any], resolved: dict[str, Any]) -> str:
    model = resolved["models"][agent["model"]]
    return (
        f"{agent_id}: {agent.get('description') or 'profile participant'}; "
        f"model={model['display_name']}; permissions={agent['permissions']}; "
        f"trust={agent['trust']}; task_kinds={','.join(agent['allowed_task_kinds'])}; "
        f"max_active={agent['max_active']}"
    )


def _available_children(
    session: dict[str, Any], resolved: dict[str, Any], caller: dict[str, Any]
) -> list[str]:
    if session.get("tainted"):
        return []
    return [
        child
        for child in caller["can_spawn"]
        if "mcp" in resolved["agents"][child].get("backends", [])
        and session.get("route_availability", {})
        .get(resolved["agents"][child]["route"], {})
        .get("available")
    ]


def _common_spawn_properties(children: list[str]) -> dict[str, Any]:
    return {
        "agent": {
            "type": "string",
            "enum": children,
            "description": "Profile participant to launch. Select by role capability.",
        },
        "mode": {
            "type": "string",
            "enum": ["read-only", "workspace-write"],
            "default": "read-only",
        },
        "cwd": {
            "type": "string",
            "description": "Optional directory inside the root session checkout.",
        },
        "write_scope": {
            "type": "array",
            "items": {"type": "string", "minLength": 1},
            "maxItems": 64,
            "description": "Required for workspace-write; smallest disjoint files/directories owned by the child.",
        },
        "attachments": {
            "type": "array",
            "items": {"type": "string", "minLength": 1},
            "maxItems": 12,
            "description": "Files inside the session root. Image files require an image-capable role and transport.",
        },
        "label": {"type": "string", "maxLength": 80},
    }


def _literal_task_schema() -> dict[str, Any]:
    path_array = {
        "type": "array",
        "items": {"type": "string", "minLength": 1},
        "minItems": 1,
        "maxItems": 32,
    }
    return {
        "oneOf": [
            {
                "type": "object",
                "properties": {
                    "operation": {"const": "locate"},
                    "needle": {"type": "string", "minLength": 1, "maxLength": 500},
                    "paths": path_array,
                    "max_results": {"type": "integer", "minimum": 1, "maximum": 200},
                },
                "required": ["operation", "needle"],
                "additionalProperties": False,
            },
            {
                "type": "object",
                "properties": {
                    "operation": {"const": "references"},
                    "symbol": {"type": "string", "minLength": 1, "maxLength": 500},
                    "paths": path_array,
                    "max_results": {"type": "integer", "minimum": 1, "maximum": 200},
                },
                "required": ["operation", "symbol"],
                "additionalProperties": False,
            },
            {
                "type": "object",
                "properties": {
                    "operation": {"const": "extract"},
                    "path": {"type": "string", "minLength": 1},
                    "start_line": {"type": "integer", "minimum": 1},
                    "end_line": {"type": "integer", "minimum": 1},
                },
                "required": ["operation", "path", "start_line", "end_line"],
                "additionalProperties": False,
            },
            {
                "type": "object",
                "properties": {
                    "operation": {"const": "summarize_supplied"},
                    "text": {"type": "string", "minLength": 1, "maxLength": 50000},
                    "max_points": {"type": "integer", "minimum": 1, "maximum": 50},
                },
                "required": ["operation", "text"],
                "additionalProperties": False,
            },
        ]
    }


def _spawn_schema(children: list[str], resolved: dict[str, Any]) -> dict[str, Any]:
    """Advertise one stable object shape; selected-role rules remain authoritative."""

    agents = [resolved["agents"][child] for child in children]
    properties = _common_spawn_properties(children)
    required = ["agent"]
    low_trust = [agent for agent in agents if agent["trust"] == "low"]
    ordinary = [agent for agent in agents if agent["trust"] != "low"]
    if low_trust:
        properties["literal_task"] = _literal_task_schema()
        properties["literal_task"]["description"] = (
            "Required only for a selected low-trust role; exact operation validation is server-side."
        )
    if ordinary:
        task_kinds = sorted({kind for agent in ordinary for kind in agent["allowed_task_kinds"]})
        properties["task_kind"] = {
            "type": "string",
            "enum": task_kinds,
            "description": "Must be allowed by the selected agent role.",
        }
        properties["task"] = {
            "type": "string",
            "minLength": min(int(agent["min_task_chars"]) for agent in ordinary),
            "maxLength": max(int(agent["max_task_chars"]) for agent in ordinary),
            "description": (
                "Self-contained bounded brief: objective, context, non-goals, deliverable, and "
                "validation. The selected role's exact length and task-kind limits are enforced "
                "server-side. Do not delegate the caller's immediate critical path."
            ),
        }
    if low_trust and not ordinary:
        required.append("literal_task")
    elif ordinary and not low_trust:
        required.extend(["task_kind", "task"])
    return {
        "type": "object",
        "properties": properties,
        "required": required,
        "additionalProperties": False,
    }


def _control_tool_definitions(
    resolved: dict[str, Any], caller: dict[str, Any]
) -> list[dict[str, Any]]:
    controls = caller.get("controls", {})
    controlled = sorted(controls) if isinstance(controls, dict) else []
    if not controlled:
        return []
    granted_actions = {
        action
        for grant in controls.values()
        if isinstance(grant, dict)
        for action in grant.get("actions", [])
    }
    role_text = ", ".join(controlled)
    base = {
        "agent_run_ref": {
            "type": "string",
            "pattern": "^ar_",
            "description": (
                "Opaque run reference from agent_list for an authorized target role: " + role_text
            ),
        }
    }
    revision = {
        "expected_revision": {
            "type": "integer",
            "minimum": 0,
            "description": "Compare-and-swap revision returned by agent_inspect/status.",
        }
    }

    def schema(properties: dict[str, Any], required: list[str]) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {**base, **properties},
            "required": ["agent_run_ref", *required],
            "additionalProperties": False,
        }

    mutation = {
        "readOnlyHint": False,
        "destructiveHint": False,
        "idempotentHint": False,
        "openWorldHint": False,
    }
    text = {"type": "string", "minLength": 1, "maxLength": 20000}
    efforts = sorted(
        {
            effort
            for role in controlled
            for effort in resolved["agents"][role].get("allowed_reasoning_efforts", [])
        }
    )
    task_kinds = sorted(
        {
            kind
            for role in controlled
            for kind in resolved["agents"][role].get("allowed_task_kinds", [])
        }
    )
    definitions: list[dict[str, Any]] = [
        {
            "name": "agent_list",
            "description": (
                "Discover authorized root, native Codex, and supervised MCP runs and their "
                "opaque run references."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
            "annotations": {"readOnlyHint": True, "openWorldHint": False},
        },
        {
            "name": "agent_inspect",
            "description": "Inspect durable and live state for an authorized agent run.",
            "inputSchema": schema({}, []),
            "annotations": {"readOnlyHint": True, "openWorldHint": False},
        },
        {
            "name": "agent_trace",
            "description": (
                "Read a paginated durable event trace. Private reasoning is filtered while "
                "completed messages and empirical tool evidence are retained."
            ),
            "inputSchema": schema(
                {
                    "cursor": {"type": "integer", "minimum": 0},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 200},
                },
                [],
            ),
            "annotations": {"readOnlyHint": True, "openWorldHint": False},
        },
        {
            "name": "agent_trace_record",
            "description": (
                "Read an oversized filtered trace record without loss. Use the record_cursor "
                "reported by agent_trace and follow next_cursor until null."
            ),
            "inputSchema": schema(
                {
                    "record_cursor": {"type": "integer", "minimum": 0},
                    "cursor": {"type": "integer", "minimum": 0},
                    "max_chars": {"type": "integer", "minimum": 500, "maximum": 500000},
                },
                ["record_cursor"],
            ),
            "annotations": {"readOnlyHint": True, "openWorldHint": False},
        },
        {
            "name": "agent_steer",
            "description": "Inject guidance into the worker's active turn.",
            "inputSchema": schema({**revision, "input": text}, ["expected_revision", "input"]),
            "annotations": mutation,
        },
        {
            "name": "agent_interrupt",
            "description": "Interrupt the active turn but retain its durable thread.",
            "inputSchema": schema(revision, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_pause",
            "description": (
                "Cold-pause a supervised worker: checkpoint evidence, retire its host, release "
                "capacity, and retain the exact thread for fresh-admission continuation."
            ),
            "inputSchema": schema(revision, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_continue",
            "description": (
                "Continue a paused, detached, or suspended run on its existing thread. "
                "Goal-mode runs may receive a larger total token budget within the compiled cap."
            ),
            "inputSchema": schema(
                {
                    **revision,
                    "input": text,
                    "goal_token_budget": {"type": "integer", "minimum": 10000},
                },
                ["expected_revision"],
            ),
            "annotations": mutation,
        },
        {
            "name": "agent_detach",
            "description": "Detach supervision while the durable run continues.",
            "inputSchema": schema(revision, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_stop",
            "description": "Fully stop the run while retaining durable history and evidence.",
            "inputSchema": schema(revision, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_finalize",
            "description": "Request evidence-only, contract-compliant finalization.",
            "inputSchema": schema({**revision, "input": text}, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_compact",
            "description": "Start Codex compaction for the durable thread.",
            "inputSchema": schema(revision, ["expected_revision"]),
            "annotations": mutation,
        },
        {
            "name": "agent_respond",
            "description": (
                "Answer a pending app-server user-input, MCP-elicitation, or approval request "
                "with the exact response object required by that request method."
            ),
            "inputSchema": schema(
                {
                    **revision,
                    "request_id": {"oneOf": [{"type": "integer"}, {"type": "string"}]},
                    "response": {"type": "object"},
                },
                ["expected_revision", "request_id", "response"],
            ),
            "annotations": mutation,
        },
        {
            "name": "agent_set_effort",
            "description": "Change reasoning effort within the role's compiled allowlist.",
            "inputSchema": schema(
                {**revision, "effort": {"type": "string", "enum": efforts}},
                ["expected_revision", "effort"],
            ),
            "annotations": mutation,
        },
        {
            "name": "agent_fork",
            "description": (
                "Fork a persisted MCP worker into a normally admitted independent job. A native "
                "fork instead inherits its role, cwd, and sandbox inside the shared root host, "
                "obeys the native-thread limit, and rejects MCP-only scope or attachment overrides."
            ),
            "inputSchema": schema(
                {
                    **revision,
                    "task_kind": {"type": "string", "enum": task_kinds},
                    "input": text,
                    "mode": {
                        "type": "string",
                        "enum": ["read-only", "workspace-write"],
                        "default": "read-only",
                    },
                    "cwd": {"type": "string"},
                    "write_scope": {
                        "type": "array",
                        "items": {"type": "string", "minLength": 1},
                        "maxItems": 64,
                    },
                    "attachments": {
                        "type": "array",
                        "items": {"type": "string", "minLength": 1},
                        "maxItems": 12,
                    },
                    "label": {"type": "string", "maxLength": 80},
                },
                ["expected_revision", "task_kind", "input"],
            ),
            "annotations": mutation,
        },
    ]
    action_by_tool = {
        "agent_list": None,
        "agent_inspect": "inspect",
        "agent_trace": "trace",
        "agent_trace_record": "trace",
        "agent_steer": "steer",
        "agent_interrupt": "interrupt",
        "agent_pause": "pause",
        "agent_continue": "continue",
        "agent_detach": "detach",
        "agent_stop": "stop",
        "agent_finalize": "finalize",
        "agent_compact": "compact",
        "agent_respond": "respond",
        "agent_set_effort": "set_effort",
        "agent_fork": "fork",
    }
    return [
        definition
        for definition in definitions
        if action_by_tool[definition["name"]] is None
        or action_by_tool[definition["name"]] in granted_actions
    ]


def tool_definitions() -> list[dict[str, Any]]:
    session, resolved, caller = _context()
    children = _available_children(session, resolved, caller)
    child_text = (
        "\n".join(
            "- " + _child_description(key, resolved["agents"][key], resolved) for key in children
        )
        or "- none"
    )
    spawn_description = (
        "Start one real asynchronous profile participant and return immediately. Continue useful "
        "non-overlapping caller work; wait only at a dependency barrier. Available children:\n"
        + child_text
    )
    tools: list[dict[str, Any]] = []
    if children:
        single_schema = _spawn_schema(children, resolved)
        tools.extend(
            [
                {
                    "name": "agent_spawn",
                    "description": spawn_description,
                    "inputSchema": single_schema,
                    "annotations": {
                        "readOnlyHint": False,
                        "destructiveHint": False,
                        "idempotentHint": False,
                        "openWorldHint": False,
                    },
                },
                {
                    "name": "agents_spawn",
                    "description": (
                        "Start multiple independent participants in one tool turn. The complete batch is "
                        "validated and admitted atomically: either every request starts or none does. "
                        "Immediately continue the caller's distinct critical-path work.\n"
                        + child_text
                    ),
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "agents": {
                                "type": "array",
                                "items": single_schema,
                                "minItems": 1,
                                "maxItems": 12,
                            }
                        },
                        "required": ["agents"],
                        "additionalProperties": False,
                    },
                    "annotations": {
                        "readOnlyHint": False,
                        "destructiveHint": False,
                        "idempotentHint": False,
                        "openWorldHint": False,
                    },
                },
            ]
        )
    management_tools = [
        {
            "name": "agent_status",
            "description": "Poll visible descendants without blocking. Omit job_ids to list recent visible jobs.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_ids": {"type": "array", "items": {"type": "string"}, "maxItems": 100},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 1000},
                },
                "additionalProperties": False,
            },
            "annotations": {
                "readOnlyHint": False,
                "destructiveHint": False,
                "idempotentHint": False,
                "openWorldHint": False,
            },
        },
        {
            "name": "agents_wait",
            "description": (
                "Wait for specified descendants only when the next action genuinely depends on them. "
                "The wait is bounded to 120 seconds and returns compact state. Pass the exact progress "
                "revision map for every requested job to return at the first durable change."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "minItems": 1,
                        "maxItems": 100,
                    },
                    "timeout_seconds": {"type": "integer", "minimum": 0, "maximum": 120},
                    "include_results": {"type": "boolean"},
                    "after_revision": {
                        "type": "object",
                        "additionalProperties": {
                            "type": "string",
                            "minLength": 64,
                            "maxLength": 64,
                        },
                    },
                },
                "required": ["job_ids"],
                "additionalProperties": False,
            },
            "annotations": {
                "readOnlyHint": False,
                "destructiveHint": False,
                "idempotentHint": False,
                "openWorldHint": False,
            },
        },
        {
            "name": "agent_result",
            "description": (
                "Read a terminal descendant result without loss. Start at cursor 0 and follow "
                "next_cursor until null. Reading changes its lifecycle from unread to read; "
                "it does not accept, reject, consume, or integrate the result."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_id": {"type": "string"},
                    "cursor": {"type": "integer", "minimum": 0},
                    "max_chars": {"type": "integer", "minimum": 500, "maximum": 500000},
                },
                "required": ["job_id"],
                "additionalProperties": False,
            },
            "annotations": {
                "readOnlyHint": False,
                "destructiveHint": False,
                "idempotentHint": False,
                "openWorldHint": False,
            },
        },
        {
            "name": "agent_result_accept",
            "description": (
                "Accept a result after reading and reviewing it. A writable result is still not "
                "integrated until agent_patch_integrate succeeds."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_id": {"type": "string"},
                    "reason": {"type": "string", "minLength": 1, "maxLength": 500},
                },
                "required": ["job_id", "reason"],
                "additionalProperties": False,
            },
            "annotations": {"readOnlyHint": False, "openWorldHint": False},
        },
        {
            "name": "agent_result_reject",
            "description": "Reject a result after reading it, with an auditable evidence-based reason.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_id": {"type": "string"},
                    "reason": {"type": "string", "minLength": 1, "maxLength": 500},
                },
                "required": ["job_id", "reason"],
                "additionalProperties": False,
            },
            "annotations": {"readOnlyHint": False, "openWorldHint": False},
        },
        {
            "name": "agent_patch_integrate",
            "description": (
                "Validate base fingerprints and git apply --check, then explicitly integrate an "
                "accepted isolated writable-worker patch into the canonical workspace."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_id": {"type": "string"},
                    "reason": {"type": "string", "minLength": 1, "maxLength": 500},
                },
                "required": ["job_id", "reason"],
                "additionalProperties": False,
            },
            "annotations": {
                "readOnlyHint": False,
                "destructiveHint": True,
                "idempotentHint": False,
                "openWorldHint": False,
            },
        },
        {
            "name": "agent_cancel",
            "description": "Cancel a stale, duplicate, mis-scoped, or no-longer-useful descendant and its child subtree.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "job_id": {"type": "string"},
                    "cascade": {"type": "boolean"},
                    "reason": {
                        "type": "string",
                        "maxLength": 500,
                        "description": "Why the job is stale, duplicate, mis-scoped, or no longer useful.",
                    },
                },
                "required": ["job_id"],
                "additionalProperties": False,
            },
            "annotations": {
                "readOnlyHint": False,
                "destructiveHint": True,
                "idempotentHint": True,
                "openWorldHint": False,
            },
        },
    ]
    if caller["can_spawn"] or caller.get("controls"):
        if session.get("tainted"):
            management_tools = [
                tool for tool in management_tools if tool["name"] != "agent_patch_integrate"
            ]
        tools.extend(management_tools)
    tools.extend(_control_tool_definitions(resolved, caller))
    allowed = agent_mcp_tool_names(
        resolved["agents"],
        root_agent=str(resolved["profile"]["root"]),
        agent_id=CALLER_AGENT,
    )
    return [tool for tool in tools if tool.get("name") in allowed]


def _spawn_kwargs(arguments: dict[str, Any]) -> dict[str, Any]:
    return {
        "session_id": SESSION_ID,
        "caller_agent": CALLER_AGENT,
        "caller_job_id": CALLER_JOB_ID,
        "caller_native": CALLER_NATIVE,
        "agent_id": arguments["agent"],
        "task_kind": arguments.get("task_kind"),
        "task": arguments.get("task"),
        "literal_task": arguments.get("literal_task"),
        "mode": arguments.get("mode", "read-only"),
        "cwd_value": arguments.get("cwd"),
        "write_scope_values": arguments.get("write_scope", []),
        "attachments": arguments.get("attachments", []),
        "label": arguments.get("label"),
    }


def _tool_call_validation_error(name: str, arguments: dict[str, Any]) -> str | None:
    definition = next(
        (item for item in tool_definitions() if item.get("name") == name),
        None,
    )
    if definition is None:
        return f"unknown tool: {name}"
    validation_errors = validate_instance(arguments, definition["inputSchema"])
    if validation_errors:
        return "invalid tool arguments: " + "; ".join(validation_errors[:12])
    return None


def _dispatch_tool(name: str, arguments: dict[str, Any]) -> Any:
    if name == "agent_spawn":
        result = spawn_job(**_spawn_kwargs(arguments))
        result["coordination_note"] = (
            "Participant started asynchronously. Continue useful non-overlapping work; "
            "do not wait unless the next action is blocked."
        )
        return result
    if name == "agents_spawn":
        return spawn_jobs(
            arguments["agents"],
            session_id=SESSION_ID,
            caller_agent=CALLER_AGENT,
            caller_job_id=CALLER_JOB_ID,
            caller_native=CALLER_NATIVE,
        )
    if name == "agent_status":
        job_ids = arguments.get("job_ids")
        return {
            "jobs": list_jobs(
                session_id=SESSION_ID,
                job_ids=job_ids,
                caller_job_id=CALLER_JOB_ID,
                caller_agent=CALLER_AGENT,
                caller_native=CALLER_NATIVE,
                limit=arguments.get("limit", 50),
            )
        }
    if name == "agents_wait":
        return wait_for_jobs(
            arguments["job_ids"],
            session_id=SESSION_ID,
            caller_job_id=CALLER_JOB_ID,
            caller_agent=CALLER_AGENT,
            caller_native=CALLER_NATIVE,
            timeout_seconds=arguments.get("timeout_seconds", 30),
            include_results=arguments.get("include_results", False),
            after_revision=arguments.get("after_revision"),
        )
    if name == "agent_result":
        return read_result(
            arguments["job_id"],
            session_id=SESSION_ID,
            caller_job_id=CALLER_JOB_ID,
            caller_agent=CALLER_AGENT,
            caller_native=CALLER_NATIVE,
            max_chars=arguments.get("max_chars"),
            cursor=arguments.get("cursor", 0),
        )
    control_context: dict[str, Any] = {
        "session_id": SESSION_ID,
        "caller_job_id": CALLER_JOB_ID,
        "caller_agent": CALLER_AGENT,
        "caller_native": CALLER_NATIVE,
    }
    if name == "agent_list":
        return {"agents": list_agent_runs(**control_context)}
    if name == "agent_inspect":
        return inspect_agent_run(arguments["agent_run_ref"], **control_context)
    if name == "agent_trace":
        return read_agent_trace(
            arguments["agent_run_ref"],
            cursor=arguments.get("cursor", 0),
            limit=arguments.get("limit", 100),
            **control_context,
        )
    if name == "agent_trace_record":
        return read_agent_trace_record(
            arguments["agent_run_ref"],
            record_cursor=arguments["record_cursor"],
            cursor=arguments.get("cursor", 0),
            max_chars=arguments.get("max_chars"),
            **control_context,
        )
    control_actions = {
        "agent_steer": "steer",
        "agent_interrupt": "interrupt",
        "agent_pause": "pause",
        "agent_continue": "continue",
        "agent_detach": "detach",
        "agent_stop": "stop",
        "agent_finalize": "finalize",
        "agent_compact": "compact",
        "agent_respond": "respond",
        "agent_set_effort": "set_effort",
        "agent_fork": "fork",
    }
    if name in control_actions:
        extras = {
            key: value
            for key, value in arguments.items()
            if key not in {"agent_run_ref", "expected_revision"}
        }
        return control_agent_run(
            arguments["agent_run_ref"],
            control_actions[name],
            expected_revision=arguments["expected_revision"],
            **control_context,
            **extras,
        )
    disposition_context: dict[str, Any] = {
        "session_id": SESSION_ID,
        "caller_job_id": CALLER_JOB_ID,
        "caller_agent": CALLER_AGENT,
        "caller_native": CALLER_NATIVE,
    }
    if name == "agent_result_accept":
        return accept_result(arguments["job_id"], arguments["reason"], **disposition_context)
    if name == "agent_result_reject":
        return reject_result(arguments["job_id"], arguments["reason"], **disposition_context)
    if name == "agent_patch_integrate":
        return integrate_patch(
            arguments["job_id"],
            arguments["reason"],
            session_id=SESSION_ID,
            caller_job_id=CALLER_JOB_ID,
            caller_agent=CALLER_AGENT,
            caller_native=CALLER_NATIVE,
        )
    if name == "agent_cancel":
        return cancel_job(
            arguments["job_id"],
            session_id=SESSION_ID,
            caller_job_id=CALLER_JOB_ID,
            caller_agent=CALLER_AGENT,
            caller_native=CALLER_NATIVE,
            cascade=arguments.get("cascade", True),
            reason=arguments.get("reason"),
        )
    raise AssertionError(f"unhandled advertised tool: {name}")


def call_tool(name: str, arguments: dict[str, Any]) -> Any:
    validation_error = _tool_call_validation_error(name, arguments)
    if validation_error is not None:
        raise ValueError(validation_error)
    return _sanitize_mcp_result(_dispatch_tool(name, arguments))


_INTERNAL_MCP_FIELDS = frozenset(
    {
        "app_server_socket_path",
        "control_socket_path",
        "events_path",
        "full_result_path",
        "partial_result_path",
        "result_path",
        "root_app_server_socket",
        "root_control_socket",
        "socket",
        "socket_path",
        "stderr_path",
        "structured_result_path",
    }
)


def _sanitize_mcp_result(value: Any, *, _path: tuple[str, ...] = ()) -> Any:
    """Remove supervisor locations without rewriting opaque model result payloads."""

    if isinstance(value, dict):
        return {
            key: (
                child
                if (not _path and key == "content")
                or (len(_path) == 2 and _path[0] == "results" and key == "preview")
                else _sanitize_mcp_result(child, _path=(*_path, key))
            )
            for key, child in value.items()
            if key not in _INTERNAL_MCP_FIELDS
        }
    if isinstance(value, list):
        return [_sanitize_mcp_result(child, _path=_path) for child in value]
    return value


def _response(request_id: Any, result: Any) -> dict[str, Any]:
    return {"jsonrpc": "2.0", "id": request_id, "result": result}


def _error(request_id: Any, code: int, message: str, data: Any = None) -> dict[str, Any]:
    value: dict[str, Any] = {"code": code, "message": message}
    if data is not None:
        value["data"] = data
    return {"jsonrpc": "2.0", "id": request_id, "error": value}


def _emit(message: dict[str, Any]) -> None:
    sys.stdout.write(
        json.dumps(message, ensure_ascii=False, separators=(",", ":"), allow_nan=False) + "\n"
    )
    sys.stdout.flush()


def _valid_request_id(value: Any) -> bool:
    return isinstance(value, str) or (isinstance(value, int) and not isinstance(value, bool))


def _invalid_request(message: dict[str, Any], detail: str) -> dict[str, Any]:
    request_id = message.get("id") if _valid_request_id(message.get("id")) else None
    return _error(request_id, -32600, "invalid request", detail)


def _validate_rpc_request(message: dict[str, Any]) -> dict[str, Any] | None:
    if message.get("jsonrpc") != "2.0":
        return _invalid_request(message, "jsonrpc must equal '2.0'")
    if not isinstance(message.get("method"), str):
        return _invalid_request(message, "method must be a string")
    if "id" in message and not _valid_request_id(message["id"]):
        return _invalid_request(message, "id must be a string or integer")
    if "params" in message and not isinstance(message["params"], dict):
        return _invalid_request(message, "params must be an object")
    return None


def _method_params(
    message: dict[str, Any], request_id: Any, *, has_id: bool
) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
    params = message.get("params", {})
    if isinstance(params, dict):
        return params, None
    if not has_id:
        return None, None
    return None, _error(request_id, -32602, "method parameters must be an object")


def _valid_client_info(value: Any) -> bool:
    return (
        isinstance(value, dict)
        and isinstance(value.get("name"), str)
        and bool(value["name"])
        and isinstance(value.get("version"), str)
        and bool(value["version"])
    )


def handle(message: dict[str, Any], state: dict[str, str]) -> dict[str, Any] | None:
    invalid = _validate_rpc_request(message)
    if invalid is not None:
        return invalid
    method = message.get("method")
    has_id = "id" in message
    request_id = message.get("id") if has_id else None
    params, params_error = _method_params(message, request_id, has_id=has_id)
    if params is None:
        return params_error
    if method == "initialize":
        if not has_id:
            return None
        if state["phase"] != "new":
            return _error(request_id, -32600, "server is already initialized")
        if (
            not isinstance(params.get("protocolVersion"), str)
            or not params["protocolVersion"]
            or not isinstance(params.get("capabilities"), dict)
            or not _valid_client_info(params.get("clientInfo"))
        ):
            return _error(request_id, -32602, "invalid initialize parameters")
        _session, resolved, caller = _context()
        children = ", ".join(caller["can_spawn"]) or "none"
        controls = caller.get("controls", {})
        controlled = (
            "; ".join(
                f"{target}={','.join(grant.get('actions', []))}"
                for target, grant in sorted(controls.items())
            )
            if isinstance(controls, dict) and controls
            else "none"
        )
        instructions = (
            f"You are profile agent {CALLER_AGENT} in {resolved['profile']['id']}. "
            f"Permitted children: {children}. Permitted control targets: {controlled}. "
        )
        if caller["can_spawn"]:
            instructions += (
                "Spawn eligible independent work early, continue your own non-overlapping "
                "critical path, and wait only at a real dependency barrier. "
            )
        else:
            instructions += "This role has no spawn authority; do not create agents. "
        if caller.get("controls"):
            instructions += (
                "Call agent_list before control. Use opaque run references and only the exact "
                "action grants shown above. MCP pause checkpoints work and retires its host, detach "
                "lets it continue live, stop is terminal, and goal continuation is token-budgeted "
                "rather than clock-estimated. "
            )
        else:
            instructions += "This role has no cross-agent control authority. "
        instructions += (
            "First obtain every requested job's progress_revision with agent_status or an initial "
            "agents_wait call. Pass that exact map to later agents_wait calls as after_revision and "
            "keep result "
            "previews disabled unless needed. Read terminal output with agent_result from cursor "
            "0 through every next_cursor until null. Never open MMO supervisor state or result "
            "files directly; use agent_result, agent_inspect, agent_trace, and agent_trace_record."
        )
        response = _response(
            request_id,
            {
                # This server supports one protocol revision. MCP requires the
                # server to return a supported revision when the client's
                # requested revision is unsupported, not echo the unknown one.
                "protocolVersion": PROTOCOL_VERSION,
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {"name": SERVER_NAME, "version": package_version()},
                "instructions": instructions,
            },
        )
        state["phase"] = "initializing"
        return response
    if method == "notifications/initialized":
        if has_id:
            return _error(request_id, -32600, "notifications/initialized must not include an id")
        if state["phase"] == "initializing":
            state["phase"] = "initialized"
        return None
    if method == "notifications/cancelled":
        if has_id:
            return _error(request_id, -32600, "notifications/cancelled must not include an id")
        return None
    if method == "ping":
        return _response(request_id, {}) if has_id else None
    if state["phase"] != "initialized":
        return _error(request_id, -32600, "server not initialized") if has_id else None
    if method == "tools/list":
        return _response(request_id, {"tools": tool_definitions()}) if has_id else None
    if method == "tools/call":
        if not has_id:
            return None
        name = params.get("name")
        arguments = params.get("arguments", {})
        if not isinstance(name, str) or not isinstance(arguments, dict):
            return _error(request_id, -32602, "invalid tools/call parameters")
        validation_error = _tool_call_validation_error(name, arguments)
        if validation_error is not None:
            return _error(request_id, -32602, "invalid params", validation_error)
        try:
            result = _sanitize_mcp_result(_dispatch_tool(name, arguments))
            return _response(
                request_id,
                {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(
                                result,
                                ensure_ascii=False,
                                separators=(",", ":"),
                                allow_nan=False,
                            ),
                        }
                    ],
                    "structuredContent": result,
                    "isError": False,
                },
            )
        except (OSError, RuntimeError, ValueError) as exc:
            return _response(
                request_id,
                {
                    "content": [{"type": "text", "text": f"{type(exc).__name__}: {exc}"}],
                    "isError": True,
                },
            )
    if not has_id:
        return None
    return _error(request_id, -32601, f"method not found: {method}")


def main() -> int:
    state = {"phase": "new"}
    for raw in sys.stdin:
        line = raw.strip()
        if not line:
            continue
        message: Any = None
        try:
            message = strict_json_loads(line)
        except (json.JSONDecodeError, ValueError) as exc:
            _emit(_error(None, -32700, "parse error", str(exc)))
            continue
        if not isinstance(message, dict):
            _emit(_error(None, -32600, "invalid request", "message must be an object"))
            continue
        try:
            result = handle(message, state)
            if result is not None:
                _emit(result)
        except Exception as exc:
            print(traceback.format_exc(), file=sys.stderr, flush=True)
            if "id" in message and _valid_request_id(message.get("id")):
                _emit(_error(message.get("id"), -32603, str(exc)))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
