← back to A2a Lab

tk_agent.py

327 lines

"""A2A bridge over the REAL tk cross-agent DM system — a2a-sdk 1.1.2.

Exposes Steve's ticket-system agent-to-agent DM log (~/.claude/tickets/events.jsonl)
as an A2A-discoverable agent, so any A2A client can query it over the standard
protocol instead of the local jsonl log directly.

READ-ONLY BY DESIGN: it only folds/reads the event log (ports the inbox + known-agents
logic from ticket-system/lib.js). It NEVER appends a dm/read event, so running it can
never pollute the shared cross-agent store. (Writing DMs over A2A would be a separate,
gated capability.)

Skills (drive via the message text):
  "agents"           -> list every known agent identity
  "inbox <agent>"    -> that agent's UNREAD inbox (add "all" for read+unread)

Run: . .venv/bin/activate && python tk_agent.py     (listens on 127.0.0.1:41242)
"""
from __future__ import annotations

import json
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

import uvicorn
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import (
    AgentCapabilities,
    AgentCard,
    AgentInterface,
    AgentSkill,
    Message,
    Part,
    Role,
)
from starlette.applications import Starlette

HOST = "127.0.0.1"
PORT = 41242
# TK_EVENTS overrides the store path (used for pollution-free testing against a temp file).
EVENTS = Path(
    os.environ.get("TK_EVENTS", str(Path.home() / ".claude" / "tickets" / "events.jsonl"))
)
LOCK = EVENTS.parent / ".lock"
# Writes are OFF unless explicitly enabled — a discovery bridge must not mutate the
# shared cross-agent log just because someone POSTs a `dm` command.
ALLOW_WRITE = os.environ.get("TK_BRIDGE_ALLOW_WRITE") == "1"
BRIDGE_FROM = os.environ.get("TK_BRIDGE_FROM", "a2a-bridge")

# ── tk store reads (faithful port of ticket-system/lib.js, READ path only) ──────
_BROADCAST = {"all", "*", "everyone"}


def _read_events() -> list[dict]:
    if not EVENTS.exists():
        return []
    out = []
    for line in EVENTS.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        try:
            out.append(json.loads(line))
        except json.JSONDecodeError:
            pass
    return out


def _messages() -> dict[str, dict]:
    m: dict[str, dict] = {}
    for ev in _read_events():
        if ev.get("type") == "dm":
            mid = ev.get("mid")
            m[mid] = {
                "mid": mid,
                "from": ev.get("from") or ev.get("agent") or "",
                "to": ev.get("to") or "",
                "text": ev.get("text") or "",
                "ticket": ev.get("ticket") or "",
                "re": ev.get("re") or "",
                "ts": ev.get("ts") or "",
                "reads": [],
            }
        elif ev.get("type") == "read":
            mm = m.get(ev.get("mid"))
            if mm and ev.get("agent") and ev["agent"] not in mm["reads"]:
                mm["reads"].append(ev["agent"])
    return m


def _thread_root(mid: str, mm: dict[str, dict]) -> str:
    cur, c, seen = mid, mm.get(mid), {mid}
    while c and c.get("re") and mm.get(c["re"]):
        if c["re"] in seen:
            break
        seen.add(c["re"])
        cur = c["re"]
        c = mm.get(cur)
    return cur


def _participants(mid: str, mm: dict[str, dict]) -> set[str]:
    root = _thread_root(mid, mm)
    s: set[str] = set()
    for m in mm.values():
        if _thread_root(m["mid"], mm) != root:
            continue
        if m["from"]:
            s.add(m["from"])
        if m["to"] and m["to"] not in _BROADCAST:
            s.add(m["to"])
    return s


def inbox(agent: str, unread_only: bool = True) -> list[dict]:
    mm = _messages()
    out = []
    for m in mm.values():
        if m["from"] == agent:
            continue
        addressed = (
            m["to"] == agent
            or m["to"] in _BROADCAST
            or agent in _participants(m["mid"], mm)
        )
        if not addressed:
            continue
        if unread_only and agent in m["reads"]:
            continue
        out.append(m)
    return sorted(out, key=lambda x: x["ts"])


def known_agents() -> list[str]:
    s: set[str] = set()
    for ev in _read_events():
        for k in ("agent", "from", "to"):
            v = ev.get(k)
            if v and v not in _BROADCAST:
                s.add(v)
    return sorted(s)


# ── tk store WRITE (gated; faithful port of lib.js withLock/nextMid/append) ─────
def _next_mid() -> str:
    mx = 0
    for ev in _read_events():
        if ev.get("type") == "dm":
            try:
                n = int(str(ev.get("mid", "")).replace("M-", ""))
            except ValueError:
                n = 0
            mx = max(mx, n)
    return "M-" + str(mx + 1).zfill(5)


