52 lines
2.2 KiB
Python
52 lines
2.2 KiB
Python
"""Which agent made this MCP call, when one app runs several sessions (Codex and Codex2) on one shared Harbor server.
|
|
|
|
An app starts one Harbor MCP server and shares it between its sessions, so the server cannot tell them apart. The
|
|
app's PreToolUse hook can: it runs before every tool call with the session's working folder. So the hook leaves a
|
|
ticket (app, tool, arguments, agent) and the server takes the matching ticket when the call arrives - oldest first,
|
|
within two minutes. Without a ticket (hooks not installed or not trusted) the server falls back to its own identity."""
|
|
import json
|
|
import time
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS session_tickets(id INTEGER PRIMARY KEY, host TEXT NOT NULL, tool TEXT NOT NULL,
|
|
args TEXT NOT NULL, agent TEXT NOT NULL, at REAL NOT NULL, used INTEGER NOT NULL DEFAULT 0);
|
|
CREATE INDEX IF NOT EXISTS session_tickets_open ON session_tickets(host, tool, used, at);
|
|
"""
|
|
MAX_AGE = 120.0
|
|
|
|
|
|
def args_key(args):
|
|
return json.dumps(args or {}, sort_keys=True, separators=(",", ":"), default=str)
|
|
|
|
|
|
def leave(harbor, host, tool, args, agent, now=None):
|
|
now = time.time() if now is None else now
|
|
|
|
def do(c):
|
|
c.execute("DELETE FROM session_tickets WHERE at < ?", (now - 10 * MAX_AGE,))
|
|
c.execute("INSERT INTO session_tickets(host, tool, args, agent, at) VALUES (?,?,?,?,?)",
|
|
(host or "", tool, args_key(args), agent, now))
|
|
ensure(harbor)
|
|
harbor._write(do)
|
|
|
|
|
|
def take(harbor, host, tool, args, now=None):
|
|
"""The agent whose hook announced this exact call, or None."""
|
|
now = time.time() if now is None else now
|
|
ensure(harbor)
|
|
|
|
def do(c):
|
|
row = c.execute("SELECT id, agent FROM session_tickets WHERE host=? AND tool=? AND args=? AND used=0 AND at>=? "
|
|
"ORDER BY at, id LIMIT 1", (host or "", tool, args_key(args), now - MAX_AGE)).fetchone()
|
|
if not row:
|
|
return None
|
|
c.execute("UPDATE session_tickets SET used=1 WHERE id=?", (row[0],))
|
|
return row[1]
|
|
return harbor._write(do)
|
|
|
|
|
|
def ensure(harbor):
|
|
if not getattr(harbor, "_session_tickets_ready", False):
|
|
harbor.conn.executescript(SCHEMA)
|
|
harbor._session_tickets_ready = True
|