feat(watch,server): thread events over t3 websocket to the gateway hook, t3_answer for pending questions
This commit is contained in:
@@ -2,17 +2,31 @@ import logging
|
||||
|
||||
from t3code_mcp.config import Settings, load_machines
|
||||
from t3code_mcp.server import Registry, build_server
|
||||
from t3code_mcp.watch import Hook, Tracker, Watcher
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
|
||||
log = logging.getLogger(__name__)
|
||||
settings = Settings()
|
||||
machines = load_machines(settings.config)
|
||||
logging.getLogger(__name__).info(
|
||||
"t3code-mcp: machines %s, listening on %s:%s",
|
||||
registry = Registry(machines, tracker=Tracker(settings.state))
|
||||
hook = (
|
||||
Hook(settings.hook_url, settings.gateway_token, registry.tracker)
|
||||
if settings.hook_url and settings.gateway_token
|
||||
else None
|
||||
)
|
||||
if hook is None:
|
||||
log.warning(
|
||||
"no T3CODE_MCP_HOOK_URL/GATEWAY_TOKEN - events are logged, not sent"
|
||||
)
|
||||
watcher = Watcher(registry, registry.tracker, hook)
|
||||
log.info(
|
||||
"t3code-mcp: machines %s, tracking %d threads, listening on %s:%s",
|
||||
list(machines),
|
||||
len(registry.tracker.threads),
|
||||
settings.host,
|
||||
settings.port,
|
||||
)
|
||||
server = build_server(Registry(machines))
|
||||
server = build_server(registry, background=watcher.run)
|
||||
server.run(transport="http", host=settings.host, port=settings.port, path="/mcp")
|
||||
|
||||
@@ -117,3 +117,6 @@ class Settings(BaseSettings):
|
||||
config: Path = Path("t3code.toml")
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8000
|
||||
state: Path = Path("t3code-mcp.json")
|
||||
hook_url: str | None = None
|
||||
gateway_token: str | None = None
|
||||
|
||||
+78
-89
@@ -8,6 +8,7 @@ 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
|
||||
@@ -18,19 +19,20 @@ from fastmcp.exceptions import ToolError
|
||||
from pydantic import Field
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from t3code_mcp import t3
|
||||
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 Tracker
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine
|
||||
|
||||
from starlette.requests import Request
|
||||
|
||||
TITLE_MAX = 80
|
||||
TEXT_MAX = 6000
|
||||
TOOLS_MAX = 12
|
||||
WAIT_MAX = 3600.0
|
||||
PENDING_EVERY = 6
|
||||
BUSY = frozenset({"running", "starting"})
|
||||
STALLED_POLLS = 3
|
||||
|
||||
|
||||
@@ -40,6 +42,7 @@ class Registry:
|
||||
transport: httpx.AsyncBaseTransport | 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:
|
||||
@@ -57,11 +60,12 @@ class Registry:
|
||||
return self.clients[name]
|
||||
|
||||
async def locate(self, thread_id: str) -> tuple[str, dict[str, Any]]:
|
||||
names = (
|
||||
[self.threads[thread_id]]
|
||||
if thread_id in self.threads
|
||||
else list(self.machines)
|
||||
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, turn_limit=1)
|
||||
if detail is not None:
|
||||
@@ -83,14 +87,6 @@ def _title(prompt: str, title: str | None) -> str:
|
||||
return text or "t3code-mcp"
|
||||
|
||||
|
||||
def _state(thread: dict[str, Any]) -> str:
|
||||
latest = thread.get("latestTurn")
|
||||
if latest:
|
||||
return latest["state"]
|
||||
messages = thread.get("messages", [])
|
||||
return "starting" if messages and messages[-1]["role"] == "user" else "idle"
|
||||
|
||||
|
||||
def _find_project(
|
||||
machine: Machine, projects: list[dict[str, Any]], key: str
|
||||
) -> dict[str, Any]:
|
||||
@@ -125,76 +121,35 @@ def _project_view(
|
||||
}
|
||||
|
||||
|
||||
def _summary(
|
||||
machine: str, detail: dict[str, Any], *, project: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
thread = detail["thread"]
|
||||
latest = thread.get("latestTurn") or {}
|
||||
session = thread.get("session") or {}
|
||||
messages = [
|
||||
{"role": m["role"], "text": m["text"][-TEXT_MAX:], "streaming": m["streaming"]}
|
||||
for m in thread.get("messages", [])
|
||||
if m.get("text")
|
||||
]
|
||||
tools = [
|
||||
{
|
||||
"summary": a["summary"],
|
||||
"detail": str((a.get("payload") or {}).get("detail", ""))[:200],
|
||||
}
|
||||
for a in thread.get("activities", [])
|
||||
if a.get("kind") == "tool.completed"
|
||||
][-TOOLS_MAX:]
|
||||
files = [
|
||||
{"path": f["path"], "kind": f["kind"], "+": f["additions"], "-": f["deletions"]}
|
||||
for c in thread.get("checkpoints", [])[-1:]
|
||||
for f in c.get("files", [])
|
||||
]
|
||||
context = next(
|
||||
(
|
||||
a["payload"]
|
||||
for a in reversed(thread.get("activities", []))
|
||||
if a.get("kind") == "context-window.updated"
|
||||
),
|
||||
None,
|
||||
)
|
||||
return {
|
||||
"thread_id": thread["id"],
|
||||
"machine": machine,
|
||||
"project": project or thread["projectId"],
|
||||
"title": thread["title"],
|
||||
"model": str(ModelSpec.from_selection(thread["modelSelection"])),
|
||||
"state": _state(thread),
|
||||
"turn_id": latest.get("turnId"),
|
||||
"session": session.get("status"),
|
||||
"error": session.get("lastError"),
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"files": files,
|
||||
"context_tokens": (context or {}).get("usedTokens"),
|
||||
}
|
||||
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()
|
||||
|
||||
|
||||
def _pending(shell: dict[str, Any], thread_id: str) -> dict[str, bool] | None:
|
||||
for thread in shell.get("threads", []):
|
||||
if thread["id"] == thread_id:
|
||||
flags = {
|
||||
"approval": bool(thread.get("hasPendingApprovals")),
|
||||
"user_input": bool(thread.get("hasPendingUserInput")),
|
||||
"plan": bool(thread.get("hasActionableProposedPlan")),
|
||||
}
|
||||
return flags if any(flags.values()) else None
|
||||
return None
|
||||
|
||||
|
||||
def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
mcp = FastMCP(
|
||||
"t3code",
|
||||
instructions=(
|
||||
"Coding threads on Бобёр's machines through T3 Code. Flow: t3_machines → "
|
||||
"t3_projects(machine) → t3_dispatch → t3_wait / t3_thread. Prompts run "
|
||||
"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). 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"])
|
||||
@@ -314,6 +269,7 @@ def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
)
|
||||
await _call(client.dispatch(t3.turn_start(tid, text, title_seed=heading)))
|
||||
registry.threads[tid] = machine
|
||||
registry.tracker.track(tid, machine, target["title"], heading)
|
||||
return {
|
||||
"thread_id": tid,
|
||||
"machine": machine,
|
||||
@@ -342,8 +298,10 @@ def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
if detail is None:
|
||||
msg = f"thread {thread_id} vanished from {machine}"
|
||||
raise ToolError(msg)
|
||||
summary = _summary(machine, detail, project=_project_title(shell, detail))
|
||||
summary["pending"] = _pending(shell, thread_id)
|
||||
summary = views.summary(
|
||||
machine, detail, project=views.project_title(shell, detail)
|
||||
)
|
||||
summary["pending"] = views.pending_in(shell, thread_id)
|
||||
return summary
|
||||
|
||||
@mcp.tool
|
||||
@@ -372,20 +330,20 @@ def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
stalled = 0
|
||||
while True:
|
||||
thread = detail["thread"]
|
||||
state = _state(thread)
|
||||
state = views.state(thread)
|
||||
session = thread.get("session") or {}
|
||||
status = session.get("status")
|
||||
if polls % PENDING_EVERY == 0:
|
||||
shell = await _call(client.shell())
|
||||
pending = _pending(shell, thread_id)
|
||||
project = _project_title(shell, detail)
|
||||
pending = views.pending_in(shell, thread_id)
|
||||
project = views.project_title(shell, detail)
|
||||
failed = status == "error" or bool(session.get("lastError"))
|
||||
stalled = stalled + 1 if state == "starting" and status == "stopped" else 0
|
||||
busy = (state in BUSY or status == "starting") and stalled < STALLED_POLLS
|
||||
done = not busy or failed or pending is not None
|
||||
timed_out = time.monotonic() - started >= timeout
|
||||
if done or timed_out:
|
||||
summary = _summary(machine, detail, project=project)
|
||||
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)
|
||||
@@ -398,6 +356,42 @@ def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
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)
|
||||
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.user_input_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")],
|
||||
@@ -421,11 +415,6 @@ def build_server(registry: Registry, *, poll: float = 5.0) -> FastMCP:
|
||||
return mcp
|
||||
|
||||
|
||||
def _project_title(shell: dict[str, Any], detail: dict[str, Any]) -> str | None:
|
||||
project_id = detail["thread"]["projectId"]
|
||||
return next((p["title"] for p in shell["projects"] if p["id"] == project_id), None)
|
||||
|
||||
|
||||
async def _call(awaitable: Any) -> Any:
|
||||
try:
|
||||
return await awaitable
|
||||
|
||||
@@ -18,6 +18,9 @@ DESCRIPTOR = "/.well-known/t3/environment"
|
||||
SHELL = "/api/orchestration/shell"
|
||||
THREADS = "/api/orchestration/threads"
|
||||
DISPATCH = "/api/orchestration/dispatch"
|
||||
WS_TICKET = "/api/auth/websocket-ticket"
|
||||
WS_PATH = "/ws"
|
||||
SUBSCRIBE_SHELL = "orchestration.subscribeShell"
|
||||
|
||||
RUNTIME_MODE = "full-access"
|
||||
INTERACTION_MODE = "default"
|
||||
@@ -101,6 +104,19 @@ def turn_interrupt(thread_id: str, turn_id: str | None = None) -> dict[str, Any]
|
||||
return command
|
||||
|
||||
|
||||
def user_input_respond(
|
||||
thread_id: str, request_id: str, answers: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "thread.user-input.respond",
|
||||
"commandId": new_id(),
|
||||
"threadId": thread_id,
|
||||
"requestId": request_id,
|
||||
"answers": answers,
|
||||
"createdAt": now_iso(),
|
||||
}
|
||||
|
||||
|
||||
class T3Client:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -142,6 +158,11 @@ class T3Client:
|
||||
async def dispatch(self, command: dict[str, Any]) -> dict[str, Any]:
|
||||
return self._json(await self._http.post(DISPATCH, json=command))
|
||||
|
||||
async def ws_url(self) -> str:
|
||||
ticket = self._json(await self._http.post(WS_TICKET))["ticket"]
|
||||
scheme = "wss" if self.url.startswith("https") else "ws"
|
||||
return f"{scheme}://{self.url.split('://', 1)[1]}{WS_PATH}?wsTicket={ticket}"
|
||||
|
||||
@staticmethod
|
||||
def _json(response: httpx.Response) -> dict[str, Any]:
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Read-model helpers shared by the tools and the watcher."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from t3code_mcp.config import ModelSpec
|
||||
|
||||
TEXT_MAX = 6000
|
||||
TOOLS_MAX = 12
|
||||
BUSY = frozenset({"running", "starting"})
|
||||
|
||||
|
||||
def state(thread: dict[str, Any]) -> str:
|
||||
latest = thread.get("latestTurn")
|
||||
if latest:
|
||||
return latest["state"]
|
||||
messages = thread.get("messages", [])
|
||||
return "starting" if messages and messages[-1]["role"] == "user" else "idle"
|
||||
|
||||
|
||||
def shell_state(shell_thread: dict[str, Any]) -> str:
|
||||
latest = shell_thread.get("latestTurn")
|
||||
return latest["state"] if latest else "starting"
|
||||
|
||||
|
||||
def pending(shell_thread: dict[str, Any]) -> dict[str, bool] | None:
|
||||
flags = {
|
||||
"approval": bool(shell_thread.get("hasPendingApprovals")),
|
||||
"user_input": bool(shell_thread.get("hasPendingUserInput")),
|
||||
"plan": bool(shell_thread.get("hasActionableProposedPlan")),
|
||||
}
|
||||
return flags if any(flags.values()) else None
|
||||
|
||||
|
||||
def pending_in(shell: dict[str, Any], thread_id: str) -> dict[str, bool] | None:
|
||||
for thread in shell.get("threads", []):
|
||||
if thread["id"] == thread_id:
|
||||
return pending(thread)
|
||||
return None
|
||||
|
||||
|
||||
def question(detail: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""The unanswered `user-input.requested` of a thread, newest first."""
|
||||
for activity in reversed(detail["thread"].get("activities", [])):
|
||||
kind = activity.get("kind")
|
||||
if kind == "user-input.resolved":
|
||||
return None
|
||||
if kind == "user-input.requested":
|
||||
payload = activity.get("payload") or {}
|
||||
return {
|
||||
"request_id": payload.get("requestId"),
|
||||
"questions": [
|
||||
{
|
||||
"id": q["id"],
|
||||
"header": q.get("header"),
|
||||
"question": q["question"],
|
||||
"options": [
|
||||
{"label": o["label"], "description": o.get("description")}
|
||||
for o in q.get("options", [])
|
||||
],
|
||||
"multi_select": bool(q.get("multiSelect")),
|
||||
}
|
||||
for q in payload.get("questions", [])
|
||||
],
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def project_title(shell: dict[str, Any], detail: dict[str, Any]) -> str | None:
|
||||
project_id = detail["thread"]["projectId"]
|
||||
return next((p["title"] for p in shell["projects"] if p["id"] == project_id), None)
|
||||
|
||||
|
||||
def summary(
|
||||
machine: str, detail: dict[str, Any], *, project: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
thread = detail["thread"]
|
||||
latest = thread.get("latestTurn") or {}
|
||||
session = thread.get("session") or {}
|
||||
messages = [
|
||||
{"role": m["role"], "text": m["text"][-TEXT_MAX:], "streaming": m["streaming"]}
|
||||
for m in thread.get("messages", [])
|
||||
if m.get("text")
|
||||
]
|
||||
tools = [
|
||||
{
|
||||
"summary": a["summary"],
|
||||
"detail": str((a.get("payload") or {}).get("detail", ""))[:200],
|
||||
}
|
||||
for a in thread.get("activities", [])
|
||||
if a.get("kind") == "tool.completed"
|
||||
][-TOOLS_MAX:]
|
||||
files = [
|
||||
{"path": f["path"], "kind": f["kind"], "+": f["additions"], "-": f["deletions"]}
|
||||
for c in thread.get("checkpoints", [])[-1:]
|
||||
for f in c.get("files", [])
|
||||
]
|
||||
context = next(
|
||||
(
|
||||
a["payload"]
|
||||
for a in reversed(thread.get("activities", []))
|
||||
if a.get("kind") == "context-window.updated"
|
||||
),
|
||||
None,
|
||||
)
|
||||
return {
|
||||
"thread_id": thread["id"],
|
||||
"machine": machine,
|
||||
"project": project or thread["projectId"],
|
||||
"title": thread["title"],
|
||||
"model": str(ModelSpec.from_selection(thread["modelSelection"])),
|
||||
"state": state(thread),
|
||||
"turn_id": latest.get("turnId"),
|
||||
"session": session.get("status"),
|
||||
"error": session.get("lastError"),
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"files": files,
|
||||
"context_tokens": (context or {}).get("usedTokens"),
|
||||
"question": question(detail),
|
||||
}
|
||||
|
||||
|
||||
def last_answer(detail: dict[str, Any]) -> str:
|
||||
for m in reversed(detail["thread"].get("messages", [])):
|
||||
if m["role"] == "assistant" and m.get("text"):
|
||||
return m["text"]
|
||||
return ""
|
||||
@@ -0,0 +1,326 @@
|
||||
"""Follow dispatched threads over T3's WebSocket RPC and push events to the gateway.
|
||||
|
||||
One `orchestration.subscribeShell` stream per machine (Effect RPC over JSON:
|
||||
`Request` → `Chunk`s acknowledged with `Ack`, `Ping` keepalive). The shell
|
||||
read-model carries everything a transition needs - `latestTurn.state`,
|
||||
`session`, `hasPendingUserInput` - so no per-thread subscription is required.
|
||||
Reconnects resume with `afterSequence`; a fresh snapshot re-derives every
|
||||
tracked thread, so transitions missed while offline are still reported.
|
||||
|
||||
Only threads started through `t3_dispatch` are tracked. Events go to the
|
||||
gateway webhook through a persisted outbox that retries until accepted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
import httpx
|
||||
import websockets
|
||||
|
||||
from t3code_mcp import t3, views
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from t3code_mcp.server import Registry
|
||||
|
||||
_log = logging.getLogger(__name__)
|
||||
|
||||
EventKind = Literal["completed", "interrupted", "failed", "question", "approval"]
|
||||
TERMINAL: dict[str, EventKind] = {
|
||||
"completed": "completed",
|
||||
"interrupted": "interrupted",
|
||||
"error": "failed",
|
||||
}
|
||||
KEEP_DAYS = 14
|
||||
PING_EVERY = 20.0
|
||||
RECONNECT_MAX = 60.0
|
||||
HOOK_RETRY_MAX = 60.0
|
||||
TEXT_MAX = 4000
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tracked:
|
||||
machine: str
|
||||
project: str
|
||||
title: str
|
||||
state: str = "starting"
|
||||
pending: str | None = None
|
||||
error: str | None = None
|
||||
updated: float = field(default_factory=time.time)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Event:
|
||||
kind: EventKind
|
||||
thread_id: str
|
||||
machine: str
|
||||
|
||||
|
||||
class Tracker:
|
||||
def __init__(self, path: Path | None = None) -> None:
|
||||
self.path = path
|
||||
self.threads: dict[str, Tracked] = {}
|
||||
self.outbox: list[dict[str, Any]] = []
|
||||
self.sequences: dict[str, int] = {}
|
||||
if path is not None and path.exists():
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
self.threads = {k: Tracked(**v) for k, v in raw.get("threads", {}).items()}
|
||||
self.outbox = raw.get("outbox", [])
|
||||
self.sequences = raw.get("sequences", {})
|
||||
|
||||
def save(self) -> None:
|
||||
if self.path is None:
|
||||
return
|
||||
cutoff = time.time() - KEEP_DAYS * 86400
|
||||
self.threads = {k: v for k, v in self.threads.items() if v.updated >= cutoff}
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload = {
|
||||
"threads": {k: asdict(v) for k, v in self.threads.items()},
|
||||
"outbox": self.outbox,
|
||||
"sequences": self.sequences,
|
||||
}
|
||||
tmp = self.path.with_suffix(".tmp")
|
||||
tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
||||
tmp.replace(self.path)
|
||||
|
||||
def track(self, thread_id: str, machine: str, project: str, title: str) -> None:
|
||||
self.threads[thread_id] = Tracked(machine=machine, project=project, title=title)
|
||||
self.save()
|
||||
|
||||
def apply(self, machine: str, shell_thread: dict[str, Any]) -> list[Event]:
|
||||
tracked = self.threads.get(shell_thread["id"])
|
||||
if tracked is None or tracked.machine != machine:
|
||||
return []
|
||||
state = views.shell_state(shell_thread)
|
||||
flags = views.pending(shell_thread) or {}
|
||||
pending = next((k for k in ("user_input", "approval") if flags.get(k)), None)
|
||||
error = (shell_thread.get("session") or {}).get("lastError")
|
||||
events: list[Event] = []
|
||||
if pending and pending != tracked.pending:
|
||||
kind: EventKind = "question" if pending == "user_input" else "approval"
|
||||
events.append(Event(kind, shell_thread["id"], machine))
|
||||
if tracked.state in views.BUSY and state in TERMINAL:
|
||||
events.append(Event(TERMINAL[state], shell_thread["id"], machine))
|
||||
elif tracked.state in views.BUSY and error and not tracked.error:
|
||||
events.append(Event("failed", shell_thread["id"], machine))
|
||||
changed = (state, pending, error) != (
|
||||
tracked.state,
|
||||
tracked.pending,
|
||||
tracked.error,
|
||||
)
|
||||
tracked.state, tracked.pending, tracked.error = state, pending, error
|
||||
if changed:
|
||||
tracked.updated = time.time()
|
||||
self.save()
|
||||
return events
|
||||
|
||||
|
||||
class Hook:
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
token: str,
|
||||
tracker: Tracker,
|
||||
*,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> None:
|
||||
self.tracker = tracker
|
||||
self._http = httpx.AsyncClient(
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
timeout=30,
|
||||
transport=transport,
|
||||
)
|
||||
self._url = url
|
||||
self._wake = asyncio.Event()
|
||||
|
||||
def send(self, payload: dict[str, Any]) -> None:
|
||||
self.tracker.outbox.append(payload)
|
||||
self.tracker.save()
|
||||
self._wake.set()
|
||||
|
||||
async def deliver_once(self) -> bool:
|
||||
payload = self.tracker.outbox[0]
|
||||
try:
|
||||
response = await self._http.post(self._url, json=payload)
|
||||
except httpx.HTTPError as exc:
|
||||
_log.warning("hook: %s unreachable: %s", self._url, exc)
|
||||
return False
|
||||
if response.is_error:
|
||||
_log.warning(
|
||||
"hook: %s answered %s: %s",
|
||||
self._url,
|
||||
response.status_code,
|
||||
response.text[:200],
|
||||
)
|
||||
return False
|
||||
self.tracker.outbox.pop(0)
|
||||
self.tracker.save()
|
||||
_log.info(
|
||||
"hook: delivered %s for %s", payload.get("event"), payload.get("thread_id")
|
||||
)
|
||||
return True
|
||||
|
||||
async def run(self) -> None:
|
||||
backoff = 1.0
|
||||
while True:
|
||||
if not self.tracker.outbox:
|
||||
self._wake.clear()
|
||||
await self._wake.wait()
|
||||
continue
|
||||
if await self.deliver_once():
|
||||
backoff = 1.0
|
||||
continue
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 2, HOOK_RETRY_MAX)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._http.aclose()
|
||||
|
||||
|
||||
class Watcher:
|
||||
def __init__(self, registry: Registry, tracker: Tracker, hook: Hook | None) -> None:
|
||||
self.registry = registry
|
||||
self.tracker = tracker
|
||||
self.hook = hook
|
||||
|
||||
async def run(self) -> None:
|
||||
tasks = [self._machine(name) for name in self.registry.machines]
|
||||
if self.hook is not None:
|
||||
tasks.append(self.hook.run())
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def _machine(self, name: str) -> None:
|
||||
backoff = 1.0
|
||||
while True:
|
||||
try:
|
||||
await self._session(name)
|
||||
backoff = 1.0
|
||||
except (
|
||||
OSError,
|
||||
websockets.WebSocketException,
|
||||
httpx.HTTPError,
|
||||
t3.T3Error,
|
||||
) as exc:
|
||||
_log.info("watch %s: %s; retry in %.0fs", name, str(exc)[:200], backoff)
|
||||
except Exception: # noqa: BLE001 - the watcher must outlive any bug
|
||||
_log.exception(
|
||||
"watch %s: unexpected error; retry in %.0fs", name, backoff
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 2, RECONNECT_MAX)
|
||||
|
||||
async def _session(self, name: str) -> None:
|
||||
client = self.registry.client(name)
|
||||
url = await client.ws_url()
|
||||
payload: dict[str, Any] = {"requestCompletionMarker": True}
|
||||
if name in self.tracker.sequences:
|
||||
payload["afterSequence"] = self.tracker.sequences[name]
|
||||
async with websockets.connect(url, max_size=None, ping_interval=None) as ws:
|
||||
request = {
|
||||
"_tag": "Request",
|
||||
"id": "1",
|
||||
"tag": t3.SUBSCRIBE_SHELL,
|
||||
"payload": payload,
|
||||
"headers": [],
|
||||
}
|
||||
await ws.send(json.dumps(request))
|
||||
_log.info(
|
||||
"watch %s: subscribed (after %s)", name, payload.get("afterSequence")
|
||||
)
|
||||
pinger = asyncio.create_task(self._ping(ws))
|
||||
try:
|
||||
async for raw in ws:
|
||||
for message in _messages(raw):
|
||||
if await self._handle(name, ws, message):
|
||||
return
|
||||
finally:
|
||||
pinger.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await pinger
|
||||
|
||||
async def _handle(self, name: str, ws: Any, message: dict[str, Any]) -> bool:
|
||||
tag = message.get("_tag")
|
||||
if tag == "Chunk":
|
||||
for item in message.get("values", []):
|
||||
await self._item(name, item)
|
||||
await ws.send(
|
||||
json.dumps({"_tag": "Ack", "requestId": message["requestId"]})
|
||||
)
|
||||
return False
|
||||
if tag == "Exit":
|
||||
_log.info(
|
||||
"watch %s: stream ended: %s", name, str(message.get("exit"))[:200]
|
||||
)
|
||||
return True
|
||||
if tag in ("Defect", "ClientProtocolError"):
|
||||
msg = f"rpc {tag}: {str(message)[:300]}"
|
||||
raise t3.T3Error(0, msg)
|
||||
return False
|
||||
|
||||
async def _item(self, name: str, item: dict[str, Any]) -> None:
|
||||
kind = item.get("kind")
|
||||
if kind == "snapshot":
|
||||
snapshot = item["snapshot"]
|
||||
self.tracker.sequences[name] = snapshot["snapshotSequence"]
|
||||
for thread in snapshot["threads"]:
|
||||
await self._events(self.tracker.apply(name, thread))
|
||||
elif kind == "thread-upserted":
|
||||
self.tracker.sequences[name] = item["sequence"]
|
||||
await self._events(self.tracker.apply(name, item["thread"]))
|
||||
elif "sequence" in item:
|
||||
self.tracker.sequences[name] = item["sequence"]
|
||||
|
||||
async def _events(self, events: list[Event]) -> None:
|
||||
for event in events:
|
||||
payload = await self.payload(event)
|
||||
_log.info("watch: %s %s on %s", event.kind, event.thread_id, event.machine)
|
||||
if self.hook is not None:
|
||||
self.hook.send(payload)
|
||||
|
||||
async def payload(self, event: Event) -> dict[str, Any]:
|
||||
tracked = self.tracker.threads[event.thread_id]
|
||||
payload: dict[str, Any] = {
|
||||
"event": event.kind,
|
||||
"thread_id": event.thread_id,
|
||||
"machine": event.machine,
|
||||
"project": tracked.project,
|
||||
"title": tracked.title,
|
||||
"state": tracked.state,
|
||||
"error": tracked.error,
|
||||
}
|
||||
try:
|
||||
detail = await self.registry.client(event.machine).thread(
|
||||
event.thread_id, turn_limit=1
|
||||
)
|
||||
except (httpx.HTTPError, t3.T3Error) as exc:
|
||||
payload["detail_error"] = str(exc)[:200]
|
||||
return payload
|
||||
if detail is None:
|
||||
return payload
|
||||
view = views.summary(event.machine, detail, project=tracked.project)
|
||||
payload.update(
|
||||
title=view["title"],
|
||||
text=views.last_answer(detail)[-TEXT_MAX:],
|
||||
files=view["files"],
|
||||
question=view["question"],
|
||||
)
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
async def _ping(ws: Any) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(PING_EVERY)
|
||||
await ws.send(json.dumps({"_tag": "Ping"}))
|
||||
|
||||
|
||||
def _messages(raw: str | bytes) -> list[dict[str, Any]]:
|
||||
decoded = json.loads(raw)
|
||||
return decoded if isinstance(decoded, list) else [decoded]
|
||||
Reference in New Issue
Block a user