472 lines
17 KiB
Python
472 lines
17 KiB
Python
"""MCP tools over one or more T3 Code servers.
|
|
|
|
`t3_wait` blocks on purpose: Claude Code has no MCP Tasks support and moves
|
|
any tool call still running after two minutes into a background task by
|
|
itself, so a plain long poll is the cheapest correct shape here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Annotated, Any
|
|
|
|
import httpx
|
|
import websockets
|
|
from fastmcp import FastMCP
|
|
from fastmcp.exceptions import ToolError
|
|
from pydantic import Field
|
|
from starlette.responses import JSONResponse
|
|
|
|
from t3code_mcp import t3, views
|
|
from t3code_mcp.config import Machine, ModelSpec
|
|
from t3code_mcp.t3 import T3Client, T3Error
|
|
from t3code_mcp.views import BUSY
|
|
from t3code_mcp.watch import Tracked, Tracker
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator, Callable, Coroutine
|
|
|
|
from starlette.requests import Request
|
|
|
|
from t3code_mcp.t3 import RpcCall
|
|
|
|
TITLE_MAX = 80
|
|
WAIT_MAX = 3600.0
|
|
|
|
|
|
@dataclass
|
|
class Registry:
|
|
machines: dict[str, Machine]
|
|
transport: httpx.AsyncBaseTransport | None = None
|
|
rpc: Callable[[str], RpcCall] | None = None
|
|
clients: dict[str, T3Client] = field(default_factory=dict)
|
|
threads: dict[str, str] = field(default_factory=dict)
|
|
tracker: Tracker = field(default_factory=Tracker)
|
|
|
|
def machine(self, name: str) -> Machine:
|
|
try:
|
|
return self.machines[name]
|
|
except KeyError:
|
|
msg = f"unknown machine {name!r}; known: {', '.join(self.machines)}"
|
|
raise ToolError(msg) from None
|
|
|
|
def client(self, name: str) -> T3Client:
|
|
machine = self.machine(name)
|
|
if name not in self.clients:
|
|
self.clients[name] = T3Client(
|
|
machine.url,
|
|
machine.token,
|
|
transport=self.transport,
|
|
rpc=self.rpc(name) if self.rpc else None,
|
|
)
|
|
return self.clients[name]
|
|
|
|
async def locate(self, thread_id: str) -> tuple[str, dict[str, Any]]:
|
|
known = self.threads.get(thread_id) or (
|
|
self.tracker.threads[thread_id].machine
|
|
if thread_id in self.tracker.threads
|
|
else None
|
|
)
|
|
names = [known] if known in self.machines else list(self.machines)
|
|
for name in names:
|
|
detail = await self.client(name).thread(thread_id)
|
|
if detail is not None:
|
|
self.threads[thread_id] = name
|
|
return name, detail
|
|
msg = f"thread {thread_id} not found on {', '.join(names)}"
|
|
raise ToolError(msg)
|
|
|
|
async def aclose(self) -> None:
|
|
for client in self.clients.values():
|
|
await client.aclose()
|
|
self.clients.clear()
|
|
|
|
|
|
def _title(prompt: str, title: str | None) -> str:
|
|
text = (title or prompt).strip().splitlines()[0].strip()
|
|
if len(text) > TITLE_MAX:
|
|
text = text[:TITLE_MAX].rsplit(" ", 1)[0] + "…"
|
|
return text or "t3code-mcp"
|
|
|
|
|
|
def _find_project(
|
|
machine: Machine, projects: list[dict[str, Any]], key: str
|
|
) -> dict[str, Any]:
|
|
allowed = [p for p in projects if machine.allows(p)]
|
|
for project in allowed:
|
|
if key in (project["id"], project["title"], project["workspaceRoot"]):
|
|
return project
|
|
names = ", ".join(p["title"] for p in allowed) or "(none allowlisted)"
|
|
msg = f"project {key!r} is not available on {machine.name}; allowed: {names}"
|
|
raise ToolError(msg)
|
|
|
|
|
|
def _project_view(
|
|
project: dict[str, Any], threads: list[dict[str, Any]]
|
|
) -> dict[str, Any]:
|
|
own = [
|
|
t
|
|
for t in threads
|
|
if t["projectId"] == project["id"] and not t.get("archivedAt")
|
|
]
|
|
running = sum(1 for t in own if views.shell_state(t) in BUSY)
|
|
default = project.get("defaultModelSelection")
|
|
return {
|
|
"id": project["id"],
|
|
"title": project["title"],
|
|
"path": project["workspaceRoot"],
|
|
"default_model": str(ModelSpec.from_selection(default)) if default else None,
|
|
"threads": len(own),
|
|
"running": running,
|
|
}
|
|
|
|
|
|
def _watch_view(thread_id: str, tracked: Tracked) -> dict[str, Any]:
|
|
return {
|
|
"thread_id": thread_id,
|
|
"machine": tracked.machine,
|
|
"project": tracked.project,
|
|
"title": tracked.title,
|
|
"state": tracked.state,
|
|
}
|
|
|
|
|
|
def build_server(
|
|
registry: Registry,
|
|
*,
|
|
poll: float = 5.0,
|
|
background: Callable[[], Coroutine[Any, Any, None]] | None = None,
|
|
) -> FastMCP:
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(_: FastMCP) -> AsyncIterator[dict[str, Any]]:
|
|
task = asyncio.create_task(background()) if background is not None else None
|
|
try:
|
|
yield {}
|
|
finally:
|
|
if task is not None:
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
await registry.aclose()
|
|
|
|
mcp = FastMCP(
|
|
"t3code",
|
|
instructions=(
|
|
"Coding threads on Бобёр's machines through T3 Code. Flow: t3_machines → "
|
|
"t3_projects(machine) → t3_dispatch, then either t3_wait (block) or just "
|
|
"carry on: the gateway injects an event when a dispatched thread finishes, "
|
|
"fails or asks a question (answer with t3_answer; t3_watch turns those "
|
|
"events off for a thread). Prompts run "
|
|
"autonomously (full-access, no approvals): say what to change, what to "
|
|
"check, what not to touch, and that the answer should end with a report."
|
|
),
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
@mcp.custom_route("/healthz", methods=["GET"])
|
|
async def healthz(_: Request) -> JSONResponse:
|
|
return JSONResponse({"status": "ok", "machines": list(registry.machines)})
|
|
|
|
@mcp.tool
|
|
async def t3_machines() -> list[dict[str, Any]]:
|
|
"""Machines with a T3 Code server: reachable now or not, project allowlist."""
|
|
result = []
|
|
for name, machine in registry.machines.items():
|
|
entry: dict[str, Any] = {
|
|
"name": name,
|
|
"url": machine.url,
|
|
"projects": machine.projects,
|
|
}
|
|
try:
|
|
descriptor = await registry.client(name).descriptor()
|
|
except (httpx.HTTPError, T3Error) as exc:
|
|
entry.update(online=False, error=str(exc)[:200])
|
|
else:
|
|
entry.update(
|
|
online=True,
|
|
label=descriptor.get("label"),
|
|
version=descriptor.get("serverVersion"),
|
|
os=(descriptor.get("platform") or {}).get("os"),
|
|
)
|
|
result.append(entry)
|
|
return result
|
|
|
|
@mcp.tool
|
|
async def t3_projects(
|
|
machine: Annotated[str, Field(description="Machine name from t3_machines")],
|
|
) -> list[dict[str, Any]]:
|
|
"""Allowlisted projects on a machine: default model, thread counts."""
|
|
spec = registry.machine(machine)
|
|
shell = await _call(registry.client(machine).shell())
|
|
return [
|
|
_project_view(p, shell["threads"])
|
|
for p in shell["projects"]
|
|
if spec.allows(p)
|
|
]
|
|
|
|
@mcp.tool
|
|
async def t3_dispatch(
|
|
machine: Annotated[str, Field(description="Machine name from t3_machines")],
|
|
project: Annotated[
|
|
str, Field(description="Project title, path or id from t3_projects")
|
|
],
|
|
prompt: Annotated[
|
|
str,
|
|
Field(
|
|
description="The task for the coding agent, complete and self-contained"
|
|
),
|
|
],
|
|
title: Annotated[
|
|
str | None,
|
|
Field(description="Thread title; default - first line of the prompt"),
|
|
] = None,
|
|
model: Annotated[
|
|
str | None,
|
|
Field(
|
|
description=(
|
|
"`instance/model`, e.g. claudeAgent/claude-opus-5; "
|
|
"default from config or the project"
|
|
)
|
|
),
|
|
] = None,
|
|
thread_id: Annotated[
|
|
str | None,
|
|
Field(description="Continue this existing thread instead of creating one"),
|
|
] = None,
|
|
watch: Annotated[ # noqa: FBT002 - MCP arguments are always named
|
|
bool,
|
|
Field(
|
|
description=(
|
|
"Report this thread back as an event; false for a thread you "
|
|
"read yourself. Every dispatch resets the flag - see t3_watch"
|
|
)
|
|
),
|
|
] = True,
|
|
) -> dict[str, Any]:
|
|
"""Start a coding thread on a machine (or follow up in one); returns at once.
|
|
|
|
The thread runs in the background: t3_wait blocks until it finishes,
|
|
t3_thread peeks without blocking.
|
|
"""
|
|
text = prompt.strip()
|
|
if not text:
|
|
msg = "prompt is empty"
|
|
raise ToolError(msg)
|
|
spec = registry.machine(machine)
|
|
client = registry.client(machine)
|
|
shell = await _call(client.shell())
|
|
target = _find_project(spec, shell["projects"], project)
|
|
if model:
|
|
selection = ModelSpec.parse(
|
|
model, spec.model.options if spec.model else None
|
|
)
|
|
elif spec.model:
|
|
selection = spec.model
|
|
elif target.get("defaultModelSelection"):
|
|
selection = ModelSpec.from_selection(target["defaultModelSelection"])
|
|
else:
|
|
msg = f"no model: set `model` in config for {machine} or pass model="
|
|
raise ToolError(msg)
|
|
heading = _title(text, title)
|
|
if thread_id:
|
|
existing = next((t for t in shell["threads"] if t["id"] == thread_id), None)
|
|
if existing is None or existing["projectId"] != target["id"]:
|
|
msg = f"thread {thread_id} is not in {target['title']} on {machine}"
|
|
raise ToolError(msg)
|
|
if views.shell_state(existing) in BUSY:
|
|
msg = f"thread {thread_id} is running; t3_wait or t3_interrupt first"
|
|
raise ToolError(msg)
|
|
tid = thread_id
|
|
await _call(
|
|
client.dispatch(t3.message_dispatch(tid, text, selection.selection()))
|
|
)
|
|
else:
|
|
tid = t3.new_id()
|
|
await _call(
|
|
client.dispatch(
|
|
t3.thread_create(tid, target["id"], heading, selection.selection())
|
|
)
|
|
)
|
|
await _call(
|
|
client.dispatch(t3.message_dispatch(tid, text, title_seed=heading))
|
|
)
|
|
registry.threads[tid] = machine
|
|
registry.tracker.track(tid, machine, target["title"], heading, watch=watch)
|
|
return {
|
|
"thread_id": tid,
|
|
"machine": machine,
|
|
"project": target["title"],
|
|
"path": target["workspaceRoot"],
|
|
"title": heading,
|
|
"model": str(selection),
|
|
"state": "running",
|
|
"watch": watch,
|
|
}
|
|
|
|
@mcp.tool
|
|
async def t3_watch(
|
|
thread_id: Annotated[
|
|
str | None,
|
|
Field(description="Thread to subscribe or unsubscribe; omit to only list"),
|
|
] = None,
|
|
enabled: Annotated[
|
|
bool | None,
|
|
Field(description="True - report the thread back again, false - stop"),
|
|
] = None,
|
|
) -> dict[str, Any]:
|
|
"""Which dispatched threads report back, and turn that on or off per thread.
|
|
|
|
t3_dispatch subscribes its thread, so its completion, failure and
|
|
questions arrive here as events. Unsubscribe a thread someone reads
|
|
directly: it keeps running, it just stops reporting. Called with no
|
|
arguments this only lists.
|
|
"""
|
|
if thread_id is not None:
|
|
if enabled is None:
|
|
msg = "pass enabled=true or enabled=false along with thread_id"
|
|
raise ToolError(msg)
|
|
try:
|
|
registry.tracker.set_watch(thread_id, enabled=enabled)
|
|
except KeyError:
|
|
msg = f"thread {thread_id} was not dispatched from here; nothing to "
|
|
raise ToolError(msg + ("watch" if enabled else "unwatch")) from None
|
|
threads = registry.tracker.threads
|
|
return {
|
|
"watched": [_watch_view(k, v) for k, v in threads.items() if v.watch],
|
|
"unwatched": [_watch_view(k, v) for k, v in threads.items() if not v.watch],
|
|
}
|
|
|
|
@mcp.tool
|
|
async def t3_thread(
|
|
thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
|
|
turns: Annotated[
|
|
int,
|
|
Field(
|
|
description="How many recent runs (user turns) to include", ge=1, le=20
|
|
),
|
|
] = 2,
|
|
) -> dict[str, Any]:
|
|
"""State of a thread without waiting.
|
|
|
|
Turn state, last messages, tools run, files changed, pending questions.
|
|
"""
|
|
machine, detail = await registry.locate(thread_id)
|
|
shell = await _call(registry.client(machine).shell())
|
|
summary = views.summary(
|
|
machine, detail, project=views.project_title(shell, detail), turns=turns
|
|
)
|
|
summary["pending"] = views.pending_detail(detail["projection"])
|
|
return summary
|
|
|
|
@mcp.tool
|
|
async def t3_wait(
|
|
thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
|
|
timeout: Annotated[ # noqa: ASYNC109 - the tool's contract, not asyncio's
|
|
float,
|
|
Field(
|
|
description="Seconds to wait before giving the current state back",
|
|
ge=5,
|
|
le=WAIT_MAX,
|
|
),
|
|
] = 900,
|
|
) -> dict[str, Any]:
|
|
"""Block until the turn finishes, asks a question, or timeout runs out.
|
|
|
|
Ends on completed / interrupted / error, on a pending approval or
|
|
user-input request, or after `timeout` seconds; returns the t3_thread view.
|
|
"""
|
|
machine, detail = await registry.locate(thread_id)
|
|
client = registry.client(machine)
|
|
started = time.monotonic()
|
|
project = views.project_title(await _call(client.shell()), detail)
|
|
while True:
|
|
projection = detail["projection"]
|
|
state = views.state(projection)
|
|
pending = views.pending_detail(projection)
|
|
stuck = views.latest_run(projection) is None and views.error(projection)
|
|
done = state not in BUSY or pending is not None or stuck
|
|
timed_out = time.monotonic() - started >= timeout
|
|
if done or timed_out:
|
|
summary = views.summary(machine, detail, project=project)
|
|
summary["pending"] = pending
|
|
summary["timed_out"] = timed_out and not done
|
|
summary["waited"] = round(time.monotonic() - started, 1)
|
|
return summary
|
|
await asyncio.sleep(poll)
|
|
refreshed = await _call(client.thread(thread_id))
|
|
if refreshed is None:
|
|
msg = f"thread {thread_id} vanished from {machine}"
|
|
raise ToolError(msg)
|
|
detail = refreshed
|
|
|
|
@mcp.tool
|
|
async def t3_answer(
|
|
thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
|
|
answers: Annotated[
|
|
dict[str, str | list[str]],
|
|
Field(
|
|
description=(
|
|
"question id → chosen option label (list for multi_select), "
|
|
"ids and options from the `question` of t3_thread / the event"
|
|
)
|
|
),
|
|
],
|
|
) -> dict[str, Any]:
|
|
"""Answer the pending question of a thread; its turn resumes."""
|
|
machine, detail = await registry.locate(thread_id)
|
|
pending = views.question(detail["projection"])
|
|
if pending is None or not pending.get("request_id"):
|
|
msg = f"thread {thread_id} has no pending question"
|
|
raise ToolError(msg)
|
|
known = {q["id"] for q in pending["questions"]}
|
|
unknown = set(answers) - known
|
|
if unknown:
|
|
msg = f"unknown question ids {sorted(unknown)}; expected {sorted(known)}"
|
|
raise ToolError(msg)
|
|
await _call(
|
|
registry.client(machine).dispatch(
|
|
t3.runtime_request_respond(
|
|
thread_id, pending["request_id"], dict(answers)
|
|
)
|
|
)
|
|
)
|
|
return {
|
|
"thread_id": thread_id,
|
|
"machine": machine,
|
|
"answered": sorted(answers),
|
|
"state": "running",
|
|
}
|
|
|
|
@mcp.tool
|
|
async def t3_interrupt(
|
|
thread_id: Annotated[str, Field(description="thread_id from t3_dispatch")],
|
|
) -> dict[str, Any]:
|
|
"""Interrupt the running turn; continue later via t3_dispatch(thread_id=...)."""
|
|
machine, detail = await registry.locate(thread_id)
|
|
latest = views.latest_run(detail["projection"])
|
|
if latest is None or latest["status"] not in BUSY:
|
|
return {
|
|
"thread_id": thread_id,
|
|
"machine": machine,
|
|
"state": latest["status"] if latest else "idle",
|
|
}
|
|
await _call(
|
|
registry.client(machine).dispatch(t3.run_interrupt(thread_id, latest["id"]))
|
|
)
|
|
return {"thread_id": thread_id, "machine": machine, "state": "interrupting"}
|
|
|
|
return mcp
|
|
|
|
|
|
async def _call(awaitable: Any) -> Any:
|
|
try:
|
|
return await awaitable
|
|
except T3Error as exc:
|
|
raise ToolError(str(exc)) from exc
|
|
except (httpx.HTTPError, websockets.WebSocketException, OSError) as exc:
|
|
msg = f"T3 unreachable: {exc}"
|
|
raise ToolError(msg) from exc
|