feat(workbench): persist operational workflows and events
This commit is contained in:
@@ -0,0 +1,259 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Hand project-scoped queued investigations to the Agent.
|
||||
|
||||
Credentials are deployment environment values, never task/session payloads.
|
||||
Run with --once for an operator-triggered tick or without it for polling.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import httpx
|
||||
from dotenv import load_dotenv
|
||||
|
||||
|
||||
logger = logging.getLogger("operational-worker")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerSettings:
|
||||
server_url: str
|
||||
agent_url: str
|
||||
project_id: str
|
||||
interval_seconds: float = 30
|
||||
max_concurrency: int = 4
|
||||
bearer: str = ""
|
||||
token_url: str = ""
|
||||
client_id: str = ""
|
||||
client_secret: str = ""
|
||||
token_refresh_seconds: float = 300
|
||||
|
||||
@classmethod
|
||||
def from_environment(cls) -> "WorkerSettings":
|
||||
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
||||
project_id = str(UUID(os.environ["OPERATIONAL_WORKER_PROJECT_ID"]))
|
||||
interval = float(os.environ.get("OPERATIONAL_WORKER_RECONCILE_SECONDS", "30"))
|
||||
if interval < 1:
|
||||
raise ValueError("worker interval must be at least one second")
|
||||
settings = cls(
|
||||
server_url=os.environ.get("OPERATIONAL_WORKER_SERVER_URL", "http://127.0.0.1:8000"),
|
||||
agent_url=os.environ.get("OPERATIONAL_WORKER_AGENT_URL", "http://127.0.0.1:8787"),
|
||||
project_id=project_id,
|
||||
interval_seconds=interval,
|
||||
max_concurrency=int(os.environ.get("OPERATIONAL_WORKER_MAX_CONCURRENCY", "4")),
|
||||
bearer=os.environ.get("OPERATIONAL_WORKER_BEARER", ""),
|
||||
token_url=os.environ.get("OPERATIONAL_WORKER_TOKEN_URL", ""),
|
||||
client_id=os.environ.get("OPERATIONAL_WORKER_CLIENT_ID", ""),
|
||||
client_secret=os.environ.get("OPERATIONAL_WORKER_CLIENT_SECRET", ""),
|
||||
token_refresh_seconds=float(os.environ.get("OPERATIONAL_WORKER_TOKEN_REFRESH_SECONDS", "300")),
|
||||
)
|
||||
if not 30 <= settings.token_refresh_seconds <= 600:
|
||||
raise ValueError("service token refresh interval must be between 30 and 600 seconds")
|
||||
if not 1 <= settings.max_concurrency <= 32:
|
||||
raise ValueError("worker concurrency must be between 1 and 32")
|
||||
if not settings.bearer and not all((settings.token_url, settings.client_id, settings.client_secret)):
|
||||
raise ValueError("configure a product service credential for the operational worker")
|
||||
return settings
|
||||
|
||||
|
||||
class CredentialProvider:
|
||||
def __init__(self, settings: WorkerSettings):
|
||||
self.settings = settings
|
||||
self._token = ""
|
||||
self._expires = 0.0
|
||||
|
||||
async def headers(self, client: httpx.AsyncClient) -> dict[str, str]:
|
||||
if self.settings.bearer:
|
||||
token = self.settings.bearer.removeprefix("Bearer ")
|
||||
else:
|
||||
if time.monotonic() >= self._expires:
|
||||
response = await client.post(self.settings.token_url, data={
|
||||
"grant_type": "client_credentials", "client_id": self.settings.client_id,
|
||||
"client_secret": self.settings.client_secret,
|
||||
})
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
self._token = payload["access_token"]
|
||||
if not isinstance(self._token, str) or not self._token:
|
||||
raise ValueError("identity provider returned no access token")
|
||||
# Product auth also restricts token age, independently of the
|
||||
# provider's exp. Do not cache a long-lived OAuth token for hours.
|
||||
lifetime = min(float(payload.get("expires_in", 60)), self.settings.token_refresh_seconds)
|
||||
self._expires = time.monotonic() + max(1, lifetime - 15)
|
||||
token = self._token
|
||||
return {"Authorization": f"Bearer {token}", "X-Project-Id": self.settings.project_id}
|
||||
|
||||
|
||||
class OperationalWorker:
|
||||
def __init__(self, settings: WorkerSettings, client: httpx.AsyncClient):
|
||||
self.settings = settings
|
||||
self.client = client
|
||||
self.credentials = CredentialProvider(settings)
|
||||
|
||||
async def tick(self) -> int:
|
||||
headers = await self.credentials.headers(self.client)
|
||||
server = self.settings.server_url.rstrip("/")
|
||||
recovery = await self.client.post(
|
||||
f"{server}/api/v1/operational-tasks/leases/recover", headers=headers, json={"limit": 100},
|
||||
)
|
||||
recovery.raise_for_status()
|
||||
active_run_ids: set[str] = set()
|
||||
for status_value in ("claimed", "running"):
|
||||
response = await self.client.get(
|
||||
f"{server}/api/v1/operational-task-runs", headers=headers,
|
||||
params={"status": status_value, "limit": self.settings.max_concurrency},
|
||||
)
|
||||
response.raise_for_status()
|
||||
active_run_ids.update(str(item["run_id"]) for item in response.json()["items"])
|
||||
available = max(0, self.settings.max_concurrency - len(active_run_ids))
|
||||
if available == 0:
|
||||
return 0
|
||||
response = await self.client.get(
|
||||
f"{server}/api/v1/operational-task-runs", headers=headers,
|
||||
params={"status": "queued", "limit": available},
|
||||
)
|
||||
response.raise_for_status()
|
||||
queued = response.json()["items"]
|
||||
|
||||
async def dispatch(run: dict) -> bool:
|
||||
if run.get("status") != "queued":
|
||||
return False
|
||||
task_id = str(UUID(run["task_id"]))
|
||||
run_id = str(UUID(run["run_id"]))
|
||||
request_id = str(UUID(run["execution_request_id"]))
|
||||
try:
|
||||
result = await self.client.post(
|
||||
f"{self.settings.agent_url.rstrip('/')}/api/v1/agent/chat/task-executions",
|
||||
headers=headers,
|
||||
json={
|
||||
"task_id": task_id,
|
||||
"run_id": run_id,
|
||||
"execution_request_id": request_id,
|
||||
},
|
||||
)
|
||||
if result.is_success:
|
||||
payload = result.json()
|
||||
if payload.get("task_id") != task_id or payload.get("run_id") != run_id:
|
||||
raise ValueError("Agent returned a different task run")
|
||||
return True
|
||||
await self._record_dispatch_failure(
|
||||
headers, run, retryable=result.status_code >= 500,
|
||||
code=f"AGENT_HTTP_{result.status_code}",
|
||||
)
|
||||
return False
|
||||
except (httpx.HTTPError, ValueError) as exc:
|
||||
await self._record_dispatch_failure(
|
||||
headers, run, retryable=True,
|
||||
code="AGENT_UNAVAILABLE", detail=type(exc).__name__,
|
||||
)
|
||||
return False
|
||||
|
||||
results = await asyncio.gather(*(dispatch(run) for run in queued))
|
||||
return sum(results)
|
||||
|
||||
async def _record_dispatch_failure(
|
||||
self, headers: dict[str, str], run: dict, *, retryable: bool,
|
||||
code: str, detail: str = "Agent rejected task run dispatch",
|
||||
) -> None:
|
||||
response = await self.client.post(
|
||||
f"{self.settings.server_url.rstrip('/')}/api/v1/operational-tasks/"
|
||||
f"{run['task_id']}/runs/{run['run_id']}/dispatch-failure",
|
||||
headers=headers,
|
||||
json={
|
||||
"request_id": str(uuid4()),
|
||||
"expected_revision": run["revision"],
|
||||
"failure_code": code,
|
||||
"failure_detail": detail,
|
||||
"retryable": retryable,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
async def listen(self, wake: asyncio.Event) -> None:
|
||||
cursor = 0
|
||||
while True:
|
||||
try:
|
||||
headers = await self.credentials.headers(self.client)
|
||||
async with self.client.stream(
|
||||
"GET",
|
||||
f"{self.settings.server_url.rstrip('/')}/api/v1/workbench/events",
|
||||
headers={**headers, "Accept": "text/event-stream"},
|
||||
params={"after": cursor},
|
||||
timeout=None,
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
event_name = ""
|
||||
reset_required = False
|
||||
async for line in response.aiter_lines():
|
||||
if line.startswith("id:"):
|
||||
cursor = int(line[3:].strip())
|
||||
elif line.startswith("event:"):
|
||||
event_name = line[6:].strip()
|
||||
elif line.startswith("data:") and event_name == "reset_required":
|
||||
snapshot = await self.client.get(
|
||||
f"{self.settings.server_url.rstrip('/')}/api/v1/workbench/sync-snapshot",
|
||||
headers=headers,
|
||||
)
|
||||
snapshot.raise_for_status()
|
||||
cursor = int(snapshot.json()["event_cursor"])
|
||||
reset_required = True
|
||||
break
|
||||
elif line.startswith("data:") and event_name == "workbench_event":
|
||||
payload = json.loads(line[5:].strip())
|
||||
if payload.get("topic") == "operational.run.queued":
|
||||
wake.set()
|
||||
elif line == "":
|
||||
event_name = ""
|
||||
if not reset_required:
|
||||
await asyncio.sleep(1)
|
||||
except (httpx.HTTPError, ValueError, KeyError):
|
||||
logger.warning("Workbench event stream disconnected; reconciliation remains active")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
async def run(once: bool) -> None:
|
||||
settings = WorkerSettings.from_environment()
|
||||
async with httpx.AsyncClient(timeout=30, follow_redirects=False, trust_env=False) as client:
|
||||
worker = OperationalWorker(settings, client)
|
||||
wake = asyncio.Event()
|
||||
listener = None if once else asyncio.create_task(worker.listen(wake))
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
accepted = await worker.tick()
|
||||
logger.info("Task dispatch tick completed: %d accepted", accepted)
|
||||
except (httpx.HTTPError, ValueError, KeyError):
|
||||
# Do not log request headers, response bodies or provider secrets.
|
||||
logger.error("Task dispatch tick failed; persisted tasks remain available for retry")
|
||||
if once:
|
||||
raise SystemExit(1) from None
|
||||
if once:
|
||||
return
|
||||
try:
|
||||
await asyncio.wait_for(wake.wait(), timeout=settings.interval_seconds)
|
||||
except TimeoutError:
|
||||
pass
|
||||
wake.clear()
|
||||
finally:
|
||||
if listener is not None:
|
||||
listener.cancel()
|
||||
await asyncio.gather(listener, return_exceptions=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--once", action="store_true")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
try:
|
||||
asyncio.run(run(args.once))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
Reference in New Issue
Block a user