def _with_lock(fn):
    EVENTS.parent.mkdir(parents=True, exist_ok=True)
    deadline = time.time() + 5.0
    while True:
        try:
            fd = os.open(str(LOCK), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        except FileExistsError:
            try:
                if time.time() - LOCK.stat().st_mtime > 10:
                    LOCK.unlink()
                    continue
            except FileNotFoundError:
                continue
            if time.time() > deadline:
                raise TimeoutError("ticket store lock timeout")
            time.sleep(0.02)
            continue
        try:
            return fn()
        finally:
            os.close(fd)
            try:
                LOCK.unlink()
            except FileNotFoundError:
                pass


def send_dm(to: str, text: str, ticket: str = "") -> dict:
    """Append a real `dm` event to the tk log (same shape lib.js writes)."""
    def _do():
        ev = {
            "ts": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "type": "dm",
            "mid": _next_mid(),
            "from": BRIDGE_FROM,
            "to": to,
            "text": text,
        }
        if ticket:
            ev["ticket"] = ticket
        with open(EVENTS, "a", encoding="utf-8") as f:
            f.write(json.dumps(ev) + "\n")
        return ev

    return _with_lock(_do)


# ── command dispatch (message text -> reply text) ───────────────────────────────
def handle(text: str) -> str:
    cmd = (text or "").strip()
    low = cmd.lower()
    if low in ("", "help"):
        w = "on" if ALLOW_WRITE else "off"
        return (
            "tk-bridge commands:  'agents'  |  'inbox <agent>' (append 'all')  |  "
            f"'dm <agent> <text>' (write is {w})"
        )
    if low == "agents":
        agents = known_agents()
        return f"{len(agents)} known agents:\n" + "\n".join(f"  • {a}" for a in agents)
    if low.startswith("inbox"):
        parts = cmd.split()
        if len(parts) < 2:
            return "usage: inbox <agent> [all]"
        agent = parts[1]
        unread_only = not (len(parts) > 2 and parts[2].lower() == "all")
        msgs = inbox(agent, unread_only=unread_only)
        scope = "unread" if unread_only else "all"
        if not msgs:
            return f"inbox({agent}) — 0 {scope} messages"
        lines = [f"inbox({agent}) — {len(msgs)} {scope} message(s):"]
        for m in msgs[-15:]:
            tkt = f" [{m['ticket']}]" if m["ticket"] else ""
            lines.append(f"  {m['mid']} {m['from']}→{m['to']}{tkt}: {m['text'][:90]}")
        return "\n".join(lines)
    if low.startswith("dm"):
        if not ALLOW_WRITE:
            return (
                "dm write is DISABLED. This bridge is read-only by default; a real "
                "cross-agent DM is a shared-log write. Start with TK_BRIDGE_ALLOW_WRITE=1 to enable."
            )
        parts = cmd.split(maxsplit=2)
        if len(parts) < 3:
            return "usage: dm <agent> <text>"
        to, body = parts[1], parts[2]
        ev = send_dm(to, body)
        return f"sent {ev['mid']} {ev['from']}→{ev['to']}: {ev['text'][:80]}"
    return f"unknown command {cmd!r} — try 'help'"


class TkBridgeExecutor(AgentExecutor):
    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        try:
            user_text = context.get_user_input() or ""
        except Exception:
            user_text = ""
        reply = Message(
            message_id=str(uuid.uuid4()),
            role=Role.ROLE_AGENT,
            parts=[Part(text=handle(user_text))],
        )
        await event_queue.enqueue_event(reply)

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        raise NotImplementedError("tk-bridge is read-only; nothing to cancel.")


def build_agent_card() -> AgentCard:
    return AgentCard(
        name="TK Bridge Agent",
        description="Read-only A2A bridge over Steve's tk cross-agent DM log. "
        "Query agent inboxes and the known-agent roster over the A2A protocol.",
        version="0.1.0",
        supported_interfaces=[
            AgentInterface(
                url=f"http://{HOST}:{PORT}/",
                protocol_binding="JSONRPC",
                protocol_version="1.0",
            )
        ],
        capabilities=AgentCapabilities(streaming=False, push_notifications=False),
        default_input_modes=["text/plain"],
        default_output_modes=["text/plain"],
        skills=[
            AgentSkill(
                id="inbox",
                name="Read agent inbox",
                description="Return an agent's unread (or all) DMs from the tk log.",
                tags=["tk", "dm", "inbox", "read-only"],
                examples=["inbox claude-mail", "inbox showroom-builder all"],
            ),
            AgentSkill(
                id="agents",
                name="List known agents",
                description="List every agent identity that has acted or been DM'd.",
                tags=["tk", "roster", "read-only"],
                examples=["agents"],
            ),
            AgentSkill(
                id="dm",
                name="Send a DM (gated)",
                description="Append a real DM to the tk log. Write is OFF unless the "
                "server runs with TK_BRIDGE_ALLOW_WRITE=1.",
                tags=["tk", "dm", "write", "gated"],
                examples=["dm vp-operations check the pg lock canary"],
            ),
        ],
    )


def build_app() -> Starlette:
    card = build_agent_card()
    handler = DefaultRequestHandler(
        agent_executor=TkBridgeExecutor(),
        task_store=InMemoryTaskStore(),
        agent_card=card,
    )
    return Starlette(
        routes=[*create_agent_card_routes(card), *create_jsonrpc_routes(handler, "/")]
    )


if __name__ == "__main__":
    print(f"A2A TK Bridge on http://{HOST}:{PORT}  (read-only over {EVENTS})")
    uvicorn.run(build_app(), host=HOST, port=PORT, log_level="warning")