Implement durable OAuth MCP bridge with local Codex worker and deployment guides

tests/conftest.py

@@ -0,0 +1,23 @@
+import pytest
+from codex_chatgpt_mcp_proxy.config import ServerConfig
+from codex_chatgpt_mcp_proxy.store import Store
+
+
+@pytest.fixture
+def config(tmp_path):
+    return ServerConfig(
+        database=str(tmp_path / "server.sqlite"),
+        mode="development",
+        worker_token="w" * 40,
+        dev_mcp_token="m" * 40,
+    )
+
+
+@pytest.fixture
+def store(config):
+    return Store(config)
+
+
+def claim(store):
+    store.register("home", "1", "test", ["demo"])
+    return store.claim("home")

tests/test_auth.py

@@ -0,0 +1,56 @@
+import time
+import jwt
+import pytest
+from cryptography.hazmat.primitives.asymmetric import rsa
+from codex_chatgpt_mcp_proxy.auth import OAuthVerifier, AuthError
+from codex_chatgpt_mcp_proxy.config import ServerConfig
+
+
+async def test_oauth_signature_issuer_audience_owner_scope_and_expiry():
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    config = ServerConfig(
+        mode="production",
+        public_url="https://bridge.example.com",
+        worker_token="w" * 40,
+        issuer="https://identity.example.com/",
+        jwks_url="https://identity.example.com/jwks",
+        audience="https://bridge.example.com",
+        owner_subject="owner",
+    )
+    config.validate()
+    verifier = OAuthVerifier(config)
+    verifier.keys = {"test": key.public_key()}
+    verifier.expires = time.time() + 300
+    claims = {
+        "iss": config.issuer,
+        "aud": config.audience,
+        "sub": "owner",
+        "scope": "bridge:access",
+        "iat": int(time.time()),
+        "exp": int(time.time()) + 60,
+    }
+
+    def token(c):
+        return jwt.encode(c, key, algorithm="RS256", headers={"kid": "test"})
+
+    await verifier.verify(token(claims))
+    for changed in (
+        {"scope": "other"},
+        {"sub": "intruder"},
+        {"aud": "different"},
+        {"iss": "different"},
+        {"exp": int(time.time()) - 60},
+    ):
+        with pytest.raises(AuthError):
+            await verifier.verify(token({**claims, **changed}))
+    with pytest.raises(AuthError):
+        await verifier.verify("garbage")
+    with pytest.raises(ValueError):
+        ServerConfig(worker_token="w" * 40).validate()
+    with pytest.raises(ValueError):
+        ServerConfig(
+            mode="development",
+            public_url="https://bridge.example.com",
+            worker_token="w" * 40,
+            dev_mcp_token="m" * 40,
+        ).validate()

tests/test_server.py

