mirror of
https://github.com/mcloh/codex-chatgpt-mcp-proxy.git
synced 2026-11-05 21:34:11 +00:00
Implement durable OAuth MCP bridge with local Codex worker and deployment guides
This commit is contained in:
Marcos Lohmann
376
src/codex_chatgpt_mcp_proxy/store.py
Normal file
376
src/codex_chatgpt_mcp_proxy/store.py
Normal file
@@ -0,0 +1,376 @@
|
||||
"""SQLite queue. BEGIN IMMEDIATE fences claims; uncertain work is never requeued."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
TERMINAL = {"succeeded", "failed", "cancelled", "timed_out"}
|
||||
ACTIVE = {
|
||||
"leased",
|
||||
"running",
|
||||
"awaiting_local_approval",
|
||||
"cancel_requested",
|
||||
"needs_reconciliation",
|
||||
}
|
||||
|
||||
|
||||
class Conflict(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
Path(config.database).parent.mkdir(parents=True, exist_ok=True)
|
||||
with self.db() as c:
|
||||
c.executescript("""
|
||||
PRAGMA journal_mode=WAL;
|
||||
CREATE TABLE IF NOT EXISTS migrations(version INTEGER PRIMARY KEY);
|
||||
INSERT OR IGNORE INTO migrations VALUES(1);
|
||||
CREATE TABLE IF NOT EXISTS tasks(
|
||||
seq INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL,
|
||||
idempotency TEXT UNIQUE NOT NULL, digest TEXT NOT NULL,
|
||||
instruction TEXT NOT NULL, workspace TEXT NOT NULL, timeout INTEGER NOT NULL,
|
||||
previous_id TEXT, state TEXT NOT NULL, worker TEXT, attempt TEXT,
|
||||
lease_until REAL, thread_id TEXT, turn_id TEXT,
|
||||
result TEXT, error TEXT, final_digest TEXT, created REAL NOT NULL, updated REAL NOT NULL);
|
||||
CREATE INDEX IF NOT EXISTS tasks_state ON tasks(state,seq);
|
||||
CREATE TABLE IF NOT EXISTS events(
|
||||
seq INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL,
|
||||
event_id TEXT NOT NULL, payload TEXT NOT NULL, created REAL NOT NULL,
|
||||
UNIQUE(task_id,event_id));
|
||||
CREATE INDEX IF NOT EXISTS event_task ON events(task_id,seq);
|
||||
INSERT OR IGNORE INTO migrations VALUES(2);
|
||||
CREATE TABLE IF NOT EXISTS history(
|
||||
seq INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL,
|
||||
state TEXT NOT NULL, created REAL NOT NULL);
|
||||
CREATE TRIGGER IF NOT EXISTS task_created AFTER INSERT ON tasks BEGIN
|
||||
INSERT INTO history(task_id,state,created) VALUES(NEW.id,NEW.state,NEW.created);
|
||||
END;
|
||||
CREATE TRIGGER IF NOT EXISTS task_transition AFTER UPDATE OF state ON tasks
|
||||
WHEN OLD.state != NEW.state BEGIN
|
||||
INSERT INTO history(task_id,state,created) VALUES(NEW.id,NEW.state,NEW.updated);
|
||||
END;
|
||||
CREATE TABLE IF NOT EXISTS workers(id TEXT PRIMARY KEY, seen REAL NOT NULL,
|
||||
protocol TEXT NOT NULL, version TEXT NOT NULL, workspaces TEXT NOT NULL);
|
||||
""")
|
||||
|
||||
@contextmanager
|
||||
def db(self):
|
||||
c = sqlite3.connect(self.config.database, timeout=10, isolation_level=None)
|
||||
c.row_factory = sqlite3.Row
|
||||
c.execute("PRAGMA busy_timeout=10000")
|
||||
try:
|
||||
yield c
|
||||
finally:
|
||||
c.close()
|
||||
|
||||
@contextmanager
|
||||
def tx(self):
|
||||
with self.db() as c:
|
||||
c.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
yield c
|
||||
c.commit()
|
||||
except BaseException:
|
||||
c.rollback()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def row(c, task_id):
|
||||
row = c.execute("SELECT * FROM tasks WHERE id=?", (task_id,)).fetchone()
|
||||
if not row:
|
||||
raise KeyError("Task not found")
|
||||
result = dict(row)
|
||||
result["history"] = [
|
||||
dict(h)
|
||||
for h in c.execute(
|
||||
"SELECT seq,state,created FROM history WHERE task_id=? ORDER BY seq", (task_id,)
|
||||
)
|
||||
]
|
||||
return result
|
||||
|
||||
def get(self, task_id):
|
||||
with self.db() as c:
|
||||
return self.row(c, task_id)
|
||||
|
||||
def submit(
|
||||
self, instruction, workspace, idempotency_key, timeout_seconds=None, previous_id=None
|
||||
):
|
||||
if workspace not in self.config.workspaces:
|
||||
raise ValueError("Unauthorized workspace")
|
||||
if not instruction.strip() or len(instruction) > 32000:
|
||||
raise ValueError("Instruction must contain 1..32000 characters")
|
||||
if not 1 <= len(idempotency_key) <= 128:
|
||||
raise ValueError("Invalid idempotency key")
|
||||
timeout = 600 if timeout_seconds is None else timeout_seconds
|
||||
if not 1 <= timeout <= self.config.max_timeout:
|
||||
raise ValueError("Timeout outside server limit")
|
||||
digest = hashlib.sha256(
|
||||
json.dumps([instruction, workspace, timeout, previous_id]).encode()
|
||||
).hexdigest()
|
||||
now = time.time()
|
||||
with self.tx() as c:
|
||||
existing = c.execute(
|
||||
"SELECT * FROM tasks WHERE idempotency=?", (idempotency_key,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
if existing["digest"] != digest:
|
||||
raise Conflict("Idempotency key reused with different payload")
|
||||
return self.row(c, existing["id"])
|
||||
thread = None
|
||||
if previous_id:
|
||||
parent = self.row(c, previous_id)
|
||||
if parent["state"] not in TERMINAL or not parent["thread_id"]:
|
||||
raise Conflict("Previous task must be terminal with a persisted thread")
|
||||
if parent["workspace"] != workspace:
|
||||
raise Conflict("Continuation cannot change workspace")
|
||||
thread = parent["thread_id"]
|
||||
if c.execute(
|
||||
"SELECT 1 FROM tasks WHERE thread_id=? AND state NOT IN ('succeeded','failed','cancelled','timed_out')",
|
||||
(thread,),
|
||||
).fetchone():
|
||||
raise Conflict("Thread already has an active or queued continuation")
|
||||
if c.execute("SELECT count(*) FROM tasks").fetchone()[0] >= self.config.max_tasks:
|
||||
raise Conflict("Task retention limit reached; run maintenance")
|
||||
task_id = str(uuid.uuid4())
|
||||
c.execute(
|
||||
"""INSERT INTO tasks(id,idempotency,digest,instruction,workspace,timeout,
|
||||
previous_id,state,thread_id,created,updated) VALUES(?,?,?,?,?,?,?,'queued',?,?,?)""",
|
||||
(
|
||||
task_id,
|
||||
idempotency_key,
|
||||
digest,
|
||||
instruction,
|
||||
workspace,
|
||||
timeout,
|
||||
previous_id,
|
||||
thread,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
return self.row(c, task_id)
|
||||
|
||||
def continue_task(self, task_id, instruction, idempotency_key, timeout_seconds=None):
|
||||
parent = self.get(task_id)
|
||||
return self.submit(
|
||||
instruction, parent["workspace"], idempotency_key, timeout_seconds, task_id
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def expire(c, now):
|
||||
c.execute(
|
||||
"""UPDATE tasks SET state='needs_reconciliation',updated=?
|
||||
WHERE state IN ('leased','running','awaiting_local_approval','cancel_requested') AND lease_until<?""",
|
||||
(now, now),
|
||||
)
|
||||
|
||||
def register(self, worker_id, protocol, version, workspaces):
|
||||
if worker_id != self.config.worker_id or protocol != "1":
|
||||
raise Conflict("Worker identity or protocol mismatch")
|
||||
if not set(workspaces).issubset(self.config.workspaces):
|
||||
raise ValueError("Worker advertised unknown workspace")
|
||||
with self.tx() as c:
|
||||
c.execute(
|
||||
"INSERT OR REPLACE INTO workers VALUES(?,?,?,?,?)",
|
||||
(worker_id, time.time(), protocol, version, json.dumps(workspaces)),
|
||||
)
|
||||
return {"protocol": "1", "lease_seconds": self.config.lease_seconds}
|
||||
|
||||
def heartbeat(self, worker_id):
|
||||
with self.tx() as c:
|
||||
self.expire(c, time.time())
|
||||
c.execute("UPDATE workers SET seen=? WHERE id=?", (time.time(), worker_id))
|
||||
return {"ok": True}
|
||||
|
||||
def claim(self, worker_id):
|
||||
with self.tx() as c:
|
||||
now = time.time()
|
||||
self.expire(c, now)
|
||||
worker = c.execute("SELECT * FROM workers WHERE id=?", (worker_id,)).fetchone()
|
||||
if not worker or now - worker["seen"] > self.config.lease_seconds:
|
||||
raise Conflict("Worker must register and heartbeat")
|
||||
if c.execute(
|
||||
"SELECT 1 FROM tasks WHERE worker=? AND state NOT IN ('succeeded','failed','cancelled','timed_out')",
|
||||
(worker_id,),
|
||||
).fetchone():
|
||||
return None
|
||||
allowed = json.loads(worker["workspaces"])
|
||||
for r in c.execute("SELECT * FROM tasks WHERE state='queued' ORDER BY seq").fetchall():
|
||||
if r["workspace"] not in allowed:
|
||||
continue
|
||||
c.execute(
|
||||
"UPDATE tasks SET state='leased',worker=?,attempt=?,lease_until=?,updated=? WHERE id=?",
|
||||
(worker_id, str(uuid.uuid4()), now + self.config.lease_seconds, now, r["id"]),
|
||||
)
|
||||
return self.row(c, r["id"])
|
||||
return None
|
||||
|
||||
def fenced(self, c, task_id, worker, attempt):
|
||||
row = self.row(c, task_id)
|
||||
if row["worker"] != worker or row["attempt"] != attempt:
|
||||
raise Conflict("Attempt ownership mismatch")
|
||||
return row
|
||||
|
||||
def renew(self, task_id, worker, attempt):
|
||||
with self.tx() as c:
|
||||
r = self.fenced(c, task_id, worker, attempt)
|
||||
if r["state"] not in TERMINAL:
|
||||
# Reconciliation remains explicit until a journaled final result arrives.
|
||||
c.execute(
|
||||
"UPDATE tasks SET lease_until=?,updated=? WHERE id=?",
|
||||
(time.time() + self.config.lease_seconds, time.time(), task_id),
|
||||
)
|
||||
return {"state": r["state"]}
|
||||
|
||||
def progress(self, task_id, worker, attempt, state, thread_id=None, turn_id=None):
|
||||
if state not in ("running", "awaiting_local_approval"):
|
||||
raise ValueError("Invalid progress state")
|
||||
with self.tx() as c:
|
||||
r = self.fenced(c, task_id, worker, attempt)
|
||||
if r["state"] in TERMINAL:
|
||||
raise Conflict("Task already terminal")
|
||||
if r["thread_id"] and thread_id and r["thread_id"] != thread_id:
|
||||
raise Conflict("Thread identity cannot change")
|
||||
if r["turn_id"] and turn_id and r["turn_id"] != turn_id:
|
||||
raise Conflict("Turn identity cannot change")
|
||||
if r["state"] in ("cancel_requested", "needs_reconciliation"):
|
||||
state = r["state"]
|
||||
c.execute(
|
||||
"UPDATE tasks SET state=?,thread_id=coalesce(?,thread_id),turn_id=coalesce(?,turn_id),updated=? WHERE id=?",
|
||||
(state, thread_id, turn_id, time.time(), task_id),
|
||||
)
|
||||
return {"ok": True}
|
||||
|
||||
def event(self, task_id, worker, attempt, event_id, payload):
|
||||
encoded = json.dumps(payload, ensure_ascii=False)
|
||||
if len(encoded.encode()) > self.config.max_event_bytes:
|
||||
raise ValueError("Event exceeds byte limit")
|
||||
with self.tx() as c:
|
||||
self.fenced(c, task_id, worker, attempt)
|
||||
old = c.execute(
|
||||
"SELECT payload FROM events WHERE task_id=? AND event_id=?", (task_id, event_id)
|
||||
).fetchone()
|
||||
if old:
|
||||
if old[0] != encoded:
|
||||
raise Conflict("Event ID reused with different payload")
|
||||
return {"ok": True}
|
||||
if (
|
||||
c.execute("SELECT count(*) FROM events WHERE task_id=?", (task_id,)).fetchone()[0]
|
||||
>= self.config.max_events
|
||||
):
|
||||
return {"ok": True, "dropped": "event_limit"}
|
||||
c.execute(
|
||||
"INSERT INTO events(task_id,event_id,payload,created) VALUES(?,?,?,?)",
|
||||
(task_id, event_id, encoded, time.time()),
|
||||
)
|
||||
return {"ok": True}
|
||||
|
||||
def finish(self, task_id, worker, attempt, state, result="", error=None):
|
||||
if state not in TERMINAL | {"needs_reconciliation"}:
|
||||
raise ValueError("Invalid final state")
|
||||
if len(result) > self.config.max_result_chars or (error and len(error) > 4000):
|
||||
raise ValueError("Result exceeds limit")
|
||||
digest = hashlib.sha256(json.dumps([state, result, error]).encode()).hexdigest()
|
||||
with self.tx() as c:
|
||||
r = self.fenced(c, task_id, worker, attempt)
|
||||
if r["final_digest"]:
|
||||
if r["final_digest"] != digest:
|
||||
raise Conflict("Conflicting final result")
|
||||
return {"ok": True}
|
||||
if r["state"] in TERMINAL:
|
||||
raise Conflict("Task terminal")
|
||||
c.execute(
|
||||
"UPDATE tasks SET state=?,result=?,error=?,final_digest=?,updated=? WHERE id=?",
|
||||
(state, result, error, digest, time.time(), task_id),
|
||||
)
|
||||
return {"ok": True}
|
||||
|
||||
def cancel(self, task_id):
|
||||
with self.tx() as c:
|
||||
r = self.row(c, task_id)
|
||||
if r["state"] not in TERMINAL:
|
||||
state = "cancelled" if r["state"] == "queued" else "cancel_requested"
|
||||
if r["state"] == "needs_reconciliation":
|
||||
raise Conflict("Reconciliation required; cannot assume execution stopped")
|
||||
c.execute(
|
||||
"UPDATE tasks SET state=?,updated=? WHERE id=?", (state, time.time(), task_id)
|
||||
)
|
||||
return self.row(c, task_id)
|
||||
|
||||
def tasks(self, after=0, limit=50, state=None):
|
||||
if limit < 1 or limit > 100 or after < 0:
|
||||
raise ValueError("Invalid pagination")
|
||||
with self.db() as c:
|
||||
rows = c.execute(
|
||||
"SELECT * FROM tasks WHERE seq>? AND (? IS NULL OR state=?) ORDER BY seq LIMIT ?",
|
||||
(after, state, state, limit),
|
||||
).fetchall()
|
||||
return {
|
||||
"items": [dict(r) for r in rows],
|
||||
"next_cursor": rows[-1]["seq"] if rows else after,
|
||||
}
|
||||
|
||||
def events(self, task_id, after=0, limit=50):
|
||||
if limit < 1 or limit > 100 or after < 0:
|
||||
raise ValueError("Invalid pagination")
|
||||
with self.db() as c:
|
||||
self.row(c, task_id)
|
||||
rows = c.execute(
|
||||
"SELECT * FROM events WHERE task_id=? AND seq>? ORDER BY seq LIMIT ?",
|
||||
(task_id, after, limit),
|
||||
).fetchall()
|
||||
return {
|
||||
"items": [{**dict(r), "payload": json.loads(r["payload"])} for r in rows],
|
||||
"next_cursor": rows[-1]["seq"] if rows else after,
|
||||
}
|
||||
|
||||
def status(self):
|
||||
with self.tx() as c:
|
||||
self.expire(c, time.time())
|
||||
worker = c.execute(
|
||||
"SELECT seen,version FROM workers WHERE id=?", (self.config.worker_id,)
|
||||
).fetchone()
|
||||
return {
|
||||
"protocol": "1",
|
||||
"worker_online": bool(
|
||||
worker and time.time() - worker["seen"] < self.config.lease_seconds
|
||||
),
|
||||
"worker_version": worker["version"] if worker else None,
|
||||
"counts": dict(
|
||||
c.execute("SELECT state,count(*) FROM tasks GROUP BY state").fetchall()
|
||||
),
|
||||
"workspaces": list(self.config.workspaces),
|
||||
}
|
||||
|
||||
def reconcile(self, task_id, worker, attempt, state):
|
||||
if state not in TERMINAL:
|
||||
raise ValueError("Reconciliation must resolve to a terminal state")
|
||||
with self.tx() as c:
|
||||
r = self.fenced(c, task_id, worker, attempt)
|
||||
if r["state"] != "needs_reconciliation":
|
||||
raise Conflict("Only uncertain tasks can be reconciled")
|
||||
c.execute(
|
||||
"UPDATE tasks SET state=?,final_digest=NULL,error=?,updated=? WHERE id=?",
|
||||
(state, "Manually reconciled locally; no execution replayed", time.time(), task_id),
|
||||
)
|
||||
return {"ok": True}
|
||||
|
||||
def cleanup(self):
|
||||
with self.tx() as c:
|
||||
cutoff = time.time() - self.config.retention_days * 86400
|
||||
# Keep idempotency tombstones and thread lineage; purge instruction/results/events.
|
||||
c.execute(
|
||||
"DELETE FROM events WHERE task_id IN (SELECT id FROM tasks WHERE updated<? AND state IN ('succeeded','failed','cancelled','timed_out'))",
|
||||
(cutoff,),
|
||||
)
|
||||
result = c.execute(
|
||||
"UPDATE tasks SET instruction='[retained metadata only]',result=NULL,error=NULL WHERE updated<? AND state IN ('succeeded','failed','cancelled','timed_out') AND instruction!='[retained metadata only]'",
|
||||
(cutoff,),
|
||||
)
|
||||
return {"purged_payloads": result.rowcount}
|
||||
Reference in New Issue
Block a user