369 lines
17 KiB
Python
369 lines
17 KiB
Python
#!/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, replace
|
|
from datetime import datetime
|
|
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
|
|
task_created_from: str | None = None
|
|
project_ids: tuple[str, ...] = ()
|
|
|
|
def project_scopes(self) -> tuple["WorkerSettings", ...]:
|
|
"""Keep credentials, sessions and cursors separate for each explicit project."""
|
|
return tuple(replace(self, project_id=project_id, project_ids=())
|
|
for project_id in self.project_ids or (self.project_id,))
|
|
|
|
@classmethod
|
|
def from_environment(cls) -> "WorkerSettings":
|
|
load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
|
|
configured_projects = os.environ.get("OPERATIONAL_WORKER_PROJECT_IDS", "").strip()
|
|
project_ids = tuple(dict.fromkeys(
|
|
str(UUID(value.strip())) for value in configured_projects.split(",")
|
|
)) if configured_projects else ()
|
|
project_id = project_ids[0] if project_ids else 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")
|
|
token_url = os.environ.get("OPERATIONAL_WORKER_TOKEN_URL", "")
|
|
if os.environ.get("ENVIRONMENT") == "development":
|
|
token_url = os.environ.get("OPERATIONAL_WORKER_DEV_TOKEN_URL") or token_url
|
|
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=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")),
|
|
task_created_from=os.environ.get("OPERATIONAL_WORKER_TASK_CREATED_FROM") or None,
|
|
project_ids=project_ids,
|
|
)
|
|
if settings.task_created_from:
|
|
boundary = datetime.fromisoformat(settings.task_created_from.replace("Z", "+00:00"))
|
|
if boundary.tzinfo is None:
|
|
raise ValueError("worker task creation boundary must include a timezone")
|
|
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)
|
|
self._agent_sessions: dict[str, str] = {}
|
|
self._message_cursor: str | None = None
|
|
|
|
async def tick(self) -> int:
|
|
headers = await self.credentials.headers(self.client)
|
|
server = self.settings.server_url.rstrip("/")
|
|
scope = {"task_type": "event_investigation"}
|
|
if self.settings.task_created_from:
|
|
scope["task_created_from"] = self.settings.task_created_from
|
|
recovery = await self.client.post(
|
|
f"{server}/api/v1/operational-tasks/leases/recover", headers=headers,
|
|
json={"limit": 100, **scope},
|
|
)
|
|
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, **scope},
|
|
)
|
|
response.raise_for_status()
|
|
for item in response.json()["items"]:
|
|
run_id = str(item["run_id"])
|
|
active_run_ids.add(run_id)
|
|
session_id = item.get("agent_session_id")
|
|
if isinstance(session_id, str) and session_id:
|
|
self._agent_sessions[run_id] = session_id
|
|
for run_id, session_id in list(self._agent_sessions.items()):
|
|
if run_id not in active_run_ids:
|
|
self._agent_sessions.pop(run_id, None)
|
|
continue
|
|
# Refresh the Agent vault with the worker's existing OAuth provider.
|
|
# Background runs never depend on an interactive browser refresh.
|
|
renewed = await self.client.post(
|
|
f"{self.settings.agent_url.rstrip('/')}/api/v1/session/{session_id}/credentials",
|
|
headers=headers, json={},
|
|
)
|
|
if renewed.status_code not in (202, 404, 409):
|
|
renewed.raise_for_status()
|
|
available = max(0, self.settings.max_concurrency - len(active_run_ids))
|
|
if available == 0:
|
|
await self._reconcile_messages(headers)
|
|
return 0
|
|
response = await self.client.get(
|
|
f"{server}/api/v1/operational-task-runs", headers=headers,
|
|
params={"status": "queued", "limit": available, **scope},
|
|
)
|
|
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/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")
|
|
session_id = payload.get("session_id")
|
|
if isinstance(session_id, str) and session_id:
|
|
self._agent_sessions[run_id] = session_id
|
|
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))
|
|
await self._reconcile_messages(headers)
|
|
return sum(results)
|
|
|
|
async def _reconcile_messages(self, headers: dict[str, str]) -> None:
|
|
# Reconcile already bound runs even when task dispatch is at capacity.
|
|
# Agent and Server share immutable message/native identities for retries.
|
|
semaphore = asyncio.Semaphore(self.settings.max_concurrency)
|
|
|
|
async def dispatch(message: dict) -> None:
|
|
async with semaphore:
|
|
message_id = str(UUID(message["message_id"]))
|
|
try:
|
|
delivered = await self.client.post(
|
|
f"{self.settings.agent_url.rstrip('/')}/api/v1/task-messages/deliver",
|
|
headers=headers, json={"message_id": message_id},
|
|
)
|
|
if delivered.status_code in {400, 403, 404, 410, 422}:
|
|
failed = await self.client.post(
|
|
f"{self.settings.server_url.rstrip('/')}/api/v1/agent-messages/{message_id}/delivery",
|
|
headers=headers, json={
|
|
"status": "failed",
|
|
"failure_detail": f"Agent rejected message delivery (HTTP {delivered.status_code})",
|
|
},
|
|
)
|
|
if failed.status_code != 409:
|
|
failed.raise_for_status()
|
|
elif not delivered.is_success:
|
|
delivered.raise_for_status()
|
|
except httpx.HTTPError:
|
|
# Avoid message content and credentials in logs. Durable
|
|
# nonterminal messages are available at the next tick.
|
|
logger.warning("Agent message delivery remains pending: %s", message_id)
|
|
|
|
# Bounded pages and a cursor retained between ticks prevent idle chats
|
|
# from occupying the first page forever. Completing a sweep starts the
|
|
# next tick from the beginning so older transient failures are retried.
|
|
for _ in range(4):
|
|
params = {"limit": 50}
|
|
if self._message_cursor:
|
|
params["cursor"] = self._message_cursor
|
|
response = await self.client.get(
|
|
f"{self.settings.server_url.rstrip('/')}/api/v1/agent-messages/pending",
|
|
headers=headers, params=params,
|
|
)
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
await asyncio.gather(*(dispatch(message) for message in payload["items"]))
|
|
self._message_cursor = payload.get("next_cursor")
|
|
if not self._message_cursor:
|
|
break
|
|
|
|
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_project(settings: WorkerSettings, client: httpx.AsyncClient, once: bool) -> None:
|
|
logger.info(
|
|
"Investigation worker scope: project=%s task_created_from=%s concurrency=%d",
|
|
settings.project_id, settings.task_created_from or "all", settings.max_concurrency,
|
|
)
|
|
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: project=%s accepted=%d", settings.project_id, accepted)
|
|
except (httpx.HTTPError, ValueError, KeyError):
|
|
logger.error("Task dispatch tick failed: project=%s; persisted tasks remain available for retry", settings.project_id)
|
|
if once:
|
|
raise
|
|
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)
|
|
|
|
|
|
async def run(once: bool) -> None:
|
|
settings = WorkerSettings.from_environment()
|
|
async with httpx.AsyncClient(timeout=30, follow_redirects=False, trust_env=False) as client:
|
|
results = await asyncio.gather(
|
|
*(run_project(scope, client, once) for scope in settings.project_scopes()),
|
|
return_exceptions=once,
|
|
)
|
|
if any(isinstance(result, Exception) for result in results):
|
|
raise SystemExit(1)
|
|
|
|
|
|
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
|