@@ -0,0 +1,102 @@
+from fastapi.testclient import TestClient
+from mcp.client.streamable_http import streamable_http_client
+from codex_chatgpt_mcp_proxy.server import create_app
+
+
+def test_api_auth_host_origin_and_isolation(config):
+    with TestClient(create_app(config), base_url=config.public_url) as c:
+        assert c.get("/healthz").json() == {"ok": True}
+        assert c.get("/healthz", headers={"host": "evil.example"}).status_code == 403
+        assert c.get("/healthz", headers={"origin": "https://evil.example"}).status_code == 403
+        for token in (None, "invalid", config.dev_mcp_token):
+            h = {"Authorization": "Bearer " + token} if token else {}
+            assert c.post("/api/worker/claim", headers=h).status_code == 401
+        assert (
+            c.post("/mcp", headers={"Authorization": "Bearer " + config.worker_token}).status_code
+            == 401
+        )
+        h = {"Authorization": "Bearer " + config.worker_token}
+        assert (
+            c.post(
+                "/api/worker/register",
+                headers=h,
+                json={
+                    "worker_id": "home",
+                    "protocol": "1",
+                    "version": "test",
+                    "workspaces": ["demo"],
+                },
+            ).status_code
+            == 200
+        )
+        assert c.post("/api/worker/claim", headers=h).json() == {"task": None}
+        assert (
+            c.post(
+                "/api/worker/register",
+                headers=h,
+                json={
+                    "worker_id": "home",
+                    "protocol": "1",
+                    "version": "test",
+                    "workspaces": ["/tmp"],
+                },
+            ).status_code
+            == 400
+        )
+        assert c.post("/api/worker/claim", headers=h, content="x" * 140000).status_code == 413
+
+
+async def test_real_mcp_transport_initialize_discover_and_validation(config):
+    import httpx2
+
+    app = create_app(config)
+    async with app.router.lifespan_context(app):
+        async with httpx2.AsyncClient(
+            transport=httpx2.ASGITransport(app=app),
+            headers={"Authorization": "Bearer " + config.dev_mcp_token},
+        ) as http:
+            async with streamable_http_client(config.public_url + "/mcp", http_client=http) as (
+                read,
+                write,
+            ):
+                from mcp import ClientSession
+
+                async with ClientSession(read, write) as session:
+                    initialization = await session.initialize()
+                    assert initialization.server_info.name == "codex-chatgpt-mcp-proxy"
+                    tools = await session.list_tools()
+                    names = {t.name for t in tools.tools}
+                    assert names == {
+                        "bridge_status",
+                        "submit_task",
+                        "get_task",
+                        "get_task_events",
+                        "continue_task",
+                        "cancel_task",
+                        "list_tasks",
+                    }
+                    invalid = await session.call_tool(
+                        "submit_task",
+                        {
+                            "instruction": "x",
+                            "workspace": "demo",
+                            "idempotency_key": "k",
+                            "timeout_seconds": 0,
+                        },
+                    )
+                    assert invalid.is_error
+                    result = await session.call_tool(
+                        "submit_task",
+                        {"instruction": "hello", "workspace": "demo", "idempotency_key": "k"},
+                    )
+                    assert not result.is_error
+                    assert result.structured_content["state"] == "queued"
+                    repeated = await session.call_tool(
+                        "submit_task",
+                        {"instruction": "hello", "workspace": "demo", "idempotency_key": "k"},
+                    )
+                    assert not repeated.is_error
+                    assert repeated.structured_content["id"] == result.structured_content["id"]
+                    for tool in tools.tools:
+                        assert tool.annotations is not None
+                        assert tool.meta["securitySchemes"][0]["type"] == "oauth2"

tests/test_store.py

