@roll-agent/octopus-agent 0.1.3 → 0.2.0

package/src/octopus_skill/octopus_run.py

@@ -1,755 +0,0 @@
-from __future__ import annotations
-
-import json
-import re
-import sys
-from dataclasses import asdict, is_dataclass
-from typing import Any, Callable
-
-from . import __version__
-from .agent import SpongeAgent, ensure_query_plan_summary_answer
-from .config import AppConfig, load_config_from_env
-from .result_renderer import RESULT_ROW_DISPLAY_LIMIT, normalize_result_format, render_result
-from .schema_cache import SchemaCacheError
-
-SERVER_NAME = "octopus-agent"
-SERVER_VERSION = __version__
-PROTOCOL_VERSION = "2024-11-05"
-
-Json = dict[str, Any]
-AgentFactory = Callable[[], SpongeAgent]
-SamplerFactory = Callable[[str, str], str]
-
-
-def main() -> None:
-    serve_stdio()
-
-
-def serve_stdio() -> None:
-    config = load_config_from_env()
-    session = StdioSession(
-        sys.stdin.buffer,
-        sys.stdout.buffer,
-        max_tokens=config.sampling.max_tokens,
-    )
-    while True:
-        message = session.read_message()
-        if message is None:
-            return
-        response = handle_request(message, sampler=session.sample)
-        if response is not None:
-            session.write_message(response)
-
-
-def handle_request(
-    message: Json,
-    agent_factory: AgentFactory | None = None,
-    sampler: SamplerFactory | None = None,
-) -> Json | None:
-    method = message.get("method")
-    if method is None:
-        return _error_response(message.get("id"), -32600, "Invalid request")
-    if "id" not in message:
-        return None
-
-    request_id = message["id"]
-    try:
-        if method == "initialize":
-            return _result_response(request_id, initialize_result())
-        if method == "ping":
-            return _result_response(request_id, {})
-        if method == "tools/list":
-            return _result_response(request_id, {"tools": tool_definitions()})
-        if method == "tools/call":
-            params = message.get("params") or {}
-            return _result_response(request_id, call_tool(params, agent_factory=agent_factory, sampler=sampler))
-        return _error_response(request_id, -32601, f"Method not found: {method}")
-    except Exception as exc:
-        return _error_response(request_id, -32603, str(exc), {"type": type(exc).__name__})
-
-
-def initialize_result() -> Json:
-    return {
-        "protocolVersion": PROTOCOL_VERSION,
-        "capabilities": {"tools": {}},
-        "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
-    }
-
-
-def tool_definitions() -> list[Json]:
-    return [
-        {
-            "name": "diagnostic_status",
-            "description": "调试专用：检查丸子Agent（Octopus Agent）配置、缓存和本地持久化状态。用户只是说使用、启动、检查、看看丸子Agent，或确认 Agent 是否可用时，只调用本工具并在返回状态后停止，不得主动调用 query_sponge 或编造示例业务问题。普通业务查询不要先调用本工具，直接调用 query_sponge。",
-            "inputSchema": {
-                "type": "object",
-                "properties": {},
-                "additionalProperties": False,
-            },
-        },
-        {
-            "name": "refresh_schema",
-            "description": "调试/运维专用：手动刷新本地 Schema Cache（表结构缓存）。只有用户明确要求刷新 schema 时才调用；普通业务查询不要先调用本工具；query_sponge 内部会先检查 schemaVersion 并按需刷新 schema cache，query_sponge 失败后也不要自动调用本工具。",
-            "inputSchema": {
-                "type": "object",
-                "required": ["traceId"],
-                "properties": {
-                    "traceId": {
-                        "type": "string",
-                        "description": "调用链路 ID，必填，会原样传给 Sponge get_schema",
-                    }
-                },
-                "additionalProperties": False,
-            },
-        },
-        {
-            "name": "query_sponge",
-            "description": "业务查询唯一入口：把自然语言问题转换为只读 SQL，并通过 Sponge MCP Server 校验和执行。只有用户给出明确业务数据查询问题时才调用本工具；用户只是说使用、启动、检查、看看丸子Agent，或确认 Agent 是否可用时，不得调用本工具，应调用 diagnostic_status 后停止，不得编造示例业务问题。本工具内部第一步会检查本地 schemaVersion，并调用 Sponge get_schema 按需刷新 schema cache；普通业务查询不得额外调用 refresh_schema。工具会先按 schema 中文语义识别表、字段、筛选条件、枚举和实体，把 search_entities 的真实命中项固化进 queryPlan，再直接生成、校验并执行 SQL，不再要求用户确认查询计划；状态枚举不单独追问，用户未提出状态条件时首次查询保持宽泛，但查询结果必须展示相关状态字段的中文含义；search_entities 只用于实体别称和模糊召回，品牌/项目等多个实体同时命中时不追问用户，queryPlan 会用同组 OR 条件承载。上层模型必须直接调用本工具；不要自行生成 SQL；不要解释 schema；不要改写用户问题。originalQuestion 必须传用户原始话术。对报告类问题必须只调用一次 query_sponge 处理完整原始问题，不要拆成多个子查询。工具返回后必须把 content[0].text 或 structuredContent.answer 原文作为最终用户可见答复完整展示；finalAnswerOnly=true（默认）时不得改写、删减、摘要、截断或追加内容，表格不得只展示前 N 条。content[0].text 只展示三段式用户可见答复（1.查询实体；2.筛选条件；3.查询结果）；structuredContent.executionSteps 会返回每一步状态和 durationMs，structuredContent.timing 会返回 totalDurationMs 与 stepDurationMs，失败时包含实际 error 原因。默认 fastMode=true、finalAnswerOnly=true，structuredContent.responseMode=verbatim；复杂查询会使用 schema context 和 LLM 生成。",
-            "inputSchema": {
-                "type": "object",
-                "required": ["question", "originalQuestion"],
-                "properties": {
-                    "question": {"type": "string", "description": "当前要处理的问题；兼容旧调用方"},
-                    "originalQuestion": {
-                        "type": "string",
-                        "description": "必填，用户原始话术。上层不得摘要、翻译、拆分或改写；即使 question 与原话一致也要原样传入。",
-                    },
-                    "fastMode": {
-                        "type": "boolean",
-                        "description": "是否启用快路径；默认 true。true 时跳过 LLM SQL 专家复核，仍保留本地规则、validate_sql 与 execute_sql 兜底；false 时仅对含聚合/子查询的高风险 SQL 启用专家复核。",
-                    },
-                    "finalAnswerOnly": {
-                        "type": "boolean",
-                        "description": "是否只返回最终业务答案；默认 true。true 时 content 与 structuredContent.answer 均为三段式 verbatim 答复：1.查询实体；2.筛选条件；3.查询结果。上层 Agent 必须原样展示给用户，不要只放在折叠步骤里，不要改写成摘要、截断表格或追加后续引导。includeSql=true 的显式调试请求仍会返回 SQL 调试信息。",
-                    },
-                    "sessionId": {"type": "string", "description": "可选会话 ID"},
-                    "resultFormat": {
-                        "type": "string",
-                        "enum": ["text", "json", "yaml", "md", "markdown", "html"],
-                        "description": "结果输出格式；不传则根据问题自动识别",
-                    },
-                    "includeSql": {
-                        "type": "boolean",
-                        "description": "是否在回参中包含 generatedSql、normalizedSql、repairedSql 等调试信息",
-                    },
-                },
-                "additionalProperties": False,
-            },
-        },
-    ]
-
-
-def call_tool(
-    params: Json,
-    agent_factory: AgentFactory | None = None,
-    sampler: SamplerFactory | None = None,
-) -> Json:
-    tool_name = params.get("name")
-    arguments = params.get("arguments") or {}
-    try:
-        if tool_name == "diagnostic_status":
-            return _tool_json(_diagnostic_status(_load_config()))
-        if tool_name == "refresh_schema":
-            agent = _make_agent(agent_factory)
-            return _tool_json(
-                _refresh_schema(
-                    agent,
-                    trace_id=_required_string(arguments, "traceId"),
-                )
-            )
-        if tool_name == "query_sponge":
-            question = _required_string(arguments, "question")
-            original_question = _required_string(arguments, "originalQuestion")
-            session_id = arguments.get("sessionId")
-            blocked_reason = _blocked_query_call_reason(
-                question=question,
-                original_question=original_question,
-            )
-            if blocked_reason is not None:
-                return _tool_blocked_query(blocked_reason)
-            agent = _make_agent(agent_factory)
-            answer_kwargs = {
-                "question": question,
-                "original_question": original_question,
-                "fast_mode": _optional_bool(arguments, "fastMode", default=True),
-                "session_id": session_id,
-                "result_format": arguments.get("resultFormat"),
-            }
-            final_answer_only = _optional_bool(arguments, "finalAnswerOnly", default=True)
-            if sampler is not None:
-                answer_kwargs["sampler"] = sampler
-            query_result = agent.query(**answer_kwargs)
-            base_question = original_question or question
-            if arguments.get("includeSql") is True:
-                payload = _query_result_payload(query_result, include_sql=True, base_question=base_question)
-                payload["answer"] = _user_visible_answer(query_result, base_question=base_question)
-                return _tool_json(payload, structured_content=payload)
-            if final_answer_only:
-                return _tool_final_answer(query_result, base_question=base_question)
-            return _tool_query_answer(query_result, final_answer_only=False, base_question=base_question)
-        return _tool_error(f"Unknown tool: {tool_name}", code="unknown_tool")
-    except Exception as exc:
-        return _tool_error(str(exc), code=type(exc).__name__)
-
-
-def _diagnostic_status(config: AppConfig) -> Json:
-    cache_exists = config.schema.cache_path.exists()
-    schema_version: str | None = None
-    cache_error: str | None = None
-    try:
-        schema_version = SpongeAgent.from_config(config).schema_cache.version()
-    except SchemaCacheError as exc:
-        cache_error = str(exc)
-
-    return {
-        "server": {"name": SERVER_NAME, "version": SERVER_VERSION},
-        "spongeMcp": {
-            "baseUrl": config.sponge_mcp.base_url,
-            "accessTokenConfigured": bool(config.sponge_mcp.access_token),
-            "timeoutMs": config.sponge_mcp.timeout_ms,
-        },
-        "schema": {
-            "cachePath": str(config.schema.cache_path),
-            "cacheExists": cache_exists,
-            "schemaVersion": schema_version,
-            "cacheError": cache_error,
-        },
-        "audit": {
-            "enabled": config.audit.enabled,
-            "logPath": str(config.audit.log_path),
-        },
-        "sampling": {
-            "maxTokens": config.sampling.max_tokens,
-        },
-        "binding": _binding_status(config),
-    }
-
-
-def _binding_status(config: AppConfig) -> Json:
-    binding = config.binding
-    return {
-        "configured": False,
-        "mode": "disabled",
-        "entitySource": "schema",
-        "dictionaryPersisted": False,
-        "shadowMode": bool(binding.shadow_mode) if binding is not None else None,
-        "enabledEntityTypes": list(binding.enabled_entity_types) if binding is not None else [],
-        "minFuzzyScore": binding.min_fuzzy_score if binding is not None else None,
-    }
-
-
-def _refresh_schema(agent: SpongeAgent, *, trace_id: str) -> Json:
-    cached_schema = agent.schema_cache.load()
-    current_version = str(cached_schema.get("schemaVersion")) if cached_schema and cached_schema.get("schemaVersion") else None
-    schema = agent.mcp_client.get_schema(
-        trace_id=trace_id,
-        schema_version=current_version,
-    )
-    response_changed = schema.get("changed")
-    if schema.get("changed") is False and cached_schema is not None:
-        schema = cached_schema
-    if "tables" in schema:
-        agent.schema_cache.save(schema)
-    schema_source = "cache" if response_changed is False and cached_schema is not None else "response"
-    agent.audit_logger.log(
-        {
-            "traceId": trace_id,
-            "tool": "get_schema",
-            "schemaVersion": schema.get("schemaVersion"),
-            "changed": response_changed,
-        }
-    )
-    return {
-        "changed": response_changed,
-        "request": {
-            "traceId": trace_id,
-            "schemaVersion": current_version,
-        },
-        "schema": {
-            "schemaVersion": schema.get("schemaVersion"),
-            "generatedAt": schema.get("generatedAt"),
-            "source": schema_source,
-        },
-        "cachePath": str(agent.schema_cache.path),
-    }
-
-
-def _load_config() -> AppConfig:
-    return load_config_from_env()
-
-
-def _make_agent(agent_factory: AgentFactory | None) -> SpongeAgent:
-    if agent_factory is not None:
-        return agent_factory()
-    return SpongeAgent.from_config(_load_config())
-
-
-def _required_string(arguments: Json, key: str) -> str:
-    value = arguments.get(key)
-    if not isinstance(value, str) or not value:
-        raise ValueError(f"{key} is required")
-    return value
-
-
-def _optional_bool(arguments: Json, key: str, *, default: bool) -> bool:
-    value = arguments.get(key)
-    if value is None:
-        return default
-    if not isinstance(value, bool):
-        raise ValueError(f"{key} must be a boolean")
-    return value
-
-
-def _required_int(arguments: Json, key: str) -> int:
-    value = arguments.get(key)
-    if not isinstance(value, int):
-        raise ValueError(f"{key} is required")
-    return value
-
-
-_PRESERVATION_SIGNAL_PATTERN = re.compile(
-    r"(今天|今日|昨天|昨日|明天|明日|本周|上周|本月|上月|最近\d+天"
-    r"|招满|已招满|在招|上架|下架|新发布|新上架"
-    r"|\d{4}年\d{1,2}月|\d{1,2}月\d{1,2}日"
-    r"|\d+日)"
-)
-_PRESERVATION_SIGNAL_EQUIVALENTS: dict[str, tuple[str, ...]] = {
-    "今天": ("今天", "今日"),
-    "今日": ("今天", "今日"),
-    "昨天": ("昨天", "昨日"),
-    "昨日": ("昨天", "昨日"),
-    "招满": ("招满", "已招满"),
-    "已招满": ("招满", "已招满"),
-}
-_REWRITTEN_SUMMARY_MARKERS = (
-    "相关的统计数据",
-    "相关统计数据",
-    "相关的统计",
-    "查询相关数据",
-    "查询相关统计",
-)
-
-
-def _extract_preservation_signals(text: str) -> list[str]:
-    return list(dict.fromkeys(match.group(0) for match in _PRESERVATION_SIGNAL_PATTERN.finditer(text)))
-
-
-def _signal_preserved(signal: str, text: str) -> bool:
-    candidates = _PRESERVATION_SIGNAL_EQUIVALENTS.get(signal, (signal,))
-    return any(candidate in text for candidate in candidates)
-
-
-def _missing_preservation_signals(*, source: str, target: str) -> list[str]:
-    missing: list[str] = []
-    for signal in _extract_preservation_signals(source):
-        if not _signal_preserved(signal, target):
-            missing.append(signal)
-    return missing
-
-
-def _looks_like_rewritten_summary(text: str) -> bool:
-    normalized = text.strip()
-    if len(normalized) > 40 or _extract_preservation_signals(normalized):
-        return False
-    if not re.search(r"(查询|统计)", normalized):
-        return False
-    return any(marker in normalized for marker in _REWRITTEN_SUMMARY_MARKERS)
-
-
-def _blocked_original_question_reason(*, question: str, original_question: str) -> str | None:
-    q = question.strip()
-    oq = original_question.strip()
-    if not q or not oq:
-        return None
-    if q == oq and _looks_like_rewritten_summary(oq):
-        return (
-            "originalQuestion 疑似上层摘要改写（例如“查询岗位相关的统计数据”），"
-            "请原样传入用户原始话术，不要摘要、翻译或改写。"
-        )
-    if q != oq:
-        missing = _missing_preservation_signals(source=q, target=oq)
-        if missing:
-            joined = "、".join(missing)
-            return (
-                f"originalQuestion 未保留 question 中的关键语义（{joined}）。"
-                "originalQuestion 必须传用户原始话术，上层不得摘要、翻译或改写。"
-            )
-    return None
-
-
-def _blocked_query_call_reason(*, question: str, original_question: str) -> str | None:
-    text = (original_question or question or "").strip()
-    if not text:
-        return "请先输入明确的业务数据查询需求，例如：查询我能看到的品牌列表。"
-    if _is_agent_meta_request(text) and not _has_explicit_business_object(text):
-        return "请先输入明确的业务数据查询需求；检查或测试丸子Agent时只允许调用 diagnostic_status。"
-    rewrite_reason = _blocked_original_question_reason(question=question, original_question=original_question)
-    if rewrite_reason is not None:
-        return rewrite_reason
-    return None
-
-
-def _is_agent_meta_request(text: str) -> bool:
-    return bool(
-        re.search(r"(丸子|octopus|agent|roll[- ]?core|工具|技能)", text, flags=re.IGNORECASE)
-        and re.search(r"(使用|启动|检查|看看|查看|测试|尝试|状态|健康|诊断|可用|功能)", text)
-    )
-
-
-def _has_explicit_business_object(text: str) -> bool:
-    return bool(
-        re.search(
-            r"(品牌|项目|岗位|门店|工单|报名|入职|候选人|城市|省份|区域|供应商|员工|薪资|招聘|数据|列表|明细|数量|多少|哪些|有哪些)",
-            text,
-        )
-    )
-
-
-def _tool_blocked_query(answer: str) -> Json:
-    payload: Json = {
-        "answer": answer,
-        "needsClarification": False,
-        "clarificationType": None,
-        "clarificationPayload": None,
-        "finalAnswerOnly": True,
-        "responseMode": "verbatim",
-        "blocked": True,
-        "blockedReason": answer,
-    }
-    return {
-        "content": _answer_content(answer),
-        "structuredContent": payload,
-    }
-
-
-def _tool_text(text: str) -> Json:
-    return {"content": [{"type": "text", "text": text}]}
-
-
-def _answer_content(answer: str) -> list[Json]:
-    return [
-        {"type": "text", "text": answer},
-    ]
-
-
-def _tool_final_answer(query_result: Any, *, base_question: str | None = None) -> Json:
-    payload = _query_result_payload(query_result, base_question=base_question)
-    answer = _user_visible_answer(query_result, base_question=base_question)
-    payload["answer"] = answer
-    payload.update(
-        {
-            "finalAnswerOnly": True,
-            "responseMode": "verbatim",
-        }
-    )
-    return {
-        "content": _answer_content(answer),
-        "structuredContent": payload,
-    }
-
-
-def _tool_query_answer(query_result: Any, *, final_answer_only: bool, base_question: str | None = None) -> Json:
-    payload = _query_result_payload(query_result, base_question=base_question)
-    payload["finalAnswerOnly"] = final_answer_only
-    payload["responseMode"] = "verbatim"
-    answer = _user_visible_answer(query_result, base_question=base_question)
-    payload["answer"] = answer
-    return {
-        "content": _answer_content(answer),
-        "structuredContent": payload,
-    }
-
-
-def _query_result_payload(query_result: Any, *, include_sql: bool = False, base_question: str | None = None) -> Json:
-    payload: Json = {
-        "answer": str(getattr(query_result, "answer", "")),
-        "needsClarification": bool(getattr(query_result, "needs_clarification", False)),
-        "clarificationType": getattr(query_result, "clarification_type", None),
-        "clarificationPayload": getattr(query_result, "clarification_payload", None),
-    }
-    query_plan = getattr(query_result, "query_plan", None)
-    if isinstance(query_plan, dict):
-        payload["queryPlan"] = _public_query_plan(query_plan)
-    execution_steps = getattr(query_result, "execution_steps", None)
-    if isinstance(execution_steps, list):
-        payload["executionSteps"] = execution_steps
-    timing = getattr(query_result, "timing", None)
-    if isinstance(timing, dict) and timing:
-        payload["timing"] = timing
-    if include_sql:
-        payload.update(
-            {
-                "generatedSql": getattr(query_result, "generated_sql", ""),
-                "normalizedSql": getattr(query_result, "normalized_sql", ""),
-                "repairedSql": getattr(query_result, "repaired_sql", None),
-                "validation": getattr(query_result, "validation", {}),
-            }
-        )
-    return _jsonable(payload)
-
-
-def _public_query_plan(query_plan: dict[str, Any]) -> Json:
-    return {
-        key: value
-        for key, value in query_plan.items()
-        if not str(key).startswith("_")
-    }
-
-
-def _user_visible_answer(query_result: Any, *, base_question: str | None = None) -> str:
-    answer = str(getattr(query_result, "answer", ""))
-    question = base_question or ""
-    result_format = getattr(query_result, "result_format", None)
-    answer = _ensure_complete_detail_answer(query_result, answer, question, result_format)
-    query_plan = getattr(query_result, "query_plan", None)
-    return _ensure_two_section_answer(answer, question, result_format, query_plan)
-
-
-def _ensure_two_section_answer(
-    answer: str,
-    question: str,
-    result_format: str | None,
-    query_plan: Any,
-) -> str:
-    if isinstance(query_plan, dict) and query_plan:
-        return ensure_query_plan_summary_answer(answer, question, result_format, query_plan)
-    if _answer_has_two_section_structure(answer, question, result_format):
-        return answer
-    return _prepend_missing_query_plan_section(answer, question, result_format)
-
-
-def _answer_has_two_section_structure(answer: str, question: str, result_format: str | None) -> bool:
-    output_format = normalize_result_format(question, result_format)
-    if output_format == "json":
-        try:
-            payload = json.loads(answer)
-        except json.JSONDecodeError:
-            return False
-        return (
-            isinstance(payload, dict)
-            and "查询实体" in payload
-            and "筛选条件" in payload
-            and "查询结果" in payload
-        )
-    if output_format == "yaml":
-        return "查询实体:" in answer and "筛选条件:" in answer and "查询结果:" in answer
-    return "1. 查询实体" in answer and "2. 筛选条件" in answer and "3. 查询结果" in answer
-
-
-def _prepend_missing_query_plan_section(answer: str, question: str, result_format: str | None) -> str:
-    output_format = normalize_result_format(question, result_format)
-    if output_format == "json":
-        try:
-            result_payload: Any = json.loads(answer)
-        except json.JSONDecodeError:
-            result_payload = answer
-        return json.dumps(
-            {
-                "查询实体": ["未返回查询计划"],
-                "筛选条件": ["未返回查询计划"],
-                "查询结果": result_payload,
-            },
-            ensure_ascii=False,
-            indent=2,
-        ) + "\n"
-    if output_format == "yaml":
-        return (
-            "查询实体:\n"
-            "  - 未返回查询计划\n"
-            "筛选条件:\n"
-            "  - 未返回查询计划\n"
-            "查询结果:\n"
-            + "\n".join(f"  {line}" if line else "" for line in answer.rstrip().splitlines())
-            + "\n"
-        )
-    if output_format == "html":
-        return (
-            "<section><p>1. 查询实体：</p><p>未返回查询计划</p></section>"
-            "<section><p>2. 筛选条件：</p><p>未返回查询计划</p></section>"
-            "<section><p>3. 查询结果：</p>"
-            + answer
-            + "</section>"
-        )
-    return (
-        "1. 查询实体：\n"
-        "未返回查询计划\n\n"
-        "2. 筛选条件：\n"
-        "未返回查询计划\n\n"
-        "3. 查询结果：\n"
-        f"{answer}"
-    )
-
-
-def _ensure_complete_detail_answer(
-    query_result: Any,
-    answer: str,
-    question: str,
-    result_format: str | None,
-) -> str:
-    raw_result = getattr(query_result, "raw_result", None)
-    if not isinstance(raw_result, dict):
-        return answer
-    rows = raw_result.get("rows")
-    if not isinstance(rows, list):
-        return answer
-    display_count = min(len(rows), RESULT_ROW_DISPLAY_LIMIT)
-    if display_count <= 0:
-        return answer
-    output_format = normalize_result_format(question, result_format)
-    if output_format in {"text", "markdown"}:
-        if _answer_contains_markdown_table_row(answer, display_count):
-            return answer
-        return render_result(question, raw_result, "markdown")
-    if _answer_contains_display_row(answer, display_count):
-        return answer
-    local_format = result_format or "markdown"
-    return render_result(question, raw_result, local_format)
-
-
-def _answer_contains_markdown_table_row(answer: str, row_number: int) -> bool:
-    escaped = re.escape(str(row_number))
-    return bool(re.search(rf"(?m)^\|\s*{escaped}\s*\|", answer))
-
-
-def _answer_contains_display_row(answer: str, row_number: int) -> bool:
-    escaped = re.escape(str(row_number))
-    patterns = (
-        rf"(?m)^\|\s*{escaped}\s*\|",
-        rf"(?m)^\s*{escaped}[\.\)、]\s+",
-        rf"<td>\s*{escaped}\s*</td>",
-        rf'"序号"\s*:\s*"?{escaped}"?',
-    )
-    return any(re.search(pattern, answer) for pattern in patterns)
-
-
-def _tool_json(data: Any, *, structured_content: Any | None = None) -> Json:
-    result = {
-        "content": [
-            {
-                "type": "text",
-                "text": json.dumps(_jsonable(data), ensure_ascii=False, sort_keys=True),
-            }
-        ]
-    }
-    if structured_content is not None:
-        result["structuredContent"] = _jsonable(structured_content)
-    return result
-
-
-def _tool_error(message: str, *, code: str) -> Json:
-    return {
-        "isError": True,
-        "content": [
-            {
-                "type": "text",
-                "text": json.dumps({"code": code, "message": message}, ensure_ascii=False, sort_keys=True),
-            }
-        ],
-    }
-
-
-def _jsonable(value: Any) -> Any:
-    if is_dataclass(value):
-        return _jsonable(asdict(value))
-    if isinstance(value, dict):
-        return {str(key): _jsonable(item) for key, item in value.items()}
-    if isinstance(value, list):
-        return [_jsonable(item) for item in value]
-    return value
-
-
-def _result_response(request_id: Any, result: Json) -> Json:
-    return {"jsonrpc": "2.0", "id": request_id, "result": result}
-
-
-def _error_response(request_id: Any, code: int, message: str, data: Json | None = None) -> Json:
-    error: Json = {"code": code, "message": message}
-    if data is not None:
-        error["data"] = data
-    return {"jsonrpc": "2.0", "id": request_id, "error": error}
-
-
-def _sampling_user_prompt(system_prompt: str, prompt: str) -> str:
-    system = (system_prompt or "").strip()
-    user = (prompt or "").strip()
-    if not system:
-        return user
-    if not user:
-        return system
-    return f"{system}\n\n---\n\n{user}"
-
-
-class StdioSession:
-    def __init__(self, input_stream: Any, output_stream: Any, *, max_tokens: int = 4096) -> None:
-        self.input_stream = input_stream
-        self.output_stream = output_stream
-        self.next_request_id = 1
-        self.max_tokens = max_tokens
-
-    def read_message(self) -> Json | None:
-        return _read_message(self.input_stream)
-
-    def write_message(self, message: Json) -> None:
-        _write_message(self.output_stream, message)
-
-    def sample(self, system_prompt: str, prompt: str) -> str:
-        request_id = f"sampling_{self.next_request_id}"
-        self.next_request_id += 1
-        # roll sampling handler currently forwards only `messages` to the LLM and
-        # ignores `systemPrompt`; embed system instructions in the user payload.
-        user_prompt = _sampling_user_prompt(system_prompt, prompt)
-        self.write_message(
-            {
-                "jsonrpc": "2.0",
-                "id": request_id,
-                "method": "sampling/createMessage",
-                "params": {
-                    "systemPrompt": system_prompt,
-                    "messages": [{"role": "user", "content": {"type": "text", "text": user_prompt}}],
-                    "temperature": 0,
-                    "maxTokens": self.max_tokens,
-                    "includeContext": "none",
-                },
-            }
-        )
-        while True:
-            message = self.read_message()
-            if message is None:
-                raise RuntimeError("Sampling failed: client closed stdio")
-            if message.get("id") == request_id:
-                if "error" in message:
-                    raise RuntimeError(f"Sampling failed: {message['error']}")
-                return _extract_sampling_text(message.get("result") or {})
-            response = handle_request(message, sampler=self.sample)
-            if response is not None:
-                self.write_message(response)
-
-
-def _extract_sampling_text(result: Json) -> str:
-    content = result.get("content")
-    if isinstance(content, dict) and content.get("type") == "text":
-        return str(content.get("text") or "")
-    if isinstance(content, list):
-        return "\n".join(
-            str(item.get("text") or "") for item in content if isinstance(item, dict) and item.get("type") == "text"
-        )
-    raise RuntimeError("Sampling response did not contain text content")
-
-
-def _read_message(stream: Any) -> Json | None:
-    while True:
-        line = stream.readline()
-        if line == b"":
-            return None
-        if line.strip():
-            return json.loads(line.decode("utf-8"))
-
-
-def _write_message(stream: Any, message: Json) -> None:
-    body = json.dumps(message, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
-    stream.write(body + b"\n")
-    stream.flush()
-
-
-if __name__ == "__main__":
-    main()