Files

1 line
21 KiB
JSON

{"repo": "carrot-peace/mimo-mcp", "files": {"server.py": {"53e2c9d1dca1f0f39196ab18b7ef40726dc78ba7": {"t": 1785674309.6110694, "n": 20228, "h": "47ffd443c188000f", "c": "#!/usr/bin/env python3\n\"\"\"MiMo Token Plan read-only MCP toolbox for Codex.\"\"\"\n\nfrom __future__ import annotations\n\nimport os\nimport re\nfrom typing import Any, Iterable, Literal\nfrom urllib.parse import urlparse\n\nimport httpx\nfrom mcp.server.fastmcp import FastMCP\n\nDEFAULT_BASE_URL = \"https://token-plan-cn.xiaomimimo.com/v1\"\nDEFAULT_MODEL = \"mimo-v2.5-pro\"\nDEFAULT_COMPLETION_BUDGET = 65536\nLARGE_COMPLETION_BUDGET = 131072\nLONG_CONTEXT_CHARS = 20000\nCOMPLEX_QUESTION_CHARS = 500\nREQUEST_TIMEOUT_SECONDS = 60.0\nVALID_TASK_TYPES = {\n \"summary\",\n \"review\",\n \"test_draft\",\n \"explain\",\n \"brainstorm\",\n \"other\",\n}\nVALID_BUDGET_MODES = {\"auto\", \"default\", \"large\"}\nAPPROVED_MIMO_HOSTS = {\n \"token-plan-cn.xiaomimimo.com\",\n \"token-plan-sgp.xiaomimimo.com\",\n \"token-plan-ams.xiaomimimo.com\",\n}\nLOCAL_HTTP_HOSTS = {\"localhost\", \"127.0.0.1\", \"::1\"}\nBASE_URL_REFUSED_ERROR = \"Refused: MIMO_BASE_URL is not an approved MiMo endpoint.\"\nBASE_URL_INVALID_ERROR = \"Error: MIMO_BASE_URL must be a valid http or https URL.\"\n\nmcp = FastMCP(\"mimo-mcp\")\n\n\nSENSITIVE_PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (\n (\"OpenAI-style API key\", re.compile(r\"(?i)(?:^|[^A-Za-z0-9_-])sk-[A-Za-z0-9_-]{12,}\")),\n (\"Token Plan API key\", re.compile(r\"(?i)(?:^|[^A-Za-z0-9_-])tp-[A-Za-z0-9_-]{8,}\")),\n (\n \"API key environment variable name\",\n re.compile(\n r\"(?i)\\b(?:OPENAI|ANTHROPIC|GEMINI|GOOGLE|MIMO|MIMO_TP|TOKEN_PLAN)\"\n r\"[_-]*(?:API[_-]*)?(?:KEY|TOKEN|SECRET)\\b\"\n ),\n ),\n (\n \"secret-like assignment\",\n re.compile(\n r\"(?im)^\\s*(?:export\\s+)?[A-Za-z_][A-Za-z0-9_]*(?:PASSWORD|PASSWD|PWD|TOKEN|SECRET|PRIVATE_KEY|ACCESS_KEY|API_KEY)\"\n r\"[A-Za-z0-9_]*\\s*[:=]\\s*['\\\"]?[^'\\\"\\s#]{6,}\"\n ),\n ),\n (\"private key block\", re.compile(r\"-----BEGIN (?:RSA |EC |OPENSSH |DSA )?PRIVATE KEY-----\")),\n (\n \".env style secret\",\n re.compile(\n r\"(?im)^\\s*[A-Za-z_][A-Za-z0-9_]*(?:KEY|TOKEN|SECRET|PASSWORD|PASSWD|PWD)\\s*=\\s*['\\\"]?[^'\\\"\\s#]{6,}\"\n ),\n ),\n (\n \"Android signing configuration\",\n re.compile(r\"(?im)^\\s*(?:storePassword|keyPassword|keyAlias|storeFile)\\s*[:=]\\s*\\S+\"),\n ),\n (\n \"Android key.properties content\",\n re.compile(\n r\"(?is)\\bstorePassword\\s*[:=].+\\bkeyPassword\\s*[:=].+\\bkeyAlias\\s*[:=].+\\bstoreFile\\s*[:=]\"\n ),\n ),\n)\n\n\ndef _env_default_model() -> str:\n return os.environ.get(\"MIMO_DEFAULT_MODEL\") or DEFAULT_MODEL\n\n\ndef _env_base_url() -> str:\n return os.environ.get(\"MIMO_BASE_URL\") or DEFAULT_BASE_URL\n\n\ndef _env_flag_enabled(name: str) -> bool:\n return os.environ.get(name) == \"1\"\n\n\ndef _normalized_base_path(path: str) -> str | None:\n if path in {\"/v1\", \"/v1/\"}:\n return \"/v1\"\n return None\n\n\ndef _chat_completions_url() -> tuple[str | None, str | None]:\n base_url = _env_base_url().strip()\n parsed = urlparse(base_url)\n if parsed.scheme not in {\"http\", \"https\"} or not parsed.netloc or not parsed.hostname:\n return None, BASE_URL_INVALID_ERROR\n try:\n port = parsed.port\n except ValueError:\n return None, BASE_URL_INVALID_ERROR\n if parsed.username or parsed.password or parsed.query or parsed.fragment:\n return None, BASE_URL_REFUSED_ERROR\n\n normalized_path = _normalized_base_path(parsed.path)\n if normalized_path is None:\n return None, BASE_URL_REFUSED_ERROR\n\n host = parsed.hostname\n if parsed.scheme == \"https\":\n official_endpoint = host in APPROVED_MIMO_HOSTS and port is None\n if official_endpoint or _env_flag_enabled(\"MIMO_ALLOW_CUSTOM_BASE_URL\"):\n netloc = parsed.netloc\n return f\"https://{netloc}{normalized_path}/chat/completions\", None\n return None, BASE_URL_REFUSED_ERROR\n\n if _env_flag_enabled(\"MIMO_ALLOW_INSECURE_LOCAL_HTTP\") and host in LOCAL_HTTP_HOSTS:\n netloc = parsed.netloc\n return f\"http://{netloc}{normalized_path}/chat/completions\", None\n return None, BASE_URL_REFUSED_ERROR\n\n\ndef _input_limit() -> int:\n raw = os.environ.get(\"MIMO_MAX_INPUT_CHARS\", \"0\").strip()\n try:\n return max(0, int(raw))\n except ValueError:\n return 0\n\n\ndef _coerce_model(model: str | None) -> str:\n if model is None:\n return _env_default_model()\n return model\n\n\ndef _env_completion_budget(name: str, default: int) -> tuple[int, bool, str | None]:\n raw = os.environ.get(name)\n if raw is None or not raw.strip():\n return default, False, None\n try:\n value = int(raw.strip())\n except ValueError:\n return default, True, f\"{name}=invalid\"\n if value <= 0:\n return default, True, f\"{name}=invalid\"\n return value, False, None\n\n\ndef _normalize_budget_mode(budget: str) -> tuple[str, bool]:\n if not isinstance(budget, str):\n return \"auto\", True\n mode = budget.strip().lower()\n if mode not in VALID_BUDGET_MODES:\n return \"auto\", True\n return mode, False\n\n\ndef _auto_budget_mode(tool: str, user_sections: list[tuple[str, str]]) -> str:\n section_map = dict(user_sections)\n if tool == \"mimo_compare\":\n return \"large\"\n if tool == \"mimo_summarize\" and len(section_map.get(\"Context\", \"\")) >= LONG_CONTEXT_CHARS:\n return \"large\"\n if tool == \"mimo_review\":\n context = section_map.get(\"Context\", \"\")\n question = section_map.get(\"Question\", \"\")\n if len(context) >= LONG_CONTEXT_CHARS or len(question) >= COMPLEX_QUESTION_CHARS:\n return \"large\"\n return \"default\"\n\n\ndef _resolve_completion_budget(\n *,\n tool: str,\n user_sections: list[tuple[str, str]],\n budget: str,\n) -> tuple[str, int, bool, list[str]]:\n requested_mode, budget_fallback_used = _normalize_budget_mode(budget)\n effective_mode = _auto_budget_mode(tool, user_sections) if requested_mode == \"auto\" else requested_mode\n default_budget, default_fallback_used, default_note = _env_completion_budget(\n \"MIMO_DEFAULT_COMPLETION_BUDGET\", DEFAULT_COMPLETION_BUDGET\n )\n large_budget, large_fallback_used, large_note = _env_completion_budget(\n \"MIMO_LARGE_COMPLETION_BUDGET\", LARGE_COMPLETION_BUDGET\n )\n env_fallbacks = [\n note\n for used, note in ((default_fallback_used, default_note), (large_fallback_used, large_note))\n if used and note\n ]\n completion_budget = large_budget if effective_mode == \"large\" else default_budget\n return requested_mode, completion_budget, budget_fallback_used, env_fallbacks\n\n\ndef _find_sensitive_categories(texts: Iterable[str]) -> list[str]:\n categories: list[str] = []\n for text in texts:\n for category, pattern in SENSITIVE_PATTERNS:\n if pattern.search(text) and category not in categories:\n categories.append(category)\n return categories\n\n\ndef _usage_report(\n *,\n tool: str,\n model: str,\n budget_mode: str,\n completion_budget: int,\n budget_fallback_used: bool,\n env_budget_fallbacks: list[str],\n input_chars: int,\n output_chars: int,\n usage: dict[str, Any] | None,\n) -> str:\n if model == \"mimo-v2.5-pro\":\n multiplier = \"usually 2x (hint only, not exact billing)\"\n elif model == \"mimo-v2.5\":\n multiplier = \"usually 1x (hint only, not exact billing)\"\n else:\n multiplier = \"unknown\"\n\n lines = [\n \"\",\n \"---\",\n \"usage report:\",\n f\"- tool: {tool}\",\n f\"- model: {model}\",\n f\"- budget_mode: {budget_mode}\",\n f\"- completion_budget: {completion_budget}\",\n f\"- budget_fallback_used: {str(budget_fallback_used).lower()}\",\n f\"- env_budget_fallback_used: {str(bool(env_budget_fallbacks)).lower()}\",\n f\"- input_chars: {input_chars}\",\n f\"- output_chars: {output_chars}\",\n f\"- estimated_credit_multiplier: {multiplier}\",\n ]\n if env_budget_fallbacks:\n lines.append(f\"- env_budget_fallbacks: {', '.join(env_budget_fallbacks)}\")\n if usage:\n lines.extend(\n [\n f\"- prompt_tokens: {usage.get('prompt_tokens')}\",\n f\"- completion_tokens: {usage.get('completion_tokens')}\",\n f\"- total_tokens: {usage.get('total_tokens')}\",\n \"- usage_source: api_usage\",\n ]\n )\n else:\n lines.append(\"- usage_source: not_reported\")\n return \"\\n\".join(lines)\n\n\ndef _empty_content_diagnostic(\n *,\n data: dict[str, Any],\n message: Any,\n finish_reason: Any,\n tool: str,\n model: str,\n budget_mode: str,\n completion_budget: int,\n budget_fallback_used: bool,\n env_budget_fallbacks: list[str],\n) -> str:\n usage = data.get(\"usage\")\n usage_value: Any = usage if isinstance(usage, dict) else None\n message_keys = sorted(message.keys()) if isinstance(message, dict) else []\n lines = [\n \"Error: MiMo API returned empty content.\",\n \"diagnostic:\",\n f\"- tool: {tool}\",\n f\"- model: {model}\",\n f\"- finish_reason: {finish_reason}\",\n f\"- message_keys: {message_keys}\",\n f\"- usage: {usage_value if usage_value is not None else 'not_reported'}\",\n f\"- budget_mode: {budget_mode}\",\n f\"- completion_budget: {completion_budget}\",\n f\"- budget_fallback_used: {str(budget_fallback_used).lower()}\",\n f\"- env_budget_fallback_used: {str(bool(env_budget_fallbacks)).lower()}\",\n \"- possible_empty_content_source: server_or_model_returned_empty_content\",\n ]\n if env_budget_fallbacks:\n lines.append(f\"- env_budget_fallbacks: {', '.join(env_budget_fallbacks)}\")\n return \"\\n\".join(lines)\n\n\ndef _build_messages(system_prompt: str, user_sections: list[tuple[str, str]]) -> list[dict[str, str]]:\n system = \"\\n\".join(\n [\n \"You are MiMo Token Plan, connected to Codex as a read-only auxiliary model.\",\n \"Only use the text explicitly provided by the caller.\",\n \"Do not claim to read local files, write files, execute commands, scan projects, or apply patches.\",\n system_prompt,\n ]\n )\n user = \"\\n\\n\".join(f\"{title}:\\n{body}\" for title, body in user_sections if body)\n return [{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": user}]\n\n\nasync def _call_mimo(\n *,\n tool: str,\n user_sections: list[tuple[str, str]],\n system_prompt: str,\n model: str | None,\n budget: str,\n) -> str:\n if not user_sections or any(not isinstance(value, str) for _, value in user_sections):\n return \"Error: all text parameters must be strings.\"\n\n required_values = [value for _, value in user_sections]\n if not any(value.strip() for value in required_values):\n return \"Error: a non-empty text parameter is required.\"\n\n model_name = _coerce_model(model)\n if not isinstance(model_name, str) or not model_name.strip():\n return \"Error: model must be a non-empty string when provided.\"\n model_name = model_name.strip()\n\n budget_mode, completion_budget, budget_fallback_used, env_budget_fallbacks = _resolve_completion_budget(\n tool=tool,\n user_sections=user_sections,\n budget=budget,\n )\n\n sensitive_categories = _find_sensitive_categories(required_values)\n if sensitive_categories:\n categories = \", \".join(sensitive_categories)\n return f\"Error: refused to call MiMo API because input may contain sensitive content. Matched rule categories: {categories}.\"\n\n input_chars = sum(len(value) for value in required_values)\n limit = _input_limit()\n if limit > 0 and input_chars > limit:\n return (\n f\"Error: input is {input_chars} characters, exceeding MIMO_MAX_INPUT_CHARS={limit}. \"\n \"Please reduce the provided context or increase MIMO_MAX_INPUT_CHARS.\"\n )\n\n api_key = os.environ.get(\"MIMO_TP_KEY\")\n if not api_key:\n return (\n \"Error: MIMO_TP_KEY is not set. Set it in your shell or with \"\n \"launchctl setenv MIMO_TP_KEY '<your-token-plan-key>' before starting Codex.\"\n )\n\n api_url, base_url_error = _chat_completions_url()\n if api_url is None:\n return base_url_error or BASE_URL_INVALID_ERROR\n\n payload = {\n \"model\": model_name,\n \"messages\": _build_messages(system_prompt, user_sections),\n \"max_completion_tokens\": completion_budget,\n }\n headers = {\"Authorization\": f\"Bearer {api_key}\", \"Content-Type\": \"application/json\"}\n\n try:\n async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT_SECONDS) as client:\n response = await client.post(api_url, headers=headers, json=payload)\n except httpx.TimeoutException:\n return f\"Error: MiMo API request timed out after {int(REQUEST_TIMEOUT_SECONDS)} seconds.\"\n except httpx.RequestError as exc:\n return f\"Error: MiMo API request failed: {exc.__class__.__name__}: {exc}\"\n\n if response.status_code in {401, 403}:\n return f\"Error: MiMo API authentication failed with HTTP {response.status_code}. Check MIMO_TP_KEY without sharing it.\"\n if response.status_code == 429:\n return \"Error: MiMo API returned HTTP 429 rate limit. Please retry later.\"\n if response.status_code < 200 or response.status_code >= 300:\n detail = response.text.strip()\n if len(detail) > 1000:\n detail = detail[:1000] + \"...\"\n return f\"Error: MiMo API returned HTTP {response.status_code}: {detail}\"\n\n try:\n data = response.json()\n except ValueError:\n return \"Error: MiMo API returned a non-JSON response.\"\n\n try:\n choice = data[\"choices\"][0]\n message = choice[\"message\"]\n content = message[\"content\"]\n except (KeyError, IndexError, TypeError):\n return \"Error: MiMo API response did not include choices[0].message.content.\"\n finish_reason = choice.get(\"finish_reason\") if isinstance(choice, dict) else None\n\n if not isinstance(content, str):\n if content is None:\n return _empty_content_diagnostic(\n data=data,\n message=message,\n finish_reason=finish_reason,\n tool=tool,\n model=model_name,\n budget_mode=budget_mode,\n completion_budget=completion_budget,\n budget_fallback_used=budget_fallback_used,\n env_budget_fallbacks=env_budget_fallbacks,\n )\n content = str(content)\n if not content.strip():\n return _empty_content_diagnostic(\n data=data,\n message=message,\n finish_reason=finish_reason,\n tool=tool,\n model=model_name,\n budget_mode=budget_mode,\n completion_budget=completion_budget,\n budget_fallback_used=budget_fallback_used,\n env_budget_fallbacks=env_budget_fallbacks,\n )\n\n usage = data.get(\"usage\")\n if not isinstance(usage, dict):\n usage = None\n return content + _usage_report(\n tool=tool,\n model=model_name,\n budget_mode=budget_mode,\n completion_budget=completion_budget,\n budget_fallback_used=budget_fallback_used,\n env_budget_fallbacks=env_budget_fallbacks,\n input_chars=input_chars,\n output_chars=len(content),\n usage=usage,\n )\n\n\n@mcp.tool()\nasync def mimo_ask(\n prompt: str,\n context: str | None = None,\n task_type: Literal[\"summary\", \"review\", \"test_draft\", \"explain\", \"brainstorm\", \"other\"] | None = None,\n model: str | None = None,\n budget: str = \"auto\",\n) -> str:\n \"\"\"Ask MiMo Token Plan using only caller-provided text.\"\"\"\n if not isinstance(prompt, str) or not prompt.strip():\n return \"Error: prompt is required and must be a non-empty string.\"\n if context is not None and not isinstance(context, str):\n return \"Error: context must be a string when provided.\"\n if task_type is not None and task_type not in VALID_TASK_TYPES:\n allowed = \", \".join(sorted(VALID_TASK_TYPES))\n return f\"Error: task_type must be one of: {allowed}.\"\n\n sections = []\n if context:\n sections.append((\"Context\", context))\n sections.append((\"Prompt\", prompt))\n task_line = f\"Task type: {task_type}.\" if task_type else \"General analysis task.\"\n return await _call_mimo(\n tool=\"mimo_ask\",\n user_sections=sections,\n system_prompt=task_line,\n model=model,\n budget=budget,\n )\n\n\n@mcp.tool()\nasync def mimo_summarize(\n context: str,\n focus: str | None = None,\n model: str | None = None,\n budget: str = \"auto\",\n) -> str:\n \"\"\"Summarize long caller-provided context.\"\"\"\n if not isinstance(context, str) or not context.strip():\n return \"Error: context is required and must be a non-empty string.\"\n if focus is not None and not isinstance(focus, str):\n return \"Error: focus must be a string when provided.\"\n sections = [(\"Context\", context)]\n if focus:\n sections.append((\"Focus\", focus))\n return await _call_mimo(\n tool=\"mimo_summarize\",\n user_sections=sections,\n system_prompt=(\n \"Summarize long code, diffs, or documents in a structured way. \"\n \"Separate facts, inferences, uncertainties, and important omissions.\"\n ),\n model=model,\n budget=budget,\n )\n\n\n@mcp.tool()\nasync def mimo_review(\n context: str,\n question: str | None = None,\n model: str | None = None,\n budget: str = \"auto\",\n) -> str:\n \"\"\"Ask MiMo for a read-only second opinion on risks and gaps.\"\"\"\n if not isinstance(context, str) or not context.strip():\n return \"Error: context is required and must be a non-empty string.\"\n if question is not None and not isinstance(question, str):\n return \"Error: question must be a string when provided.\"\n sections = [(\"Context\", context)]\n if question:\n sections.append((\"Question\", question))\n return await _call_mimo(\n tool=\"mimo_review\",\n user_sections=sections,\n system_prompt=(\n \"Provide a second opinion focused only on risks, edge cases, potential bugs, and missing tests. \"\n \"Do not ask to modify code and do not give a final implementation decision; Codex decides.\"\n ),\n model=model,\n budget=budget,\n )\n\n\n@mcp.tool()\nasync def mimo_test_draft(\n context: str,\n test_goal: str | None = None,\n model: str | None = None,\n budget: str = \"auto\",\n) -> str:\n \"\"\"Generate test ideas or draft tests from caller-provided context.\"\"\"\n if not isinstance(context, str) or not context.strip():\n return \"Error: context is required and must be a non-empty string.\"\n if test_goal is not None and not isinstance(test_goal, str):\n return \"Error: test_goal must be a string when provided.\"\n sections = [(\"Context\", context)]\n if test_goal:\n sections.append((\"Test goal\", test_goal))\n return await _call_mimo(\n tool=\"mimo_test_draft\",\n user_sections=sections,\n system_prompt=(\n \"Draft test scenarios, boundary cases, pseudocode, or test sketches. \"\n \"Do not assume project APIs that were not provided. Do not write files.\"\n ),\n model=model,\n budget=budget,\n )\n\n\n@mcp.tool()\nasync def mimo_compare(\n options: str,\n criteria: str | None = None,\n model: str | None = None,\n budget: str = \"auto\",\n) -> str:\n \"\"\"Compare caller-provided options without unsupported claims.\"\"\"\n if not isinstance(options, str) or not options.strip():\n return \"Error: options is required and must be a non-empty string.\"\n if criteria is not None and not isinstance(criteria, str):\n return \"Error: criteria must be a string when provided.\"\n sections = [(\"Options\", options)]\n if criteria:\n sections.append((\"Criteria\", criteria))\n return await _call_mimo(\n tool=\"mimo_compare\",\n user_sections=sections,\n system_prompt=(\n \"Compare options by complexity, risk, maintainability, and test cost. \"\n \"Mark uncertainty clearly and avoid claims not supported by the provided evidence.\"\n ),\n model=model,\n budget=budget,\n )\n\n\nif __name__ == \"__main__\":\n mcp.run()\n"}}}}