@@ -0,0 +1,116 @@
+import concurrent.futures
+import time
+import pytest
+from codex_chatgpt_mcp_proxy.store import Store, Conflict
+from conftest import claim
+
+
+def test_idempotency_and_payload_conflict(store):
+    first = store.submit("hello", "demo", "key")
+    assert store.submit("hello", "demo", "key")["id"] == first["id"]
+    with pytest.raises(Conflict):
+        store.submit("different", "demo", "key")
+    with pytest.raises(ValueError):
+        store.submit("hello", "/tmp", "another")
+    with pytest.raises(ValueError):
+        store.submit("hello", "demo", "another", 4000)
+
+
+def test_concurrent_claim_is_atomic(store):
+    store.submit("hello", "demo", "key")
+    store.register("home", "1", "test", ["demo"])
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+        rows = list(pool.map(lambda _: store.claim("home"), range(8)))
+    assert sum(r is not None for r in rows) == 1
+
+
+def test_renew_expiry_fencing_and_final_retry(store):
+    task = store.submit("hello", "demo", "key")
+    r = claim(store)
+    with pytest.raises(Conflict):
+        store.renew(task["id"], "home", "wrong")
+    store.progress(task["id"], "home", r["attempt"], "running", "thread", "turn")
+    assert store.renew(task["id"], "home", r["attempt"])["state"] == "running"
+    with store.tx() as c:
+        c.execute("UPDATE tasks SET lease_until=?", (time.time() - 1,))
+    assert store.claim("home") is None
+    assert store.get(task["id"])["state"] == "needs_reconciliation"
+    store.finish(task["id"], "home", r["attempt"], "succeeded", "done")
+    store.finish(task["id"], "home", r["attempt"], "succeeded", "done")
+    with pytest.raises(Conflict):
+        store.finish(task["id"], "home", r["attempt"], "failed", "different")
+    with pytest.raises(Conflict):
+        store.progress(task["id"], "home", r["attempt"], "running")
+    assert store.cancel(task["id"])["state"] == "succeeded"
+
+
+def test_continuation_thread_serialization(store):
+    task = store.submit("hello", "demo", "key")
+    r = claim(store)
+    with pytest.raises(Conflict):
+        store.continue_task(task["id"], "continue", "key2")
+    store.progress(task["id"], "home", r["attempt"], "running", "thread", "turn")
+    store.finish(task["id"], "home", r["attempt"], "succeeded", "done")
+    second = store.continue_task(task["id"], "continue", "key2")
+    assert second["thread_id"] == "thread"
+    with pytest.raises(Conflict):
+        store.continue_task(task["id"], "again", "key3")
+    assert claim(store)["id"] == second["id"]
+
+
+def test_events_persistence_and_pagination(store, config):
+    task = store.submit("hello", "demo", "key")
+    r = claim(store)
+    for i in range(4):
+        store.event(task["id"], "home", r["attempt"], str(i), {"i": i})
+    store.event(task["id"], "home", r["attempt"], "0", {"i": 0})
+    with pytest.raises(Conflict):
+        store.event(task["id"], "home", r["attempt"], "0", {"i": 99})
+    restarted = Store(config)
+    page = restarted.events(task["id"], limit=2)
+    assert len(page["items"]) == 2
+    page2 = restarted.events(task["id"], after=page["next_cursor"])
+    assert [r["payload"]["i"] for r in page2["items"]] == [2, 3]
+    assert restarted.get(task["id"])["attempt"] == r["attempt"]
+
+
+def test_cancel_queued_running_and_terminal(store):
+    queued = store.submit("hello", "demo", "key")
+    assert store.cancel(queued["id"])["state"] == "cancelled"
+    task = store.submit("hello", "demo", "key2")
+    r = claim(store)
+    assert store.cancel(task["id"])["state"] == "cancel_requested"
+    store.progress(task["id"], "home", r["attempt"], "running")
+    assert store.get(task["id"])["state"] == "cancel_requested"
+    store.finish(task["id"], "home", r["attempt"], "cancelled")
+    assert store.cancel(task["id"])["state"] == "cancelled"
+
+
+def test_bounds_and_retention(store):
+    task = store.submit("hello", "demo", "key")
+    r = claim(store)
+    with pytest.raises(ValueError):
+        store.event(task["id"], "home", r["attempt"], "large", {"v": "x" * 20000})
+    store.finish(task["id"], "home", r["attempt"], "succeeded", "done")
+    with store.tx() as c:
+        c.execute("UPDATE tasks SET updated=?", (time.time() - 86400 * 31,))
+    assert store.cleanup()["purged_payloads"] == 1
+    assert store.get(task["id"])["result"] is None
+
+
+def test_manual_reconciliation_audit_and_release(store):
+    task = store.submit("hello", "demo", "k")
+    leased = claim(store)
+    store.finish(task["id"], "home", leased["attempt"], "needs_reconciliation", error="uncertain")
+    store.reconcile(task["id"], "home", leased["attempt"], "failed")
+    assert store.get(task["id"])["state"] == "failed"
+    assert [h["state"] for h in store.get(task["id"])["history"]] == [
+        "queued",
+        "leased",
+        "needs_reconciliation",
+        "failed",
+    ]
+    with pytest.raises(Conflict):
+        store.reconcile(task["id"], "home", leased["attempt"], "failed")
+    store.submit("next", "demo", "k2")
+    assert claim(store) is not None

