feat(workbench): persist operational workflows and events

tests/unit/test_operational_worker.py

@@ -0,0 +1,141 @@
+import json
+from uuid import uuid4
+
+import httpx
+import pytest
+
+from scripts.run_operational_worker import CredentialProvider, OperationalWorker, WorkerSettings
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"
+
+
+@pytest.mark.anyio
+async def test_persisted_queue_retries_with_backend_execution_identity():
+    task_id, run_id, execution_id = str(uuid4()), str(uuid4()), str(uuid4())
+    submitted = []
+    def handler(request):
+        assert request.headers["Authorization"] == "Bearer TEST_ONLY"
+        assert not request.url.path.endswith("schedules/dispatch-due")
+        if request.url.path.endswith("leases/recover"):
+            return httpx.Response(200, json={"items": []})
+        if request.method == "GET":
+            if request.url.params.get("status") != "queued":
+                return httpx.Response(200, json={"items": []})
+            return httpx.Response(200, json={"items": [{
+                "task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
+                "status": "queued", "revision": 1,
+            }]})
+        body = json.loads(request.content)
+        submitted.append(body)
+        return httpx.Response(202, json={"task_id": task_id, "run_id": run_id, "session_id": "execution-session"})
+    settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        worker = OperationalWorker(settings, client)
+        assert await worker.tick() == 1
+        assert await worker.tick() == 1
+    assert submitted[0] == submitted[1]
+    assert submitted[0] == {
+        "task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
+    }
+    assert "credential" not in submitted[0]
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("run_status", ["claimed", "running", "failed"])
+async def test_only_queued_runs_are_submitted(run_status):
+    def handler(request):
+        assert not request.url.path.endswith("schedules/dispatch-due")
+        if request.method == "GET":
+            requested = request.url.params.get("status")
+            if requested == run_status and run_status in {"claimed", "running"}:
+                return httpx.Response(200, json={"items": [{"run_id": str(uuid4()), "status": run_status}]})
+            return httpx.Response(200, json={"items": []})
+        assert request.url.host == "server.test"
+        return httpx.Response(200, json={"items": []})
+    settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        assert await OperationalWorker(settings, client).tick() == 0
+
+
+@pytest.mark.anyio
+async def test_bounded_queue_continues_after_one_agent_handoff_fails():
+    runs = [{"task_id": str(uuid4()), "run_id": str(uuid4()),
+             "execution_request_id": str(uuid4()), "status": "queued",
+             "revision": 1} for _ in range(2)]
+    limits, submitted, failures = [], [], []
+
+    def handler(request):
+        if request.method == "GET":
+            if request.url.params.get("status") != "queued":
+                return httpx.Response(200, json={"items": []})
+            limits.append(int(request.url.params["limit"]))
+            return httpx.Response(200, json={"items": runs})
+        if request.url.host == "server.test":
+            if request.url.path.endswith("leases/recover"):
+                return httpx.Response(200, json={"items": []})
+            failures.append(json.loads(request.content))
+            return httpx.Response(200, json={})
+        body = json.loads(request.content)
+        submitted.append(body)
+        if body["task_id"] == runs[0]["task_id"]:
+            return httpx.Response(503, json={"detail": "temporary failure"})
+        return httpx.Response(202, json={"task_id": body["task_id"], "run_id": body["run_id"]})
+
+    settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        assert await OperationalWorker(settings, client).tick() == 1
+    assert limits == [settings.max_concurrency]
+    assert [body["task_id"] for body in submitted] == [run["task_id"] for run in runs]
+    assert [body["execution_request_id"] for body in submitted] == [run["execution_request_id"] for run in runs]
+    assert failures[0]["retryable"] is True
+
+
+@pytest.mark.anyio
+async def test_service_credentials_are_cached_and_refreshed_without_entering_task_payload(monkeypatch):
+    clock = [100.0]
+    monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
+    requests = []
+
+    def handler(request):
+        requests.append(request)
+        assert request.url == "https://identity.test/token"
+        assert "authorization" not in request.headers
+        return httpx.Response(200, json={"access_token": f"TEST_ONLY_{len(requests)}", "expires_in": 60})
+
+    settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
+                              token_url="https://identity.test/token", client_id="test-service",
+                              client_secret="TEST_ONLY_SECRET")
+    provider = CredentialProvider(settings)
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        first = await provider.headers(client)
+        clock[0] = 140.0
+        assert await provider.headers(client) == first
+        assert len(requests) == 1
+        clock[0] = 146.0
+        refreshed = await provider.headers(client)
+    assert len(requests) == 2
+    assert first["Authorization"] == "Bearer TEST_ONLY_1"
+    assert refreshed == {"Authorization": "Bearer TEST_ONLY_2", "X-Project-Id": settings.project_id}
+
+
+@pytest.mark.anyio
+async def test_long_lived_oauth_token_is_refreshed_before_product_token_age_limit(monkeypatch):
+    clock, requests = [0.0], []
+    monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
+
+    def handler(request):
+        requests.append(request)
+        return httpx.Response(200, json={"access_token": "TEST_ONLY", "expires_in": 3600})
+
+    settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
+                              token_url="https://identity.test/token", client_id="test-service",
+                              client_secret="TEST_ONLY_SECRET")
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        provider = CredentialProvider(settings)
+        await provider.headers(client)
+        clock[0] = 301.0
+        await provider.headers(client)
+    assert len(requests) == 2