349 lines
13 KiB
Python
349 lines
13 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Model Context Protocol (MCP) server command (artemis mcp)."""
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import sys
|
|
import threading
|
|
import tomllib
|
|
from typing import Annotated
|
|
|
|
from mcp_server.base import mcp as agent_mcp
|
|
import mcp_server.tools # noqa: F401
|
|
from mcp_server.utils import env_utils
|
|
from artemis.mcp.adb_server import mcp as adb_mcp
|
|
from artemis.runtime import shutdown_awake_service, start_awake_service
|
|
from artemis.utils.logger import get_logger
|
|
from rich.console import Console
|
|
from rich.syntax import Syntax
|
|
import typer
|
|
|
|
logger = get_logger(__name__)
|
|
console = Console()
|
|
|
|
ARTEMIS_MCP_TOOLS = (
|
|
"mobile_run_task",
|
|
"mobile_manage_task",
|
|
"mobile_get_device_state",
|
|
"mobile_inspect_trace",
|
|
"mobile_diagnose",
|
|
"mobile_workspace",
|
|
"mobile_config",
|
|
"mobile_task_file",
|
|
"mobile_batch",
|
|
"mobile_search_history",
|
|
"mobile_export_report",
|
|
)
|
|
|
|
|
|
def _get_config_snippet(client: str, python_exe: str, project_root: str) -> dict:
|
|
"""Generate the supported Codex MCP configuration."""
|
|
if client != "codex":
|
|
raise ValueError("Only Codex MCP configuration is supported")
|
|
return {
|
|
"mcp_servers": {
|
|
"artemis": {
|
|
"command": python_exe,
|
|
"args": ["-m", "mcp_server"],
|
|
"cwd": project_root,
|
|
"env": {
|
|
"PYTHONUNBUFFERED": "1",
|
|
"PYTHONPATH": project_root,
|
|
"ARTEMIS_DESKTOP_NOTIFY": "true",
|
|
},
|
|
"enabled": True,
|
|
"required": True,
|
|
"startup_timeout_sec": 120,
|
|
"enabled_tools": list(ARTEMIS_MCP_TOOLS),
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
def _codex_toml_block(server_config: dict) -> str:
|
|
"""Renders a managed Codex MCP server block using TOML-compatible JSON strings."""
|
|
command = json.dumps(str(server_config["command"]), ensure_ascii=False)
|
|
args = json.dumps(server_config.get("args", []), ensure_ascii=False)
|
|
cwd = json.dumps(str(server_config["cwd"]), ensure_ascii=False)
|
|
env = server_config.get("env", {})
|
|
lines = [
|
|
"# BEGIN ARTEMIS MCP CONFIG",
|
|
"[mcp_servers.artemis]",
|
|
f"command = {command}",
|
|
f"args = {args}",
|
|
f"cwd = {cwd}",
|
|
]
|
|
for key in ("enabled", "required"):
|
|
if key in server_config:
|
|
lines.append(f"{key} = {'true' if server_config[key] else 'false'}")
|
|
if "startup_timeout_sec" in server_config:
|
|
lines.append(f"startup_timeout_sec = {int(server_config['startup_timeout_sec'])}")
|
|
if "enabled_tools" in server_config:
|
|
lines.append(
|
|
f"enabled_tools = {json.dumps(server_config['enabled_tools'], ensure_ascii=False)}"
|
|
)
|
|
if env:
|
|
lines.extend(["", "[mcp_servers.artemis.env]"])
|
|
for key, value in env.items():
|
|
lines.append(f"{key} = {json.dumps(str(value), ensure_ascii=False)}")
|
|
lines.append("# END ARTEMIS MCP CONFIG")
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def _merge_codex_toml(file_path: Path, server_config: dict) -> bool:
|
|
"""Adds or updates Artemis in Codex config.toml while preserving unrelated settings."""
|
|
try:
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
existing = file_path.read_text(encoding="utf-8") if file_path.exists() else ""
|
|
managed_block = _codex_toml_block(server_config)
|
|
begin_marker = "# BEGIN ARTEMIS MCP CONFIG"
|
|
end_marker = "# END ARTEMIS MCP CONFIG"
|
|
|
|
if begin_marker in existing and end_marker in existing:
|
|
managed_pattern = re.compile(
|
|
rf"(?ms)^{re.escape(begin_marker)}\n.*?^{re.escape(end_marker)}\n?"
|
|
)
|
|
# Windows paths contain backslashes; a replacement string would
|
|
# interpret them as regex escapes on a second installation.
|
|
updated = managed_pattern.sub(lambda _: managed_block, existing, count=1)
|
|
else:
|
|
# Remove an older, unmanaged Artemis table before adding the managed block.
|
|
kept_lines: list[str] = []
|
|
skipping_artemis = False
|
|
header_pattern = re.compile(r"^\s*\[{1,2}\s*(.+?)\s*\]{1,2}\s*(?:#.*)?$")
|
|
for line in existing.splitlines(keepends=True):
|
|
header_match = header_pattern.match(line.rstrip("\r\n"))
|
|
if header_match:
|
|
normalized = re.sub(r"[\s\"']", "", header_match.group(1)).lower()
|
|
skipping_artemis = normalized == "mcp_servers.artemis" or normalized.startswith(
|
|
"mcp_servers.artemis."
|
|
)
|
|
if not skipping_artemis:
|
|
kept_lines.append(line)
|
|
preserved = "".join(kept_lines).rstrip()
|
|
updated = f"{preserved}\n\n{managed_block}" if preserved else managed_block
|
|
|
|
# Never replace a user's config with invalid TOML.
|
|
tomllib.loads(updated)
|
|
if updated != existing:
|
|
file_path.write_text(updated, encoding="utf-8")
|
|
return True
|
|
except Exception as e:
|
|
if file_path.exists():
|
|
try:
|
|
backup_path = file_path.with_suffix(file_path.suffix + ".bak")
|
|
backup_path.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8")
|
|
logger.warning(
|
|
f"Could not update {file_path}; backup created at {backup_path}: {e}"
|
|
)
|
|
except (OSError, UnicodeDecodeError) as backup_exc:
|
|
logger.warning(
|
|
f"Could not update {file_path} ({e}); backup also failed: {backup_exc}"
|
|
)
|
|
else:
|
|
logger.warning(f"Could not create Codex MCP config file {file_path}: {e}")
|
|
return False
|
|
|
|
|
|
def _inject_rules_block(file_path: Path, content: str) -> bool:
|
|
"""Injects or updates a managed ARTEMIS block in a global rules file without overwriting user content."""
|
|
try:
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
begin_marker = "<!-- BEGIN ARTEMIS MOBILE TESTING RULES -->"
|
|
end_marker = "<!-- END ARTEMIS MOBILE TESTING RULES -->"
|
|
new_block = f"{begin_marker}\n{content.strip()}\n{end_marker}\n"
|
|
|
|
existing = ""
|
|
if file_path.exists():
|
|
try:
|
|
existing = file_path.read_text(encoding="utf-8")
|
|
except Exception:
|
|
existing = ""
|
|
|
|
if begin_marker in existing and end_marker in existing:
|
|
prefix = existing.split(begin_marker)[0]
|
|
suffix = existing.split(end_marker)[1]
|
|
updated = f"{prefix.rstrip()}\n\n{new_block}\n{suffix.lstrip()}".strip() + "\n"
|
|
else:
|
|
updated = f"{existing.rstrip()}\n\n{new_block}".strip() + "\n"
|
|
|
|
if existing == updated:
|
|
return True
|
|
|
|
file_path.write_text(updated, encoding="utf-8")
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"Could not inject rules block into {file_path}: {e}")
|
|
return False
|
|
|
|
|
|
def install_rules(client: str, project_root: str) -> list[str]:
|
|
"""Install testing rules only into the Codex instructions file."""
|
|
if client != "codex":
|
|
raise ValueError("Only Codex MCP configuration is supported")
|
|
rules_src = Path(project_root) / "mcp_server" / "rules.md"
|
|
if not rules_src.exists():
|
|
rules_src = Path(__file__).resolve().parents[4] / "mcp_server" / "rules.md"
|
|
try:
|
|
raw_rules = rules_src.read_text(encoding="utf-8").strip()
|
|
except OSError:
|
|
logger.warning(f"Could not read rules file {rules_src}")
|
|
return []
|
|
codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
|
|
override_file = codex_home / "AGENTS.override.md"
|
|
agents_file = (
|
|
override_file
|
|
if override_file.exists() and override_file.read_text(encoding="utf-8").strip()
|
|
else codex_home / "AGENTS.md"
|
|
)
|
|
return [str(agents_file)] if _inject_rules_block(agents_file, raw_rules) else []
|
|
|
|
|
|
def install_mcp_config(client: str, python_exe: str, project_root: str) -> list[str]:
|
|
"""Install ARTEMIS configuration and rules only for Codex."""
|
|
server_cfg = _get_config_snippet(client, python_exe, project_root)["mcp_servers"]["artemis"]
|
|
codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
|
|
codex_path = codex_home / "config.toml"
|
|
if not _merge_codex_toml(codex_path, server_cfg):
|
|
raise typer.Exit(1)
|
|
return [str(codex_path), *install_rules(client, project_root)]
|
|
|
|
|
|
def mcp_command(
|
|
server_type: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--type",
|
|
"-t",
|
|
help="Type of MCP server to start: 'agent' (default, universal IDE mobile agent), 'adb' (raw adb), 'xml' (xml fuzzy search).",
|
|
),
|
|
] = "agent",
|
|
transport: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--transport",
|
|
help="MCP transport protocol ('stdio' or 'sse').",
|
|
),
|
|
] = "stdio",
|
|
host: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--host",
|
|
help="Host address when running with SSE transport.",
|
|
),
|
|
] = "127.0.0.1",
|
|
port: Annotated[
|
|
int,
|
|
typer.Option(
|
|
"--port",
|
|
"-p",
|
|
help="Port number when running with SSE transport.",
|
|
),
|
|
] = 8001,
|
|
install_config: Annotated[
|
|
str | None,
|
|
typer.Option(
|
|
"--install",
|
|
"-i",
|
|
help="Install ARTEMIS MCP configuration and testing rules for Codex: --install codex.",
|
|
),
|
|
] = None,
|
|
generate_config: Annotated[
|
|
str | None,
|
|
typer.Option(
|
|
"--generate-config",
|
|
"-g",
|
|
help="Output ready-to-use Codex TOML configuration: --generate-config codex.",
|
|
),
|
|
] = None,
|
|
) -> None:
|
|
"""Launch or configure the ARTEMIS Model Context Protocol (MCP) server."""
|
|
project_root = env_utils.get_project_root()
|
|
python_exe = env_utils.resolve_python_executable(project_root)
|
|
|
|
if install_config:
|
|
client = install_config.lower()
|
|
if client != "codex":
|
|
console.print(
|
|
"[bold red]目前仅支持 Codex MCP 配置，请使用 --install codex。[/bold red]"
|
|
)
|
|
raise typer.Exit(1)
|
|
installed_paths = install_mcp_config(client, python_exe, project_root)
|
|
console.print(
|
|
"[bold green]✔ Successfully installed ARTEMIS MCP server configuration & rules to:[/bold green]"
|
|
)
|
|
for path in installed_paths:
|
|
console.print(f" • [cyan]{path}[/cyan]")
|
|
console.print(
|
|
"\n[dim]Please restart or reload your IDE window to activate the Artemis MCP tools.[/dim]"
|
|
)
|
|
raise typer.Exit(0)
|
|
|
|
if generate_config:
|
|
client = generate_config.lower()
|
|
if client != "codex":
|
|
console.print(
|
|
"[bold red]目前仅支持 Codex MCP 配置，请使用 --generate-config codex。[/bold red]"
|
|
)
|
|
raise typer.Exit(1)
|
|
server_cfg = _get_config_snippet(client, python_exe, project_root)["mcp_servers"]["artemis"]
|
|
json_str = _codex_toml_block(server_cfg)
|
|
syntax_language = "toml"
|
|
|
|
syntax = Syntax(json_str, syntax_language, theme="monokai", line_numbers=False)
|
|
console.print(f"[bold cyan]MCP Configuration for {client.upper()}:[/bold cyan]")
|
|
console.print(syntax)
|
|
raise typer.Exit(0)
|
|
|
|
threading.Thread(target=start_awake_service, daemon=True, name="artemis-awake-init").start()
|
|
try:
|
|
st = server_type.lower()
|
|
if st in ("agent", "mobile", "artemis", "default"):
|
|
logger.info(f"Starting Artemis Mobile Agent MCP Server over {transport}...")
|
|
if transport.lower() == "sse":
|
|
agent_mcp.run(transport="sse", host=host, port=port)
|
|
else:
|
|
agent_mcp.run(transport="stdio")
|
|
elif st == "adb":
|
|
if transport.lower() == "sse":
|
|
logger.info(f"Starting Artemis ADB MCP Server over {transport}...")
|
|
adb_mcp.run(transport="sse", host=host, port=port)
|
|
else:
|
|
# stdio carries the MCP JSON-RPC stream: redirect logging and detach
|
|
# from the parent DataEngine before any output can corrupt it.
|
|
from artemis.mcp.adb_server import configure_stdio_mode
|
|
|
|
configure_stdio_mode()
|
|
logger.info("Starting Artemis ADB MCP Server over stdio...")
|
|
adb_mcp.run(transport="stdio")
|
|
elif st == "xml":
|
|
from artemis.mcp.xml_search_server import mcp as xml_mcp
|
|
|
|
logger.info(f"Starting Artemis XML Fuzzy Search MCP Server over {transport}...")
|
|
if transport.lower() == "sse":
|
|
xml_mcp.run(transport="sse", host=host, port=port)
|
|
else:
|
|
xml_mcp.run(transport="stdio")
|
|
else:
|
|
logger.error(
|
|
f"Unsupported MCP server type: {server_type}. Use 'agent', 'adb', or 'xml'."
|
|
)
|
|
raise typer.Exit(1)
|
|
finally:
|
|
shutdown_awake_service()
|