tests/test_worker.py

@@ -0,0 +1,226 @@
+import asyncio
+import httpx
+import pytest
+from codex_chatgpt_mcp_proxy.config import WorkerConfig
+from codex_chatgpt_mcp_proxy.journal import Journal
+from codex_chatgpt_mcp_proxy.server import create_app
+from codex_chatgpt_mcp_proxy.worker import Worker
+from conftest import claim
+
+
+class FakeCodex:
+    starts = 0
+    threads = []
+    delay = 0.05
+    disconnect = False
+    approval_needed = False
+
+    def __init__(self, config, approval):
+        self.approval = approval
+        self.notifications = asyncio.Queue()
+        self.job = None
+
+    async def start(self, cwd):
+        pass
+
+    async def thread(self, cwd, thread_id=None):
+        self.threads.append(thread_id)
+        return thread_id or "fake-thread"
+
+    async def turn(self, thread_id, instruction):
+        type(self).starts += 1
+
+        async def complete():
+            await asyncio.sleep(self.delay)
+            if self.approval_needed:
+                decision = await self.approval(
+                    {
+                        "id": 1,
+                        "method": "item/commandExecution/requestApproval",
+                        "params": {"command": "harmless"},
+                    }
+                )
+                assert decision in ("accept", "decline", "cancel")
+            if self.disconnect:
+                await self.notifications.put({"method": "bridge/disconnected", "params": {}})
+            else:
+                await self.notifications.put(
+                    {
+                        "method": "item/completed",
+                        "params": {
+                            "threadId": thread_id,
+                            "turnId": "fake-turn",
+                            "item": {"type": "agentMessage", "text": "done"},
+                        },
+                    }
+                )
+                await self.notifications.put(
+                    {
+                        "method": "turn/completed",
+                        "params": {
+                            "threadId": thread_id,
+                            "turn": {"id": "fake-turn", "status": "completed"},
+                        },
+                    }
+                )
+
+        self.job = asyncio.create_task(complete())
+        return "fake-turn"
+
+    async def interrupt(self, thread_id, turn_id):
+        if self.job:
+            self.job.cancel()
+        await self.notifications.put(
+            {
+                "method": "turn/completed",
+                "params": {"threadId": thread_id, "turn": {"id": turn_id, "status": "interrupted"}},
+            }
+        )
+
+    async def close(self):
+        if self.job:
+            self.job.cancel()
+            await asyncio.gather(self.job, return_exceptions=True)
+
+
+@pytest.fixture
+def worker_setup(config, tmp_path):
+    demo = tmp_path / "demo"
+    demo.mkdir()
+    app = create_app(config)
+    wc = WorkerConfig(
+        url=config.public_url,
+        token=config.worker_token,
+        state=str(tmp_path / "worker.sqlite"),
+        workspaces={"demo": str(demo)},
+        poll=0.1,
+    )
+    http = httpx.AsyncClient(
+        transport=httpx.ASGITransport(app=app),
+        base_url=config.public_url,
+        headers={"Authorization": "Bearer " + config.worker_token},
+    )
+    for n, v in [
+        ("starts", 0),
+        ("threads", []),
+        ("delay", 0.05),
+        ("disconnect", False),
+        ("approval_needed", False),
+    ]:
+        setattr(FakeCodex, n, v)
+    return app.state.store, Worker(wc, FakeCodex, http)
+
+
+async def run_task(store, worker, timeout=None):
+    row = store.submit("hello", "demo", "key", timeout)
+    task = claim(store)
+    worker.journal.save(task)
+    await worker.execute(task)
+    await worker.flush()
+    return store.get(row["id"])
+
+
+async def test_worker_success_and_continuation(worker_setup):
+    store, worker = worker_setup
+    task = await run_task(store, worker)
+    assert task["state"] == "succeeded" and task["result"] == "done"
+    second = store.continue_task(task["id"], "remember", "key2")
+    await worker.execute(claim(store))
+    await worker.flush()
+    assert store.get(second["id"])["state"] == "succeeded"
+    assert FakeCodex.threads == [None, "fake-thread"]
+    await worker.http.aclose()
+
+
+async def test_worker_timeout_and_cancel(worker_setup):
+    store, worker = worker_setup
+    FakeCodex.delay = 10
+    task = await run_task(store, worker, 1)
+    assert task["state"] == "timed_out"
+    row = store.submit("cancel", "demo", "key2")
+    task = claim(store)
+    store.cancel(row["id"])
+    await worker.execute(task)
+    await worker.flush()
+    assert store.get(row["id"])["state"] == "cancelled"
+    await worker.http.aclose()
+
+
+async def test_codex_disconnect_requires_reconciliation(worker_setup):
+    store, worker = worker_setup
+    FakeCodex.disconnect = True
+    task = await run_task(store, worker)
+    assert task["state"] == "needs_reconciliation"
+    assert claim(store) is None
+    await worker.http.aclose()
+
+
+async def test_network_loss_durable_final_and_restart_no_reexecution(worker_setup):
+    store, worker = worker_setup
+    row = store.submit("hello", "demo", "key")
+    task = claim(store)
+    worker.journal.save(task)
+    original = worker.http
+
+    class Offline:
+        async def post(self, *args, **kwargs):
+            raise httpx.ConnectError("offline")
+
+    worker.http = Offline()
+    await worker.execute(task)
+    with pytest.raises(httpx.ConnectError):
+        await worker.flush()
+    assert worker.journal.pending()
+    worker.http = original
+    await worker.recover()
+    assert store.get(row["id"])["state"] == "succeeded"
+    assert FakeCodex.starts == 1 and worker.journal.tasks() == []
+    await worker.http.aclose()
+
+
+async def test_uncertain_worker_restart_never_replays(worker_setup):
+    store, worker = worker_setup
+    row = store.submit("hello", "demo", "key")
+    task = claim(store)
+    task["turn_start_intent"] = True
+    worker.journal.save(task)
+    await worker.recover()
+    assert store.get(row["id"])["state"] == "needs_reconciliation"
+    assert FakeCodex.starts == 0
+    await worker.http.aclose()
+
+
+async def test_local_approval_only(worker_setup):
+    store, worker = worker_setup
+    FakeCodex.approval_needed = True
+    job = asyncio.create_task(run_task(store, worker))
+    for _ in range(100):
+        approvals = worker.journal.approvals()
+        if approvals:
+            break
+        await asyncio.sleep(0.02)
+    assert approvals
+    restarted = Journal(worker.config.state)
+    restarted.decide(approvals[0]["id"], "decline")
+    assert (await job)["state"] == "succeeded"
+    assert restarted.approvals() == []
+    assert "harmless" not in str(store.events(store.tasks()["items"][0]["id"])["items"])
+    await worker.http.aclose()
+
+
+def test_workspace_symlinks_and_retarget(tmp_path):
+    demo = tmp_path / "demo"
+    demo.mkdir()
+    alias = tmp_path / "alias"
+    alias.symlink_to(demo, target_is_directory=True)
+    config = WorkerConfig(token="w" * 40, workspaces={"demo": str(alias)})
+    with pytest.raises(ValueError):
+        config.validate()
+    config.workspaces = {"demo": str(demo)}
+    config.validate()
+    demo.rmdir()
+    demo.symlink_to(tmp_path, target_is_directory=True)
+    with pytest.raises(ValueError):
+        config.resolve_workspace("demo")
+    with pytest.raises(ValueError):
+        config.resolve_workspace("../escape")