diff --git a/coding-agent/src/coding_agent/environment/archive.py b/coding-agent/src/coding_agent/environment/archive.py new file mode 100644 index 000000000..d44463d4e --- /dev/null +++ b/coding-agent/src/coding_agent/environment/archive.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import base64 +import gzip +import io +import os +import shutil +import tarfile +import tempfile +import zlib +from pathlib import Path +from typing import Protocol, override + + +_PRESERVED_ROOT_NAMES = frozenset({".git"}) +_EXCLUDED_DIRECTORY_NAMES = frozenset({"__pycache__"}) +_EXCLUDED_FILE_SUFFIXES = (".pyc", ".pyo") +_MAX_WORKSPACE_ARCHIVE_BYTES = 8 * 1024 * 1024 +_MAX_WORKSPACE_ARCHIVE_BASE64_CHARS = 4 * ((_MAX_WORKSPACE_ARCHIVE_BYTES + 2) // 3) +_MAX_WORKSPACE_TAR_STREAM_BYTES = 16 * 1024 * 1024 +_MAX_WORKSPACE_ARCHIVE_MEMBERS = 4096 + + +def create_workspace_archive_base64(workspace_root: Path) -> str: + root = workspace_root.expanduser().resolve() + if not root.is_dir(): + raise ValueError(f"workspace root does not exist: {root}") + + buffer = io.BytesIO() + total_input_bytes = 0 + with tarfile.open(fileobj=buffer, mode="w:gz") as archive: + for path in sorted(root.rglob("*")): + relative = path.relative_to(root) + if should_exclude_workspace_archive_path(relative): + continue + if path.is_symlink(): + raise ValueError( + f"workspace archive does not support symlinks: {relative}" + ) + if not path.is_file(): + continue + stat_result = path.stat() + total_input_bytes += stat_result.st_size + _raise_if_archive_too_large(total_input_bytes) + data = path.read_bytes() + info = tarfile.TarInfo(name=relative.as_posix()) + info.size = len(data) + info.mode = stat_result.st_mode & 0o777 + info.mtime = int(stat_result.st_mtime) + archive.addfile(info, io.BytesIO(data)) + archive_bytes = buffer.getvalue() + _raise_if_archive_too_large(len(archive_bytes)) + return base64.b64encode(archive_bytes).decode("ascii") + + +def should_exclude_workspace_archive_path(relative: Path) -> bool: + if not relative.parts: + return True + if relative.parts[0] in _PRESERVED_ROOT_NAMES: + return True + if any(part in _EXCLUDED_DIRECTORY_NAMES for part in relative.parts): + return True + return relative.name.endswith(_EXCLUDED_FILE_SUFFIXES) + + +def extract_workspace_archive_base64(workspace_root: Path, archive_base64: str) -> None: + root = workspace_root.expanduser().resolve() + if not root.exists(): + root.mkdir(parents=True, exist_ok=True) + if not root.is_dir(): + raise ValueError(f"workspace root is not a directory: {root}") + + _raise_if_archive_base64_too_large(archive_base64) + try: + archive_bytes = base64.b64decode(archive_base64.encode("ascii"), validate=True) + except ValueError as exc: + raise ValueError("workspace archive must be valid base64") from exc + _raise_if_archive_too_large(len(archive_bytes)) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir).resolve() + try: + _extract_workspace_archive_to_temp(temp_root, archive_bytes) + except (tarfile.TarError, gzip.BadGzipFile, EOFError, zlib.error) as exc: + raise ValueError("workspace archive is not a valid tar.gz") from exc + + _clear_extract_target(root) + _copy_extracted_workspace(temp_root, root) + + +def _extract_workspace_archive_to_temp(temp_root: Path, archive_bytes: bytes) -> None: + archive_buffer = io.BytesIO(archive_bytes) + gzip_stream = gzip.GzipFile(fileobj=archive_buffer, mode="rb") + counted_stream = _CountingReader( + gzip_stream, + max_bytes=_MAX_WORKSPACE_TAR_STREAM_BYTES, + ) + extracted_bytes = 0 + member_count = 0 + with tarfile.open(fileobj=counted_stream, mode="r|") as archive: + for member in archive: + member_count += 1 + if member_count > _MAX_WORKSPACE_ARCHIVE_MEMBERS: + raise ValueError("workspace archive contains too many members") + _reject_preserved_root_member(member.name) + target_path = _safe_archive_target(temp_root, member.name) + if member.isdir(): + target_path.mkdir(parents=True, exist_ok=True) + target_path.chmod(_member_mode(member)) + continue + if not member.isfile(): + raise ValueError( + f"workspace archive only supports regular files and directories: {member.name}" + ) + if member.size < 0: + raise ValueError("workspace archive is not a valid tar.gz") + extracted_bytes += member.size + _raise_if_archive_too_large(extracted_bytes) + file_obj = archive.extractfile(member) + if file_obj is None: + raise ValueError( + f"workspace archive member is unreadable: {member.name}" + ) + _copy_archive_file(file_obj, target_path, member) + + +def _copy_archive_file( + file_obj: _ReadableBinary, target_path: Path, member: tarfile.TarInfo +) -> None: + target_path.parent.mkdir(parents=True, exist_ok=True) + copied_bytes = 0 + with target_path.open("wb") as target_file: + while True: + chunk = file_obj.read(64 * 1024) + if not chunk: + break + copied_bytes += len(chunk) + if copied_bytes > member.size: + raise ValueError("workspace archive is not a valid tar.gz") + _ = target_file.write(chunk) + if copied_bytes != member.size: + raise ValueError("workspace archive is not a valid tar.gz") + target_path.chmod(_member_mode(member)) + os.utime(target_path, (member.mtime, member.mtime)) + + +def _copy_extracted_workspace(temp_root: Path, workspace_root: Path) -> None: + for source_path in sorted(temp_root.rglob("*")): + relative = source_path.relative_to(temp_root) + target_path = workspace_root / relative + if source_path.is_dir(): + target_path.mkdir(parents=True, exist_ok=True) + target_path.chmod(source_path.stat().st_mode & 0o777) + continue + target_path.parent.mkdir(parents=True, exist_ok=True) + _ = shutil.copy2(source_path, target_path) + + +def _clear_extract_target(workspace_root: Path) -> None: + for child in workspace_root.iterdir(): + if child.name in _PRESERVED_ROOT_NAMES: + continue + if child.is_dir() and not child.is_symlink(): + shutil.rmtree(child) + continue + child.unlink() + + +def _member_mode(member: tarfile.TarInfo) -> int: + mode = member.mode & 0o777 + if mode != 0: + return mode + if member.isdir(): + return 0o755 + return 0o644 + + +def _safe_archive_target(workspace_root: Path, member_name: str) -> Path: + member_path = Path(member_name) + if member_path.is_absolute() or ".." in member_path.parts: + raise ValueError( + f"workspace archive member escapes workspace root: {member_name}" + ) + if not member_path.parts: + raise ValueError("workspace archive member name must be non-empty") + candidate = (workspace_root / member_path).resolve() + try: + _ = candidate.relative_to(workspace_root) + except ValueError as exc: + raise ValueError( + f"workspace archive member escapes workspace root: {member_name}" + ) from exc + return candidate + + +def _reject_preserved_root_member(member_name: str) -> None: + member_path = Path(member_name) + if member_path.parts and member_path.parts[0] in _PRESERVED_ROOT_NAMES: + raise ValueError( + f"workspace archive must not contain preserved root entry: {member_name}" + ) + + +def _raise_if_archive_too_large(size_bytes: int) -> None: + if size_bytes > _MAX_WORKSPACE_ARCHIVE_BYTES: + raise ValueError("workspace archive exceeds 8 MiB limit") + + +def _raise_if_archive_base64_too_large(archive_base64: str) -> None: + if len(archive_base64) > _MAX_WORKSPACE_ARCHIVE_BASE64_CHARS: + raise ValueError("workspace archive exceeds 8 MiB limit") + + +class _ReadableBinary(Protocol): + def read(self, size: int = -1, /) -> bytes: ... + + +class _CountingReader(io.BufferedIOBase): + def __init__(self, stream: _ReadableBinary, *, max_bytes: int) -> None: + self._stream: _ReadableBinary = stream + self._max_bytes: int = max_bytes + self._bytes_read: int = 0 + + @override + def readable(self) -> bool: + return True + + @override + def read(self, size: int | None = -1) -> bytes: + if size is None: + size = -1 + chunk = self._stream.read(size) + self._bytes_read += len(chunk) + if self._bytes_read > self._max_bytes: + raise ValueError("workspace archive exceeds 8 MiB limit") + return chunk diff --git a/coding-agent/src/coding_agent/environment/docker_workspace_provider.py b/coding-agent/src/coding_agent/environment/docker_workspace_provider.py index 6fe03c081..90d2c553b 100644 --- a/coding-agent/src/coding_agent/environment/docker_workspace_provider.py +++ b/coding-agent/src/coding_agent/environment/docker_workspace_provider.py @@ -38,7 +38,7 @@ WorkspaceStatus, register_workspace_provider, ) -from ..workspace_archive import ( +from .archive import ( create_workspace_archive_base64, extract_workspace_archive_base64, should_exclude_workspace_archive_path, diff --git a/coding-agent/src/coding_agent/executors/external.py b/coding-agent/src/coding_agent/executors/external.py new file mode 100644 index 000000000..237c8192d --- /dev/null +++ b/coding-agent/src/coding_agent/executors/external.py @@ -0,0 +1,1509 @@ +"""Generic external executor records and registry. + +External executors consume already-authorized Bee execution plans. This module +defines the product-layer model and durable run store; concrete executor +adapters are added in later goals. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import datetime +import hashlib +from inspect import isawaitable +from pathlib import Path +from typing import Final, Protocol + +from agentkit.storage.pg import AsyncPGPool, PGPool +from coding_agent.action_safety.approval_routing import ActionApprovalRoute +from coding_agent.action_safety.command_policy import CommandPolicyDecision +from coding_agent.bee_command_bridge import ( + BeeCommandIntentPlan, + BeeNodeCompletionEvidence, + is_authorized_bee_command_plan, +) +from coding_agent.topic_store import JSONObject, JSONValue + +ExecutorKind = str +ExecutorRunStatus = str +WorkspaceResolver = Callable[[str], Path | None] + +_EXECUTOR_KINDS: Final[frozenset[str]] = frozenset({ + "local", + "docker", + "kubernetes_job", + "argo_workflow", + "fixture", +}) +_EXECUTOR_RUN_STATUSES: Final[frozenset[str]] = frozenset({ + "planned", + "submitted", + "running", + "succeeded", + "failed", + "cancelled", +}) +_MAX_SAFE_LABEL_CHARS: Final[int] = 128 +_MAX_DISPLAY_TEXT_CHARS: Final[int] = 256 +_MAX_METADATA_STRING_CHARS: Final[int] = 256 +_FORBIDDEN_METADATA_KEY_PARTS: Final[frozenset[str]] = frozenset({ + "api_key", + "apikey", + "args", + "argv", + "bearer", + "cmd", + "command", + "commands", + "command_output", + "content", + "credential", + "credentials", + "env", + "environment", + "exec", + "key", + "log", + "message", + "password", + "pod_name", + "prompt", + "result", + "script", + "secret", + "shell", + "stderr", + "stdout", + "text", + "token", + "workflow_name", + "job_name", +}) +_FORBIDDEN_LABEL_VALUE_PARTS: Final[frozenset[str]] = frozenset({ + "api_key", + "apikey", + "bearer", + "command_output", + "content", + "credential", + "credentials", + "env", + "environment", + "key", + "message", + "password", + "prompt", + "result", + "secret", + "stderr", + "stdout", + "text", + "token", +}) +_FORBIDDEN_REF_PARTS: Final[frozenset[str]] = frozenset({ + "job_name", + "jobs/", + "kubeconfig", + "password", + "pod_name", + "pods/", + "secret", + "token", + "workflow_name", + "workflows/", +}) +_RAW_EXECUTION_TEXT_MARKERS: Final[tuple[str, ...]] = ( + "command:", + "command=", + "raw log", + "raw_log", + "rm -rf", + "stderr:", + "stderr=", + "stdout:", + "stdout=", + "traceback", + "uv run", + "python -c", + "pytest ", + "pytest -q", +) +_SECRET_VALUE_MARKERS: Final[tuple[str, ...]] = ( + "-----begin ", + "akia", + "bearer ", + "gho_", + "ghp_", + "kubeconfig", + "password=", + "private key", + "secret=", + "sk-", + "token=", +) +_EXECUTOR_PLAN_AUTHORIZATION_TOKEN: Final[object] = object() + + +class ExternalExecutor(Protocol): + @property + def kind(self) -> ExecutorKind: ... + + def capability(self) -> ExecutorCapability: ... + + async def submit(self, plan: ExecutorPlan) -> ExecutorResult: ... + + +class DockerCapabilityClient(Protocol): + def available(self) -> bool: ... + + +class KubernetesCapabilityClient(Protocol): + def available(self) -> bool: ... + + +class ArgoWorkflowCapabilityClient(Protocol): + def available(self) -> bool: ... + + +class ExecutorRunStore(Protocol): + async def create_executor_run( + self, + record: ExecutorRunRecord, + ) -> ExecutorRunRecord: ... + + async def load_executor_run( + self, + executor_run_id: str, + ) -> ExecutorRunRecord | None: ... + + async def list_executor_runs( + self, + *, + task_id: str | None = None, + node_id: str | None = None, + executor_kind: ExecutorKind | None = None, + status: ExecutorRunStatus | None = None, + limit: int = 100, + ) -> list[ExecutorRunRecord]: ... + + async def update_executor_run_status( + self, + executor_run_id: str, + *, + status: ExecutorRunStatus, + submitted_at: datetime | None, + started_at: datetime | None, + finished_at: datetime | None, + error_type: str | None, + error_message: str | None, + metadata: JSONObject, + ) -> ExecutorRunRecord: ... + + async def attach_executor_result( + self, + executor_run_id: str, + *, + result: ExecutorResult, + ) -> ExecutorRunRecord: ... + + +@dataclass(frozen=True) +class ExecutorCapability: + executor_kind: ExecutorKind + enabled: bool + available: bool + status: str + reason: str | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_executor_kind(self.executor_kind) + if not isinstance(self.enabled, bool): + raise TypeError("enabled must be a bool") + if not isinstance(self.available, bool): + raise TypeError("available must be a bool") + _require_optional_label("status", self.status) + _require_optional_display_text("reason", self.reason) + _require_json_object("metadata", self.metadata) + + +@dataclass(frozen=True) +class ExecutorPlan: + executor_kind: ExecutorKind + task_id: str + node_id: str + workspace_ref: str + requested_at: datetime + executor_run_id: str | None = None + approved_workspace_hash: str | None = None + command_category: str = "unknown" + command_profile: str = "unknown" + launch_id: str | None = None + topic_id: str | None = None + timeout_seconds: int | None = None + validation_label: str | None = None + metadata: JSONObject = field(default_factory=dict) + authorization_token: object | None = field(default=None, repr=False, compare=False) + authorization_signature: str | None = field( + default=None, init=False, repr=False, compare=False + ) + + def __post_init__(self) -> None: + _require_executor_kind(self.executor_kind) + _require_non_empty("task_id", self.task_id) + _require_non_empty("node_id", self.node_id) + _require_optional_label("workspace_ref", self.workspace_ref) + _require_datetime("requested_at", self.requested_at) + _require_optional_id("executor_run_id", self.executor_run_id) + _require_optional_id("approved_workspace_hash", self.approved_workspace_hash) + _require_optional_label("command_category", self.command_category) + _require_optional_label("command_profile", self.command_profile) + _require_optional_id("launch_id", self.launch_id) + _require_optional_id("topic_id", self.topic_id) + if self.timeout_seconds is not None: + _require_positive_int("timeout_seconds", self.timeout_seconds) + _require_optional_label("validation_label", self.validation_label) + _require_json_object("metadata", self.metadata) + + +@dataclass(frozen=True) +class ExecutorEvidence: + evidence_kind: str + evidence_ref: str + summary: str | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_optional_label("evidence_kind", self.evidence_kind) + _require_safe_ref("evidence_ref", self.evidence_ref) + _require_optional_display_text("summary", self.summary) + _require_json_object("metadata", self.metadata) + + def to_safe_dict(self) -> JSONObject: + data: JSONObject = { + "evidence_kind": self.evidence_kind, + "evidence_ref": self.evidence_ref, + "metadata": dict(self.metadata), + } + if self.summary is not None: + data["summary"] = self.summary + return data + + +@dataclass(frozen=True) +class ExecutorResult: + status: ExecutorRunStatus + sanitized_summary: str | None = None + evidence: tuple[ExecutorEvidence, ...] = () + finished_at: datetime | None = None + error_type: str | None = None + error_message: str | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_status("executor result status", self.status) + _require_optional_display_text("sanitized_summary", self.sanitized_summary) + _require_optional_datetime("finished_at", self.finished_at) + _require_optional_label("error_type", self.error_type) + _require_optional_display_text("error_message", self.error_message) + if not isinstance(self.evidence, tuple): + raise TypeError("evidence must be a tuple") + for item in self.evidence: + if not isinstance(item, ExecutorEvidence): + raise TypeError("evidence items must be ExecutorEvidence") + _require_json_object("metadata", self.metadata) + + def evidence_as_json(self) -> list[JSONObject]: + return [item.to_safe_dict() for item in self.evidence] + + +LocalExecutorRunner = Callable[ + [ExecutorPlan, Path], + ExecutorResult | Awaitable[ExecutorResult], +] + + +@dataclass(frozen=True) +class ExecutorRunRecord: + executor_run_id: str + executor_kind: ExecutorKind + task_id: str + node_id: str + status: ExecutorRunStatus + requested_at: datetime + launch_id: str | None = None + topic_id: str | None = None + submitted_at: datetime | None = None + started_at: datetime | None = None + finished_at: datetime | None = None + error_type: str | None = None + error_message: str | None = None + sanitized_summary: str | None = None + evidence: tuple[ExecutorEvidence, ...] = () + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_non_empty("executor_run_id", self.executor_run_id) + _require_executor_kind(self.executor_kind) + _require_non_empty("task_id", self.task_id) + _require_non_empty("node_id", self.node_id) + _require_status("executor run status", self.status) + _require_datetime("requested_at", self.requested_at) + _require_optional_id("launch_id", self.launch_id) + _require_optional_id("topic_id", self.topic_id) + _require_optional_datetime("submitted_at", self.submitted_at) + _require_optional_datetime("started_at", self.started_at) + _require_optional_datetime("finished_at", self.finished_at) + _require_optional_label("error_type", self.error_type) + _require_optional_display_text("error_message", self.error_message) + _require_optional_display_text("sanitized_summary", self.sanitized_summary) + if not isinstance(self.evidence, tuple): + raise TypeError("evidence must be a tuple") + for item in self.evidence: + if not isinstance(item, ExecutorEvidence): + raise TypeError("evidence items must be ExecutorEvidence") + _require_json_object("metadata", self.metadata) + + +class ExecutorRegistry: + def __init__(self) -> None: + self._executors: dict[ExecutorKind, ExternalExecutor] = {} + + def register(self, executor: ExternalExecutor) -> None: + _require_executor_kind(executor.kind) + if executor.kind in self._executors: + raise ValueError(f"executor already registered: {executor.kind}") + self._executors[executor.kind] = executor + + def resolve(self, executor_kind: ExecutorKind) -> ExternalExecutor: + _require_executor_kind(executor_kind) + executor = self._executors.get(executor_kind) + if executor is None: + raise KeyError(f"executor not registered: {executor_kind}") + return executor + + def list_kinds(self) -> tuple[ExecutorKind, ...]: + return tuple(sorted(self._executors)) + + +class LocalExecutorAdapter: + """Local executor adapter over already-approved Bee command bridge plans.""" + + kind: Final[ExecutorKind] = "local" + + def __init__( + self, + *, + store: ExecutorRunStore, + runner: LocalExecutorRunner | None = None, + workspace_resolver: WorkspaceResolver | None = None, + now: Callable[[], datetime], + executor_run_id_factory: Callable[[ExecutorPlan], str] | None = None, + ) -> None: + self._store = store + self._runner = runner + self._workspace_resolver = workspace_resolver + self._now = now + self._executor_run_id_factory = ( + executor_run_id_factory or _default_executor_run_id + ) + + def capability(self) -> ExecutorCapability: + available = self._runner is not None and self._workspace_resolver is not None + return ExecutorCapability( + executor_kind=self.kind, + enabled=True, + available=available, + status="available" if available else "unavailable", + reason=None if available else "local executor dependencies not configured", + ) + + async def submit(self, plan: ExecutorPlan) -> ExecutorResult: + _require_local_executor_plan(plan) + workspace_root = self._resolve_workspace(plan) + if workspace_root is None: + raise ValueError( + f"local executor workspace not found: {plan.workspace_ref}" + ) + _require_approved_workspace_binding(plan, workspace_root) + if self._runner is None: + raise ValueError("local executor runner is not configured") + executor_run_id = plan.executor_run_id or self._executor_run_id_factory(plan) + requested_at = plan.requested_at + await self._store.create_executor_run( + ExecutorRunRecord( + executor_run_id=executor_run_id, + executor_kind=self.kind, + task_id=plan.task_id, + node_id=plan.node_id, + launch_id=plan.launch_id, + topic_id=plan.topic_id, + status="planned", + requested_at=requested_at, + metadata=_executor_run_metadata(plan), + ) + ) + started_at = self._now() + await self._store.update_executor_run_status( + executor_run_id, + status="running", + submitted_at=started_at, + started_at=started_at, + finished_at=None, + error_type=None, + error_message=None, + metadata={"phase": "local_runner_started"}, + ) + try: + result_or_awaitable = self._runner(plan, workspace_root) + result = ( + await result_or_awaitable + if isawaitable(result_or_awaitable) + else result_or_awaitable + ) + result = _sanitize_local_executor_result(result) + except ValueError as exc: + if _is_safety_validation_error(exc): + raise + result = ExecutorResult( + status="failed", + sanitized_summary="Local executor failed with sanitized error", + finished_at=self._now(), + error_type=exc.__class__.__name__, + error_message="local executor runner failed", + metadata={"phase": "local_runner_failed"}, + ) + except Exception as exc: + result = ExecutorResult( + status="failed", + sanitized_summary="Local executor failed with sanitized error", + finished_at=self._now(), + error_type=exc.__class__.__name__, + error_message="local executor runner failed", + metadata={"phase": "local_runner_failed"}, + ) + await self._store.attach_executor_result(executor_run_id, result=result) + return result + + def _resolve_workspace(self, plan: ExecutorPlan) -> Path | None: + if self._workspace_resolver is None: + return None + workspace_root = self._workspace_resolver(plan.workspace_ref) + if workspace_root is None: + return None + resolved = Path(workspace_root).expanduser().resolve() + if not resolved.is_dir(): + return None + return resolved + + +class DockerExecutorAdapter: + """Optional Docker executor adapter in disabled/dry-run mode.""" + + kind: Final[ExecutorKind] = "docker" + + def __init__( + self, + *, + enabled: bool = False, + capability_client: DockerCapabilityClient | None = None, + ) -> None: + self._enabled = enabled + self._capability_client = capability_client + + def capability(self) -> ExecutorCapability: + if not self._enabled: + return ExecutorCapability( + executor_kind=self.kind, + enabled=False, + available=False, + status="disabled", + reason="docker executor disabled", + ) + available = ( + self._capability_client.available() + if self._capability_client is not None + else False + ) + return ExecutorCapability( + executor_kind=self.kind, + enabled=True, + available=available, + status="available" if available else "unavailable", + reason=None if available else "docker capability unavailable", + ) + + async def submit(self, plan: ExecutorPlan) -> ExecutorResult: + raise RuntimeError("docker executor execution is deferred; use dry_run_render") + + def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: + _require_authorized_executor_plan(plan, expected_kind=self.kind) + return { + "executor_kind": self.kind, + "mode": "dry_run", + "task_ref": _safe_ref_token(plan.task_id), + "node_ref": _safe_ref_token(plan.node_id), + "workspace_ref": _safe_ref_token(plan.workspace_ref), + "intent_category": plan.command_category, + "intent_profile": plan.command_profile, + "timeout_seconds": plan.timeout_seconds or 0, + } + + +class KubernetesJobExecutorAdapter: + """Kubernetes Job adapter in disabled/dry-run/fake-status mode.""" + + kind: Final[ExecutorKind] = "kubernetes_job" + + def __init__( + self, + *, + enabled: bool = False, + capability_client: KubernetesCapabilityClient | None = None, + ) -> None: + self._enabled = enabled + self._capability_client = capability_client + + def capability(self) -> ExecutorCapability: + if not self._enabled: + return ExecutorCapability( + executor_kind=self.kind, + enabled=False, + available=False, + status="disabled", + reason="kubernetes job executor disabled", + ) + available = ( + self._capability_client.available() + if self._capability_client is not None + else False + ) + return ExecutorCapability( + executor_kind=self.kind, + enabled=True, + available=available, + status="available" if available else "unavailable", + reason=None if available else "kubernetes capability unavailable", + ) + + async def submit(self, plan: ExecutorPlan) -> ExecutorResult: + raise RuntimeError( + "kubernetes job executor execution is deferred; use dry_run_render" + ) + + def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: + _require_authorized_executor_plan(plan, expected_kind=self.kind) + return { + "apiVersion": "batch/v1", + "kind": "Job", + "metadata": { + "generateName": "coding-agent-executor-", + "labels": { + "app": "coding-agent", + "executor_kind": self.kind, + "intent_category": plan.command_category, + "intent_profile": plan.command_profile, + }, + }, + "spec": { + "backoffLimit": 0, + "template": { + "spec": { + "restartPolicy": "Never", + "containers": [ + { + "name": "executor", + "image": "coding-agent-executor:dry-run", + "args": ["dry-run"], + } + ], + } + }, + }, + "safe_refs": { + "task_ref": _safe_ref_token(plan.task_id), + "node_ref": _safe_ref_token(plan.node_id), + "workspace_ref": _safe_ref_token(plan.workspace_ref), + }, + } + + def import_status(self, plan: ExecutorPlan, status: JSONObject) -> ExecutorResult: + _require_authorized_executor_plan(plan, expected_kind=self.kind) + executor_status = _kubernetes_job_status(status) + return ExecutorResult( + status=executor_status, + sanitized_summary=f"Kubernetes Job {executor_status}", + evidence=( + ExecutorEvidence( + evidence_kind="sanitized_artifact", + evidence_ref="evidence/kubernetes-job-status", + ), + ), + metadata={"phase": "kubernetes_status_imported"}, + ) + + +class ArgoWorkflowExecutorAdapter: + """Argo Workflows adapter in disabled/dry-run/fake-status mode.""" + + kind: Final[ExecutorKind] = "argo_workflow" + + def __init__( + self, + *, + enabled: bool = False, + capability_client: ArgoWorkflowCapabilityClient | None = None, + ) -> None: + self._enabled = enabled + self._capability_client = capability_client + + def capability(self) -> ExecutorCapability: + if not self._enabled: + return ExecutorCapability( + executor_kind=self.kind, + enabled=False, + available=False, + status="disabled", + reason="argo workflow executor disabled", + ) + available = ( + self._capability_client.available() + if self._capability_client is not None + else False + ) + return ExecutorCapability( + executor_kind=self.kind, + enabled=True, + available=available, + status="available" if available else "unavailable", + reason=None if available else "argo workflow capability unavailable", + ) + + async def submit(self, plan: ExecutorPlan) -> ExecutorResult: + raise RuntimeError( + "argo workflow executor execution is deferred; use dry_run_render" + ) + + def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: + _require_authorized_executor_plan(plan, expected_kind=self.kind) + return { + "apiVersion": "argoproj.io/v1alpha1", + "kind": "Workflow", + "metadata": { + "generateName": "coding-agent-executor-", + "labels": { + "app": "coding-agent", + "executor_kind": self.kind, + "intent_category": plan.command_category, + "intent_profile": plan.command_profile, + }, + }, + "spec": { + "entrypoint": "executor", + "templates": [ + { + "name": "executor", + "container": { + "image": "coding-agent-executor:dry-run", + "args": ["dry-run"], + }, + } + ], + }, + "safe_refs": { + "task_ref": _safe_ref_token(plan.task_id), + "node_ref": _safe_ref_token(plan.node_id), + "workspace_ref": _safe_ref_token(plan.workspace_ref), + }, + } + + def import_status(self, plan: ExecutorPlan, status: JSONObject) -> ExecutorResult: + _require_authorized_executor_plan(plan, expected_kind=self.kind) + executor_status = _argo_workflow_status(status) + return ExecutorResult( + status=executor_status, + sanitized_summary=f"Argo Workflow {executor_status}", + evidence=( + ExecutorEvidence( + evidence_kind="sanitized_artifact", + evidence_ref="evidence/argo-workflow-status", + ), + ), + metadata={"phase": "argo_workflow_status_imported"}, + ) + + +def build_local_executor_plan_from_bee_command_plan( + *, + command_plan: BeeCommandIntentPlan, + task_id: str, + workspace_ref: str, + approved_workspace_root: Path | str, + requested_at: datetime, + executor_run_id: str | None = None, + launch_id: str | None = None, + topic_id: str | None = None, + timeout_seconds: int | None = None, +) -> ExecutorPlan: + if command_plan.status != "ready": + raise ValueError("local executor requires a ready Bee command plan") + if not is_authorized_bee_command_plan(command_plan): + raise ValueError("local executor requires authorized Bee command plan") + if command_plan.policy is None: + raise ValueError("local executor requires command policy verdict") + if command_plan.policy.decision != CommandPolicyDecision.ALLOW: + raise ValueError("local executor requires allowed command policy") + if command_plan.approval_route is None: + raise ValueError("local executor requires approval route") + if command_plan.approval_route.route != ActionApprovalRoute.ALLOW: + raise ValueError("local executor requires allow approval route") + intent = command_plan.resolution.intent + if intent is None: + raise ValueError("local executor requires resolved command intent") + return _authorize_executor_plan( + ExecutorPlan( + executor_kind="local", + executor_run_id=executor_run_id, + approved_workspace_hash=_workspace_hash( + Path(approved_workspace_root).expanduser().resolve() + ), + task_id=task_id, + node_id=command_plan.resolution.node_id, + workspace_ref=workspace_ref, + requested_at=requested_at, + command_category=intent.category, + command_profile=intent.profile, + launch_id=launch_id, + topic_id=topic_id, + timeout_seconds=timeout_seconds or command_plan.policy.timeout_seconds, + validation_label=intent.validation_label or intent.name, + metadata={ + "policy_decision": command_plan.policy.decision.value, + "approval_route": command_plan.approval_route.route.value, + "template_id": command_plan.resolution.template_id, + "runtime_kind": command_plan.policy.environment_kind, + }, + authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, + ) + ) + + +def build_docker_executor_plan_from_local_plan( + local_plan: ExecutorPlan, + *, + executor_run_id: str | None = None, +) -> ExecutorPlan: + _require_authorized_executor_plan(local_plan, expected_kind="local") + return _authorize_executor_plan( + ExecutorPlan( + executor_kind="docker", + executor_run_id=executor_run_id, + approved_workspace_hash=local_plan.approved_workspace_hash, + task_id=local_plan.task_id, + node_id=local_plan.node_id, + workspace_ref=local_plan.workspace_ref, + requested_at=local_plan.requested_at, + command_category=local_plan.command_category, + command_profile=local_plan.command_profile, + launch_id=local_plan.launch_id, + topic_id=local_plan.topic_id, + timeout_seconds=local_plan.timeout_seconds, + validation_label=local_plan.validation_label, + metadata={ + "policy_decision": "allow", + "approval_route": "allow", + "source_kind": "local", + }, + authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, + ) + ) + + +def build_kubernetes_job_executor_plan_from_local_plan( + local_plan: ExecutorPlan, + *, + executor_run_id: str | None = None, +) -> ExecutorPlan: + _require_authorized_executor_plan(local_plan, expected_kind="local") + return _authorize_executor_plan( + ExecutorPlan( + executor_kind="kubernetes_job", + executor_run_id=executor_run_id, + approved_workspace_hash=local_plan.approved_workspace_hash, + task_id=local_plan.task_id, + node_id=local_plan.node_id, + workspace_ref=local_plan.workspace_ref, + requested_at=local_plan.requested_at, + command_category=local_plan.command_category, + command_profile=local_plan.command_profile, + launch_id=local_plan.launch_id, + topic_id=local_plan.topic_id, + timeout_seconds=local_plan.timeout_seconds, + validation_label=local_plan.validation_label, + metadata={ + "policy_decision": "allow", + "approval_route": "allow", + "source_kind": "local", + }, + authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, + ) + ) + + +def build_argo_workflow_executor_plan_from_local_plan( + local_plan: ExecutorPlan, + *, + executor_run_id: str | None = None, +) -> ExecutorPlan: + _require_authorized_executor_plan(local_plan, expected_kind="local") + return _authorize_executor_plan( + ExecutorPlan( + executor_kind="argo_workflow", + executor_run_id=executor_run_id, + approved_workspace_hash=local_plan.approved_workspace_hash, + task_id=local_plan.task_id, + node_id=local_plan.node_id, + workspace_ref=local_plan.workspace_ref, + requested_at=local_plan.requested_at, + command_category=local_plan.command_category, + command_profile=local_plan.command_profile, + launch_id=local_plan.launch_id, + topic_id=local_plan.topic_id, + timeout_seconds=local_plan.timeout_seconds, + validation_label=local_plan.validation_label, + metadata={ + "policy_decision": "allow", + "approval_route": "allow", + "source_kind": "local", + }, + authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, + ) + ) + + +def executor_result_completion_evidence( + result: ExecutorResult, +) -> tuple[BeeNodeCompletionEvidence, ...]: + status = "passed" if result.status == "succeeded" else "failed" + return tuple( + BeeNodeCompletionEvidence( + evidence_kind=item.evidence_kind, + evidence_ref=item.evidence_ref, + status=status, + ) + for item in result.evidence + ) + + +def _sanitize_local_executor_result(result: ExecutorResult) -> ExecutorResult: + return ExecutorResult( + status=result.status, + sanitized_summary=_local_executor_status_summary(result.status), + evidence=tuple( + evidence for item in result.evidence if (evidence := _safe_evidence(item)) + ), + finished_at=result.finished_at, + error_type=result.error_type, + error_message=_local_executor_error_summary(result), + metadata={"phase": "local_runner_result"}, + ) + + +def _local_executor_status_summary(status: ExecutorRunStatus) -> str: + if status == "succeeded": + return "Local executor succeeded" + if status == "failed": + return "Local executor failed" + if status == "cancelled": + return "Local executor cancelled" + return "Local executor status recorded" + + +def _local_executor_error_summary(result: ExecutorResult) -> str | None: + if result.status == "failed": + return "local executor returned failed status" + return None + + +def _safe_evidence(item: ExecutorEvidence) -> ExecutorEvidence | None: + try: + return ExecutorEvidence( + evidence_kind=item.evidence_kind, + evidence_ref=item.evidence_ref, + ) + except ValueError: + return None + + +class PGExecutorRunStore: + _CREATE_SCHEMA_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS executor_runs ( + executor_run_id TEXT PRIMARY KEY, + executor_kind TEXT NOT NULL, + task_id TEXT NOT NULL, + node_id TEXT NOT NULL, + launch_id TEXT, + topic_id TEXT, + status TEXT NOT NULL, + requested_at TIMESTAMPTZ NOT NULL, + submitted_at TIMESTAMPTZ, + started_at TIMESTAMPTZ, + finished_at TIMESTAMPTZ, + error_type TEXT, + error_message TEXT, + sanitized_summary TEXT, + evidence JSONB NOT NULL DEFAULT '[]'::jsonb, + metadata JSONB NOT NULL DEFAULT '{}'::jsonb + ); + + CREATE INDEX IF NOT EXISTS executor_runs_task_node_idx + ON executor_runs (task_id, node_id, requested_at, executor_run_id); + + CREATE INDEX IF NOT EXISTS executor_runs_kind_status_idx + ON executor_runs (executor_kind, status, requested_at, executor_run_id); + """ + _INSERT_RUN_SQL: Final[str] = """ + INSERT INTO executor_runs ( + executor_run_id, + executor_kind, + task_id, + node_id, + launch_id, + topic_id, + status, + requested_at, + submitted_at, + started_at, + finished_at, + error_type, + error_message, + sanitized_summary, + evidence, + metadata + ) + VALUES ( + $1, $2, $3, $4, $5, $6, $7, $8, + $9, $10, $11, $12, $13, $14, $15::jsonb, $16::jsonb + ) + ON CONFLICT (executor_run_id) DO NOTHING + RETURNING * + """ + _SELECT_RUN_SQL: Final[str] = """ + SELECT * FROM executor_runs WHERE executor_run_id = $1 + """ + _LIST_RUNS_SQL: Final[str] = """ + SELECT * FROM executor_runs + WHERE ($1::text IS NULL OR task_id = $1) + AND ($2::text IS NULL OR node_id = $2) + AND ($3::text IS NULL OR executor_kind = $3) + AND ($4::text IS NULL OR status = $4) + ORDER BY requested_at, executor_run_id + LIMIT $5 + """ + _UPDATE_RUN_STATUS_SQL: Final[str] = """ + UPDATE executor_runs + SET status = $2, + submitted_at = $3, + started_at = $4, + finished_at = $5, + error_type = $6, + error_message = $7, + metadata = $8::jsonb + WHERE executor_run_id = $1 + RETURNING * + """ + _ATTACH_RESULT_SQL: Final[str] = """ + UPDATE executor_runs + SET status = $2, + finished_at = $3, + error_type = $4, + error_message = $5, + sanitized_summary = $6, + evidence = $7::jsonb, + metadata = $8::jsonb + WHERE executor_run_id = $1 + RETURNING * + """ + + def __init__(self, *, pool: PGPool) -> None: + self._pool = pool + self._schema_ready = False + + async def _ensure_schema(self) -> AsyncPGPool: + pool = await self._pool.get_pool() + if not self._schema_ready: + _ = await pool.execute(self._CREATE_SCHEMA_SQL) + self._schema_ready = True + return pool + + async def create_executor_run( + self, + record: ExecutorRunRecord, + ) -> ExecutorRunRecord: + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._INSERT_RUN_SQL, + record.executor_run_id, + record.executor_kind, + record.task_id, + record.node_id, + record.launch_id, + record.topic_id, + record.status, + record.requested_at, + record.submitted_at, + record.started_at, + record.finished_at, + record.error_type, + record.error_message, + record.sanitized_summary, + [item.to_safe_dict() for item in record.evidence], + record.metadata, + ) + if row is None: + raise ValueError(f"executor run already exists: {record.executor_run_id}") + return _run_from_row(row) + + async def load_executor_run( + self, + executor_run_id: str, + ) -> ExecutorRunRecord | None: + _require_non_empty("executor_run_id", executor_run_id) + pool = await self._ensure_schema() + row = await pool.fetchrow(self._SELECT_RUN_SQL, executor_run_id) + return None if row is None else _run_from_row(row) + + async def list_executor_runs( + self, + *, + task_id: str | None = None, + node_id: str | None = None, + executor_kind: ExecutorKind | None = None, + status: ExecutorRunStatus | None = None, + limit: int = 100, + ) -> list[ExecutorRunRecord]: + _require_optional_id("task_id", task_id) + _require_optional_id("node_id", node_id) + if executor_kind is not None: + _require_executor_kind(executor_kind) + if status is not None: + _require_status("executor run status", status) + _require_positive_int("limit", limit) + pool = await self._ensure_schema() + rows = await pool.fetch( + self._LIST_RUNS_SQL, + task_id, + node_id, + executor_kind, + status, + limit, + ) + return [_run_from_row(row) for row in rows] + + async def update_executor_run_status( + self, + executor_run_id: str, + *, + status: ExecutorRunStatus, + submitted_at: datetime | None, + started_at: datetime | None, + finished_at: datetime | None, + error_type: str | None, + error_message: str | None, + metadata: JSONObject, + ) -> ExecutorRunRecord: + _require_non_empty("executor_run_id", executor_run_id) + _require_status("executor run status", status) + _require_optional_datetime("submitted_at", submitted_at) + _require_optional_datetime("started_at", started_at) + _require_optional_datetime("finished_at", finished_at) + _require_optional_label("error_type", error_type) + _require_optional_display_text("error_message", error_message) + _require_json_object("metadata", metadata) + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._UPDATE_RUN_STATUS_SQL, + executor_run_id, + status, + submitted_at, + started_at, + finished_at, + error_type, + error_message, + metadata, + ) + if row is None: + raise KeyError(f"executor run not found: {executor_run_id}") + return _run_from_row(row) + + async def attach_executor_result( + self, + executor_run_id: str, + *, + result: ExecutorResult, + ) -> ExecutorRunRecord: + _require_non_empty("executor_run_id", executor_run_id) + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._ATTACH_RESULT_SQL, + executor_run_id, + result.status, + result.finished_at, + result.error_type, + result.error_message, + result.sanitized_summary, + result.evidence_as_json(), + result.metadata, + ) + if row is None: + raise KeyError(f"executor run not found: {executor_run_id}") + return _run_from_row(row) + + +def _run_from_row(row: dict[str, object]) -> ExecutorRunRecord: + return ExecutorRunRecord( + executor_run_id=_required_str(row, "executor_run_id", context="executor run"), + executor_kind=_required_str(row, "executor_kind", context="executor run"), + task_id=_required_str(row, "task_id", context="executor run"), + node_id=_required_str(row, "node_id", context="executor run"), + launch_id=_optional_str(row, "launch_id", context="executor run"), + topic_id=_optional_str(row, "topic_id", context="executor run"), + status=_required_str(row, "status", context="executor run"), + requested_at=_required_datetime(row, "requested_at", context="executor run"), + submitted_at=_optional_datetime(row, "submitted_at", context="executor run"), + started_at=_optional_datetime(row, "started_at", context="executor run"), + finished_at=_optional_datetime(row, "finished_at", context="executor run"), + error_type=_optional_str(row, "error_type", context="executor run"), + error_message=_optional_str(row, "error_message", context="executor run"), + sanitized_summary=_optional_str( + row, "sanitized_summary", context="executor run" + ), + evidence=_required_evidence_tuple(row, "evidence", context="executor run"), + metadata=_required_json_object(row, "metadata", context="executor run"), + ) + + +def _required_evidence_tuple( + row: dict[str, object], + key: str, + *, + context: str, +) -> tuple[ExecutorEvidence, ...]: + value = row.get(key) + if not isinstance(value, list): + raise TypeError(f"postgres {context} must include list {key}") + return tuple(_evidence_from_json(item, context=context) for item in value) + + +def _evidence_from_json(value: object, *, context: str) -> ExecutorEvidence: + if not isinstance(value, dict): + raise TypeError(f"postgres {context} evidence must include objects") + evidence_kind = value.get("evidence_kind") + evidence_ref = value.get("evidence_ref") + summary = value.get("summary") + metadata = value.get("metadata", {}) + if not isinstance(evidence_kind, str): + raise TypeError(f"postgres {context} evidence missing evidence_kind") + if not isinstance(evidence_ref, str): + raise TypeError(f"postgres {context} evidence missing evidence_ref") + if summary is not None and not isinstance(summary, str): + raise TypeError(f"postgres {context} evidence summary must be a string") + if not isinstance(metadata, dict): + raise TypeError(f"postgres {context} evidence metadata must be an object") + return ExecutorEvidence( + evidence_kind=evidence_kind, + evidence_ref=evidence_ref, + summary=summary, + metadata=metadata, + ) + + +def _require_executor_kind(value: str) -> None: + if value not in _EXECUTOR_KINDS: + raise ValueError(f"executor kind must be one of {sorted(_EXECUTOR_KINDS)}") + + +def _require_local_executor_plan(plan: ExecutorPlan) -> None: + _require_authorized_executor_plan(plan, expected_kind="local") + if plan.approved_workspace_hash is None: + raise ValueError("local executor requires approved workspace binding") + + +def _require_authorized_executor_plan( + plan: ExecutorPlan, + *, + expected_kind: ExecutorKind, +) -> None: + if plan.executor_kind != expected_kind: + raise ValueError( + f"{expected_kind} executor requires executor_kind={expected_kind!r}" + ) + if plan.authorization_token is not _EXECUTOR_PLAN_AUTHORIZATION_TOKEN: + raise ValueError(f"{expected_kind} executor requires authorized executor plan") + if plan.authorization_signature != _executor_plan_signature(plan): + raise ValueError(f"{expected_kind} executor requires signed executor plan") + if plan.metadata.get("policy_decision") != "allow": + raise ValueError(f"{expected_kind} executor requires allowed command policy") + if plan.metadata.get("approval_route") != "allow": + raise ValueError(f"{expected_kind} executor requires allow approval route") + + +def _require_approved_workspace_binding( + plan: ExecutorPlan, workspace_root: Path +) -> None: + if plan.approved_workspace_hash != _workspace_hash(workspace_root): + raise ValueError("local executor workspace binding mismatch") + + +def _executor_run_metadata(plan: ExecutorPlan) -> JSONObject: + return { + "phase": "local_executor_planned", + "intent_category": plan.command_category, + "intent_profile": plan.command_profile, + } + + +def _default_executor_run_id(plan: ExecutorPlan) -> str: + return f"executor-{plan.task_id}-{plan.node_id}" + + +def _authorize_executor_plan(plan: ExecutorPlan) -> ExecutorPlan: + object.__setattr__( + plan, + "authorization_signature", + _executor_plan_signature(plan), + ) + return plan + + +def _executor_plan_signature(plan: ExecutorPlan) -> str: + metadata_items = tuple( + sorted((key, str(value)) for key, value in plan.metadata.items()) + ) + parts = ( + plan.executor_kind, + plan.executor_run_id or "", + plan.approved_workspace_hash or "", + plan.task_id, + plan.node_id, + plan.workspace_ref, + plan.command_category, + plan.command_profile, + plan.launch_id or "", + plan.topic_id or "", + str(plan.timeout_seconds or ""), + plan.validation_label or "", + repr(metadata_items), + ) + return hashlib.sha256("\x1f".join(parts).encode("utf-8")).hexdigest() + + +def _workspace_hash(path: Path) -> str: + return hashlib.sha256(str(path).encode("utf-8")).hexdigest()[:16] + + +def _safe_ref_token(value: str) -> str: + return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16] + + +def _kubernetes_job_status(status: JSONObject) -> ExecutorRunStatus: + if status.get("succeeded") is True or status.get("condition") == "complete": + return "succeeded" + if status.get("failed") is True or status.get("condition") == "failed": + return "failed" + return "running" + + +def _argo_workflow_status(status: JSONObject) -> ExecutorRunStatus: + phase = status.get("phase") + if phase == "Succeeded" or status.get("succeeded") is True: + return "succeeded" + if phase in {"Failed", "Error"} or status.get("failed") is True: + return "failed" + return "running" + + +def _is_safety_validation_error(exc: ValueError) -> bool: + message = str(exc) + return ( + "raw execution text" in message + or "secret-shaped" in message + or "forbidden metadata key" in message + or "sensitive reference text" in message + or "must not contain whitespace" in message + ) + + +def _require_non_empty(field_name: str, value: str) -> None: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field_name} must be a non-empty string") + + +def _require_optional_id(field_name: str, value: str | None) -> None: + if value is None: + return + _require_non_empty(field_name, value) + + +def _require_optional_label(field_name: str, value: str | None) -> None: + if value is None: + return + _require_non_empty(field_name, value) + if len(value) > _MAX_SAFE_LABEL_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_SAFE_LABEL_CHARS} characters" + ) + folded = value.casefold() + if any(part in folded for part in _FORBIDDEN_LABEL_VALUE_PARTS): + raise ValueError(f"{field_name} must not contain sensitive label text") + _reject_secret_shaped_value(field_name, value) + + +def _require_safe_ref(field_name: str, value: str) -> None: + _require_non_empty(field_name, value) + if len(value) > _MAX_DISPLAY_TEXT_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" + ) + if any(character.isspace() for character in value): + raise ValueError(f"{field_name} must not contain whitespace") + _reject_secret_shaped_value(field_name, value) + folded = value.casefold() + if any(part in folded for part in _FORBIDDEN_REF_PARTS): + raise ValueError(f"{field_name} must not contain sensitive reference text") + _reject_raw_execution_text(field_name, value) + + +def _require_status(field_name: str, value: str) -> None: + if value not in _EXECUTOR_RUN_STATUSES: + raise ValueError( + f"{field_name} must be one of {sorted(_EXECUTOR_RUN_STATUSES)}" + ) + + +def _require_datetime(field_name: str, value: datetime) -> None: + if not isinstance(value, datetime): + raise TypeError(f"{field_name} must be a datetime") + + +def _require_optional_datetime(field_name: str, value: datetime | None) -> None: + if value is not None: + _require_datetime(field_name, value) + + +def _require_optional_display_text(field_name: str, value: str | None) -> None: + if value is None: + return + _require_non_empty(field_name, value) + if len(value) > _MAX_DISPLAY_TEXT_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" + ) + _reject_secret_shaped_value(field_name, value) + _reject_raw_execution_text(field_name, value) + + +def _require_positive_int(field_name: str, value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{field_name} must be a positive int") + + +def _require_json_object(field_name: str, value: JSONObject) -> None: + if not isinstance(value, dict): + raise TypeError(f"{field_name} must be a JSON object") + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{field_name} keys must be strings") + key_folded = key.casefold() + if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): + raise ValueError(f"{field_name} contains forbidden metadata key: {key}") + _require_json_value(f"{field_name}.{key}", item) + + +def _require_json_value(field_name: str, value: JSONValue) -> None: + if value is None or isinstance(value, str | int | float | bool): + if isinstance(value, str): + if len(value) > _MAX_METADATA_STRING_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_METADATA_STRING_CHARS} characters" + ) + _reject_secret_shaped_value(field_name, value) + _reject_raw_execution_text(field_name, value) + return + if isinstance(value, list): + for index, item in enumerate(value): + _require_json_value(f"{field_name}[{index}]", item) + return + if isinstance(value, dict): + _require_json_object(field_name, value) + return + raise TypeError(f"{field_name} must be JSON-safe") + + +def _reject_secret_shaped_value(field_name: str, value: str) -> None: + folded = value.casefold() + if any(marker in folded for marker in _SECRET_VALUE_MARKERS): + raise ValueError(f"{field_name} must not contain secret-shaped values") + + +def _reject_raw_execution_text(field_name: str, value: str) -> None: + folded = value.casefold() + if any(marker in folded for marker in _RAW_EXECUTION_TEXT_MARKERS): + raise ValueError(f"{field_name} must not contain raw execution text") + + +def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: + if row is None: + raise RuntimeError(f"postgres did not return a row for {context}") + return row + + +def _required_str(row: dict[str, object], key: str, *, context: str) -> str: + value = row.get(key) + if not isinstance(value, str): + raise TypeError(f"postgres {context} must include string {key}") + return value + + +def _optional_str(row: dict[str, object], key: str, *, context: str) -> str | None: + value = row.get(key) + if value is None: + return None + if not isinstance(value, str): + raise TypeError(f"postgres {context} must include string or None {key}") + return value + + +def _required_datetime(row: dict[str, object], key: str, *, context: str) -> datetime: + value = row.get(key) + if not isinstance(value, datetime): + raise TypeError(f"postgres {context} must include datetime {key}") + return value + + +def _optional_datetime( + row: dict[str, object], + key: str, + *, + context: str, +) -> datetime | None: + value = row.get(key) + if value is None: + return None + if not isinstance(value, datetime): + raise TypeError(f"postgres {context} must include datetime or None {key}") + return value + + +def _required_json_object( + row: dict[str, object], + key: str, + *, + context: str, +) -> JSONObject: + value = row.get(key) + if not isinstance(value, dict): + raise TypeError(f"postgres {context} must include dict {key}") + _require_json_object(key, value) + return value diff --git a/coding-agent/src/coding_agent/external_executor.py b/coding-agent/src/coding_agent/external_executor.py index 237c8192d..eb9432951 100644 --- a/coding-agent/src/coding_agent/external_executor.py +++ b/coding-agent/src/coding_agent/external_executor.py @@ -1,1509 +1,4 @@ -"""Generic external executor records and registry. - -External executors consume already-authorized Bee execution plans. This module -defines the product-layer model and durable run store; concrete executor -adapters are added in later goals. -""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from dataclasses import dataclass, field -from datetime import datetime -import hashlib -from inspect import isawaitable -from pathlib import Path -from typing import Final, Protocol - -from agentkit.storage.pg import AsyncPGPool, PGPool -from coding_agent.action_safety.approval_routing import ActionApprovalRoute -from coding_agent.action_safety.command_policy import CommandPolicyDecision -from coding_agent.bee_command_bridge import ( - BeeCommandIntentPlan, - BeeNodeCompletionEvidence, - is_authorized_bee_command_plan, -) -from coding_agent.topic_store import JSONObject, JSONValue - -ExecutorKind = str -ExecutorRunStatus = str -WorkspaceResolver = Callable[[str], Path | None] - -_EXECUTOR_KINDS: Final[frozenset[str]] = frozenset({ - "local", - "docker", - "kubernetes_job", - "argo_workflow", - "fixture", -}) -_EXECUTOR_RUN_STATUSES: Final[frozenset[str]] = frozenset({ - "planned", - "submitted", - "running", - "succeeded", - "failed", - "cancelled", -}) -_MAX_SAFE_LABEL_CHARS: Final[int] = 128 -_MAX_DISPLAY_TEXT_CHARS: Final[int] = 256 -_MAX_METADATA_STRING_CHARS: Final[int] = 256 -_FORBIDDEN_METADATA_KEY_PARTS: Final[frozenset[str]] = frozenset({ - "api_key", - "apikey", - "args", - "argv", - "bearer", - "cmd", - "command", - "commands", - "command_output", - "content", - "credential", - "credentials", - "env", - "environment", - "exec", - "key", - "log", - "message", - "password", - "pod_name", - "prompt", - "result", - "script", - "secret", - "shell", - "stderr", - "stdout", - "text", - "token", - "workflow_name", - "job_name", -}) -_FORBIDDEN_LABEL_VALUE_PARTS: Final[frozenset[str]] = frozenset({ - "api_key", - "apikey", - "bearer", - "command_output", - "content", - "credential", - "credentials", - "env", - "environment", - "key", - "message", - "password", - "prompt", - "result", - "secret", - "stderr", - "stdout", - "text", - "token", -}) -_FORBIDDEN_REF_PARTS: Final[frozenset[str]] = frozenset({ - "job_name", - "jobs/", - "kubeconfig", - "password", - "pod_name", - "pods/", - "secret", - "token", - "workflow_name", - "workflows/", -}) -_RAW_EXECUTION_TEXT_MARKERS: Final[tuple[str, ...]] = ( - "command:", - "command=", - "raw log", - "raw_log", - "rm -rf", - "stderr:", - "stderr=", - "stdout:", - "stdout=", - "traceback", - "uv run", - "python -c", - "pytest ", - "pytest -q", -) -_SECRET_VALUE_MARKERS: Final[tuple[str, ...]] = ( - "-----begin ", - "akia", - "bearer ", - "gho_", - "ghp_", - "kubeconfig", - "password=", - "private key", - "secret=", - "sk-", - "token=", -) -_EXECUTOR_PLAN_AUTHORIZATION_TOKEN: Final[object] = object() - - -class ExternalExecutor(Protocol): - @property - def kind(self) -> ExecutorKind: ... - - def capability(self) -> ExecutorCapability: ... - - async def submit(self, plan: ExecutorPlan) -> ExecutorResult: ... - - -class DockerCapabilityClient(Protocol): - def available(self) -> bool: ... - - -class KubernetesCapabilityClient(Protocol): - def available(self) -> bool: ... - - -class ArgoWorkflowCapabilityClient(Protocol): - def available(self) -> bool: ... - - -class ExecutorRunStore(Protocol): - async def create_executor_run( - self, - record: ExecutorRunRecord, - ) -> ExecutorRunRecord: ... - - async def load_executor_run( - self, - executor_run_id: str, - ) -> ExecutorRunRecord | None: ... - - async def list_executor_runs( - self, - *, - task_id: str | None = None, - node_id: str | None = None, - executor_kind: ExecutorKind | None = None, - status: ExecutorRunStatus | None = None, - limit: int = 100, - ) -> list[ExecutorRunRecord]: ... - - async def update_executor_run_status( - self, - executor_run_id: str, - *, - status: ExecutorRunStatus, - submitted_at: datetime | None, - started_at: datetime | None, - finished_at: datetime | None, - error_type: str | None, - error_message: str | None, - metadata: JSONObject, - ) -> ExecutorRunRecord: ... - - async def attach_executor_result( - self, - executor_run_id: str, - *, - result: ExecutorResult, - ) -> ExecutorRunRecord: ... - - -@dataclass(frozen=True) -class ExecutorCapability: - executor_kind: ExecutorKind - enabled: bool - available: bool - status: str - reason: str | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_executor_kind(self.executor_kind) - if not isinstance(self.enabled, bool): - raise TypeError("enabled must be a bool") - if not isinstance(self.available, bool): - raise TypeError("available must be a bool") - _require_optional_label("status", self.status) - _require_optional_display_text("reason", self.reason) - _require_json_object("metadata", self.metadata) - - -@dataclass(frozen=True) -class ExecutorPlan: - executor_kind: ExecutorKind - task_id: str - node_id: str - workspace_ref: str - requested_at: datetime - executor_run_id: str | None = None - approved_workspace_hash: str | None = None - command_category: str = "unknown" - command_profile: str = "unknown" - launch_id: str | None = None - topic_id: str | None = None - timeout_seconds: int | None = None - validation_label: str | None = None - metadata: JSONObject = field(default_factory=dict) - authorization_token: object | None = field(default=None, repr=False, compare=False) - authorization_signature: str | None = field( - default=None, init=False, repr=False, compare=False - ) - - def __post_init__(self) -> None: - _require_executor_kind(self.executor_kind) - _require_non_empty("task_id", self.task_id) - _require_non_empty("node_id", self.node_id) - _require_optional_label("workspace_ref", self.workspace_ref) - _require_datetime("requested_at", self.requested_at) - _require_optional_id("executor_run_id", self.executor_run_id) - _require_optional_id("approved_workspace_hash", self.approved_workspace_hash) - _require_optional_label("command_category", self.command_category) - _require_optional_label("command_profile", self.command_profile) - _require_optional_id("launch_id", self.launch_id) - _require_optional_id("topic_id", self.topic_id) - if self.timeout_seconds is not None: - _require_positive_int("timeout_seconds", self.timeout_seconds) - _require_optional_label("validation_label", self.validation_label) - _require_json_object("metadata", self.metadata) - - -@dataclass(frozen=True) -class ExecutorEvidence: - evidence_kind: str - evidence_ref: str - summary: str | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_optional_label("evidence_kind", self.evidence_kind) - _require_safe_ref("evidence_ref", self.evidence_ref) - _require_optional_display_text("summary", self.summary) - _require_json_object("metadata", self.metadata) - - def to_safe_dict(self) -> JSONObject: - data: JSONObject = { - "evidence_kind": self.evidence_kind, - "evidence_ref": self.evidence_ref, - "metadata": dict(self.metadata), - } - if self.summary is not None: - data["summary"] = self.summary - return data - - -@dataclass(frozen=True) -class ExecutorResult: - status: ExecutorRunStatus - sanitized_summary: str | None = None - evidence: tuple[ExecutorEvidence, ...] = () - finished_at: datetime | None = None - error_type: str | None = None - error_message: str | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_status("executor result status", self.status) - _require_optional_display_text("sanitized_summary", self.sanitized_summary) - _require_optional_datetime("finished_at", self.finished_at) - _require_optional_label("error_type", self.error_type) - _require_optional_display_text("error_message", self.error_message) - if not isinstance(self.evidence, tuple): - raise TypeError("evidence must be a tuple") - for item in self.evidence: - if not isinstance(item, ExecutorEvidence): - raise TypeError("evidence items must be ExecutorEvidence") - _require_json_object("metadata", self.metadata) - - def evidence_as_json(self) -> list[JSONObject]: - return [item.to_safe_dict() for item in self.evidence] - - -LocalExecutorRunner = Callable[ - [ExecutorPlan, Path], - ExecutorResult | Awaitable[ExecutorResult], -] - - -@dataclass(frozen=True) -class ExecutorRunRecord: - executor_run_id: str - executor_kind: ExecutorKind - task_id: str - node_id: str - status: ExecutorRunStatus - requested_at: datetime - launch_id: str | None = None - topic_id: str | None = None - submitted_at: datetime | None = None - started_at: datetime | None = None - finished_at: datetime | None = None - error_type: str | None = None - error_message: str | None = None - sanitized_summary: str | None = None - evidence: tuple[ExecutorEvidence, ...] = () - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_non_empty("executor_run_id", self.executor_run_id) - _require_executor_kind(self.executor_kind) - _require_non_empty("task_id", self.task_id) - _require_non_empty("node_id", self.node_id) - _require_status("executor run status", self.status) - _require_datetime("requested_at", self.requested_at) - _require_optional_id("launch_id", self.launch_id) - _require_optional_id("topic_id", self.topic_id) - _require_optional_datetime("submitted_at", self.submitted_at) - _require_optional_datetime("started_at", self.started_at) - _require_optional_datetime("finished_at", self.finished_at) - _require_optional_label("error_type", self.error_type) - _require_optional_display_text("error_message", self.error_message) - _require_optional_display_text("sanitized_summary", self.sanitized_summary) - if not isinstance(self.evidence, tuple): - raise TypeError("evidence must be a tuple") - for item in self.evidence: - if not isinstance(item, ExecutorEvidence): - raise TypeError("evidence items must be ExecutorEvidence") - _require_json_object("metadata", self.metadata) - - -class ExecutorRegistry: - def __init__(self) -> None: - self._executors: dict[ExecutorKind, ExternalExecutor] = {} - - def register(self, executor: ExternalExecutor) -> None: - _require_executor_kind(executor.kind) - if executor.kind in self._executors: - raise ValueError(f"executor already registered: {executor.kind}") - self._executors[executor.kind] = executor - - def resolve(self, executor_kind: ExecutorKind) -> ExternalExecutor: - _require_executor_kind(executor_kind) - executor = self._executors.get(executor_kind) - if executor is None: - raise KeyError(f"executor not registered: {executor_kind}") - return executor - - def list_kinds(self) -> tuple[ExecutorKind, ...]: - return tuple(sorted(self._executors)) - - -class LocalExecutorAdapter: - """Local executor adapter over already-approved Bee command bridge plans.""" - - kind: Final[ExecutorKind] = "local" - - def __init__( - self, - *, - store: ExecutorRunStore, - runner: LocalExecutorRunner | None = None, - workspace_resolver: WorkspaceResolver | None = None, - now: Callable[[], datetime], - executor_run_id_factory: Callable[[ExecutorPlan], str] | None = None, - ) -> None: - self._store = store - self._runner = runner - self._workspace_resolver = workspace_resolver - self._now = now - self._executor_run_id_factory = ( - executor_run_id_factory or _default_executor_run_id - ) - - def capability(self) -> ExecutorCapability: - available = self._runner is not None and self._workspace_resolver is not None - return ExecutorCapability( - executor_kind=self.kind, - enabled=True, - available=available, - status="available" if available else "unavailable", - reason=None if available else "local executor dependencies not configured", - ) - - async def submit(self, plan: ExecutorPlan) -> ExecutorResult: - _require_local_executor_plan(plan) - workspace_root = self._resolve_workspace(plan) - if workspace_root is None: - raise ValueError( - f"local executor workspace not found: {plan.workspace_ref}" - ) - _require_approved_workspace_binding(plan, workspace_root) - if self._runner is None: - raise ValueError("local executor runner is not configured") - executor_run_id = plan.executor_run_id or self._executor_run_id_factory(plan) - requested_at = plan.requested_at - await self._store.create_executor_run( - ExecutorRunRecord( - executor_run_id=executor_run_id, - executor_kind=self.kind, - task_id=plan.task_id, - node_id=plan.node_id, - launch_id=plan.launch_id, - topic_id=plan.topic_id, - status="planned", - requested_at=requested_at, - metadata=_executor_run_metadata(plan), - ) - ) - started_at = self._now() - await self._store.update_executor_run_status( - executor_run_id, - status="running", - submitted_at=started_at, - started_at=started_at, - finished_at=None, - error_type=None, - error_message=None, - metadata={"phase": "local_runner_started"}, - ) - try: - result_or_awaitable = self._runner(plan, workspace_root) - result = ( - await result_or_awaitable - if isawaitable(result_or_awaitable) - else result_or_awaitable - ) - result = _sanitize_local_executor_result(result) - except ValueError as exc: - if _is_safety_validation_error(exc): - raise - result = ExecutorResult( - status="failed", - sanitized_summary="Local executor failed with sanitized error", - finished_at=self._now(), - error_type=exc.__class__.__name__, - error_message="local executor runner failed", - metadata={"phase": "local_runner_failed"}, - ) - except Exception as exc: - result = ExecutorResult( - status="failed", - sanitized_summary="Local executor failed with sanitized error", - finished_at=self._now(), - error_type=exc.__class__.__name__, - error_message="local executor runner failed", - metadata={"phase": "local_runner_failed"}, - ) - await self._store.attach_executor_result(executor_run_id, result=result) - return result - - def _resolve_workspace(self, plan: ExecutorPlan) -> Path | None: - if self._workspace_resolver is None: - return None - workspace_root = self._workspace_resolver(plan.workspace_ref) - if workspace_root is None: - return None - resolved = Path(workspace_root).expanduser().resolve() - if not resolved.is_dir(): - return None - return resolved - - -class DockerExecutorAdapter: - """Optional Docker executor adapter in disabled/dry-run mode.""" - - kind: Final[ExecutorKind] = "docker" - - def __init__( - self, - *, - enabled: bool = False, - capability_client: DockerCapabilityClient | None = None, - ) -> None: - self._enabled = enabled - self._capability_client = capability_client - - def capability(self) -> ExecutorCapability: - if not self._enabled: - return ExecutorCapability( - executor_kind=self.kind, - enabled=False, - available=False, - status="disabled", - reason="docker executor disabled", - ) - available = ( - self._capability_client.available() - if self._capability_client is not None - else False - ) - return ExecutorCapability( - executor_kind=self.kind, - enabled=True, - available=available, - status="available" if available else "unavailable", - reason=None if available else "docker capability unavailable", - ) - - async def submit(self, plan: ExecutorPlan) -> ExecutorResult: - raise RuntimeError("docker executor execution is deferred; use dry_run_render") - - def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: - _require_authorized_executor_plan(plan, expected_kind=self.kind) - return { - "executor_kind": self.kind, - "mode": "dry_run", - "task_ref": _safe_ref_token(plan.task_id), - "node_ref": _safe_ref_token(plan.node_id), - "workspace_ref": _safe_ref_token(plan.workspace_ref), - "intent_category": plan.command_category, - "intent_profile": plan.command_profile, - "timeout_seconds": plan.timeout_seconds or 0, - } - - -class KubernetesJobExecutorAdapter: - """Kubernetes Job adapter in disabled/dry-run/fake-status mode.""" - - kind: Final[ExecutorKind] = "kubernetes_job" - - def __init__( - self, - *, - enabled: bool = False, - capability_client: KubernetesCapabilityClient | None = None, - ) -> None: - self._enabled = enabled - self._capability_client = capability_client - - def capability(self) -> ExecutorCapability: - if not self._enabled: - return ExecutorCapability( - executor_kind=self.kind, - enabled=False, - available=False, - status="disabled", - reason="kubernetes job executor disabled", - ) - available = ( - self._capability_client.available() - if self._capability_client is not None - else False - ) - return ExecutorCapability( - executor_kind=self.kind, - enabled=True, - available=available, - status="available" if available else "unavailable", - reason=None if available else "kubernetes capability unavailable", - ) - - async def submit(self, plan: ExecutorPlan) -> ExecutorResult: - raise RuntimeError( - "kubernetes job executor execution is deferred; use dry_run_render" - ) - - def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: - _require_authorized_executor_plan(plan, expected_kind=self.kind) - return { - "apiVersion": "batch/v1", - "kind": "Job", - "metadata": { - "generateName": "coding-agent-executor-", - "labels": { - "app": "coding-agent", - "executor_kind": self.kind, - "intent_category": plan.command_category, - "intent_profile": plan.command_profile, - }, - }, - "spec": { - "backoffLimit": 0, - "template": { - "spec": { - "restartPolicy": "Never", - "containers": [ - { - "name": "executor", - "image": "coding-agent-executor:dry-run", - "args": ["dry-run"], - } - ], - } - }, - }, - "safe_refs": { - "task_ref": _safe_ref_token(plan.task_id), - "node_ref": _safe_ref_token(plan.node_id), - "workspace_ref": _safe_ref_token(plan.workspace_ref), - }, - } - - def import_status(self, plan: ExecutorPlan, status: JSONObject) -> ExecutorResult: - _require_authorized_executor_plan(plan, expected_kind=self.kind) - executor_status = _kubernetes_job_status(status) - return ExecutorResult( - status=executor_status, - sanitized_summary=f"Kubernetes Job {executor_status}", - evidence=( - ExecutorEvidence( - evidence_kind="sanitized_artifact", - evidence_ref="evidence/kubernetes-job-status", - ), - ), - metadata={"phase": "kubernetes_status_imported"}, - ) - - -class ArgoWorkflowExecutorAdapter: - """Argo Workflows adapter in disabled/dry-run/fake-status mode.""" - - kind: Final[ExecutorKind] = "argo_workflow" - - def __init__( - self, - *, - enabled: bool = False, - capability_client: ArgoWorkflowCapabilityClient | None = None, - ) -> None: - self._enabled = enabled - self._capability_client = capability_client - - def capability(self) -> ExecutorCapability: - if not self._enabled: - return ExecutorCapability( - executor_kind=self.kind, - enabled=False, - available=False, - status="disabled", - reason="argo workflow executor disabled", - ) - available = ( - self._capability_client.available() - if self._capability_client is not None - else False - ) - return ExecutorCapability( - executor_kind=self.kind, - enabled=True, - available=available, - status="available" if available else "unavailable", - reason=None if available else "argo workflow capability unavailable", - ) - - async def submit(self, plan: ExecutorPlan) -> ExecutorResult: - raise RuntimeError( - "argo workflow executor execution is deferred; use dry_run_render" - ) - - def dry_run_render(self, plan: ExecutorPlan) -> JSONObject: - _require_authorized_executor_plan(plan, expected_kind=self.kind) - return { - "apiVersion": "argoproj.io/v1alpha1", - "kind": "Workflow", - "metadata": { - "generateName": "coding-agent-executor-", - "labels": { - "app": "coding-agent", - "executor_kind": self.kind, - "intent_category": plan.command_category, - "intent_profile": plan.command_profile, - }, - }, - "spec": { - "entrypoint": "executor", - "templates": [ - { - "name": "executor", - "container": { - "image": "coding-agent-executor:dry-run", - "args": ["dry-run"], - }, - } - ], - }, - "safe_refs": { - "task_ref": _safe_ref_token(plan.task_id), - "node_ref": _safe_ref_token(plan.node_id), - "workspace_ref": _safe_ref_token(plan.workspace_ref), - }, - } - - def import_status(self, plan: ExecutorPlan, status: JSONObject) -> ExecutorResult: - _require_authorized_executor_plan(plan, expected_kind=self.kind) - executor_status = _argo_workflow_status(status) - return ExecutorResult( - status=executor_status, - sanitized_summary=f"Argo Workflow {executor_status}", - evidence=( - ExecutorEvidence( - evidence_kind="sanitized_artifact", - evidence_ref="evidence/argo-workflow-status", - ), - ), - metadata={"phase": "argo_workflow_status_imported"}, - ) - - -def build_local_executor_plan_from_bee_command_plan( - *, - command_plan: BeeCommandIntentPlan, - task_id: str, - workspace_ref: str, - approved_workspace_root: Path | str, - requested_at: datetime, - executor_run_id: str | None = None, - launch_id: str | None = None, - topic_id: str | None = None, - timeout_seconds: int | None = None, -) -> ExecutorPlan: - if command_plan.status != "ready": - raise ValueError("local executor requires a ready Bee command plan") - if not is_authorized_bee_command_plan(command_plan): - raise ValueError("local executor requires authorized Bee command plan") - if command_plan.policy is None: - raise ValueError("local executor requires command policy verdict") - if command_plan.policy.decision != CommandPolicyDecision.ALLOW: - raise ValueError("local executor requires allowed command policy") - if command_plan.approval_route is None: - raise ValueError("local executor requires approval route") - if command_plan.approval_route.route != ActionApprovalRoute.ALLOW: - raise ValueError("local executor requires allow approval route") - intent = command_plan.resolution.intent - if intent is None: - raise ValueError("local executor requires resolved command intent") - return _authorize_executor_plan( - ExecutorPlan( - executor_kind="local", - executor_run_id=executor_run_id, - approved_workspace_hash=_workspace_hash( - Path(approved_workspace_root).expanduser().resolve() - ), - task_id=task_id, - node_id=command_plan.resolution.node_id, - workspace_ref=workspace_ref, - requested_at=requested_at, - command_category=intent.category, - command_profile=intent.profile, - launch_id=launch_id, - topic_id=topic_id, - timeout_seconds=timeout_seconds or command_plan.policy.timeout_seconds, - validation_label=intent.validation_label or intent.name, - metadata={ - "policy_decision": command_plan.policy.decision.value, - "approval_route": command_plan.approval_route.route.value, - "template_id": command_plan.resolution.template_id, - "runtime_kind": command_plan.policy.environment_kind, - }, - authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, - ) - ) - - -def build_docker_executor_plan_from_local_plan( - local_plan: ExecutorPlan, - *, - executor_run_id: str | None = None, -) -> ExecutorPlan: - _require_authorized_executor_plan(local_plan, expected_kind="local") - return _authorize_executor_plan( - ExecutorPlan( - executor_kind="docker", - executor_run_id=executor_run_id, - approved_workspace_hash=local_plan.approved_workspace_hash, - task_id=local_plan.task_id, - node_id=local_plan.node_id, - workspace_ref=local_plan.workspace_ref, - requested_at=local_plan.requested_at, - command_category=local_plan.command_category, - command_profile=local_plan.command_profile, - launch_id=local_plan.launch_id, - topic_id=local_plan.topic_id, - timeout_seconds=local_plan.timeout_seconds, - validation_label=local_plan.validation_label, - metadata={ - "policy_decision": "allow", - "approval_route": "allow", - "source_kind": "local", - }, - authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, - ) - ) - - -def build_kubernetes_job_executor_plan_from_local_plan( - local_plan: ExecutorPlan, - *, - executor_run_id: str | None = None, -) -> ExecutorPlan: - _require_authorized_executor_plan(local_plan, expected_kind="local") - return _authorize_executor_plan( - ExecutorPlan( - executor_kind="kubernetes_job", - executor_run_id=executor_run_id, - approved_workspace_hash=local_plan.approved_workspace_hash, - task_id=local_plan.task_id, - node_id=local_plan.node_id, - workspace_ref=local_plan.workspace_ref, - requested_at=local_plan.requested_at, - command_category=local_plan.command_category, - command_profile=local_plan.command_profile, - launch_id=local_plan.launch_id, - topic_id=local_plan.topic_id, - timeout_seconds=local_plan.timeout_seconds, - validation_label=local_plan.validation_label, - metadata={ - "policy_decision": "allow", - "approval_route": "allow", - "source_kind": "local", - }, - authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, - ) - ) - - -def build_argo_workflow_executor_plan_from_local_plan( - local_plan: ExecutorPlan, - *, - executor_run_id: str | None = None, -) -> ExecutorPlan: - _require_authorized_executor_plan(local_plan, expected_kind="local") - return _authorize_executor_plan( - ExecutorPlan( - executor_kind="argo_workflow", - executor_run_id=executor_run_id, - approved_workspace_hash=local_plan.approved_workspace_hash, - task_id=local_plan.task_id, - node_id=local_plan.node_id, - workspace_ref=local_plan.workspace_ref, - requested_at=local_plan.requested_at, - command_category=local_plan.command_category, - command_profile=local_plan.command_profile, - launch_id=local_plan.launch_id, - topic_id=local_plan.topic_id, - timeout_seconds=local_plan.timeout_seconds, - validation_label=local_plan.validation_label, - metadata={ - "policy_decision": "allow", - "approval_route": "allow", - "source_kind": "local", - }, - authorization_token=_EXECUTOR_PLAN_AUTHORIZATION_TOKEN, - ) - ) - - -def executor_result_completion_evidence( - result: ExecutorResult, -) -> tuple[BeeNodeCompletionEvidence, ...]: - status = "passed" if result.status == "succeeded" else "failed" - return tuple( - BeeNodeCompletionEvidence( - evidence_kind=item.evidence_kind, - evidence_ref=item.evidence_ref, - status=status, - ) - for item in result.evidence - ) - - -def _sanitize_local_executor_result(result: ExecutorResult) -> ExecutorResult: - return ExecutorResult( - status=result.status, - sanitized_summary=_local_executor_status_summary(result.status), - evidence=tuple( - evidence for item in result.evidence if (evidence := _safe_evidence(item)) - ), - finished_at=result.finished_at, - error_type=result.error_type, - error_message=_local_executor_error_summary(result), - metadata={"phase": "local_runner_result"}, - ) - - -def _local_executor_status_summary(status: ExecutorRunStatus) -> str: - if status == "succeeded": - return "Local executor succeeded" - if status == "failed": - return "Local executor failed" - if status == "cancelled": - return "Local executor cancelled" - return "Local executor status recorded" - - -def _local_executor_error_summary(result: ExecutorResult) -> str | None: - if result.status == "failed": - return "local executor returned failed status" - return None - - -def _safe_evidence(item: ExecutorEvidence) -> ExecutorEvidence | None: - try: - return ExecutorEvidence( - evidence_kind=item.evidence_kind, - evidence_ref=item.evidence_ref, - ) - except ValueError: - return None - - -class PGExecutorRunStore: - _CREATE_SCHEMA_SQL: Final[str] = """ - CREATE TABLE IF NOT EXISTS executor_runs ( - executor_run_id TEXT PRIMARY KEY, - executor_kind TEXT NOT NULL, - task_id TEXT NOT NULL, - node_id TEXT NOT NULL, - launch_id TEXT, - topic_id TEXT, - status TEXT NOT NULL, - requested_at TIMESTAMPTZ NOT NULL, - submitted_at TIMESTAMPTZ, - started_at TIMESTAMPTZ, - finished_at TIMESTAMPTZ, - error_type TEXT, - error_message TEXT, - sanitized_summary TEXT, - evidence JSONB NOT NULL DEFAULT '[]'::jsonb, - metadata JSONB NOT NULL DEFAULT '{}'::jsonb - ); - - CREATE INDEX IF NOT EXISTS executor_runs_task_node_idx - ON executor_runs (task_id, node_id, requested_at, executor_run_id); - - CREATE INDEX IF NOT EXISTS executor_runs_kind_status_idx - ON executor_runs (executor_kind, status, requested_at, executor_run_id); - """ - _INSERT_RUN_SQL: Final[str] = """ - INSERT INTO executor_runs ( - executor_run_id, - executor_kind, - task_id, - node_id, - launch_id, - topic_id, - status, - requested_at, - submitted_at, - started_at, - finished_at, - error_type, - error_message, - sanitized_summary, - evidence, - metadata - ) - VALUES ( - $1, $2, $3, $4, $5, $6, $7, $8, - $9, $10, $11, $12, $13, $14, $15::jsonb, $16::jsonb - ) - ON CONFLICT (executor_run_id) DO NOTHING - RETURNING * - """ - _SELECT_RUN_SQL: Final[str] = """ - SELECT * FROM executor_runs WHERE executor_run_id = $1 - """ - _LIST_RUNS_SQL: Final[str] = """ - SELECT * FROM executor_runs - WHERE ($1::text IS NULL OR task_id = $1) - AND ($2::text IS NULL OR node_id = $2) - AND ($3::text IS NULL OR executor_kind = $3) - AND ($4::text IS NULL OR status = $4) - ORDER BY requested_at, executor_run_id - LIMIT $5 - """ - _UPDATE_RUN_STATUS_SQL: Final[str] = """ - UPDATE executor_runs - SET status = $2, - submitted_at = $3, - started_at = $4, - finished_at = $5, - error_type = $6, - error_message = $7, - metadata = $8::jsonb - WHERE executor_run_id = $1 - RETURNING * - """ - _ATTACH_RESULT_SQL: Final[str] = """ - UPDATE executor_runs - SET status = $2, - finished_at = $3, - error_type = $4, - error_message = $5, - sanitized_summary = $6, - evidence = $7::jsonb, - metadata = $8::jsonb - WHERE executor_run_id = $1 - RETURNING * - """ - - def __init__(self, *, pool: PGPool) -> None: - self._pool = pool - self._schema_ready = False - - async def _ensure_schema(self) -> AsyncPGPool: - pool = await self._pool.get_pool() - if not self._schema_ready: - _ = await pool.execute(self._CREATE_SCHEMA_SQL) - self._schema_ready = True - return pool - - async def create_executor_run( - self, - record: ExecutorRunRecord, - ) -> ExecutorRunRecord: - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._INSERT_RUN_SQL, - record.executor_run_id, - record.executor_kind, - record.task_id, - record.node_id, - record.launch_id, - record.topic_id, - record.status, - record.requested_at, - record.submitted_at, - record.started_at, - record.finished_at, - record.error_type, - record.error_message, - record.sanitized_summary, - [item.to_safe_dict() for item in record.evidence], - record.metadata, - ) - if row is None: - raise ValueError(f"executor run already exists: {record.executor_run_id}") - return _run_from_row(row) - - async def load_executor_run( - self, - executor_run_id: str, - ) -> ExecutorRunRecord | None: - _require_non_empty("executor_run_id", executor_run_id) - pool = await self._ensure_schema() - row = await pool.fetchrow(self._SELECT_RUN_SQL, executor_run_id) - return None if row is None else _run_from_row(row) - - async def list_executor_runs( - self, - *, - task_id: str | None = None, - node_id: str | None = None, - executor_kind: ExecutorKind | None = None, - status: ExecutorRunStatus | None = None, - limit: int = 100, - ) -> list[ExecutorRunRecord]: - _require_optional_id("task_id", task_id) - _require_optional_id("node_id", node_id) - if executor_kind is not None: - _require_executor_kind(executor_kind) - if status is not None: - _require_status("executor run status", status) - _require_positive_int("limit", limit) - pool = await self._ensure_schema() - rows = await pool.fetch( - self._LIST_RUNS_SQL, - task_id, - node_id, - executor_kind, - status, - limit, - ) - return [_run_from_row(row) for row in rows] - - async def update_executor_run_status( - self, - executor_run_id: str, - *, - status: ExecutorRunStatus, - submitted_at: datetime | None, - started_at: datetime | None, - finished_at: datetime | None, - error_type: str | None, - error_message: str | None, - metadata: JSONObject, - ) -> ExecutorRunRecord: - _require_non_empty("executor_run_id", executor_run_id) - _require_status("executor run status", status) - _require_optional_datetime("submitted_at", submitted_at) - _require_optional_datetime("started_at", started_at) - _require_optional_datetime("finished_at", finished_at) - _require_optional_label("error_type", error_type) - _require_optional_display_text("error_message", error_message) - _require_json_object("metadata", metadata) - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._UPDATE_RUN_STATUS_SQL, - executor_run_id, - status, - submitted_at, - started_at, - finished_at, - error_type, - error_message, - metadata, - ) - if row is None: - raise KeyError(f"executor run not found: {executor_run_id}") - return _run_from_row(row) - - async def attach_executor_result( - self, - executor_run_id: str, - *, - result: ExecutorResult, - ) -> ExecutorRunRecord: - _require_non_empty("executor_run_id", executor_run_id) - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._ATTACH_RESULT_SQL, - executor_run_id, - result.status, - result.finished_at, - result.error_type, - result.error_message, - result.sanitized_summary, - result.evidence_as_json(), - result.metadata, - ) - if row is None: - raise KeyError(f"executor run not found: {executor_run_id}") - return _run_from_row(row) - - -def _run_from_row(row: dict[str, object]) -> ExecutorRunRecord: - return ExecutorRunRecord( - executor_run_id=_required_str(row, "executor_run_id", context="executor run"), - executor_kind=_required_str(row, "executor_kind", context="executor run"), - task_id=_required_str(row, "task_id", context="executor run"), - node_id=_required_str(row, "node_id", context="executor run"), - launch_id=_optional_str(row, "launch_id", context="executor run"), - topic_id=_optional_str(row, "topic_id", context="executor run"), - status=_required_str(row, "status", context="executor run"), - requested_at=_required_datetime(row, "requested_at", context="executor run"), - submitted_at=_optional_datetime(row, "submitted_at", context="executor run"), - started_at=_optional_datetime(row, "started_at", context="executor run"), - finished_at=_optional_datetime(row, "finished_at", context="executor run"), - error_type=_optional_str(row, "error_type", context="executor run"), - error_message=_optional_str(row, "error_message", context="executor run"), - sanitized_summary=_optional_str( - row, "sanitized_summary", context="executor run" - ), - evidence=_required_evidence_tuple(row, "evidence", context="executor run"), - metadata=_required_json_object(row, "metadata", context="executor run"), - ) - - -def _required_evidence_tuple( - row: dict[str, object], - key: str, - *, - context: str, -) -> tuple[ExecutorEvidence, ...]: - value = row.get(key) - if not isinstance(value, list): - raise TypeError(f"postgres {context} must include list {key}") - return tuple(_evidence_from_json(item, context=context) for item in value) - - -def _evidence_from_json(value: object, *, context: str) -> ExecutorEvidence: - if not isinstance(value, dict): - raise TypeError(f"postgres {context} evidence must include objects") - evidence_kind = value.get("evidence_kind") - evidence_ref = value.get("evidence_ref") - summary = value.get("summary") - metadata = value.get("metadata", {}) - if not isinstance(evidence_kind, str): - raise TypeError(f"postgres {context} evidence missing evidence_kind") - if not isinstance(evidence_ref, str): - raise TypeError(f"postgres {context} evidence missing evidence_ref") - if summary is not None and not isinstance(summary, str): - raise TypeError(f"postgres {context} evidence summary must be a string") - if not isinstance(metadata, dict): - raise TypeError(f"postgres {context} evidence metadata must be an object") - return ExecutorEvidence( - evidence_kind=evidence_kind, - evidence_ref=evidence_ref, - summary=summary, - metadata=metadata, - ) - - -def _require_executor_kind(value: str) -> None: - if value not in _EXECUTOR_KINDS: - raise ValueError(f"executor kind must be one of {sorted(_EXECUTOR_KINDS)}") - - -def _require_local_executor_plan(plan: ExecutorPlan) -> None: - _require_authorized_executor_plan(plan, expected_kind="local") - if plan.approved_workspace_hash is None: - raise ValueError("local executor requires approved workspace binding") - - -def _require_authorized_executor_plan( - plan: ExecutorPlan, - *, - expected_kind: ExecutorKind, -) -> None: - if plan.executor_kind != expected_kind: - raise ValueError( - f"{expected_kind} executor requires executor_kind={expected_kind!r}" - ) - if plan.authorization_token is not _EXECUTOR_PLAN_AUTHORIZATION_TOKEN: - raise ValueError(f"{expected_kind} executor requires authorized executor plan") - if plan.authorization_signature != _executor_plan_signature(plan): - raise ValueError(f"{expected_kind} executor requires signed executor plan") - if plan.metadata.get("policy_decision") != "allow": - raise ValueError(f"{expected_kind} executor requires allowed command policy") - if plan.metadata.get("approval_route") != "allow": - raise ValueError(f"{expected_kind} executor requires allow approval route") - - -def _require_approved_workspace_binding( - plan: ExecutorPlan, workspace_root: Path -) -> None: - if plan.approved_workspace_hash != _workspace_hash(workspace_root): - raise ValueError("local executor workspace binding mismatch") - - -def _executor_run_metadata(plan: ExecutorPlan) -> JSONObject: - return { - "phase": "local_executor_planned", - "intent_category": plan.command_category, - "intent_profile": plan.command_profile, - } - - -def _default_executor_run_id(plan: ExecutorPlan) -> str: - return f"executor-{plan.task_id}-{plan.node_id}" - - -def _authorize_executor_plan(plan: ExecutorPlan) -> ExecutorPlan: - object.__setattr__( - plan, - "authorization_signature", - _executor_plan_signature(plan), - ) - return plan - - -def _executor_plan_signature(plan: ExecutorPlan) -> str: - metadata_items = tuple( - sorted((key, str(value)) for key, value in plan.metadata.items()) - ) - parts = ( - plan.executor_kind, - plan.executor_run_id or "", - plan.approved_workspace_hash or "", - plan.task_id, - plan.node_id, - plan.workspace_ref, - plan.command_category, - plan.command_profile, - plan.launch_id or "", - plan.topic_id or "", - str(plan.timeout_seconds or ""), - plan.validation_label or "", - repr(metadata_items), - ) - return hashlib.sha256("\x1f".join(parts).encode("utf-8")).hexdigest() - - -def _workspace_hash(path: Path) -> str: - return hashlib.sha256(str(path).encode("utf-8")).hexdigest()[:16] - - -def _safe_ref_token(value: str) -> str: - return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16] - - -def _kubernetes_job_status(status: JSONObject) -> ExecutorRunStatus: - if status.get("succeeded") is True or status.get("condition") == "complete": - return "succeeded" - if status.get("failed") is True or status.get("condition") == "failed": - return "failed" - return "running" - - -def _argo_workflow_status(status: JSONObject) -> ExecutorRunStatus: - phase = status.get("phase") - if phase == "Succeeded" or status.get("succeeded") is True: - return "succeeded" - if phase in {"Failed", "Error"} or status.get("failed") is True: - return "failed" - return "running" - - -def _is_safety_validation_error(exc: ValueError) -> bool: - message = str(exc) - return ( - "raw execution text" in message - or "secret-shaped" in message - or "forbidden metadata key" in message - or "sensitive reference text" in message - or "must not contain whitespace" in message - ) - - -def _require_non_empty(field_name: str, value: str) -> None: - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"{field_name} must be a non-empty string") - - -def _require_optional_id(field_name: str, value: str | None) -> None: - if value is None: - return - _require_non_empty(field_name, value) - - -def _require_optional_label(field_name: str, value: str | None) -> None: - if value is None: - return - _require_non_empty(field_name, value) - if len(value) > _MAX_SAFE_LABEL_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_SAFE_LABEL_CHARS} characters" - ) - folded = value.casefold() - if any(part in folded for part in _FORBIDDEN_LABEL_VALUE_PARTS): - raise ValueError(f"{field_name} must not contain sensitive label text") - _reject_secret_shaped_value(field_name, value) - - -def _require_safe_ref(field_name: str, value: str) -> None: - _require_non_empty(field_name, value) - if len(value) > _MAX_DISPLAY_TEXT_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" - ) - if any(character.isspace() for character in value): - raise ValueError(f"{field_name} must not contain whitespace") - _reject_secret_shaped_value(field_name, value) - folded = value.casefold() - if any(part in folded for part in _FORBIDDEN_REF_PARTS): - raise ValueError(f"{field_name} must not contain sensitive reference text") - _reject_raw_execution_text(field_name, value) - - -def _require_status(field_name: str, value: str) -> None: - if value not in _EXECUTOR_RUN_STATUSES: - raise ValueError( - f"{field_name} must be one of {sorted(_EXECUTOR_RUN_STATUSES)}" - ) - - -def _require_datetime(field_name: str, value: datetime) -> None: - if not isinstance(value, datetime): - raise TypeError(f"{field_name} must be a datetime") - - -def _require_optional_datetime(field_name: str, value: datetime | None) -> None: - if value is not None: - _require_datetime(field_name, value) - - -def _require_optional_display_text(field_name: str, value: str | None) -> None: - if value is None: - return - _require_non_empty(field_name, value) - if len(value) > _MAX_DISPLAY_TEXT_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" - ) - _reject_secret_shaped_value(field_name, value) - _reject_raw_execution_text(field_name, value) - - -def _require_positive_int(field_name: str, value: int) -> None: - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError(f"{field_name} must be a positive int") - - -def _require_json_object(field_name: str, value: JSONObject) -> None: - if not isinstance(value, dict): - raise TypeError(f"{field_name} must be a JSON object") - for key, item in value.items(): - if not isinstance(key, str): - raise TypeError(f"{field_name} keys must be strings") - key_folded = key.casefold() - if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): - raise ValueError(f"{field_name} contains forbidden metadata key: {key}") - _require_json_value(f"{field_name}.{key}", item) - - -def _require_json_value(field_name: str, value: JSONValue) -> None: - if value is None or isinstance(value, str | int | float | bool): - if isinstance(value, str): - if len(value) > _MAX_METADATA_STRING_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_METADATA_STRING_CHARS} characters" - ) - _reject_secret_shaped_value(field_name, value) - _reject_raw_execution_text(field_name, value) - return - if isinstance(value, list): - for index, item in enumerate(value): - _require_json_value(f"{field_name}[{index}]", item) - return - if isinstance(value, dict): - _require_json_object(field_name, value) - return - raise TypeError(f"{field_name} must be JSON-safe") - - -def _reject_secret_shaped_value(field_name: str, value: str) -> None: - folded = value.casefold() - if any(marker in folded for marker in _SECRET_VALUE_MARKERS): - raise ValueError(f"{field_name} must not contain secret-shaped values") - - -def _reject_raw_execution_text(field_name: str, value: str) -> None: - folded = value.casefold() - if any(marker in folded for marker in _RAW_EXECUTION_TEXT_MARKERS): - raise ValueError(f"{field_name} must not contain raw execution text") - - -def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: - if row is None: - raise RuntimeError(f"postgres did not return a row for {context}") - return row - - -def _required_str(row: dict[str, object], key: str, *, context: str) -> str: - value = row.get(key) - if not isinstance(value, str): - raise TypeError(f"postgres {context} must include string {key}") - return value - - -def _optional_str(row: dict[str, object], key: str, *, context: str) -> str | None: - value = row.get(key) - if value is None: - return None - if not isinstance(value, str): - raise TypeError(f"postgres {context} must include string or None {key}") - return value - - -def _required_datetime(row: dict[str, object], key: str, *, context: str) -> datetime: - value = row.get(key) - if not isinstance(value, datetime): - raise TypeError(f"postgres {context} must include datetime {key}") - return value - - -def _optional_datetime( - row: dict[str, object], - key: str, - *, - context: str, -) -> datetime | None: - value = row.get(key) - if value is None: - return None - if not isinstance(value, datetime): - raise TypeError(f"postgres {context} must include datetime or None {key}") - return value - - -def _required_json_object( - row: dict[str, object], - key: str, - *, - context: str, -) -> JSONObject: - value = row.get(key) - if not isinstance(value, dict): - raise TypeError(f"postgres {context} must include dict {key}") - _require_json_object(key, value) - return value +"""Deprecated import location. Moved to coding_agent.executors.external.""" +import sys as _sys +from coding_agent.executors import external as _moved +_sys.modules[__name__] = _moved diff --git a/coding-agent/src/coding_agent/local_durable_store.py b/coding-agent/src/coding_agent/local_durable_store.py index e80b378f6..447c04bc7 100644 --- a/coding-agent/src/coding_agent/local_durable_store.py +++ b/coding-agent/src/coding_agent/local_durable_store.py @@ -1,1284 +1,4 @@ -from __future__ import annotations - -import json -import sqlite3 -import threading -from collections.abc import Callable, Mapping -from datetime import UTC, datetime -from pathlib import Path -from typing import Any, Final, cast - -from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot -from agentkit.storage.protocols import TapeInfo, TapeSearchResult -from agentkit.storage.sqlite import ( - SQLiteCheckpointStore, - SQLiteTapeStore, - _checkpoint_snapshot_from_payload, - _checkpoint_snapshot_to_payload, - _entry_anchor_type, - _entry_nested_str, - _json_object_from_text, - _optional_entry_str, -) - -from coding_agent.runtime_store import ( - AgentRunRecord, - AgentInteractionRecord, - RuntimeEventRecord, - RunMessageSnapshotRecord, - SQLiteRuntimeStore, - _agent_run_from_sqlite_row, - _agent_run_sqlite_values, - _datetime_to_json, - _interaction_from_sqlite_row, - _interaction_sqlite_values, - _json_to_sql, - _runtime_event_from_sqlite_row, -) -from coding_agent.server.stores.session_owner_store import ( - OwnerAuthority, - SQLiteSessionOwnerStore, - SessionOwnershipConflictError, - SessionOwnershipConflictReason, - _datetime_from_sqlite_text, - _datetime_to_sqlite_text, -) -from coding_agent.server.stores.session_store import SQLiteSessionStore, SessionPayload - - -class SQLiteLocalDurableStore: - """SQLite-local protected mutation facade. - - This facade is intentionally local to the SQLite bundle path. It does not replace - the generic store protocols; it provides the transaction shape required by local - durable fencing: owner epoch check, target ownership check, and mutation in the - same SQLite transaction. - """ - - _SESSION_TAPES_SCHEMA_SQL: Final[str] = """ - CREATE TABLE IF NOT EXISTS session_tapes ( - session_id TEXT PRIMARY KEY, - tape_id TEXT UNIQUE NOT NULL, - created_at TEXT NOT NULL, - updated_at TEXT NOT NULL - ); - CREATE INDEX IF NOT EXISTS session_tapes_tape_id_idx - ON session_tapes (tape_id); - """ - - def __init__(self, path: Path) -> None: - self._path = path - self._lock = threading.Lock() - self._path.parent.mkdir(parents=True, exist_ok=True) - self._owner_store = SQLiteSessionOwnerStore(path) - with self._connect() as connection: - connection.executescript(SQLiteSessionStore._CREATE_TABLE_SQL) - connection.executescript(SQLiteTapeStore._CREATE_SCHEMA_SQL) - connection.executescript(SQLiteCheckpointStore._CREATE_SCHEMA_SQL) - connection.executescript(SQLiteRuntimeStore._CREATE_SCHEMA_SQL) - connection.executescript(self._SESSION_TAPES_SCHEMA_SQL) - - def _connect(self) -> sqlite3.Connection: - connection = sqlite3.connect(self._path, timeout=5.0) - connection.row_factory = sqlite3.Row - connection.execute("PRAGMA busy_timeout = 5000") - return connection - - async def acquire_owner( - self, - session_id: str, - owner_id: str, - lease_seconds: float = 30.0, - ) -> OwnerAuthority: - return await self._owner_store.acquire_authority( - session_id, - owner_id, - lease_seconds=lease_seconds, - ) - - async def renew_owner( - self, - authority: OwnerAuthority, - lease_seconds: float = 30.0, - ) -> OwnerAuthority: - return await self._owner_store.renew_authority( - authority, - lease_seconds=lease_seconds, - ) - - async def save_session( - self, - authority: OwnerAuthority, - payload: SessionPayload, - ) -> None: - _require_json_object("session payload", payload) - payload_id = payload.get("id") - if payload_id != authority.session_id: - raise SessionOwnershipConflictError( - "session payload belongs to another owner" - ) - payload_session_id = payload.get("session_id") - if ( - payload_session_id is not None - and payload_session_id != authority.session_id - ): - raise SessionOwnershipConflictError( - "session payload belongs to another owner" - ) - tape_id = payload.get("tape_id") - if tape_id is not None and not isinstance(tape_id, str): - raise TypeError("session payload tape_id must be a string") - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - if tape_id: - self._bind_tape(connection, authority.session_id, tape_id) - connection.execute( - """ - INSERT INTO agent_http_sessions (session_id, payload, updated_at) - VALUES (?, ?, CURRENT_TIMESTAMP) - ON CONFLICT(session_id) - DO UPDATE SET payload = excluded.payload, - updated_at = CURRENT_TIMESTAMP - """, - (authority.session_id, json.dumps(payload, sort_keys=True)), - ) - - async def delete_session(self, authority: OwnerAuthority) -> None: - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - connection.execute( - "DELETE FROM agent_http_sessions WHERE session_id = ?", - (authority.session_id,), - ) - connection.execute( - "DELETE FROM session_tapes WHERE session_id = ?", - (authority.session_id,), - ) - - def load_session(self, session_id: str) -> SessionPayload | None: - _require_non_empty("session_id", session_id) - with self._lock, self._connect() as connection: - row = connection.execute( - "SELECT payload FROM agent_http_sessions WHERE session_id = ?", - (session_id,), - ).fetchone() - if row is None: - return None - payload = json.loads(str(row["payload"])) - if not isinstance(payload, dict): - raise TypeError("sqlite session payload must be a JSON object") - return cast(SessionPayload, payload) - - async def append_tape_entries( - self, - authority: OwnerAuthority, - tape_id: str, - entries: list[dict[str, Any]], - ) -> None: - _require_non_empty("tape_id", tape_id) - if not entries: - return - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_tape_belongs_to_session( - connection, - tape_id=tape_id, - session_id=authority.session_id, - ) - row = connection.execute( - "SELECT COALESCE(MAX(seq), -1) AS max_seq FROM tape_entries WHERE tape_id = ?", - (tape_id,), - ).fetchone() - max_seq = _row_required_int(row, "max_seq", context="tape max seq") - values: list[tuple[object, ...]] = [] - now = datetime.now(UTC).isoformat() - for offset, entry in enumerate(entries, start=1): - _require_json_object("tape entry", entry) - values.append( - ( - tape_id, - max_seq + offset, - json.dumps(entry, sort_keys=True), - _optional_entry_str(entry, "kind"), - _entry_nested_str(entry, "run_id"), - _entry_nested_str(entry, "tool_call_id"), - _entry_anchor_type(entry), - now, - ) - ) - connection.executemany( - """ - INSERT INTO tape_entries ( - tape_id, seq, entry_json, kind, run_id, tool_call_id, - anchor_type, created_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, - values, - ) - - async def load_tape(self, tape_id: str) -> list[dict[str, Any]]: - _require_non_empty("tape_id", tape_id) - with self._lock, self._connect() as connection: - rows = connection.execute( - """ - SELECT entry_json - FROM tape_entries - WHERE tape_id = ? - ORDER BY seq - """, - (tape_id,), - ).fetchall() - return [ - _json_object_from_text(row["entry_json"], context="tape entry") - for row in rows - ] - - async def truncate_tape( - self, - authority: OwnerAuthority, - tape_id: str, - keep: int, - ) -> None: - _require_non_empty("tape_id", tape_id) - if keep < 0: - raise ValueError("keep must be >= 0") - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_tape_belongs_to_session( - connection, - tape_id=tape_id, - session_id=authority.session_id, - ) - connection.execute( - "DELETE FROM tape_entries WHERE tape_id = ? AND seq >= ?", - (tape_id, keep), - ) - - def session_id_for_tape(self, tape_id: str) -> str | None: - _require_non_empty("tape_id", tape_id) - with self._lock, self._connect() as connection: - row = connection.execute( - "SELECT session_id FROM session_tapes WHERE tape_id = ?", - (tape_id,), - ).fetchone() - if row is None: - return None - session_id = row["session_id"] - if not isinstance(session_id, str): - raise TypeError("session_tapes session_id must be text") - return session_id - - async def create_agent_run( - self, - authority: OwnerAuthority, - record: AgentRunRecord, - ) -> AgentRunRecord: - if record.session_id != authority.session_id: - raise SessionOwnershipConflictError("agent run belongs to another session") - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - existing = connection.execute( - "SELECT session_id FROM agent_runs WHERE run_id = ?", - (record.run_id,), - ).fetchone() - if existing is not None and existing["session_id"] != authority.session_id: - raise SessionOwnershipConflictError( - "agent run target belongs to another session" - ) - connection.execute( - """ - INSERT INTO agent_runs ( - run_id, session_id, tape_id, parent_run_id, agent_id, status, - started_at, ended_at, metadata, result, error - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(run_id) - DO UPDATE SET - tape_id = excluded.tape_id, - parent_run_id = excluded.parent_run_id, - agent_id = excluded.agent_id, - status = excluded.status, - started_at = excluded.started_at, - ended_at = excluded.ended_at, - metadata = excluded.metadata, - result = excluded.result, - error = excluded.error - """, - _agent_run_sqlite_values(record), - ) - return record - - async def update_agent_run( - self, - authority: OwnerAuthority, - run_id: str, - *, - status: str, - ended_at: datetime | None, - metadata: dict[str, Any], - result: dict[str, Any], - error: str | None, - ) -> AgentRunRecord: - _require_non_empty("run_id", run_id) - _require_non_empty("status", status) - _require_json_object("metadata", metadata) - _require_json_object("result", result) - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_run_belongs_to_session( - connection, - run_id=run_id, - session_id=authority.session_id, - ) - connection.execute( - """ - UPDATE agent_runs - SET status = ?, - ended_at = ?, - metadata = ?, - result = ?, - error = ? - WHERE run_id = ? - """, - ( - status, - None if ended_at is None else _datetime_to_json(ended_at), - _json_to_sql(cast(dict[str, Any], metadata)), - _json_to_sql(cast(dict[str, Any], result)), - error, - run_id, - ), - ) - row = connection.execute( - "SELECT * FROM agent_runs WHERE run_id = ?", - (run_id,), - ).fetchone() - if row is None: - raise KeyError(f"agent run not found after update: {run_id}") - return _agent_run_from_sqlite_row(row) - - async def append_runtime_event( - self, - authority: OwnerAuthority, - record: RuntimeEventRecord, - ) -> RuntimeEventRecord: - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_run_belongs_to_session( - connection, - run_id=record.run_id, - session_id=authority.session_id, - ) - existing = connection.execute( - """ - SELECT runtime_events.* - FROM runtime_events - JOIN agent_runs ON agent_runs.run_id = runtime_events.run_id - WHERE runtime_events.event_id = ? - """, - (record.event_id,), - ).fetchone() - if existing is not None: - if existing["run_id"] != record.run_id: - raise SessionOwnershipConflictError( - "runtime event target belongs to another run" - ) - return _runtime_event_from_sqlite_row(existing) - connection.execute( - """ - INSERT INTO runtime_events ( - event_id, run_id, event_kind, payload, created_at - ) - VALUES (?, ?, ?, ?, ?) - """, - ( - record.event_id, - record.run_id, - record.event_kind, - _json_to_sql(record.payload), - _datetime_to_json(record.created_at), - ), - ) - row = connection.execute( - "SELECT * FROM runtime_events WHERE event_id = ?", - (record.event_id,), - ).fetchone() - if row is None: - raise RuntimeError("sqlite runtime event insert returned no row") - return _runtime_event_from_sqlite_row(row) - - async def replay_runtime_events( - self, - run_id: str, - *, - after_sequence: int = 0, - limit: int = 1000, - ) -> list[RuntimeEventRecord]: - _require_non_empty("run_id", run_id) - if after_sequence < 0: - raise ValueError("after_sequence must be >= 0") - if limit <= 0: - raise ValueError("limit must be positive") - with self._lock, self._connect() as connection: - rows = connection.execute( - """ - SELECT * FROM runtime_events - WHERE run_id = ? AND sequence > ? - ORDER BY sequence - LIMIT ? - """, - (run_id, after_sequence, limit), - ).fetchall() - return [_runtime_event_from_sqlite_row(row) for row in rows] - - async def save_message_snapshot( - self, - authority: OwnerAuthority, - record: RunMessageSnapshotRecord, - ) -> RunMessageSnapshotRecord: - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_run_belongs_to_session( - connection, - run_id=record.run_id, - session_id=authority.session_id, - ) - existing = connection.execute( - """ - SELECT run_message_snapshots.run_id, agent_runs.session_id - FROM run_message_snapshots - JOIN agent_runs ON agent_runs.run_id = run_message_snapshots.run_id - WHERE run_message_snapshots.snapshot_id = ? - """, - (record.snapshot_id,), - ).fetchone() - if existing is not None and ( - existing["run_id"] != record.run_id - or existing["session_id"] != authority.session_id - ): - raise SessionOwnershipConflictError( - "message snapshot target belongs to another session" - ) - connection.execute( - """ - INSERT INTO run_message_snapshots ( - snapshot_id, run_id, messages, metadata, created_at - ) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT(snapshot_id) - DO UPDATE SET - messages = excluded.messages, - metadata = excluded.metadata, - created_at = excluded.created_at - """, - ( - record.snapshot_id, - record.run_id, - _json_to_sql(cast(list[Any], record.messages)), - _json_to_sql(record.metadata), - _datetime_to_json(record.created_at), - ), - ) - return record - - async def create_agent_interaction( - self, - authority: OwnerAuthority, - record: AgentInteractionRecord, - ) -> AgentInteractionRecord: - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_run_belongs_to_session( - connection, - run_id=record.run_id, - session_id=authority.session_id, - ) - existing = connection.execute( - """ - SELECT agent_interactions.run_id, agent_runs.session_id - FROM agent_interactions - JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id - WHERE agent_interactions.interaction_id = ? - """, - (record.interaction_id,), - ).fetchone() - if existing is not None: - if ( - existing["run_id"] != record.run_id - or existing["session_id"] != authority.session_id - ): - raise SessionOwnershipConflictError( - "interaction target belongs to another session" - ) - row = connection.execute( - "SELECT * FROM agent_interactions WHERE interaction_id = ?", - (record.interaction_id,), - ).fetchone() - if row is None: - raise RuntimeError( - "sqlite agent interaction lookup returned no row" - ) - return _interaction_from_sqlite_row(row) - connection.execute( - """ - INSERT INTO agent_interactions ( - interaction_id, run_id, interaction_kind, status, - request_payload, response_payload, metadata, - created_at, resolved_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - _interaction_sqlite_values(record), - ) - row = connection.execute( - "SELECT * FROM agent_interactions WHERE interaction_id = ?", - (record.interaction_id,), - ).fetchone() - if row is None: - raise RuntimeError("sqlite agent interaction insert returned no row") - return _interaction_from_sqlite_row(row) - - async def resolve_agent_interaction( - self, - authority: OwnerAuthority, - interaction_id: str, - *, - status: str, - response_payload: dict[str, Any], - resolved_at: datetime, - ) -> AgentInteractionRecord: - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - existing = connection.execute( - """ - SELECT agent_interactions.* - FROM agent_interactions - JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id - WHERE agent_interactions.interaction_id = ? - AND agent_runs.session_id = ? - """, - (interaction_id, authority.session_id), - ).fetchone() - if existing is None: - raise SessionOwnershipConflictError( - "interaction target belongs to another session" - ) - if existing["resolved_at"] is None: - connection.execute( - """ - UPDATE agent_interactions - SET status = ?, - response_payload = ?, - resolved_at = ? - WHERE interaction_id = ? - """, - ( - status, - _json_to_sql(response_payload), - _datetime_to_json(resolved_at), - interaction_id, - ), - ) - row = connection.execute( - "SELECT * FROM agent_interactions WHERE interaction_id = ?", - (interaction_id,), - ).fetchone() - if row is None: - raise RuntimeError("sqlite agent interaction resolve returned no row") - return _interaction_from_sqlite_row(row) - - async def claim_attached_executor_run( - self, - authorities: Mapping[str, OwnerAuthority], - *, - session_id: str | None, - executor_kind: str, - claim_metadata: dict[str, Any], - ) -> AgentRunRecord | None: - _require_non_empty("executor_kind", executor_kind) - _require_json_object("claim_metadata", claim_metadata) - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - filter_sql = "" - values: list[object] = [] - if session_id is not None: - filter_sql = "AND session_id = ?" - values.append(session_id) - rows = connection.execute( - f""" - SELECT * FROM agent_runs - WHERE status IN ('requested', 'expired') - {filter_sql} - ORDER BY started_at, run_id - """, - values, - ).fetchall() - candidates = [_agent_run_from_sqlite_row(row) for row in rows] - candidates = [ - run - for run in candidates - if run.metadata.get("executor_ref_kind") - in {"external_worker", "local_attached"} - and run.metadata.get("executor_kind") == executor_kind - ] - selected: AgentRunRecord | None = None - authority: OwnerAuthority | None = None - for candidate in candidates: - candidate_authority = authorities.get(candidate.session_id) - if candidate_authority is None: - continue - try: - self._assert_authority(connection, candidate_authority) - except SessionOwnershipConflictError: - continue - selected = candidate - authority = candidate_authority - break - if selected is None or authority is None: - return None - metadata = cast(dict[str, Any], {**selected.metadata, **claim_metadata}) - connection.execute( - """ - UPDATE agent_runs - SET status = 'claimed', - metadata = ? - WHERE run_id = ? - AND session_id = ? - """, - ( - _json_to_sql(metadata), - selected.run_id, - authority.session_id, - ), - ) - row = connection.execute( - "SELECT * FROM agent_runs WHERE run_id = ?", - (selected.run_id,), - ).fetchone() - if row is None: - raise KeyError(f"agent run not found after claim: {selected.run_id}") - return _agent_run_from_sqlite_row(row) - - def session_id_for_run(self, run_id: str) -> str | None: - _require_non_empty("run_id", run_id) - with self._lock, self._connect() as connection: - row = connection.execute( - "SELECT session_id FROM agent_runs WHERE run_id = ?", - (run_id,), - ).fetchone() - if row is None: - return None - session_id = row["session_id"] - if not isinstance(session_id, str): - raise TypeError("agent run session_id must be text") - return session_id - - async def save_checkpoint( - self, - authority: OwnerAuthority, - snapshot: CheckpointSnapshot, - ) -> None: - meta = snapshot.meta - if meta.session_id != authority.session_id: - raise SessionOwnershipConflictError("checkpoint belongs to another session") - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_tape_belongs_to_session( - connection, - tape_id=meta.tape_id, - session_id=authority.session_id, - ) - existing = connection.execute( - """ - SELECT session_id, tape_id - FROM checkpoints - WHERE checkpoint_id = ? - """, - (meta.checkpoint_id,), - ).fetchone() - if existing is not None and ( - existing["session_id"] != authority.session_id - or existing["tape_id"] != meta.tape_id - ): - raise SessionOwnershipConflictError( - "checkpoint target belongs to another session" - ) - payload = _checkpoint_snapshot_to_payload(snapshot) - connection.execute( - """ - INSERT INTO checkpoints ( - checkpoint_id, tape_id, session_id, entry_count, window_start, - label, snapshot_json, created_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(checkpoint_id) - DO UPDATE SET - entry_count = excluded.entry_count, - window_start = excluded.window_start, - label = excluded.label, - snapshot_json = excluded.snapshot_json, - created_at = excluded.created_at - """, - ( - meta.checkpoint_id, - meta.tape_id, - meta.session_id, - meta.entry_count, - meta.window_start, - meta.label, - json.dumps(payload, sort_keys=True), - meta.created_at.isoformat(), - ), - ) - - async def load_checkpoint(self, checkpoint_id: str) -> CheckpointSnapshot | None: - _require_non_empty("checkpoint_id", checkpoint_id) - with self._lock, self._connect() as connection: - row = connection.execute( - "SELECT snapshot_json FROM checkpoints WHERE checkpoint_id = ?", - (checkpoint_id,), - ).fetchone() - if row is None: - return None - return _checkpoint_snapshot_from_payload( - _json_object_from_text(row["snapshot_json"], context="checkpoint snapshot") - ) - - async def delete_checkpoint( - self, - authority: OwnerAuthority, - checkpoint_id: str, - ) -> None: - _require_non_empty("checkpoint_id", checkpoint_id) - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - row = connection.execute( - "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", - (checkpoint_id,), - ).fetchone() - if row is None: - return - if row["session_id"] != authority.session_id: - raise SessionOwnershipConflictError( - "checkpoint target belongs to another session" - ) - connection.execute( - "DELETE FROM checkpoints WHERE checkpoint_id = ?", - (checkpoint_id,), - ) - - async def restore_checkpoint_state( - self, - authority: OwnerAuthority, - snapshot: CheckpointSnapshot, - session_payload: SessionPayload, - ) -> None: - meta = snapshot.meta - if meta.session_id != authority.session_id: - raise SessionOwnershipConflictError("checkpoint belongs to another session") - if session_payload.get("id") != authority.session_id: - raise SessionOwnershipConflictError( - "session payload belongs to another owner" - ) - if session_payload.get("tape_id") != meta.tape_id: - raise SessionOwnershipConflictError( - "checkpoint restore session payload has mismatched tape id" - ) - with self._lock, self._connect() as connection: - connection.execute("BEGIN IMMEDIATE") - self._assert_authority(connection, authority) - self._assert_tape_belongs_to_session( - connection, - tape_id=meta.tape_id, - session_id=authority.session_id, - ) - checkpoint_row = connection.execute( - """ - SELECT session_id, tape_id - FROM checkpoints - WHERE checkpoint_id = ? - """, - (meta.checkpoint_id,), - ).fetchone() - if checkpoint_row is None: - raise KeyError(f"checkpoint not found: {meta.checkpoint_id}") - if ( - checkpoint_row["session_id"] != authority.session_id - or checkpoint_row["tape_id"] != meta.tape_id - ): - raise SessionOwnershipConflictError( - "checkpoint target belongs to another session" - ) - connection.execute( - "DELETE FROM tape_entries WHERE tape_id = ?", - (meta.tape_id,), - ) - now = datetime.now(UTC).isoformat() - values: list[tuple[object, ...]] = [] - for seq, entry in enumerate(snapshot.tape_entries): - _require_json_object("tape entry", entry) - values.append( - ( - meta.tape_id, - seq, - json.dumps(entry, sort_keys=True), - _optional_entry_str(entry, "kind"), - _entry_nested_str(entry, "run_id"), - _entry_nested_str(entry, "tool_call_id"), - _entry_anchor_type(entry), - now, - ) - ) - connection.executemany( - """ - INSERT INTO tape_entries ( - tape_id, seq, entry_json, kind, run_id, tool_call_id, - anchor_type, created_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, - values, - ) - connection.execute( - """ - INSERT INTO agent_http_sessions (session_id, payload, updated_at) - VALUES (?, ?, CURRENT_TIMESTAMP) - ON CONFLICT(session_id) - DO UPDATE SET payload = excluded.payload, - updated_at = CURRENT_TIMESTAMP - """, - (authority.session_id, json.dumps(session_payload, sort_keys=True)), - ) - connection.execute( - """ - DELETE FROM checkpoints - WHERE tape_id = ? - AND session_id = ? - AND entry_count > ? - """, - (meta.tape_id, authority.session_id, meta.entry_count), - ) - - def session_id_for_checkpoint(self, checkpoint_id: str) -> str | None: - _require_non_empty("checkpoint_id", checkpoint_id) - with self._lock, self._connect() as connection: - row = connection.execute( - "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", - (checkpoint_id,), - ).fetchone() - if row is None: - return None - session_id = row["session_id"] - if not isinstance(session_id, str): - raise TypeError("checkpoint session_id must be text") - return session_id - - def _assert_authority( - self, - connection: sqlite3.Connection, - authority: OwnerAuthority, - ) -> None: - row = connection.execute( - """ - SELECT owner_id, lease_expires_at, fencing_token - FROM session_owners - WHERE session_id = ? - """, - (authority.session_id,), - ).fetchone() - if row is None: - raise SessionOwnershipConflictError( - "session owner lease is missing", - reason=SessionOwnershipConflictReason.MISSING_OWNER, - ) - lease_expires_at = _datetime_from_sqlite_text(row["lease_expires_at"]) - if lease_expires_at <= datetime.now(UTC): - raise SessionOwnershipConflictError( - "session owner lease has expired", - reason=SessionOwnershipConflictReason.EXPIRED_LEASE, - ) - if ( - row["owner_id"] != authority.owner_id - or row["fencing_token"] != authority.epoch - ): - raise SessionOwnershipConflictError("stale owner or fencing token rejected") - - def _bind_tape( - self, - connection: sqlite3.Connection, - session_id: str, - tape_id: str, - ) -> None: - now = _datetime_to_sqlite_text(datetime.now(UTC)) - session_row = connection.execute( - "SELECT tape_id FROM session_tapes WHERE session_id = ?", - (session_id,), - ).fetchone() - if session_row is not None and session_row["tape_id"] != tape_id: - raise SessionOwnershipConflictError("session tape target cannot be rebound") - tape_row = connection.execute( - "SELECT session_id FROM session_tapes WHERE tape_id = ?", - (tape_id,), - ).fetchone() - if tape_row is not None and tape_row["session_id"] != session_id: - raise SessionOwnershipConflictError( - "tape target belongs to another session" - ) - connection.execute( - """ - INSERT INTO session_tapes (session_id, tape_id, created_at, updated_at) - VALUES (?, ?, ?, ?) - ON CONFLICT(session_id) - DO UPDATE SET - tape_id = excluded.tape_id, - updated_at = excluded.updated_at - """, - (session_id, tape_id, now, now), - ) - - def _assert_tape_belongs_to_session( - self, - connection: sqlite3.Connection, - *, - tape_id: str, - session_id: str, - ) -> None: - row = connection.execute( - """ - SELECT session_id - FROM session_tapes - WHERE tape_id = ? - """, - (tape_id,), - ).fetchone() - if row is None: - raise SessionOwnershipConflictError("tape target is not bound to a session") - if row["session_id"] != session_id: - raise SessionOwnershipConflictError( - "tape target belongs to another session" - ) - - def _assert_run_belongs_to_session( - self, - connection: sqlite3.Connection, - *, - run_id: str, - session_id: str, - ) -> None: - row = connection.execute( - "SELECT session_id FROM agent_runs WHERE run_id = ?", - (run_id,), - ).fetchone() - if row is None: - raise KeyError(f"agent run not found: {run_id}") - if row["session_id"] != session_id: - raise SessionOwnershipConflictError("agent run belongs to another session") - - -def _require_non_empty(field_name: str, value: str) -> None: - if not isinstance(value, str) or not value: - raise ValueError(f"{field_name} must be non-empty") - - -def _require_json_object(field_name: str, value: object) -> None: - if not isinstance(value, dict): - raise TypeError(f"{field_name} must be a JSON object") - - -def _row_required_int(row: sqlite3.Row | None, key: str, *, context: str) -> int: - if row is None: - raise TypeError(f"{context} row is missing") - value = row[key] - if not isinstance(value, int): - raise TypeError(f"{context} {key} must be an int") - return value - - -class FencedSQLiteTapeStore: - def __init__( - self, - *, - durable_store: SQLiteLocalDurableStore, - path: Path, - authority_for_session: Callable[[str], OwnerAuthority], - ) -> None: - self._durable_store = durable_store - self._delegate = SQLiteTapeStore(path) - self._authority_for_session = authority_for_session - - async def save(self, tape_id: str, entries: list[dict[str, Any]]) -> None: - session_id = self._require_session_id_for_tape(tape_id) - await self._durable_store.append_tape_entries( - self._authority_for_session(session_id), - tape_id, - entries, - ) - - async def load(self, tape_id: str) -> list[dict[str, Any]]: - return await self._delegate.load(tape_id) - - async def list_ids(self) -> list[str]: - return await self._delegate.list_ids() - - async def truncate(self, tape_id: str, keep: int) -> None: - session_id = self._require_session_id_for_tape(tape_id) - await self._durable_store.truncate_tape( - self._authority_for_session(session_id), - tape_id, - keep, - ) - - async def info(self, tape_id: str) -> TapeInfo | None: - return await self._delegate.info(tape_id) - - async def search( - self, - *, - tape_id: str | None = None, - kind: str | None = None, - run_id: str | None = None, - tool_call_id: str | None = None, - anchor_type: str | None = None, - limit: int = 100, - ) -> list[TapeSearchResult]: - return await self._delegate.search( - tape_id=tape_id, - kind=kind, - run_id=run_id, - tool_call_id=tool_call_id, - anchor_type=anchor_type, - limit=limit, - ) - - def _require_session_id_for_tape(self, tape_id: str) -> str: - session_id = self._durable_store.session_id_for_tape(tape_id) - if session_id is None: - raise SessionOwnershipConflictError("tape target is not bound to a session") - return session_id - - -class FencedSQLiteCheckpointStore: - def __init__( - self, - *, - durable_store: SQLiteLocalDurableStore, - path: Path, - authority_for_session: Callable[[str], OwnerAuthority], - ) -> None: - self._durable_store = durable_store - self._delegate = SQLiteCheckpointStore(path) - self._authority_for_session = authority_for_session - - async def save(self, snapshot: CheckpointSnapshot) -> None: - session_id = snapshot.meta.session_id - if session_id is None: - raise SessionOwnershipConflictError( - "checkpoint target is not bound to a session" - ) - await self._durable_store.save_checkpoint( - self._authority_for_session(session_id), - snapshot, - ) - - async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: - return await self._delegate.load(checkpoint_id) - - async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: - return await self._delegate.list_by_tape(tape_id) - - async def delete(self, checkpoint_id: str) -> None: - session_id = self._durable_store.session_id_for_checkpoint(checkpoint_id) - if session_id is None: - return - await self._durable_store.delete_checkpoint( - self._authority_for_session(session_id), - checkpoint_id, - ) - - -class FencedSQLiteRuntimeStore: - def __init__( - self, - *, - durable_store: SQLiteLocalDurableStore, - path: Path, - authority_for_session: Callable[[str], OwnerAuthority], - authorities: Callable[[], Mapping[str, OwnerAuthority]], - ) -> None: - self._durable_store = durable_store - self._delegate = SQLiteRuntimeStore(path) - self._authority_for_session = authority_for_session - self._authorities = authorities - - async def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: - return await self._durable_store.create_agent_run( - self._authority_for_session(record.session_id), - record, - ) - - async def update_agent_run( - self, - run_id: str, - *, - status: str, - ended_at: datetime | None, - metadata: dict[str, Any], - result: dict[str, Any], - error: str | None, - ) -> AgentRunRecord: - session_id = self._require_session_id_for_run(run_id) - return await self._durable_store.update_agent_run( - self._authority_for_session(session_id), - run_id, - status=status, - ended_at=ended_at, - metadata=metadata, - result=result, - error=error, - ) - - async def load_agent_run(self, run_id: str) -> AgentRunRecord | None: - return await self._delegate.load_agent_run(run_id) - - async def list_agent_runs(self, session_id: str) -> list[AgentRunRecord]: - return await self._delegate.list_agent_runs(session_id) - - async def claim_attached_executor_run( - self, - *, - session_id: str | None, - executor_kind: str, - claim_metadata: dict[str, Any], - ) -> AgentRunRecord | None: - return await self._durable_store.claim_attached_executor_run( - self._authorities(), - session_id=session_id, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - - async def claim_external_worker_run( - self, - *, - session_id: str | None, - executor_kind: str, - claim_metadata: dict[str, Any], - ) -> AgentRunRecord | None: - return await self._durable_store.claim_attached_executor_run( - self._authorities(), - session_id=session_id, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - - async def append_runtime_event( - self, - record: RuntimeEventRecord, - ) -> RuntimeEventRecord: - session_id = self._require_session_id_for_run(record.run_id) - return await self._durable_store.append_runtime_event( - self._authority_for_session(session_id), - record, - ) - - async def load_runtime_event( - self, - event_id: str, - ) -> RuntimeEventRecord | None: - return await self._delegate.load_runtime_event(event_id) - - async def replay_runtime_events( - self, - run_id: str, - *, - after_sequence: int = 0, - limit: int = 1000, - ) -> list[RuntimeEventRecord]: - return await self._delegate.replay_runtime_events( - run_id, - after_sequence=after_sequence, - limit=limit, - ) - - async def save_message_snapshot( - self, - record: RunMessageSnapshotRecord, - ) -> RunMessageSnapshotRecord: - session_id = self._require_session_id_for_run(record.run_id) - return await self._durable_store.save_message_snapshot( - self._authority_for_session(session_id), - record, - ) - - async def load_message_snapshot( - self, - snapshot_id: str, - ) -> RunMessageSnapshotRecord | None: - return await self._delegate.load_message_snapshot(snapshot_id) - - async def list_message_snapshots( - self, - run_id: str, - ) -> list[RunMessageSnapshotRecord]: - return await self._delegate.list_message_snapshots(run_id) - - async def create_agent_interaction( - self, - record: AgentInteractionRecord, - ) -> AgentInteractionRecord: - session_id = self._require_session_id_for_run(record.run_id) - return await self._durable_store.create_agent_interaction( - self._authority_for_session(session_id), - record, - ) - - async def load_agent_interaction( - self, - interaction_id: str, - ) -> AgentInteractionRecord | None: - return await self._delegate.load_agent_interaction(interaction_id) - - async def list_agent_interactions( - self, - run_id: str, - ) -> list[AgentInteractionRecord]: - return await self._delegate.list_agent_interactions(run_id) - - async def resolve_agent_interaction( - self, - interaction_id: str, - *, - status: str, - response_payload: dict[str, Any], - resolved_at: datetime, - ) -> AgentInteractionRecord: - interaction = await self._delegate.load_agent_interaction(interaction_id) - if interaction is None: - raise KeyError(f"agent interaction not found: {interaction_id}") - session_id = self._require_session_id_for_run(interaction.run_id) - return await self._durable_store.resolve_agent_interaction( - self._authority_for_session(session_id), - interaction_id, - status=status, - response_payload=response_payload, - resolved_at=resolved_at, - ) - - def _require_session_id_for_run(self, run_id: str) -> str: - session_id = self._durable_store.session_id_for_run(run_id) - if session_id is None: - raise KeyError(f"agent run not found: {run_id}") - return session_id +"""Deprecated import location. Moved to coding_agent.stores.durable_local.""" +import sys as _sys +from coding_agent.stores import durable_local as _moved +_sys.modules[__name__] = _moved diff --git a/coding-agent/src/coding_agent/pg_durable_store.py b/coding-agent/src/coding_agent/pg_durable_store.py index 07f7b83ef..70f010e98 100644 --- a/coding-agent/src/coding_agent/pg_durable_store.py +++ b/coding-agent/src/coding_agent/pg_durable_store.py @@ -1,1110 +1,4 @@ -from __future__ import annotations - -from collections.abc import Callable, Mapping -from datetime import datetime -import json -from typing import Any, cast - -from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot -from agentkit.storage.pg import ( - PGCheckpointStore, - PGPool, - PGSessionOwnerStore, - PGTapeStore, -) -from agentkit.storage.protocols import TapeInfo, TapeSearchResult - -from coding_agent.runtime_store import ( - AgentInteractionRecord, - AgentRunRecord, - JSONObject, - PGRuntimeStore, - RunMessageSnapshotRecord, - RuntimeEventRecord, - _agent_run_from_row, - _interaction_from_row, - _message_snapshot_from_row, - _runtime_event_from_row, -) -from coding_agent.server.stores.session_owner_store import ( - OwnerAuthority, - SessionOwnershipConflictError, - SessionOwnershipConflictReason, -) -from coding_agent.server.stores.session_store import PGSessionMetadataStore - - -class PGDurableStore: - """PostgreSQL protected write facade for durable owner fencing.""" - - _UPSERT_SESSION_SQL = """ - INSERT INTO agent_http_sessions (session_id, payload) - VALUES ($1, $2::jsonb) - ON CONFLICT (session_id) - DO UPDATE SET payload = EXCLUDED.payload, updated_at = NOW() - """ - _DELETE_SESSION_SQL = "DELETE FROM agent_http_sessions WHERE session_id = $1" - _SELECT_SESSION_FOR_UPDATE_SQL = """ - SELECT payload - FROM agent_http_sessions - WHERE session_id = $1 - FOR UPDATE - """ - _SELECT_SESSION_BY_TAPE_SQL = """ - SELECT session_id - FROM session_tapes - WHERE tape_id = $1 - LIMIT 1 - """ - _CREATE_SESSION_TAPES_SQL = """ - CREATE TABLE IF NOT EXISTS session_tapes ( - session_id TEXT PRIMARY KEY, - tape_id TEXT UNIQUE NOT NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() - ); - CREATE INDEX IF NOT EXISTS session_tapes_tape_id_idx - ON session_tapes (tape_id); - """ - _SELECT_SESSION_TAPE_FOR_UPDATE_SQL = """ - SELECT tape_id - FROM session_tapes - WHERE session_id = $1 - FOR UPDATE - """ - _SELECT_TAPE_SESSION_FOR_UPDATE_SQL = """ - SELECT session_id - FROM session_tapes - WHERE tape_id = $1 - FOR UPDATE - """ - _UPSERT_SESSION_TAPE_SQL = """ - INSERT INTO session_tapes (session_id, tape_id) - VALUES ($1, $2) - ON CONFLICT DO NOTHING - """ - _SELECT_OWNER_FOR_UPDATE_SQL = """ - SELECT owner_id, lease_expires_at, fencing_token - FROM session_owners - WHERE session_id = $1 - FOR UPDATE - """ - _SELECT_RUN_SESSION_FOR_UPDATE_SQL = """ - SELECT session_id - FROM agent_runs - WHERE run_id = $1 - FOR UPDATE - """ - _SELECT_INTERACTION_RUN_FOR_UPDATE_SQL = """ - SELECT run_id - FROM agent_interactions - WHERE interaction_id = $1 - FOR UPDATE - """ - _SELECT_CHECKPOINT_META_FOR_UPDATE_SQL = """ - SELECT meta - FROM agent_checkpoints - WHERE checkpoint_id = $1 - FOR UPDATE - """ - _UPSERT_OWNED_RUN_SQL = """ - INSERT INTO agent_runs ( - run_id, - session_id, - tape_id, - parent_run_id, - agent_id, - status, - started_at, - ended_at, - metadata, - result, - error - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::jsonb, $10::jsonb, $11) - ON CONFLICT (run_id) - DO UPDATE SET - tape_id = EXCLUDED.tape_id, - parent_run_id = EXCLUDED.parent_run_id, - agent_id = EXCLUDED.agent_id, - status = EXCLUDED.status, - started_at = EXCLUDED.started_at, - ended_at = EXCLUDED.ended_at, - metadata = EXCLUDED.metadata, - result = EXCLUDED.result, - error = EXCLUDED.error, - updated_at = NOW() - WHERE agent_runs.session_id = $2 - RETURNING * - """ - _INSERT_OWNED_EVENT_SQL = """ - WITH inserted AS ( - INSERT INTO runtime_events ( - event_id, - run_id, - event_kind, - payload, - created_at - ) - VALUES ($1, $2, $3, $4::jsonb, $5) - ON CONFLICT (event_id) DO NOTHING - RETURNING * - ) - SELECT * FROM inserted - UNION ALL - SELECT runtime_events.* - FROM runtime_events - WHERE runtime_events.event_id = $1 - AND runtime_events.run_id = $2 - AND NOT EXISTS (SELECT 1 FROM inserted) - """ - _UPSERT_OWNED_MESSAGE_SNAPSHOT_SQL = """ - INSERT INTO run_message_snapshots ( - snapshot_id, - run_id, - messages, - metadata, - created_at - ) - VALUES ($1, $2, $3::jsonb, $4::jsonb, $5) - ON CONFLICT (snapshot_id) - DO UPDATE SET - messages = EXCLUDED.messages, - metadata = EXCLUDED.metadata, - created_at = EXCLUDED.created_at, - updated_at = NOW() - WHERE run_message_snapshots.run_id = $2 - RETURNING * - """ - _INSERT_OWNED_INTERACTION_SQL = """ - WITH inserted AS ( - INSERT INTO agent_interactions ( - interaction_id, - run_id, - interaction_kind, - status, - request_payload, - response_payload, - metadata, - created_at, - resolved_at - ) - VALUES ($1, $2, $3, $4, $5::jsonb, $6::jsonb, $7::jsonb, $8, $9) - ON CONFLICT (interaction_id) DO NOTHING - RETURNING * - ) - SELECT * FROM inserted - UNION ALL - SELECT agent_interactions.* - FROM agent_interactions - WHERE agent_interactions.interaction_id = $1 - AND agent_interactions.run_id = $2 - AND NOT EXISTS (SELECT 1 FROM inserted) - """ - _UPSERT_OWNED_CHECKPOINT_SQL = """ - INSERT INTO agent_checkpoints ( - checkpoint_id, - tape_id, - meta, - entries, - plugin_states, - extra - ) - VALUES ($1, $2, $3::jsonb, $4::jsonb, $5::jsonb, $6::jsonb) - ON CONFLICT (checkpoint_id) DO UPDATE SET - tape_id = EXCLUDED.tape_id, - meta = EXCLUDED.meta, - entries = EXCLUDED.entries, - plugin_states = EXCLUDED.plugin_states, - extra = EXCLUDED.extra - WHERE agent_checkpoints.meta->>'session_id' = $7 - RETURNING checkpoint_id - """ - _INSERT_TAPE_SQL = PGTapeStore._INSERT_SQL - _TRUNCATE_TAPE_SQL = PGTapeStore._TRUNCATE_SQL - _DELETE_CHECKPOINT_SQL = PGCheckpointStore._DELETE_SQL - _DELETE_NEWER_CHECKPOINTS_SQL = """ - DELETE FROM agent_checkpoints - WHERE tape_id = $1 - AND meta->>'session_id' = $2 - AND (meta->>'entry_count')::int > $3 - """ - _DELETE_SESSION_TAPE_SQL = "DELETE FROM session_tapes WHERE session_id = $1" - - def __init__(self, *, pool: PGPool) -> None: - self._pool = pool - self._schema_ready = False - - async def _ensure_schema(self) -> None: - pool = await self._pool.get_pool() - if self._schema_ready: - return - _ = await pool.execute(PGSessionOwnerStore._CREATE_TABLE_SQL) - _ = await pool.execute(PGSessionMetadataStore._CREATE_TABLE_SQL) - _ = await pool.execute(self._CREATE_SESSION_TAPES_SQL) - _ = await pool.execute(PGTapeStore._CREATE_TABLE_SQL) - _ = await pool.execute(PGCheckpointStore._CREATE_TABLE_SQL) - _ = await pool.execute(PGRuntimeStore._CREATE_SCHEMA_SQL) - self._schema_ready = True - - async def _with_transaction(self, body: Callable[[Any], Any]) -> Any: - await self._ensure_schema() - connection = await self._pool.acquire() - try: - _ = await connection.execute("BEGIN") - result = await body(connection) - _ = await connection.execute("COMMIT") - return result - except BaseException: - try: - _ = await connection.execute("ROLLBACK") - except Exception: - pass - raise - finally: - await self._pool.release(connection) - - async def session_id_for_tape(self, tape_id: str) -> str | None: - await self._ensure_schema() - pool = await self._pool.get_pool() - row = await pool.fetchrow(self._SELECT_SESSION_BY_TAPE_SQL, tape_id) - if row is None: - return None - return _required_str(dict(row), "session_id") - - async def save_session( - self, - authority: OwnerAuthority, - payload: dict[str, Any], - ) -> None: - _require_payload_session(authority, payload) - tape_id = payload.get("tape_id") - if tape_id is not None and not isinstance(tape_id, str): - raise TypeError("session payload tape_id must be a string") - - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - if tape_id: - await self._bind_tape(connection, authority.session_id, tape_id) - _ = await connection.fetchrow( - self._SELECT_SESSION_FOR_UPDATE_SQL, - authority.session_id, - ) - await connection.execute( - self._UPSERT_SESSION_SQL, - authority.session_id, - payload, - ) - - await self._with_transaction(body) - - async def delete_session(self, authority: OwnerAuthority) -> None: - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - _ = await connection.fetchrow( - self._SELECT_SESSION_FOR_UPDATE_SQL, - authority.session_id, - ) - await connection.execute(self._DELETE_SESSION_SQL, authority.session_id) - await connection.execute( - self._DELETE_SESSION_TAPE_SQL, authority.session_id - ) - - await self._with_transaction(body) - - async def append_tape_entries( - self, - authority: OwnerAuthority, - tape_id: str, - entries: list[dict[str, Any]], - ) -> None: - if not entries: - return - - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - await self._require_stable_tape(connection, authority, tape_id) - payload_values = [json.dumps(entry) for entry in entries] - await connection.execute(self._INSERT_TAPE_SQL, tape_id, payload_values) - - await self._with_transaction(body) - - async def truncate_tape( - self, - authority: OwnerAuthority, - tape_id: str, - keep: int, - ) -> None: - if keep < 0: - raise ValueError("keep must be >= 0") - - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - await self._require_stable_tape(connection, authority, tape_id) - await connection.execute(self._TRUNCATE_TAPE_SQL, tape_id, keep) - - await self._with_transaction(body) - - async def create_agent_run( - self, - authority: OwnerAuthority, - record: AgentRunRecord, - ) -> AgentRunRecord: - if record.session_id != authority.session_id: - raise SessionOwnershipConflictError("run target belongs to another owner") - if record.tape_id is None: - raise SessionOwnershipConflictError("run target is not bound to a tape") - - async def body(connection: Any) -> AgentRunRecord: - await self._require_owner(connection, authority) - await self._require_stable_tape(connection, authority, record.tape_id) - row = await connection.fetchrow( - self._UPSERT_OWNED_RUN_SQL, - record.run_id, - record.session_id, - record.tape_id, - record.parent_run_id, - record.agent_id, - record.status, - record.started_at, - record.ended_at, - record.metadata, - record.result, - record.error, - ) - return _agent_run_from_row( - _required_owned_row(row, "run target belongs to another owner") - ) - - return cast(AgentRunRecord, await self._with_transaction(body)) - - async def update_agent_run( - self, - authority: OwnerAuthority, - run_id: str, - *, - status: str, - ended_at: datetime | None, - metadata: JSONObject, - result: JSONObject, - error: str | None, - ) -> AgentRunRecord: - async def body(connection: Any) -> AgentRunRecord: - await self._require_owner(connection, authority) - await self._require_run_owner(connection, authority, run_id) - row = await connection.fetchrow( - PGRuntimeStore._UPDATE_RUN_SQL, - run_id, - status, - ended_at, - metadata, - result, - error, - ) - return _agent_run_from_row(_required_row(row, "agent run update")) - - return cast(AgentRunRecord, await self._with_transaction(body)) - - async def append_runtime_event( - self, - authority: OwnerAuthority, - record: RuntimeEventRecord, - ) -> RuntimeEventRecord: - async def body(connection: Any) -> RuntimeEventRecord: - await self._require_owner(connection, authority) - await self._require_run_owner(connection, authority, record.run_id) - row = await connection.fetchrow( - self._INSERT_OWNED_EVENT_SQL, - record.event_id, - record.run_id, - record.event_kind, - record.payload, - record.created_at, - ) - return _runtime_event_from_row( - _required_owned_row( - row, - "runtime event target belongs to another owner", - ) - ) - - return cast(RuntimeEventRecord, await self._with_transaction(body)) - - async def save_message_snapshot( - self, - authority: OwnerAuthority, - record: RunMessageSnapshotRecord, - ) -> RunMessageSnapshotRecord: - async def body(connection: Any) -> RunMessageSnapshotRecord: - await self._require_owner(connection, authority) - await self._require_run_owner(connection, authority, record.run_id) - row = await connection.fetchrow( - self._UPSERT_OWNED_MESSAGE_SNAPSHOT_SQL, - record.snapshot_id, - record.run_id, - record.messages, - record.metadata, - record.created_at, - ) - return _message_snapshot_from_row( - _required_owned_row( - row, - "message snapshot target belongs to another owner", - ) - ) - - return cast(RunMessageSnapshotRecord, await self._with_transaction(body)) - - async def create_agent_interaction( - self, - authority: OwnerAuthority, - record: AgentInteractionRecord, - ) -> AgentInteractionRecord: - async def body(connection: Any) -> AgentInteractionRecord: - await self._require_owner(connection, authority) - await self._require_run_owner(connection, authority, record.run_id) - row = await connection.fetchrow( - self._INSERT_OWNED_INTERACTION_SQL, - record.interaction_id, - record.run_id, - record.interaction_kind, - record.status, - record.request_payload, - record.response_payload, - record.metadata, - record.created_at, - record.resolved_at, - ) - return _interaction_from_row( - _required_owned_row(row, "interaction target belongs to another owner") - ) - - return cast(AgentInteractionRecord, await self._with_transaction(body)) - - async def resolve_agent_interaction( - self, - authority: OwnerAuthority, - interaction_id: str, - *, - status: str, - response_payload: JSONObject, - resolved_at: datetime, - ) -> AgentInteractionRecord: - async def body(connection: Any) -> AgentInteractionRecord: - await self._require_owner(connection, authority) - existing = await connection.fetchrow( - self._SELECT_INTERACTION_RUN_FOR_UPDATE_SQL, - interaction_id, - ) - if existing is None: - raise KeyError(f"agent interaction not found: {interaction_id}") - await self._require_run_owner( - connection, - authority, - _required_str(dict(existing), "run_id"), - ) - row = await connection.fetchrow( - PGRuntimeStore._RESOLVE_INTERACTION_SQL, - interaction_id, - status, - response_payload, - resolved_at, - ) - return _interaction_from_row( - _required_row(row, "agent interaction resolve") - ) - - return cast(AgentInteractionRecord, await self._with_transaction(body)) - - async def claim_attached_executor_run( - self, - authorities: Mapping[str, OwnerAuthority], - *, - session_id: str | None, - executor_kind: str, - claim_metadata: JSONObject, - ) -> AgentRunRecord | None: - if session_id is not None: - authority = authorities.get(session_id) - if authority is None: - return None - return await self._claim_attached_executor_run_for_authority( - authority, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - for authority in sorted(authorities.values(), key=lambda item: item.session_id): - claimed = await self._claim_attached_executor_run_for_authority( - authority, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - if claimed is not None: - return claimed - return None - - async def _claim_attached_executor_run_for_authority( - self, - authority: OwnerAuthority, - *, - executor_kind: str, - claim_metadata: JSONObject, - ) -> AgentRunRecord | None: - async def body(connection: Any) -> AgentRunRecord | None: - await self._require_owner(connection, authority) - row = await connection.fetchrow( - PGRuntimeStore._CLAIM_ATTACHED_EXECUTOR_RUN_SQL, - authority.session_id, - executor_kind, - claim_metadata, - ) - if row is None: - return None - return _agent_run_from_row(row) - - return cast(AgentRunRecord | None, await self._with_transaction(body)) - - async def save_checkpoint( - self, - authority: OwnerAuthority, - snapshot: CheckpointSnapshot, - ) -> None: - meta = snapshot.meta - if meta.session_id != authority.session_id: - raise SessionOwnershipConflictError( - "checkpoint target belongs to another owner" - ) - - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - await self._require_stable_tape(connection, authority, meta.tape_id) - row = await connection.fetchrow( - self._UPSERT_OWNED_CHECKPOINT_SQL, - meta.checkpoint_id, - meta.tape_id, - _checkpoint_meta_payload(meta), - list(snapshot.tape_entries), - snapshot.plugin_states, - snapshot.extra, - authority.session_id, - ) - _required_owned_row(row, "checkpoint target belongs to another owner") - - await self._with_transaction(body) - - async def delete_checkpoint( - self, - authority: OwnerAuthority, - checkpoint_id: str, - ) -> None: - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - await self._require_checkpoint_owner(connection, authority, checkpoint_id) - await connection.execute(self._DELETE_CHECKPOINT_SQL, checkpoint_id) - - await self._with_transaction(body) - - async def restore_checkpoint_state( - self, - authority: OwnerAuthority, - snapshot: CheckpointSnapshot, - session_payload: dict[str, Any], - ) -> None: - _require_payload_session(authority, session_payload) - meta = snapshot.meta - if meta.session_id != authority.session_id: - raise SessionOwnershipConflictError( - "checkpoint target belongs to another owner" - ) - if session_payload.get("tape_id") != meta.tape_id: - raise SessionOwnershipConflictError( - "checkpoint restore session payload has mismatched tape id" - ) - - async def body(connection: Any) -> None: - await self._require_owner(connection, authority) - await self._require_stable_tape(connection, authority, meta.tape_id) - await self._require_checkpoint_owner( - connection, - authority, - meta.checkpoint_id, - ) - await connection.execute(self._TRUNCATE_TAPE_SQL, meta.tape_id, 0) - if snapshot.tape_entries: - payload_values = [json.dumps(entry) for entry in snapshot.tape_entries] - await connection.execute( - self._INSERT_TAPE_SQL, meta.tape_id, payload_values - ) - await connection.execute( - self._UPSERT_SESSION_SQL, - authority.session_id, - session_payload, - ) - await connection.execute( - self._DELETE_NEWER_CHECKPOINTS_SQL, - meta.tape_id, - authority.session_id, - meta.entry_count, - ) - - await self._with_transaction(body) - - async def _require_owner(self, connection: Any, authority: OwnerAuthority) -> None: - row = await connection.fetchrow( - self._SELECT_OWNER_FOR_UPDATE_SQL, - authority.session_id, - ) - if row is None: - raise SessionOwnershipConflictError( - "session owner lease is missing", - reason=SessionOwnershipConflictReason.MISSING_OWNER, - ) - row_dict = dict(row) - lease_expires_at = row_dict.get("lease_expires_at") - if not isinstance(lease_expires_at, datetime): - raise TypeError("postgres owner row missing datetime lease_expires_at") - now = datetime.now(lease_expires_at.tzinfo) - if lease_expires_at <= now: - raise SessionOwnershipConflictError( - "session owner lease has expired", - reason=SessionOwnershipConflictReason.EXPIRED_LEASE, - ) - if ( - row_dict.get("owner_id") != authority.owner_id - or row_dict.get("fencing_token") != authority.epoch - ): - raise SessionOwnershipConflictError("stale owner or fencing token rejected") - - async def _require_stable_tape( - self, - connection: Any, - authority: OwnerAuthority, - tape_id: str, - ) -> None: - row = await connection.fetchrow( - self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, - tape_id, - ) - if row is None: - raise SessionOwnershipConflictError("tape target is not bound to a session") - if dict(row).get("session_id") != authority.session_id: - raise SessionOwnershipConflictError("tape target belongs to another owner") - - async def _bind_tape( - self, - connection: Any, - session_id: str, - tape_id: str, - ) -> None: - await connection.execute(self._UPSERT_SESSION_TAPE_SQL, session_id, tape_id) - session_row = await connection.fetchrow( - self._SELECT_SESSION_TAPE_FOR_UPDATE_SQL, - session_id, - ) - if session_row is None: - tape_row = await connection.fetchrow( - self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, - tape_id, - ) - if tape_row is not None and dict(tape_row).get("session_id") != session_id: - raise SessionOwnershipConflictError( - "tape target belongs to another session" - ) - raise SessionOwnershipConflictError("session tape target is not bound") - if dict(session_row).get("tape_id") != tape_id: - raise SessionOwnershipConflictError("session tape target cannot be rebound") - tape_row = await connection.fetchrow( - self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, - tape_id, - ) - if tape_row is None: - raise SessionOwnershipConflictError("tape target is not bound to a session") - if dict(tape_row).get("session_id") != session_id: - raise SessionOwnershipConflictError( - "tape target belongs to another session" - ) - - async def _session_id_for_run( - self, - connection: Any, - run_id: str, - *, - for_update: bool, - ) -> str | None: - del for_update - row = await connection.fetchrow(self._SELECT_RUN_SESSION_FOR_UPDATE_SQL, run_id) - if row is None: - return None - return _required_str(dict(row), "session_id") - - async def _require_run_owner( - self, - connection: Any, - authority: OwnerAuthority, - run_id: str, - ) -> None: - session_id = await self._session_id_for_run( - connection, - run_id, - for_update=True, - ) - if session_id is None: - raise KeyError(f"agent run not found: {run_id}") - if session_id != authority.session_id: - raise SessionOwnershipConflictError("run target belongs to another owner") - - async def _require_checkpoint_owner( - self, - connection: Any, - authority: OwnerAuthority, - checkpoint_id: str, - ) -> None: - row = await connection.fetchrow( - self._SELECT_CHECKPOINT_META_FOR_UPDATE_SQL, - checkpoint_id, - ) - if row is None: - raise KeyError(f"checkpoint not found: {checkpoint_id}") - meta = _required_dict(dict(row), "meta") - if meta.get("session_id") != authority.session_id: - raise SessionOwnershipConflictError( - "checkpoint target belongs to another owner" - ) - - -class FencedPGTapeStore: - def __init__( - self, - *, - durable_store: PGDurableStore, - pool: PGPool, - authority_for_session: Callable[[str], OwnerAuthority], - ) -> None: - self._durable_store = durable_store - self._delegate = PGTapeStore(pool=pool) - self._authority_for_session = authority_for_session - - async def save(self, tape_id: str, entries: list[dict[str, Any]]) -> None: - session_id = await self._require_session_id_for_tape_async(tape_id) - await self._durable_store.append_tape_entries( - self._authority_for_session(session_id), - tape_id, - entries, - ) - - async def load(self, tape_id: str) -> list[dict[str, Any]]: - return await self._delegate.load(tape_id) - - async def list_ids(self) -> list[str]: - return await self._delegate.list_ids() - - async def truncate(self, tape_id: str, keep: int) -> None: - session_id = await self._require_session_id_for_tape_async(tape_id) - await self._durable_store.truncate_tape( - self._authority_for_session(session_id), - tape_id, - keep, - ) - - async def info(self, tape_id: str) -> TapeInfo | None: - return await self._delegate.info(tape_id) - - async def search( - self, - *, - tape_id: str | None = None, - kind: str | None = None, - run_id: str | None = None, - tool_call_id: str | None = None, - anchor_type: str | None = None, - limit: int = 100, - ) -> list[TapeSearchResult]: - return await self._delegate.search( - tape_id=tape_id, - kind=kind, - run_id=run_id, - tool_call_id=tool_call_id, - anchor_type=anchor_type, - limit=limit, - ) - - async def _require_session_id_for_tape_async(self, tape_id: str) -> str: - session_id = await self._durable_store.session_id_for_tape(tape_id) - if session_id is None: - raise SessionOwnershipConflictError("tape target is not bound to a session") - return session_id - - -class FencedPGCheckpointStore: - def __init__( - self, - *, - durable_store: PGDurableStore, - pool: PGPool, - authority_for_session: Callable[[str], OwnerAuthority], - ) -> None: - self._durable_store = durable_store - self._delegate = PGCheckpointStore(pool=pool) - self._authority_for_session = authority_for_session - - async def save(self, snapshot: CheckpointSnapshot) -> None: - session_id = snapshot.meta.session_id - if session_id is None: - raise SessionOwnershipConflictError( - "checkpoint target is not bound to a session" - ) - await self._durable_store.save_checkpoint( - self._authority_for_session(session_id), - snapshot, - ) - - async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: - return await self._delegate.load(checkpoint_id) - - async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: - return await self._delegate.list_by_tape(tape_id) - - async def delete(self, checkpoint_id: str) -> None: - snapshot = await self._delegate.load(checkpoint_id) - if snapshot is None: - return - session_id = snapshot.meta.session_id - if session_id is None: - raise SessionOwnershipConflictError( - "checkpoint target is not bound to a session" - ) - await self._durable_store.delete_checkpoint( - self._authority_for_session(session_id), - checkpoint_id, - ) - - -class FencedPGRuntimeStore: - def __init__( - self, - *, - durable_store: PGDurableStore, - pool: PGPool, - authority_for_session: Callable[[str], OwnerAuthority], - authorities: Callable[[], Mapping[str, OwnerAuthority]], - ) -> None: - self._durable_store = durable_store - self._delegate = PGRuntimeStore(pool=pool) - self._authority_for_session = authority_for_session - self._authorities = authorities - - async def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: - return await self._durable_store.create_agent_run( - self._authority_for_session(record.session_id), - record, - ) - - async def update_agent_run( - self, - run_id: str, - *, - status: str, - ended_at: datetime | None, - metadata: JSONObject, - result: JSONObject, - error: str | None, - ) -> AgentRunRecord: - existing = await self._delegate.load_agent_run(run_id) - if existing is None: - raise KeyError(f"agent run not found: {run_id}") - return await self._durable_store.update_agent_run( - self._authority_for_session(existing.session_id), - run_id, - status=status, - ended_at=ended_at, - metadata=metadata, - result=result, - error=error, - ) - - async def load_agent_run(self, run_id: str) -> AgentRunRecord | None: - return await self._delegate.load_agent_run(run_id) - - async def list_agent_runs(self, session_id: str) -> list[AgentRunRecord]: - return await self._delegate.list_agent_runs(session_id) - - async def claim_attached_executor_run( - self, - *, - session_id: str | None, - executor_kind: str, - claim_metadata: JSONObject, - ) -> AgentRunRecord | None: - return await self._durable_store.claim_attached_executor_run( - self._authorities(), - session_id=session_id, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - - async def claim_external_worker_run( - self, - *, - session_id: str | None, - executor_kind: str, - claim_metadata: JSONObject, - ) -> AgentRunRecord | None: - return await self.claim_attached_executor_run( - session_id=session_id, - executor_kind=executor_kind, - claim_metadata=claim_metadata, - ) - - async def append_runtime_event( - self, - record: RuntimeEventRecord, - ) -> RuntimeEventRecord: - existing = await self._delegate.load_agent_run(record.run_id) - if existing is None: - raise KeyError(f"agent run not found: {record.run_id}") - return await self._durable_store.append_runtime_event( - self._authority_for_session(existing.session_id), - record, - ) - - async def replay_runtime_events( - self, - run_id: str, - *, - after_sequence: int = 0, - limit: int = 1000, - ) -> list[RuntimeEventRecord]: - return await self._delegate.replay_runtime_events( - run_id, - after_sequence=after_sequence, - limit=limit, - ) - - async def load_runtime_event( - self, - event_id: str, - ) -> RuntimeEventRecord | None: - return await self._delegate.load_runtime_event(event_id) - - async def save_message_snapshot( - self, - record: RunMessageSnapshotRecord, - ) -> RunMessageSnapshotRecord: - existing = await self._delegate.load_agent_run(record.run_id) - if existing is None: - raise KeyError(f"agent run not found: {record.run_id}") - return await self._durable_store.save_message_snapshot( - self._authority_for_session(existing.session_id), - record, - ) - - async def load_message_snapshot( - self, - snapshot_id: str, - ) -> RunMessageSnapshotRecord | None: - return await self._delegate.load_message_snapshot(snapshot_id) - - async def list_message_snapshots( - self, - run_id: str, - ) -> list[RunMessageSnapshotRecord]: - return await self._delegate.list_message_snapshots(run_id) - - async def create_agent_interaction( - self, - record: AgentInteractionRecord, - ) -> AgentInteractionRecord: - existing = await self._delegate.load_agent_run(record.run_id) - if existing is None: - raise KeyError(f"agent run not found: {record.run_id}") - return await self._durable_store.create_agent_interaction( - self._authority_for_session(existing.session_id), - record, - ) - - async def resolve_agent_interaction( - self, - interaction_id: str, - *, - status: str, - response_payload: JSONObject, - resolved_at: datetime, - ) -> AgentInteractionRecord: - existing = await self._delegate.load_agent_interaction(interaction_id) - if existing is None: - raise KeyError(f"agent interaction not found: {interaction_id}") - run = await self._delegate.load_agent_run(existing.run_id) - if run is None: - raise KeyError(f"agent run not found: {existing.run_id}") - return await self._durable_store.resolve_agent_interaction( - self._authority_for_session(run.session_id), - interaction_id, - status=status, - response_payload=response_payload, - resolved_at=resolved_at, - ) - - async def load_agent_interaction( - self, - interaction_id: str, - ) -> AgentInteractionRecord | None: - return await self._delegate.load_agent_interaction(interaction_id) - - async def list_agent_interactions( - self, - run_id: str, - ) -> list[AgentInteractionRecord]: - return await self._delegate.list_agent_interactions(run_id) - - -def _require_payload_session( - authority: OwnerAuthority, - payload: dict[str, Any], -) -> None: - payload_id = payload.get("id") - if payload_id != authority.session_id: - raise SessionOwnershipConflictError("session payload belongs to another owner") - payload_session_id = payload.get("session_id") - if payload_session_id is not None and payload_session_id != authority.session_id: - raise SessionOwnershipConflictError("session payload belongs to another owner") - - -def _checkpoint_meta_payload(meta: CheckpointMeta) -> dict[str, Any]: - return { - "checkpoint_id": meta.checkpoint_id, - "tape_id": meta.tape_id, - "session_id": meta.session_id, - "entry_count": meta.entry_count, - "window_start": meta.window_start, - "created_at": meta.created_at.isoformat(), - "label": meta.label, - } - - -def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: - if row is None: - raise RuntimeError(f"postgres {context} returned no row") - return row - - -def _required_owned_row( - row: dict[str, object] | None, - conflict_message: str, -) -> dict[str, object]: - if row is None: - raise SessionOwnershipConflictError(conflict_message) - return row - - -def _required_str(row: dict[str, object], key: str) -> str: - value = row.get(key) - if not isinstance(value, str): - raise TypeError(f"postgres row must include string {key}") - return value - - -def _required_dict(row: dict[str, object], key: str) -> dict[str, object]: - value = row.get(key) - if not isinstance(value, dict): - raise TypeError(f"postgres row must include dict {key}") - return cast(dict[str, object], value) +"""Deprecated import location. Moved to coding_agent.stores.durable_pg.""" +import sys as _sys +from coding_agent.stores import durable_pg as _moved +_sys.modules[__name__] = _moved diff --git a/coding-agent/src/coding_agent/runs/scheduled.py b/coding-agent/src/coding_agent/runs/scheduled.py new file mode 100644 index 000000000..11e73e19c --- /dev/null +++ b/coding-agent/src/coding_agent/runs/scheduled.py @@ -0,0 +1,1192 @@ +"""Durable scheduled run and proactive signal records.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from typing import Final, Protocol + +from agentkit.storage.pg import AsyncPGPool, PGPool +from agentkit.tape.tape import Tape +from coding_agent.topic_lifecycle import TopicLifecycle +from coding_agent.topic_store import JSONObject +from coding_agent.topic_store import TopicRecord + +ScheduleStatus = str +SignalStatus = str +TriggerStatus = str + +_SCHEDULE_STATUSES: Final[frozenset[str]] = frozenset( + {"active", "paused", "disabled", "completed"} +) +_SIGNAL_STATUSES: Final[frozenset[str]] = frozenset( + {"new", "planned", "ignored", "consumed"} +) +_TRIGGER_STATUSES: Final[frozenset[str]] = frozenset( + {"planned", "launched", "skipped", "failed"} +) +_MAX_DISPLAY_TEXT_CHARS: Final[int] = 256 +_MAX_METADATA_STRING_CHARS: Final[int] = 256 +_MAX_SAFE_LABEL_CHARS: Final[int] = 128 +_FORBIDDEN_METADATA_KEY_PARTS: Final[frozenset[str]] = frozenset( + { + "command_output", + "content", + "env", + "message", + "prompt", + "result", + "secret", + "stderr", + "stdout", + "text", + } +) +_FORBIDDEN_BEE_LAUNCH_KEY_PARTS: Final[frozenset[str]] = frozenset( + { + "api_key", + "apikey", + "args", + "argv", + "bearer", + "cmd", + "command", + "commands", + "command_output", + "content", + "credential", + "credentials", + "env", + "environment", + "exec", + "executor", + "message", + "password", + "prompt", + "result", + "script", + "secret", + "shell", + "stderr", + "stdout", + "text", + "token", + } +) +_SECRET_VALUE_MARKERS: Final[tuple[str, ...]] = ( + "-----begin ", + "akia", + "password=", + "secret=", + "sk-", + "token=", +) + + +@dataclass(frozen=True) +class ScheduleRecord: + schedule_id: str + session_id: str + kind: str + status: ScheduleStatus + cadence: str + created_at: datetime + updated_at: datetime + topic_id: str | None = None + owner: str | None = None + title: str | None = None + next_due_at: datetime | None = None + last_triggered_at: datetime | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_non_empty("schedule_id", self.schedule_id) + _require_non_empty("session_id", self.session_id) + _require_safe_label("kind", self.kind) + _require_safe_label("cadence", self.cadence) + _require_status("schedule status", self.status, _SCHEDULE_STATUSES) + _require_datetime("created_at", self.created_at) + _require_datetime("updated_at", self.updated_at) + if self.topic_id is not None: + _require_non_empty("topic_id", self.topic_id) + _require_optional_display_text("owner", self.owner) + _require_optional_display_text("title", self.title) + if self.next_due_at is not None: + _require_datetime("next_due_at", self.next_due_at) + if self.last_triggered_at is not None: + _require_datetime("last_triggered_at", self.last_triggered_at) + _require_json_object("metadata", self.metadata) + if "bee_launch" in self.metadata: + _require_bee_launch_metadata_object( + "metadata.bee_launch", + self.metadata["bee_launch"], + ) + + +@dataclass(frozen=True) +class ScheduleTriggerRecord: + trigger_id: str + schedule_id: str + status: TriggerStatus + due_at: datetime + planned_at: datetime + topic_id: str | None = None + signal_id: str | None = None + run_id: str | None = None + reason: str | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_non_empty("trigger_id", self.trigger_id) + _require_non_empty("schedule_id", self.schedule_id) + _require_status("trigger status", self.status, _TRIGGER_STATUSES) + _require_datetime("due_at", self.due_at) + _require_datetime("planned_at", self.planned_at) + if self.topic_id is not None: + _require_non_empty("topic_id", self.topic_id) + if self.signal_id is not None: + _require_non_empty("signal_id", self.signal_id) + if self.run_id is not None: + _require_non_empty("run_id", self.run_id) + _require_optional_display_text("reason", self.reason) + _require_json_object("metadata", self.metadata) + + +@dataclass(frozen=True) +class ProactiveSignalRecord: + signal_id: str + dedupe_key: str + kind: str + status: SignalStatus + observed_at: datetime + summary: str + session_id: str | None = None + topic_id: str | None = None + cooldown_until: datetime | None = None + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_non_empty("signal_id", self.signal_id) + _require_safe_label("dedupe_key", self.dedupe_key) + _require_safe_label("kind", self.kind) + _require_status("signal status", self.status, _SIGNAL_STATUSES) + _require_datetime("observed_at", self.observed_at) + _require_optional_display_text("summary", self.summary) + if self.session_id is not None: + _require_non_empty("session_id", self.session_id) + if self.topic_id is not None: + _require_non_empty("topic_id", self.topic_id) + if self.cooldown_until is not None: + _require_datetime("cooldown_until", self.cooldown_until) + _require_json_object("metadata", self.metadata) + if "bee_launch" in self.metadata: + _require_bee_launch_metadata_object( + "metadata.bee_launch", + self.metadata["bee_launch"], + ) + + +@dataclass(frozen=True) +class ScheduledLaunchIntent: + trigger_id: str + schedule_id: str + session_id: str + topic_id: str | None + reason: str + due_at: datetime + planned_at: datetime + metadata: JSONObject = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_non_empty("trigger_id", self.trigger_id) + _require_non_empty("schedule_id", self.schedule_id) + _require_non_empty("session_id", self.session_id) + if self.topic_id is not None: + _require_non_empty("topic_id", self.topic_id) + _require_safe_label("reason", self.reason) + _require_datetime("due_at", self.due_at) + _require_datetime("planned_at", self.planned_at) + _require_json_object("metadata", self.metadata) + + +@dataclass(frozen=True) +class PreparedScheduledRun: + trigger_id: str + schedule_id: str + session_id: str + topic_id: str + run_metadata: JSONObject + + def __post_init__(self) -> None: + _require_non_empty("trigger_id", self.trigger_id) + _require_non_empty("schedule_id", self.schedule_id) + _require_non_empty("session_id", self.session_id) + _require_non_empty("topic_id", self.topic_id) + _require_json_object("run_metadata", self.run_metadata) + + +class ScheduledRunPlannerStore(Protocol): + async def list_due_schedules( + self, + *, + due_at_or_before: datetime, + limit: int, + ) -> list[ScheduleRecord]: ... + + async def update_schedule_status( + self, + schedule_id: str, + *, + status: ScheduleStatus, + next_due_at: datetime | None, + last_triggered_at: datetime | None, + updated_at: datetime, + metadata: JSONObject, + ) -> ScheduleRecord: ... + + async def record_trigger( + self, + record: ScheduleTriggerRecord, + ) -> ScheduleTriggerRecord: ... + + +class ProactiveSignalPlannerStore(Protocol): + async def list_signals( + self, + *, + status: SignalStatus | None = None, + session_id: str | None = None, + topic_id: str | None = None, + limit: int = 100, + ) -> list[ProactiveSignalRecord]: ... + + async def update_signal_status( + self, + signal_id: str, + *, + status: SignalStatus, + cooldown_until: datetime | None, + metadata: JSONObject, + ) -> ProactiveSignalRecord: ... + + async def record_trigger( + self, + record: ScheduleTriggerRecord, + ) -> ScheduleTriggerRecord: ... + + +class ScheduledRunPlanner: + def __init__( + self, + *, + store: ScheduledRunPlannerStore, + trigger_id_factory: Callable[[ScheduleRecord, datetime], str] | None = None, + ) -> None: + self._store = store + self._trigger_id_factory = trigger_id_factory or _default_trigger_id + + async def plan_due_schedules( + self, + *, + now: datetime, + max_due: int, + ) -> list[ScheduledLaunchIntent]: + _require_datetime("now", now) + _require_positive_int("max_due", max_due) + schedules = await self._store.list_due_schedules( + due_at_or_before=now, + limit=max_due, + ) + intents: list[ScheduledLaunchIntent] = [] + for schedule in schedules: + if len(intents) >= max_due: + break + if schedule.next_due_at is None or schedule.next_due_at > now: + continue + trigger_id = self._trigger_id_factory(schedule, now) + trigger = await self._store.record_trigger( + ScheduleTriggerRecord( + trigger_id=trigger_id, + schedule_id=schedule.schedule_id, + topic_id=schedule.topic_id, + status="planned", + due_at=schedule.next_due_at, + planned_at=now, + reason="schedule_due", + metadata={ + "trigger_kind": "schedule", + "schedule_kind": schedule.kind, + **_schedule_bee_launch_metadata(schedule), + }, + ) + ) + next_due_at = _next_due_at(schedule, now) + await self._store.update_schedule_status( + schedule.schedule_id, + status="active", + next_due_at=next_due_at, + last_triggered_at=now, + updated_at=now, + metadata={ + **schedule.metadata, + "last_trigger_status": "planned", + }, + ) + intents.append( + ScheduledLaunchIntent( + trigger_id=trigger.trigger_id, + schedule_id=schedule.schedule_id, + session_id=schedule.session_id, + topic_id=schedule.topic_id, + reason="schedule_due", + due_at=trigger.due_at, + planned_at=trigger.planned_at, + metadata={ + "schedule_kind": schedule.kind, + **_schedule_bee_launch_metadata(schedule), + }, + ) + ) + return intents + + +class ProactiveSignalPlanner: + def __init__( + self, + *, + store: ProactiveSignalPlannerStore, + cooldown: timedelta = timedelta(minutes=30), + ) -> None: + if cooldown.total_seconds() <= 0: + raise ValueError("cooldown must be positive") + self._store = store + self._cooldown = cooldown + + async def plan_new_signals( + self, + *, + now: datetime, + max_signals: int, + ) -> list[ScheduledLaunchIntent]: + _require_datetime("now", now) + _require_positive_int("max_signals", max_signals) + signals = await self._store.list_signals(status="new", limit=max_signals) + intents: list[ScheduledLaunchIntent] = [] + for signal in signals: + if len(intents) >= max_signals: + break + if signal.cooldown_until is not None and signal.cooldown_until > now: + await self._store.update_signal_status( + signal.signal_id, + status="ignored", + cooldown_until=signal.cooldown_until, + metadata={ + **signal.metadata, + "skip_reason": "cooldown", + }, + ) + continue + session_id = _required_signal_session(signal) + cooldown_until = now + self._cooldown + updated = await self._store.update_signal_status( + signal.signal_id, + status="planned", + cooldown_until=cooldown_until, + metadata={ + **signal.metadata, + "planned_reason": "proactive_signal", + }, + ) + trigger = await self._store.record_trigger( + ScheduleTriggerRecord( + trigger_id=f"signal-trigger-{updated.signal_id}", + schedule_id=f"signal:{updated.signal_id}", + signal_id=updated.signal_id, + topic_id=updated.topic_id, + status="planned", + due_at=updated.observed_at, + planned_at=now, + reason="proactive_signal", + metadata={ + "trigger_kind": "proactive_signal", + "signal_kind": updated.kind, + **_signal_bee_launch_metadata(updated), + }, + ) + ) + intents.append( + ScheduledLaunchIntent( + trigger_id=trigger.trigger_id, + schedule_id=trigger.schedule_id, + session_id=session_id, + topic_id=updated.topic_id, + reason="proactive_signal", + due_at=trigger.due_at, + planned_at=trigger.planned_at, + metadata={ + "signal_kind": updated.kind, + **_signal_bee_launch_metadata(updated), + }, + ) + ) + return intents + + +class ScheduledTopicStore(Protocol): + async def load_topic(self, topic_id: str) -> TopicRecord | None: ... + + async def find_open_topic( + self, + *, + session_id: str, + tape_id: str, + ) -> TopicRecord | None: ... + + +class ScheduledRunLaunchPreparer: + def __init__( + self, + *, + topic_store: ScheduledTopicStore, + topic_lifecycle: TopicLifecycle, + ) -> None: + self._topic_store = topic_store + self._topic_lifecycle = topic_lifecycle + + async def prepare( + self, + *, + intent: ScheduledLaunchIntent, + tape: Tape, + ) -> PreparedScheduledRun: + topic = await self._resolve_topic(intent=intent, tape=tape) + metadata: JSONObject = { + "scheduled_run": { + "schedule_id": intent.schedule_id, + "trigger_id": intent.trigger_id, + "reason": intent.reason, + "planned_at": intent.planned_at.isoformat(), + }, + "topic": { + "topic_id": topic.topic_id, + "tape_id": topic.tape_id, + "session_id": topic.session_id, + "kind": topic.kind, + "status": topic.status, + "topic_initial_seq": topic.topic_initial_seq, + "topic_finalized_seq": topic.topic_finalized_seq, + }, + } + return PreparedScheduledRun( + trigger_id=intent.trigger_id, + schedule_id=intent.schedule_id, + session_id=intent.session_id, + topic_id=topic.topic_id, + run_metadata=metadata, + ) + + async def _resolve_topic( + self, + *, + intent: ScheduledLaunchIntent, + tape: Tape, + ) -> TopicRecord: + if intent.topic_id is not None: + topic = await self._topic_store.load_topic(intent.topic_id) + if topic is not None and topic.status == "open": + _require_topic_matches_intent(topic=topic, intent=intent, tape=tape) + return topic + raise ValueError("scheduled topic is not open") + open_topic = await self._topic_store.find_open_topic( + session_id=intent.session_id, + tape_id=tape.tape_id, + ) + if open_topic is not None: + return open_topic + return await self._create_scheduled_topic(intent=intent, tape=tape) + + async def _create_scheduled_topic( + self, + *, + intent: ScheduledLaunchIntent, + tape: Tape, + ) -> TopicRecord: + metadata: JSONObject = { + "source": "scheduled_run", + "schedule_kind": _safe_metadata_label(intent.metadata.get("schedule_kind")), + } + return await self._topic_lifecycle.create_topic( + tape=tape, + session_id=intent.session_id, + kind="coding", + title="Scheduled run", + owner=None, + metadata=metadata, + ) + + +class PGScheduledRunStore: + _CREATE_SCHEMA_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS scheduled_runs ( + schedule_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + topic_id TEXT, + kind TEXT NOT NULL, + status TEXT NOT NULL, + cadence TEXT NOT NULL, + owner TEXT, + title TEXT, + next_due_at TIMESTAMPTZ, + last_triggered_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL, + updated_at TIMESTAMPTZ NOT NULL, + metadata JSONB NOT NULL DEFAULT '{}'::jsonb + ); + + CREATE INDEX IF NOT EXISTS scheduled_runs_session_status_due_idx + ON scheduled_runs (session_id, status, next_due_at, schedule_id); + + CREATE INDEX IF NOT EXISTS scheduled_runs_topic_status_due_idx + ON scheduled_runs (topic_id, status, next_due_at, schedule_id); + + CREATE TABLE IF NOT EXISTS scheduled_run_triggers ( + trigger_id TEXT PRIMARY KEY, + schedule_id TEXT NOT NULL, + signal_id TEXT, + topic_id TEXT, + run_id TEXT, + status TEXT NOT NULL, + due_at TIMESTAMPTZ NOT NULL, + planned_at TIMESTAMPTZ NOT NULL, + reason TEXT, + metadata JSONB NOT NULL DEFAULT '{}'::jsonb + ); + + CREATE INDEX IF NOT EXISTS scheduled_run_triggers_schedule_due_idx + ON scheduled_run_triggers (schedule_id, due_at, trigger_id); + + CREATE TABLE IF NOT EXISTS proactive_signals ( + signal_id TEXT PRIMARY KEY, + dedupe_key TEXT NOT NULL UNIQUE, + session_id TEXT, + topic_id TEXT, + kind TEXT NOT NULL, + status TEXT NOT NULL, + observed_at TIMESTAMPTZ NOT NULL, + cooldown_until TIMESTAMPTZ, + summary TEXT NOT NULL, + metadata JSONB NOT NULL DEFAULT '{}'::jsonb + ); + + CREATE INDEX IF NOT EXISTS proactive_signals_status_observed_idx + ON proactive_signals (status, observed_at, signal_id); + """ + _UPSERT_SCHEDULE_SQL: Final[str] = """ + INSERT INTO scheduled_runs ( + schedule_id, + session_id, + topic_id, + kind, + status, + cadence, + owner, + title, + next_due_at, + last_triggered_at, + created_at, + updated_at, + metadata + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13::jsonb) + ON CONFLICT (schedule_id) + DO UPDATE SET + session_id = EXCLUDED.session_id, + topic_id = EXCLUDED.topic_id, + kind = EXCLUDED.kind, + status = EXCLUDED.status, + cadence = EXCLUDED.cadence, + owner = EXCLUDED.owner, + title = EXCLUDED.title, + next_due_at = EXCLUDED.next_due_at, + last_triggered_at = EXCLUDED.last_triggered_at, + updated_at = EXCLUDED.updated_at, + metadata = EXCLUDED.metadata + RETURNING * + """ + _SELECT_SCHEDULE_SQL: Final[str] = """ + SELECT * FROM scheduled_runs WHERE schedule_id = $1 + """ + _LIST_SCHEDULES_SQL: Final[str] = """ + SELECT * FROM scheduled_runs + WHERE ($1::text IS NULL OR session_id = $1) + AND ($2::text IS NULL OR topic_id = $2) + AND ($3::text IS NULL OR status = $3) + ORDER BY COALESCE(next_due_at, updated_at), schedule_id + LIMIT $4 + """ + _LIST_DUE_SCHEDULES_SQL: Final[str] = """ + SELECT * FROM scheduled_runs + WHERE status = 'active' + AND next_due_at IS NOT NULL + AND next_due_at <= $1 + ORDER BY next_due_at, schedule_id + LIMIT $2 + """ + _UPDATE_SCHEDULE_STATUS_SQL: Final[str] = """ + UPDATE scheduled_runs + SET status = $2, + next_due_at = $3, + last_triggered_at = $4, + updated_at = $5, + metadata = $6::jsonb + WHERE schedule_id = $1 + RETURNING * + """ + _INSERT_TRIGGER_SQL: Final[str] = """ + INSERT INTO scheduled_run_triggers ( + trigger_id, + schedule_id, + signal_id, + topic_id, + run_id, + status, + due_at, + planned_at, + reason, + metadata + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb) + ON CONFLICT (trigger_id) + DO UPDATE SET + schedule_id = EXCLUDED.schedule_id, + signal_id = EXCLUDED.signal_id, + topic_id = EXCLUDED.topic_id, + run_id = EXCLUDED.run_id, + status = EXCLUDED.status, + due_at = EXCLUDED.due_at, + planned_at = EXCLUDED.planned_at, + reason = EXCLUDED.reason, + metadata = EXCLUDED.metadata + RETURNING * + """ + _LIST_TRIGGERS_SQL: Final[str] = """ + SELECT * FROM scheduled_run_triggers + WHERE schedule_id = $1 + ORDER BY due_at, trigger_id + LIMIT $2 + """ + _UPSERT_SIGNAL_SQL: Final[str] = """ + INSERT INTO proactive_signals ( + signal_id, + dedupe_key, + session_id, + topic_id, + kind, + status, + observed_at, + cooldown_until, + summary, + metadata + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb) + ON CONFLICT (dedupe_key) + DO UPDATE SET + status = proactive_signals.status, + metadata = proactive_signals.metadata + RETURNING * + """ + _SELECT_SIGNAL_SQL: Final[str] = """ + SELECT * FROM proactive_signals WHERE signal_id = $1 + """ + _LIST_SIGNALS_SQL: Final[str] = """ + SELECT * FROM proactive_signals + WHERE ($1::text IS NULL OR status = $1) + AND ($2::text IS NULL OR session_id = $2) + AND ($3::text IS NULL OR topic_id = $3) + ORDER BY observed_at, signal_id + LIMIT $4 + """ + _UPDATE_SIGNAL_STATUS_SQL: Final[str] = """ + UPDATE proactive_signals + SET status = $2, + cooldown_until = $3, + metadata = $4::jsonb + WHERE signal_id = $1 + RETURNING * + """ + + def __init__(self, *, pool: PGPool) -> None: + self._pool = pool + self._schema_ready = False + + async def _ensure_schema(self) -> AsyncPGPool: + pool = await self._pool.get_pool() + if not self._schema_ready: + _ = await pool.execute(self._CREATE_SCHEMA_SQL) + self._schema_ready = True + return pool + + async def create_schedule(self, record: ScheduleRecord) -> ScheduleRecord: + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._UPSERT_SCHEDULE_SQL, + record.schedule_id, + record.session_id, + record.topic_id, + record.kind, + record.status, + record.cadence, + record.owner, + record.title, + record.next_due_at, + record.last_triggered_at, + record.created_at, + record.updated_at, + record.metadata, + ) + return _schedule_from_row(_required_row(row, "schedule upsert")) + + async def load_schedule(self, schedule_id: str) -> ScheduleRecord | None: + _require_non_empty("schedule_id", schedule_id) + pool = await self._ensure_schema() + row = await pool.fetchrow(self._SELECT_SCHEDULE_SQL, schedule_id) + return None if row is None else _schedule_from_row(row) + + async def list_schedules( + self, + *, + session_id: str | None = None, + topic_id: str | None = None, + status: ScheduleStatus | None = None, + limit: int = 100, + ) -> list[ScheduleRecord]: + if session_id is not None: + _require_non_empty("session_id", session_id) + if topic_id is not None: + _require_non_empty("topic_id", topic_id) + if status is not None: + _require_status("schedule status", status, _SCHEDULE_STATUSES) + _require_positive_int("limit", limit) + pool = await self._ensure_schema() + rows = await pool.fetch( + self._LIST_SCHEDULES_SQL, session_id, topic_id, status, limit + ) + return [_schedule_from_row(row) for row in rows] + + async def list_due_schedules( + self, + *, + due_at_or_before: datetime, + limit: int = 100, + ) -> list[ScheduleRecord]: + _require_datetime("due_at_or_before", due_at_or_before) + _require_positive_int("limit", limit) + pool = await self._ensure_schema() + rows = await pool.fetch(self._LIST_DUE_SCHEDULES_SQL, due_at_or_before, limit) + return [_schedule_from_row(row) for row in rows] + + async def update_schedule_status( + self, + schedule_id: str, + *, + status: ScheduleStatus, + next_due_at: datetime | None, + last_triggered_at: datetime | None, + updated_at: datetime, + metadata: JSONObject, + ) -> ScheduleRecord: + _require_non_empty("schedule_id", schedule_id) + _require_status("schedule status", status, _SCHEDULE_STATUSES) + if next_due_at is not None: + _require_datetime("next_due_at", next_due_at) + if last_triggered_at is not None: + _require_datetime("last_triggered_at", last_triggered_at) + _require_datetime("updated_at", updated_at) + _require_json_object("metadata", metadata) + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._UPDATE_SCHEDULE_STATUS_SQL, + schedule_id, + status, + next_due_at, + last_triggered_at, + updated_at, + metadata, + ) + if row is None: + raise KeyError(f"schedule not found: {schedule_id}") + return _schedule_from_row(row) + + async def record_trigger( + self, + record: ScheduleTriggerRecord, + ) -> ScheduleTriggerRecord: + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._INSERT_TRIGGER_SQL, + record.trigger_id, + record.schedule_id, + record.signal_id, + record.topic_id, + record.run_id, + record.status, + record.due_at, + record.planned_at, + record.reason, + record.metadata, + ) + return _trigger_from_row(_required_row(row, "schedule trigger upsert")) + + async def list_triggers( + self, + schedule_id: str, + *, + limit: int = 100, + ) -> list[ScheduleTriggerRecord]: + _require_non_empty("schedule_id", schedule_id) + _require_positive_int("limit", limit) + pool = await self._ensure_schema() + rows = await pool.fetch(self._LIST_TRIGGERS_SQL, schedule_id, limit) + return [_trigger_from_row(row) for row in rows] + + async def record_signal( + self, + record: ProactiveSignalRecord, + ) -> ProactiveSignalRecord: + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._UPSERT_SIGNAL_SQL, + record.signal_id, + record.dedupe_key, + record.session_id, + record.topic_id, + record.kind, + record.status, + record.observed_at, + record.cooldown_until, + record.summary, + record.metadata, + ) + return _signal_from_row(_required_row(row, "proactive signal upsert")) + + async def load_signal(self, signal_id: str) -> ProactiveSignalRecord | None: + _require_non_empty("signal_id", signal_id) + pool = await self._ensure_schema() + row = await pool.fetchrow(self._SELECT_SIGNAL_SQL, signal_id) + return None if row is None else _signal_from_row(row) + + async def list_signals( + self, + *, + status: SignalStatus | None = None, + session_id: str | None = None, + topic_id: str | None = None, + limit: int = 100, + ) -> list[ProactiveSignalRecord]: + if status is not None: + _require_status("signal status", status, _SIGNAL_STATUSES) + if session_id is not None: + _require_non_empty("session_id", session_id) + if topic_id is not None: + _require_non_empty("topic_id", topic_id) + _require_positive_int("limit", limit) + pool = await self._ensure_schema() + rows = await pool.fetch( + self._LIST_SIGNALS_SQL, status, session_id, topic_id, limit + ) + return [_signal_from_row(row) for row in rows] + + async def update_signal_status( + self, + signal_id: str, + *, + status: SignalStatus, + cooldown_until: datetime | None, + metadata: JSONObject, + ) -> ProactiveSignalRecord: + _require_non_empty("signal_id", signal_id) + _require_status("signal status", status, _SIGNAL_STATUSES) + if cooldown_until is not None: + _require_datetime("cooldown_until", cooldown_until) + _require_json_object("metadata", metadata) + pool = await self._ensure_schema() + row = await pool.fetchrow( + self._UPDATE_SIGNAL_STATUS_SQL, + signal_id, + status, + cooldown_until, + metadata, + ) + if row is None: + raise KeyError(f"proactive signal not found: {signal_id}") + return _signal_from_row(row) + + +def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: + if row is None: + raise RuntimeError(f"postgres {context} returned no row") + return row + + +def _schedule_from_row(row: dict[str, object]) -> ScheduleRecord: + return ScheduleRecord( + schedule_id=_required_str(row, "schedule_id", context="schedule row"), + session_id=_required_str(row, "session_id", context="schedule row"), + topic_id=_optional_str(row, "topic_id", context="schedule row"), + kind=_required_str(row, "kind", context="schedule row"), + status=_required_str(row, "status", context="schedule row"), + cadence=_required_str(row, "cadence", context="schedule row"), + owner=_optional_str(row, "owner", context="schedule row"), + title=_optional_str(row, "title", context="schedule row"), + next_due_at=_optional_datetime(row, "next_due_at", context="schedule row"), + last_triggered_at=_optional_datetime( + row, "last_triggered_at", context="schedule row" + ), + created_at=_required_datetime(row, "created_at", context="schedule row"), + updated_at=_required_datetime(row, "updated_at", context="schedule row"), + metadata=_required_json_object(row, "metadata", context="schedule row"), + ) + + +def _trigger_from_row(row: dict[str, object]) -> ScheduleTriggerRecord: + return ScheduleTriggerRecord( + trigger_id=_required_str(row, "trigger_id", context="trigger row"), + schedule_id=_required_str(row, "schedule_id", context="trigger row"), + signal_id=_optional_str(row, "signal_id", context="trigger row"), + topic_id=_optional_str(row, "topic_id", context="trigger row"), + run_id=_optional_str(row, "run_id", context="trigger row"), + status=_required_str(row, "status", context="trigger row"), + due_at=_required_datetime(row, "due_at", context="trigger row"), + planned_at=_required_datetime(row, "planned_at", context="trigger row"), + reason=_optional_str(row, "reason", context="trigger row"), + metadata=_required_json_object(row, "metadata", context="trigger row"), + ) + + +def _signal_from_row(row: dict[str, object]) -> ProactiveSignalRecord: + return ProactiveSignalRecord( + signal_id=_required_str(row, "signal_id", context="signal row"), + dedupe_key=_required_str(row, "dedupe_key", context="signal row"), + session_id=_optional_str(row, "session_id", context="signal row"), + topic_id=_optional_str(row, "topic_id", context="signal row"), + kind=_required_str(row, "kind", context="signal row"), + status=_required_str(row, "status", context="signal row"), + observed_at=_required_datetime(row, "observed_at", context="signal row"), + cooldown_until=_optional_datetime(row, "cooldown_until", context="signal row"), + summary=_required_str(row, "summary", context="signal row"), + metadata=_required_json_object(row, "metadata", context="signal row"), + ) + + +def _require_non_empty(field_name: str, value: str) -> None: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field_name} must be a non-empty string") + + +def _require_safe_label(field_name: str, value: str) -> None: + _require_non_empty(field_name, value) + if len(value) > _MAX_SAFE_LABEL_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_SAFE_LABEL_CHARS} characters" + ) + key_folded = field_name.casefold() + value_folded = value.casefold() + if any(part in value_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): + raise ValueError(f"{field_name} must not contain sensitive label text") + if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): + raise ValueError(f"{field_name} must not be sensitive label text") + _reject_secret_shaped_value(field_name, value) + + +def _require_status(field_name: str, value: str, allowed: frozenset[str]) -> None: + if value not in allowed: + raise ValueError(f"{field_name} must be one of {sorted(allowed)}") + + +def _require_optional_display_text(field_name: str, value: str | None) -> None: + if value is None: + return + _require_non_empty(field_name, value) + if len(value) > _MAX_DISPLAY_TEXT_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" + ) + _reject_secret_shaped_value(field_name, value) + + +def _require_datetime(field_name: str, value: datetime) -> None: + if not isinstance(value, datetime): + raise TypeError(f"{field_name} must be a datetime") + + +def _require_positive_int(field_name: str, value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{field_name} must be a positive int") + + +def _require_json_object(field_name: str, value: JSONObject) -> None: + if not isinstance(value, dict): + raise TypeError(f"{field_name} must be a JSON object") + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{field_name} keys must be strings") + key_folded = key.casefold() + if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): + raise ValueError(f"{field_name} contains forbidden metadata key: {key}") + _require_json_value(f"{field_name}.{key}", item) + + +def _require_json_value(field_name: str, value: object) -> None: + if value is None or isinstance(value, str | int | float | bool): + if isinstance(value, str): + if len(value) > _MAX_METADATA_STRING_CHARS: + raise ValueError( + f"{field_name} must be at most {_MAX_METADATA_STRING_CHARS} characters" + ) + _reject_secret_shaped_value(field_name, value) + return + if isinstance(value, list): + for index, item in enumerate(value): + _require_json_value(f"{field_name}[{index}]", item) + return + if isinstance(value, dict): + _require_json_object(field_name, value) + return + raise TypeError(f"{field_name} must be JSON-safe") + + +def _reject_secret_shaped_value(field_name: str, value: str) -> None: + folded = value.casefold() + if any( + marker in folded for marker in _SECRET_VALUE_MARKERS if marker != "sk-" + ) or folded.startswith("sk-"): + raise ValueError(f"{field_name} must not contain secret-shaped values") + + +def _default_trigger_id(schedule: ScheduleRecord, now: datetime) -> str: + safe_time = now.isoformat().replace(":", "").replace("+", "_") + return f"trigger-{schedule.schedule_id}-{safe_time}" + + +def _next_due_at(schedule: ScheduleRecord, now: datetime) -> datetime | None: + if schedule.cadence == "once": + return None + if schedule.cadence == "hourly": + return now + timedelta(hours=1) + if schedule.cadence == "daily": + return now + timedelta(days=1) + return None + + +def _schedule_bee_launch_metadata(schedule: ScheduleRecord) -> JSONObject: + value = schedule.metadata.get("bee_launch") + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("schedule.metadata.bee_launch must be an object") + bee_launch = dict(value) + _require_bee_launch_metadata_object("schedule.metadata.bee_launch", bee_launch) + return {"bee_launch": bee_launch} + + +def _signal_bee_launch_metadata(signal: ProactiveSignalRecord) -> JSONObject: + value = signal.metadata.get("bee_launch") + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("signal.metadata.bee_launch must be an object") + bee_launch = dict(value) + _require_bee_launch_metadata_object("signal.metadata.bee_launch", bee_launch) + return {"bee_launch": bee_launch} + + +def _require_bee_launch_metadata_object( + field_name: str, + value: JSONObject, +) -> None: + _require_json_object(field_name, value) + for key, item in value.items(): + key_folded = key.casefold() + if any(part in key_folded for part in _FORBIDDEN_BEE_LAUNCH_KEY_PARTS): + raise ValueError( + f"{field_name} contains forbidden Bee launch metadata key: {key}" + ) + if isinstance(item, dict): + _require_bee_launch_metadata_object(f"{field_name}.{key}", item) + elif isinstance(item, list): + for index, nested in enumerate(item): + if isinstance(nested, dict): + _require_bee_launch_metadata_object( + f"{field_name}.{key}[{index}]", + nested, + ) + + +def _require_topic_matches_intent( + *, + topic: TopicRecord, + intent: ScheduledLaunchIntent, + tape: Tape, +) -> None: + if topic.session_id != intent.session_id: + raise ValueError("scheduled topic session does not match launch intent") + if topic.tape_id != tape.tape_id: + raise ValueError("scheduled topic tape does not match launch tape") + + +def _safe_metadata_label(value: object) -> str: + if not isinstance(value, str): + return "unknown" + _require_safe_label("metadata_label", value) + return value + + +def _required_signal_session(signal: ProactiveSignalRecord) -> str: + if signal.session_id is None: + raise ValueError("proactive signal requires session_id for launch planning") + return signal.session_id + + +def _required_str(row: dict[str, object], key: str, *, context: str) -> str: + value = row.get(key) + if not isinstance(value, str): + raise TypeError(f"postgres {context} must include string {key}") + return value + + +def _optional_str(row: dict[str, object], key: str, *, context: str) -> str | None: + value = row.get(key) + if value is None: + return None + if not isinstance(value, str): + raise TypeError(f"postgres {context} must include string or None {key}") + return value + + +def _required_datetime(row: dict[str, object], key: str, *, context: str) -> datetime: + value = row.get(key) + if not isinstance(value, datetime): + raise TypeError(f"postgres {context} must include datetime {key}") + return value + + +def _optional_datetime( + row: dict[str, object], key: str, *, context: str +) -> datetime | None: + value = row.get(key) + if value is None: + return None + if not isinstance(value, datetime): + raise TypeError(f"postgres {context} must include datetime or None {key}") + return value + + +def _required_json_object( + row: dict[str, object], key: str, *, context: str +) -> JSONObject: + value = row.get(key) + if not isinstance(value, dict): + raise TypeError(f"postgres {context} must include dict {key}") + _require_json_object(key, value) + return value diff --git a/coding-agent/src/coding_agent/scheduled_runs.py b/coding-agent/src/coding_agent/scheduled_runs.py index 11e73e19c..22beee6ed 100644 --- a/coding-agent/src/coding_agent/scheduled_runs.py +++ b/coding-agent/src/coding_agent/scheduled_runs.py @@ -1,1192 +1,4 @@ -"""Durable scheduled run and proactive signal records.""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass, field -from datetime import datetime, timedelta -from typing import Final, Protocol - -from agentkit.storage.pg import AsyncPGPool, PGPool -from agentkit.tape.tape import Tape -from coding_agent.topic_lifecycle import TopicLifecycle -from coding_agent.topic_store import JSONObject -from coding_agent.topic_store import TopicRecord - -ScheduleStatus = str -SignalStatus = str -TriggerStatus = str - -_SCHEDULE_STATUSES: Final[frozenset[str]] = frozenset( - {"active", "paused", "disabled", "completed"} -) -_SIGNAL_STATUSES: Final[frozenset[str]] = frozenset( - {"new", "planned", "ignored", "consumed"} -) -_TRIGGER_STATUSES: Final[frozenset[str]] = frozenset( - {"planned", "launched", "skipped", "failed"} -) -_MAX_DISPLAY_TEXT_CHARS: Final[int] = 256 -_MAX_METADATA_STRING_CHARS: Final[int] = 256 -_MAX_SAFE_LABEL_CHARS: Final[int] = 128 -_FORBIDDEN_METADATA_KEY_PARTS: Final[frozenset[str]] = frozenset( - { - "command_output", - "content", - "env", - "message", - "prompt", - "result", - "secret", - "stderr", - "stdout", - "text", - } -) -_FORBIDDEN_BEE_LAUNCH_KEY_PARTS: Final[frozenset[str]] = frozenset( - { - "api_key", - "apikey", - "args", - "argv", - "bearer", - "cmd", - "command", - "commands", - "command_output", - "content", - "credential", - "credentials", - "env", - "environment", - "exec", - "executor", - "message", - "password", - "prompt", - "result", - "script", - "secret", - "shell", - "stderr", - "stdout", - "text", - "token", - } -) -_SECRET_VALUE_MARKERS: Final[tuple[str, ...]] = ( - "-----begin ", - "akia", - "password=", - "secret=", - "sk-", - "token=", -) - - -@dataclass(frozen=True) -class ScheduleRecord: - schedule_id: str - session_id: str - kind: str - status: ScheduleStatus - cadence: str - created_at: datetime - updated_at: datetime - topic_id: str | None = None - owner: str | None = None - title: str | None = None - next_due_at: datetime | None = None - last_triggered_at: datetime | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_non_empty("schedule_id", self.schedule_id) - _require_non_empty("session_id", self.session_id) - _require_safe_label("kind", self.kind) - _require_safe_label("cadence", self.cadence) - _require_status("schedule status", self.status, _SCHEDULE_STATUSES) - _require_datetime("created_at", self.created_at) - _require_datetime("updated_at", self.updated_at) - if self.topic_id is not None: - _require_non_empty("topic_id", self.topic_id) - _require_optional_display_text("owner", self.owner) - _require_optional_display_text("title", self.title) - if self.next_due_at is not None: - _require_datetime("next_due_at", self.next_due_at) - if self.last_triggered_at is not None: - _require_datetime("last_triggered_at", self.last_triggered_at) - _require_json_object("metadata", self.metadata) - if "bee_launch" in self.metadata: - _require_bee_launch_metadata_object( - "metadata.bee_launch", - self.metadata["bee_launch"], - ) - - -@dataclass(frozen=True) -class ScheduleTriggerRecord: - trigger_id: str - schedule_id: str - status: TriggerStatus - due_at: datetime - planned_at: datetime - topic_id: str | None = None - signal_id: str | None = None - run_id: str | None = None - reason: str | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_non_empty("trigger_id", self.trigger_id) - _require_non_empty("schedule_id", self.schedule_id) - _require_status("trigger status", self.status, _TRIGGER_STATUSES) - _require_datetime("due_at", self.due_at) - _require_datetime("planned_at", self.planned_at) - if self.topic_id is not None: - _require_non_empty("topic_id", self.topic_id) - if self.signal_id is not None: - _require_non_empty("signal_id", self.signal_id) - if self.run_id is not None: - _require_non_empty("run_id", self.run_id) - _require_optional_display_text("reason", self.reason) - _require_json_object("metadata", self.metadata) - - -@dataclass(frozen=True) -class ProactiveSignalRecord: - signal_id: str - dedupe_key: str - kind: str - status: SignalStatus - observed_at: datetime - summary: str - session_id: str | None = None - topic_id: str | None = None - cooldown_until: datetime | None = None - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_non_empty("signal_id", self.signal_id) - _require_safe_label("dedupe_key", self.dedupe_key) - _require_safe_label("kind", self.kind) - _require_status("signal status", self.status, _SIGNAL_STATUSES) - _require_datetime("observed_at", self.observed_at) - _require_optional_display_text("summary", self.summary) - if self.session_id is not None: - _require_non_empty("session_id", self.session_id) - if self.topic_id is not None: - _require_non_empty("topic_id", self.topic_id) - if self.cooldown_until is not None: - _require_datetime("cooldown_until", self.cooldown_until) - _require_json_object("metadata", self.metadata) - if "bee_launch" in self.metadata: - _require_bee_launch_metadata_object( - "metadata.bee_launch", - self.metadata["bee_launch"], - ) - - -@dataclass(frozen=True) -class ScheduledLaunchIntent: - trigger_id: str - schedule_id: str - session_id: str - topic_id: str | None - reason: str - due_at: datetime - planned_at: datetime - metadata: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - _require_non_empty("trigger_id", self.trigger_id) - _require_non_empty("schedule_id", self.schedule_id) - _require_non_empty("session_id", self.session_id) - if self.topic_id is not None: - _require_non_empty("topic_id", self.topic_id) - _require_safe_label("reason", self.reason) - _require_datetime("due_at", self.due_at) - _require_datetime("planned_at", self.planned_at) - _require_json_object("metadata", self.metadata) - - -@dataclass(frozen=True) -class PreparedScheduledRun: - trigger_id: str - schedule_id: str - session_id: str - topic_id: str - run_metadata: JSONObject - - def __post_init__(self) -> None: - _require_non_empty("trigger_id", self.trigger_id) - _require_non_empty("schedule_id", self.schedule_id) - _require_non_empty("session_id", self.session_id) - _require_non_empty("topic_id", self.topic_id) - _require_json_object("run_metadata", self.run_metadata) - - -class ScheduledRunPlannerStore(Protocol): - async def list_due_schedules( - self, - *, - due_at_or_before: datetime, - limit: int, - ) -> list[ScheduleRecord]: ... - - async def update_schedule_status( - self, - schedule_id: str, - *, - status: ScheduleStatus, - next_due_at: datetime | None, - last_triggered_at: datetime | None, - updated_at: datetime, - metadata: JSONObject, - ) -> ScheduleRecord: ... - - async def record_trigger( - self, - record: ScheduleTriggerRecord, - ) -> ScheduleTriggerRecord: ... - - -class ProactiveSignalPlannerStore(Protocol): - async def list_signals( - self, - *, - status: SignalStatus | None = None, - session_id: str | None = None, - topic_id: str | None = None, - limit: int = 100, - ) -> list[ProactiveSignalRecord]: ... - - async def update_signal_status( - self, - signal_id: str, - *, - status: SignalStatus, - cooldown_until: datetime | None, - metadata: JSONObject, - ) -> ProactiveSignalRecord: ... - - async def record_trigger( - self, - record: ScheduleTriggerRecord, - ) -> ScheduleTriggerRecord: ... - - -class ScheduledRunPlanner: - def __init__( - self, - *, - store: ScheduledRunPlannerStore, - trigger_id_factory: Callable[[ScheduleRecord, datetime], str] | None = None, - ) -> None: - self._store = store - self._trigger_id_factory = trigger_id_factory or _default_trigger_id - - async def plan_due_schedules( - self, - *, - now: datetime, - max_due: int, - ) -> list[ScheduledLaunchIntent]: - _require_datetime("now", now) - _require_positive_int("max_due", max_due) - schedules = await self._store.list_due_schedules( - due_at_or_before=now, - limit=max_due, - ) - intents: list[ScheduledLaunchIntent] = [] - for schedule in schedules: - if len(intents) >= max_due: - break - if schedule.next_due_at is None or schedule.next_due_at > now: - continue - trigger_id = self._trigger_id_factory(schedule, now) - trigger = await self._store.record_trigger( - ScheduleTriggerRecord( - trigger_id=trigger_id, - schedule_id=schedule.schedule_id, - topic_id=schedule.topic_id, - status="planned", - due_at=schedule.next_due_at, - planned_at=now, - reason="schedule_due", - metadata={ - "trigger_kind": "schedule", - "schedule_kind": schedule.kind, - **_schedule_bee_launch_metadata(schedule), - }, - ) - ) - next_due_at = _next_due_at(schedule, now) - await self._store.update_schedule_status( - schedule.schedule_id, - status="active", - next_due_at=next_due_at, - last_triggered_at=now, - updated_at=now, - metadata={ - **schedule.metadata, - "last_trigger_status": "planned", - }, - ) - intents.append( - ScheduledLaunchIntent( - trigger_id=trigger.trigger_id, - schedule_id=schedule.schedule_id, - session_id=schedule.session_id, - topic_id=schedule.topic_id, - reason="schedule_due", - due_at=trigger.due_at, - planned_at=trigger.planned_at, - metadata={ - "schedule_kind": schedule.kind, - **_schedule_bee_launch_metadata(schedule), - }, - ) - ) - return intents - - -class ProactiveSignalPlanner: - def __init__( - self, - *, - store: ProactiveSignalPlannerStore, - cooldown: timedelta = timedelta(minutes=30), - ) -> None: - if cooldown.total_seconds() <= 0: - raise ValueError("cooldown must be positive") - self._store = store - self._cooldown = cooldown - - async def plan_new_signals( - self, - *, - now: datetime, - max_signals: int, - ) -> list[ScheduledLaunchIntent]: - _require_datetime("now", now) - _require_positive_int("max_signals", max_signals) - signals = await self._store.list_signals(status="new", limit=max_signals) - intents: list[ScheduledLaunchIntent] = [] - for signal in signals: - if len(intents) >= max_signals: - break - if signal.cooldown_until is not None and signal.cooldown_until > now: - await self._store.update_signal_status( - signal.signal_id, - status="ignored", - cooldown_until=signal.cooldown_until, - metadata={ - **signal.metadata, - "skip_reason": "cooldown", - }, - ) - continue - session_id = _required_signal_session(signal) - cooldown_until = now + self._cooldown - updated = await self._store.update_signal_status( - signal.signal_id, - status="planned", - cooldown_until=cooldown_until, - metadata={ - **signal.metadata, - "planned_reason": "proactive_signal", - }, - ) - trigger = await self._store.record_trigger( - ScheduleTriggerRecord( - trigger_id=f"signal-trigger-{updated.signal_id}", - schedule_id=f"signal:{updated.signal_id}", - signal_id=updated.signal_id, - topic_id=updated.topic_id, - status="planned", - due_at=updated.observed_at, - planned_at=now, - reason="proactive_signal", - metadata={ - "trigger_kind": "proactive_signal", - "signal_kind": updated.kind, - **_signal_bee_launch_metadata(updated), - }, - ) - ) - intents.append( - ScheduledLaunchIntent( - trigger_id=trigger.trigger_id, - schedule_id=trigger.schedule_id, - session_id=session_id, - topic_id=updated.topic_id, - reason="proactive_signal", - due_at=trigger.due_at, - planned_at=trigger.planned_at, - metadata={ - "signal_kind": updated.kind, - **_signal_bee_launch_metadata(updated), - }, - ) - ) - return intents - - -class ScheduledTopicStore(Protocol): - async def load_topic(self, topic_id: str) -> TopicRecord | None: ... - - async def find_open_topic( - self, - *, - session_id: str, - tape_id: str, - ) -> TopicRecord | None: ... - - -class ScheduledRunLaunchPreparer: - def __init__( - self, - *, - topic_store: ScheduledTopicStore, - topic_lifecycle: TopicLifecycle, - ) -> None: - self._topic_store = topic_store - self._topic_lifecycle = topic_lifecycle - - async def prepare( - self, - *, - intent: ScheduledLaunchIntent, - tape: Tape, - ) -> PreparedScheduledRun: - topic = await self._resolve_topic(intent=intent, tape=tape) - metadata: JSONObject = { - "scheduled_run": { - "schedule_id": intent.schedule_id, - "trigger_id": intent.trigger_id, - "reason": intent.reason, - "planned_at": intent.planned_at.isoformat(), - }, - "topic": { - "topic_id": topic.topic_id, - "tape_id": topic.tape_id, - "session_id": topic.session_id, - "kind": topic.kind, - "status": topic.status, - "topic_initial_seq": topic.topic_initial_seq, - "topic_finalized_seq": topic.topic_finalized_seq, - }, - } - return PreparedScheduledRun( - trigger_id=intent.trigger_id, - schedule_id=intent.schedule_id, - session_id=intent.session_id, - topic_id=topic.topic_id, - run_metadata=metadata, - ) - - async def _resolve_topic( - self, - *, - intent: ScheduledLaunchIntent, - tape: Tape, - ) -> TopicRecord: - if intent.topic_id is not None: - topic = await self._topic_store.load_topic(intent.topic_id) - if topic is not None and topic.status == "open": - _require_topic_matches_intent(topic=topic, intent=intent, tape=tape) - return topic - raise ValueError("scheduled topic is not open") - open_topic = await self._topic_store.find_open_topic( - session_id=intent.session_id, - tape_id=tape.tape_id, - ) - if open_topic is not None: - return open_topic - return await self._create_scheduled_topic(intent=intent, tape=tape) - - async def _create_scheduled_topic( - self, - *, - intent: ScheduledLaunchIntent, - tape: Tape, - ) -> TopicRecord: - metadata: JSONObject = { - "source": "scheduled_run", - "schedule_kind": _safe_metadata_label(intent.metadata.get("schedule_kind")), - } - return await self._topic_lifecycle.create_topic( - tape=tape, - session_id=intent.session_id, - kind="coding", - title="Scheduled run", - owner=None, - metadata=metadata, - ) - - -class PGScheduledRunStore: - _CREATE_SCHEMA_SQL: Final[str] = """ - CREATE TABLE IF NOT EXISTS scheduled_runs ( - schedule_id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - topic_id TEXT, - kind TEXT NOT NULL, - status TEXT NOT NULL, - cadence TEXT NOT NULL, - owner TEXT, - title TEXT, - next_due_at TIMESTAMPTZ, - last_triggered_at TIMESTAMPTZ, - created_at TIMESTAMPTZ NOT NULL, - updated_at TIMESTAMPTZ NOT NULL, - metadata JSONB NOT NULL DEFAULT '{}'::jsonb - ); - - CREATE INDEX IF NOT EXISTS scheduled_runs_session_status_due_idx - ON scheduled_runs (session_id, status, next_due_at, schedule_id); - - CREATE INDEX IF NOT EXISTS scheduled_runs_topic_status_due_idx - ON scheduled_runs (topic_id, status, next_due_at, schedule_id); - - CREATE TABLE IF NOT EXISTS scheduled_run_triggers ( - trigger_id TEXT PRIMARY KEY, - schedule_id TEXT NOT NULL, - signal_id TEXT, - topic_id TEXT, - run_id TEXT, - status TEXT NOT NULL, - due_at TIMESTAMPTZ NOT NULL, - planned_at TIMESTAMPTZ NOT NULL, - reason TEXT, - metadata JSONB NOT NULL DEFAULT '{}'::jsonb - ); - - CREATE INDEX IF NOT EXISTS scheduled_run_triggers_schedule_due_idx - ON scheduled_run_triggers (schedule_id, due_at, trigger_id); - - CREATE TABLE IF NOT EXISTS proactive_signals ( - signal_id TEXT PRIMARY KEY, - dedupe_key TEXT NOT NULL UNIQUE, - session_id TEXT, - topic_id TEXT, - kind TEXT NOT NULL, - status TEXT NOT NULL, - observed_at TIMESTAMPTZ NOT NULL, - cooldown_until TIMESTAMPTZ, - summary TEXT NOT NULL, - metadata JSONB NOT NULL DEFAULT '{}'::jsonb - ); - - CREATE INDEX IF NOT EXISTS proactive_signals_status_observed_idx - ON proactive_signals (status, observed_at, signal_id); - """ - _UPSERT_SCHEDULE_SQL: Final[str] = """ - INSERT INTO scheduled_runs ( - schedule_id, - session_id, - topic_id, - kind, - status, - cadence, - owner, - title, - next_due_at, - last_triggered_at, - created_at, - updated_at, - metadata - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13::jsonb) - ON CONFLICT (schedule_id) - DO UPDATE SET - session_id = EXCLUDED.session_id, - topic_id = EXCLUDED.topic_id, - kind = EXCLUDED.kind, - status = EXCLUDED.status, - cadence = EXCLUDED.cadence, - owner = EXCLUDED.owner, - title = EXCLUDED.title, - next_due_at = EXCLUDED.next_due_at, - last_triggered_at = EXCLUDED.last_triggered_at, - updated_at = EXCLUDED.updated_at, - metadata = EXCLUDED.metadata - RETURNING * - """ - _SELECT_SCHEDULE_SQL: Final[str] = """ - SELECT * FROM scheduled_runs WHERE schedule_id = $1 - """ - _LIST_SCHEDULES_SQL: Final[str] = """ - SELECT * FROM scheduled_runs - WHERE ($1::text IS NULL OR session_id = $1) - AND ($2::text IS NULL OR topic_id = $2) - AND ($3::text IS NULL OR status = $3) - ORDER BY COALESCE(next_due_at, updated_at), schedule_id - LIMIT $4 - """ - _LIST_DUE_SCHEDULES_SQL: Final[str] = """ - SELECT * FROM scheduled_runs - WHERE status = 'active' - AND next_due_at IS NOT NULL - AND next_due_at <= $1 - ORDER BY next_due_at, schedule_id - LIMIT $2 - """ - _UPDATE_SCHEDULE_STATUS_SQL: Final[str] = """ - UPDATE scheduled_runs - SET status = $2, - next_due_at = $3, - last_triggered_at = $4, - updated_at = $5, - metadata = $6::jsonb - WHERE schedule_id = $1 - RETURNING * - """ - _INSERT_TRIGGER_SQL: Final[str] = """ - INSERT INTO scheduled_run_triggers ( - trigger_id, - schedule_id, - signal_id, - topic_id, - run_id, - status, - due_at, - planned_at, - reason, - metadata - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb) - ON CONFLICT (trigger_id) - DO UPDATE SET - schedule_id = EXCLUDED.schedule_id, - signal_id = EXCLUDED.signal_id, - topic_id = EXCLUDED.topic_id, - run_id = EXCLUDED.run_id, - status = EXCLUDED.status, - due_at = EXCLUDED.due_at, - planned_at = EXCLUDED.planned_at, - reason = EXCLUDED.reason, - metadata = EXCLUDED.metadata - RETURNING * - """ - _LIST_TRIGGERS_SQL: Final[str] = """ - SELECT * FROM scheduled_run_triggers - WHERE schedule_id = $1 - ORDER BY due_at, trigger_id - LIMIT $2 - """ - _UPSERT_SIGNAL_SQL: Final[str] = """ - INSERT INTO proactive_signals ( - signal_id, - dedupe_key, - session_id, - topic_id, - kind, - status, - observed_at, - cooldown_until, - summary, - metadata - ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb) - ON CONFLICT (dedupe_key) - DO UPDATE SET - status = proactive_signals.status, - metadata = proactive_signals.metadata - RETURNING * - """ - _SELECT_SIGNAL_SQL: Final[str] = """ - SELECT * FROM proactive_signals WHERE signal_id = $1 - """ - _LIST_SIGNALS_SQL: Final[str] = """ - SELECT * FROM proactive_signals - WHERE ($1::text IS NULL OR status = $1) - AND ($2::text IS NULL OR session_id = $2) - AND ($3::text IS NULL OR topic_id = $3) - ORDER BY observed_at, signal_id - LIMIT $4 - """ - _UPDATE_SIGNAL_STATUS_SQL: Final[str] = """ - UPDATE proactive_signals - SET status = $2, - cooldown_until = $3, - metadata = $4::jsonb - WHERE signal_id = $1 - RETURNING * - """ - - def __init__(self, *, pool: PGPool) -> None: - self._pool = pool - self._schema_ready = False - - async def _ensure_schema(self) -> AsyncPGPool: - pool = await self._pool.get_pool() - if not self._schema_ready: - _ = await pool.execute(self._CREATE_SCHEMA_SQL) - self._schema_ready = True - return pool - - async def create_schedule(self, record: ScheduleRecord) -> ScheduleRecord: - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._UPSERT_SCHEDULE_SQL, - record.schedule_id, - record.session_id, - record.topic_id, - record.kind, - record.status, - record.cadence, - record.owner, - record.title, - record.next_due_at, - record.last_triggered_at, - record.created_at, - record.updated_at, - record.metadata, - ) - return _schedule_from_row(_required_row(row, "schedule upsert")) - - async def load_schedule(self, schedule_id: str) -> ScheduleRecord | None: - _require_non_empty("schedule_id", schedule_id) - pool = await self._ensure_schema() - row = await pool.fetchrow(self._SELECT_SCHEDULE_SQL, schedule_id) - return None if row is None else _schedule_from_row(row) - - async def list_schedules( - self, - *, - session_id: str | None = None, - topic_id: str | None = None, - status: ScheduleStatus | None = None, - limit: int = 100, - ) -> list[ScheduleRecord]: - if session_id is not None: - _require_non_empty("session_id", session_id) - if topic_id is not None: - _require_non_empty("topic_id", topic_id) - if status is not None: - _require_status("schedule status", status, _SCHEDULE_STATUSES) - _require_positive_int("limit", limit) - pool = await self._ensure_schema() - rows = await pool.fetch( - self._LIST_SCHEDULES_SQL, session_id, topic_id, status, limit - ) - return [_schedule_from_row(row) for row in rows] - - async def list_due_schedules( - self, - *, - due_at_or_before: datetime, - limit: int = 100, - ) -> list[ScheduleRecord]: - _require_datetime("due_at_or_before", due_at_or_before) - _require_positive_int("limit", limit) - pool = await self._ensure_schema() - rows = await pool.fetch(self._LIST_DUE_SCHEDULES_SQL, due_at_or_before, limit) - return [_schedule_from_row(row) for row in rows] - - async def update_schedule_status( - self, - schedule_id: str, - *, - status: ScheduleStatus, - next_due_at: datetime | None, - last_triggered_at: datetime | None, - updated_at: datetime, - metadata: JSONObject, - ) -> ScheduleRecord: - _require_non_empty("schedule_id", schedule_id) - _require_status("schedule status", status, _SCHEDULE_STATUSES) - if next_due_at is not None: - _require_datetime("next_due_at", next_due_at) - if last_triggered_at is not None: - _require_datetime("last_triggered_at", last_triggered_at) - _require_datetime("updated_at", updated_at) - _require_json_object("metadata", metadata) - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._UPDATE_SCHEDULE_STATUS_SQL, - schedule_id, - status, - next_due_at, - last_triggered_at, - updated_at, - metadata, - ) - if row is None: - raise KeyError(f"schedule not found: {schedule_id}") - return _schedule_from_row(row) - - async def record_trigger( - self, - record: ScheduleTriggerRecord, - ) -> ScheduleTriggerRecord: - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._INSERT_TRIGGER_SQL, - record.trigger_id, - record.schedule_id, - record.signal_id, - record.topic_id, - record.run_id, - record.status, - record.due_at, - record.planned_at, - record.reason, - record.metadata, - ) - return _trigger_from_row(_required_row(row, "schedule trigger upsert")) - - async def list_triggers( - self, - schedule_id: str, - *, - limit: int = 100, - ) -> list[ScheduleTriggerRecord]: - _require_non_empty("schedule_id", schedule_id) - _require_positive_int("limit", limit) - pool = await self._ensure_schema() - rows = await pool.fetch(self._LIST_TRIGGERS_SQL, schedule_id, limit) - return [_trigger_from_row(row) for row in rows] - - async def record_signal( - self, - record: ProactiveSignalRecord, - ) -> ProactiveSignalRecord: - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._UPSERT_SIGNAL_SQL, - record.signal_id, - record.dedupe_key, - record.session_id, - record.topic_id, - record.kind, - record.status, - record.observed_at, - record.cooldown_until, - record.summary, - record.metadata, - ) - return _signal_from_row(_required_row(row, "proactive signal upsert")) - - async def load_signal(self, signal_id: str) -> ProactiveSignalRecord | None: - _require_non_empty("signal_id", signal_id) - pool = await self._ensure_schema() - row = await pool.fetchrow(self._SELECT_SIGNAL_SQL, signal_id) - return None if row is None else _signal_from_row(row) - - async def list_signals( - self, - *, - status: SignalStatus | None = None, - session_id: str | None = None, - topic_id: str | None = None, - limit: int = 100, - ) -> list[ProactiveSignalRecord]: - if status is not None: - _require_status("signal status", status, _SIGNAL_STATUSES) - if session_id is not None: - _require_non_empty("session_id", session_id) - if topic_id is not None: - _require_non_empty("topic_id", topic_id) - _require_positive_int("limit", limit) - pool = await self._ensure_schema() - rows = await pool.fetch( - self._LIST_SIGNALS_SQL, status, session_id, topic_id, limit - ) - return [_signal_from_row(row) for row in rows] - - async def update_signal_status( - self, - signal_id: str, - *, - status: SignalStatus, - cooldown_until: datetime | None, - metadata: JSONObject, - ) -> ProactiveSignalRecord: - _require_non_empty("signal_id", signal_id) - _require_status("signal status", status, _SIGNAL_STATUSES) - if cooldown_until is not None: - _require_datetime("cooldown_until", cooldown_until) - _require_json_object("metadata", metadata) - pool = await self._ensure_schema() - row = await pool.fetchrow( - self._UPDATE_SIGNAL_STATUS_SQL, - signal_id, - status, - cooldown_until, - metadata, - ) - if row is None: - raise KeyError(f"proactive signal not found: {signal_id}") - return _signal_from_row(row) - - -def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: - if row is None: - raise RuntimeError(f"postgres {context} returned no row") - return row - - -def _schedule_from_row(row: dict[str, object]) -> ScheduleRecord: - return ScheduleRecord( - schedule_id=_required_str(row, "schedule_id", context="schedule row"), - session_id=_required_str(row, "session_id", context="schedule row"), - topic_id=_optional_str(row, "topic_id", context="schedule row"), - kind=_required_str(row, "kind", context="schedule row"), - status=_required_str(row, "status", context="schedule row"), - cadence=_required_str(row, "cadence", context="schedule row"), - owner=_optional_str(row, "owner", context="schedule row"), - title=_optional_str(row, "title", context="schedule row"), - next_due_at=_optional_datetime(row, "next_due_at", context="schedule row"), - last_triggered_at=_optional_datetime( - row, "last_triggered_at", context="schedule row" - ), - created_at=_required_datetime(row, "created_at", context="schedule row"), - updated_at=_required_datetime(row, "updated_at", context="schedule row"), - metadata=_required_json_object(row, "metadata", context="schedule row"), - ) - - -def _trigger_from_row(row: dict[str, object]) -> ScheduleTriggerRecord: - return ScheduleTriggerRecord( - trigger_id=_required_str(row, "trigger_id", context="trigger row"), - schedule_id=_required_str(row, "schedule_id", context="trigger row"), - signal_id=_optional_str(row, "signal_id", context="trigger row"), - topic_id=_optional_str(row, "topic_id", context="trigger row"), - run_id=_optional_str(row, "run_id", context="trigger row"), - status=_required_str(row, "status", context="trigger row"), - due_at=_required_datetime(row, "due_at", context="trigger row"), - planned_at=_required_datetime(row, "planned_at", context="trigger row"), - reason=_optional_str(row, "reason", context="trigger row"), - metadata=_required_json_object(row, "metadata", context="trigger row"), - ) - - -def _signal_from_row(row: dict[str, object]) -> ProactiveSignalRecord: - return ProactiveSignalRecord( - signal_id=_required_str(row, "signal_id", context="signal row"), - dedupe_key=_required_str(row, "dedupe_key", context="signal row"), - session_id=_optional_str(row, "session_id", context="signal row"), - topic_id=_optional_str(row, "topic_id", context="signal row"), - kind=_required_str(row, "kind", context="signal row"), - status=_required_str(row, "status", context="signal row"), - observed_at=_required_datetime(row, "observed_at", context="signal row"), - cooldown_until=_optional_datetime(row, "cooldown_until", context="signal row"), - summary=_required_str(row, "summary", context="signal row"), - metadata=_required_json_object(row, "metadata", context="signal row"), - ) - - -def _require_non_empty(field_name: str, value: str) -> None: - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"{field_name} must be a non-empty string") - - -def _require_safe_label(field_name: str, value: str) -> None: - _require_non_empty(field_name, value) - if len(value) > _MAX_SAFE_LABEL_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_SAFE_LABEL_CHARS} characters" - ) - key_folded = field_name.casefold() - value_folded = value.casefold() - if any(part in value_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): - raise ValueError(f"{field_name} must not contain sensitive label text") - if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): - raise ValueError(f"{field_name} must not be sensitive label text") - _reject_secret_shaped_value(field_name, value) - - -def _require_status(field_name: str, value: str, allowed: frozenset[str]) -> None: - if value not in allowed: - raise ValueError(f"{field_name} must be one of {sorted(allowed)}") - - -def _require_optional_display_text(field_name: str, value: str | None) -> None: - if value is None: - return - _require_non_empty(field_name, value) - if len(value) > _MAX_DISPLAY_TEXT_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_DISPLAY_TEXT_CHARS} characters" - ) - _reject_secret_shaped_value(field_name, value) - - -def _require_datetime(field_name: str, value: datetime) -> None: - if not isinstance(value, datetime): - raise TypeError(f"{field_name} must be a datetime") - - -def _require_positive_int(field_name: str, value: int) -> None: - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError(f"{field_name} must be a positive int") - - -def _require_json_object(field_name: str, value: JSONObject) -> None: - if not isinstance(value, dict): - raise TypeError(f"{field_name} must be a JSON object") - for key, item in value.items(): - if not isinstance(key, str): - raise TypeError(f"{field_name} keys must be strings") - key_folded = key.casefold() - if any(part in key_folded for part in _FORBIDDEN_METADATA_KEY_PARTS): - raise ValueError(f"{field_name} contains forbidden metadata key: {key}") - _require_json_value(f"{field_name}.{key}", item) - - -def _require_json_value(field_name: str, value: object) -> None: - if value is None or isinstance(value, str | int | float | bool): - if isinstance(value, str): - if len(value) > _MAX_METADATA_STRING_CHARS: - raise ValueError( - f"{field_name} must be at most {_MAX_METADATA_STRING_CHARS} characters" - ) - _reject_secret_shaped_value(field_name, value) - return - if isinstance(value, list): - for index, item in enumerate(value): - _require_json_value(f"{field_name}[{index}]", item) - return - if isinstance(value, dict): - _require_json_object(field_name, value) - return - raise TypeError(f"{field_name} must be JSON-safe") - - -def _reject_secret_shaped_value(field_name: str, value: str) -> None: - folded = value.casefold() - if any( - marker in folded for marker in _SECRET_VALUE_MARKERS if marker != "sk-" - ) or folded.startswith("sk-"): - raise ValueError(f"{field_name} must not contain secret-shaped values") - - -def _default_trigger_id(schedule: ScheduleRecord, now: datetime) -> str: - safe_time = now.isoformat().replace(":", "").replace("+", "_") - return f"trigger-{schedule.schedule_id}-{safe_time}" - - -def _next_due_at(schedule: ScheduleRecord, now: datetime) -> datetime | None: - if schedule.cadence == "once": - return None - if schedule.cadence == "hourly": - return now + timedelta(hours=1) - if schedule.cadence == "daily": - return now + timedelta(days=1) - return None - - -def _schedule_bee_launch_metadata(schedule: ScheduleRecord) -> JSONObject: - value = schedule.metadata.get("bee_launch") - if value is None: - return {} - if not isinstance(value, dict): - raise TypeError("schedule.metadata.bee_launch must be an object") - bee_launch = dict(value) - _require_bee_launch_metadata_object("schedule.metadata.bee_launch", bee_launch) - return {"bee_launch": bee_launch} - - -def _signal_bee_launch_metadata(signal: ProactiveSignalRecord) -> JSONObject: - value = signal.metadata.get("bee_launch") - if value is None: - return {} - if not isinstance(value, dict): - raise TypeError("signal.metadata.bee_launch must be an object") - bee_launch = dict(value) - _require_bee_launch_metadata_object("signal.metadata.bee_launch", bee_launch) - return {"bee_launch": bee_launch} - - -def _require_bee_launch_metadata_object( - field_name: str, - value: JSONObject, -) -> None: - _require_json_object(field_name, value) - for key, item in value.items(): - key_folded = key.casefold() - if any(part in key_folded for part in _FORBIDDEN_BEE_LAUNCH_KEY_PARTS): - raise ValueError( - f"{field_name} contains forbidden Bee launch metadata key: {key}" - ) - if isinstance(item, dict): - _require_bee_launch_metadata_object(f"{field_name}.{key}", item) - elif isinstance(item, list): - for index, nested in enumerate(item): - if isinstance(nested, dict): - _require_bee_launch_metadata_object( - f"{field_name}.{key}[{index}]", - nested, - ) - - -def _require_topic_matches_intent( - *, - topic: TopicRecord, - intent: ScheduledLaunchIntent, - tape: Tape, -) -> None: - if topic.session_id != intent.session_id: - raise ValueError("scheduled topic session does not match launch intent") - if topic.tape_id != tape.tape_id: - raise ValueError("scheduled topic tape does not match launch tape") - - -def _safe_metadata_label(value: object) -> str: - if not isinstance(value, str): - return "unknown" - _require_safe_label("metadata_label", value) - return value - - -def _required_signal_session(signal: ProactiveSignalRecord) -> str: - if signal.session_id is None: - raise ValueError("proactive signal requires session_id for launch planning") - return signal.session_id - - -def _required_str(row: dict[str, object], key: str, *, context: str) -> str: - value = row.get(key) - if not isinstance(value, str): - raise TypeError(f"postgres {context} must include string {key}") - return value - - -def _optional_str(row: dict[str, object], key: str, *, context: str) -> str | None: - value = row.get(key) - if value is None: - return None - if not isinstance(value, str): - raise TypeError(f"postgres {context} must include string or None {key}") - return value - - -def _required_datetime(row: dict[str, object], key: str, *, context: str) -> datetime: - value = row.get(key) - if not isinstance(value, datetime): - raise TypeError(f"postgres {context} must include datetime {key}") - return value - - -def _optional_datetime( - row: dict[str, object], key: str, *, context: str -) -> datetime | None: - value = row.get(key) - if value is None: - return None - if not isinstance(value, datetime): - raise TypeError(f"postgres {context} must include datetime or None {key}") - return value - - -def _required_json_object( - row: dict[str, object], key: str, *, context: str -) -> JSONObject: - value = row.get(key) - if not isinstance(value, dict): - raise TypeError(f"postgres {context} must include dict {key}") - _require_json_object(key, value) - return value +"""Deprecated import location. Moved to coding_agent.runs.scheduled.""" +import sys as _sys +from coding_agent.runs import scheduled as _moved +_sys.modules[__name__] = _moved diff --git a/coding-agent/src/coding_agent/storage_migration.py b/coding-agent/src/coding_agent/storage_migration.py index 63ff3c5ac..4dbfb8571 100644 --- a/coding-agent/src/coding_agent/storage_migration.py +++ b/coding-agent/src/coding_agent/storage_migration.py @@ -1,166 +1,4 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any, cast - -from agentkit.storage.checkpoint_fs import FSCheckpointStore -from agentkit.storage.sqlite import SQLiteCheckpointStore, SQLiteTapeStore -from coding_agent.local_storage import local_sqlite_path - - -@dataclass(frozen=True) -class StoreMigrationReport: - scanned: int = 0 - migrated: int = 0 - skipped: int = 0 - - -@dataclass(frozen=True) -class LegacySQLiteMigrationReport: - tapes: StoreMigrationReport - checkpoints: StoreMigrationReport - - -async def migrate_legacy_storage_to_sqlite( - data_dir: Path, - *, - tapes_dir: Path | None = None, - checkpoints_dir: Path | None = None, - tape_sqlite_path: Path | None = None, - checkpoint_sqlite_path: Path | None = None, - replace_tapes: bool = False, - dry_run: bool = False, -) -> LegacySQLiteMigrationReport: - resolved_data_dir = Path(data_dir) - return LegacySQLiteMigrationReport( - tapes=await migrate_jsonl_tapes_to_sqlite( - tapes_dir or resolved_data_dir / "tapes", - tape_sqlite_path or local_sqlite_path(resolved_data_dir), - replace=replace_tapes, - dry_run=dry_run, - ), - checkpoints=await migrate_fs_checkpoints_to_sqlite( - checkpoints_dir or resolved_data_dir / "checkpoints", - checkpoint_sqlite_path or local_sqlite_path(resolved_data_dir), - dry_run=dry_run, - ), - ) - - -async def migrate_jsonl_tapes_to_sqlite( - source_dir: Path, - sqlite_path: Path, - *, - replace: bool = False, - dry_run: bool = False, -) -> StoreMigrationReport: - source_dir = Path(source_dir) - sqlite_path = Path(sqlite_path) - if not source_dir.exists(): - return StoreMigrationReport() - if not source_dir.is_dir(): - raise NotADirectoryError(str(source_dir)) - - store = ( - None if dry_run and not sqlite_path.exists() else SQLiteTapeStore(sqlite_path) - ) - scanned = 0 - migrated = 0 - skipped = 0 - for path in sorted(source_dir.glob("*.jsonl")): - scanned += 1 - tape_id = path.stem - entries = _read_jsonl_objects(path) - existing = [] if store is None else await store.load(tape_id) - if not entries and not existing: - skipped += 1 - continue - if existing: - if existing == entries: - skipped += 1 - continue - if not replace: - raise ValueError( - f"tape {tape_id!r} already exists with different entries" - ) - if not dry_run: - if store is None: - store = SQLiteTapeStore(sqlite_path) - await store.truncate(tape_id, 0) - if not dry_run: - if store is None: - store = SQLiteTapeStore(sqlite_path) - await store.save(tape_id, entries) - migrated += 1 - return StoreMigrationReport(scanned=scanned, migrated=migrated, skipped=skipped) - - -async def migrate_fs_checkpoints_to_sqlite( - source_dir: Path, - sqlite_path: Path, - *, - dry_run: bool = False, -) -> StoreMigrationReport: - source_dir = Path(source_dir) - sqlite_path = Path(sqlite_path) - if not source_dir.exists(): - return StoreMigrationReport() - if not source_dir.is_dir(): - raise NotADirectoryError(str(source_dir)) - - source_store = FSCheckpointStore(source_dir) - target_store = ( - None - if dry_run and not sqlite_path.exists() - else SQLiteCheckpointStore(sqlite_path) - ) - scanned = 0 - migrated = 0 - skipped = 0 - for meta_path in sorted(source_dir.glob("*.meta.json")): - checkpoint_id = meta_path.name.removesuffix(".meta.json") - scanned += 1 - snapshot = await source_store.load(checkpoint_id) - if snapshot is None: - raise RuntimeError( - f"checkpoint metadata {meta_path} did not load checkpoint snapshot" - ) - existing = ( - None if target_store is None else await target_store.load(checkpoint_id) - ) - if existing == snapshot: - skipped += 1 - continue - if not dry_run: - if target_store is None: - target_store = SQLiteCheckpointStore(sqlite_path) - await target_store.save(snapshot) - migrated += 1 - return StoreMigrationReport(scanned=scanned, migrated=migrated, skipped=skipped) - - -def _read_jsonl_objects(path: Path) -> list[dict[str, Any]]: - entries: list[dict[str, Any]] = [] - with path.open(encoding="utf-8") as handle: - for line_number, line in enumerate(handle, start=1): - stripped = line.strip() - if not stripped: - continue - loaded = json.loads(stripped) - if not isinstance(loaded, dict): - raise TypeError( - f"{path}:{line_number} must contain a JSON object tape entry" - ) - entries.append(cast(dict[str, Any], loaded)) - return entries - - -__all__ = [ - "LegacySQLiteMigrationReport", - "StoreMigrationReport", - "migrate_fs_checkpoints_to_sqlite", - "migrate_jsonl_tapes_to_sqlite", - "migrate_legacy_storage_to_sqlite", -] +"""Deprecated import location. Moved to coding_agent.stores.migration.""" +import sys as _sys +from coding_agent.stores import migration as _moved +_sys.modules[__name__] = _moved diff --git a/coding-agent/src/coding_agent/stores/durable_local.py b/coding-agent/src/coding_agent/stores/durable_local.py new file mode 100644 index 000000000..e80b378f6 --- /dev/null +++ b/coding-agent/src/coding_agent/stores/durable_local.py @@ -0,0 +1,1284 @@ +from __future__ import annotations + +import json +import sqlite3 +import threading +from collections.abc import Callable, Mapping +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Final, cast + +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot +from agentkit.storage.protocols import TapeInfo, TapeSearchResult +from agentkit.storage.sqlite import ( + SQLiteCheckpointStore, + SQLiteTapeStore, + _checkpoint_snapshot_from_payload, + _checkpoint_snapshot_to_payload, + _entry_anchor_type, + _entry_nested_str, + _json_object_from_text, + _optional_entry_str, +) + +from coding_agent.runtime_store import ( + AgentRunRecord, + AgentInteractionRecord, + RuntimeEventRecord, + RunMessageSnapshotRecord, + SQLiteRuntimeStore, + _agent_run_from_sqlite_row, + _agent_run_sqlite_values, + _datetime_to_json, + _interaction_from_sqlite_row, + _interaction_sqlite_values, + _json_to_sql, + _runtime_event_from_sqlite_row, +) +from coding_agent.server.stores.session_owner_store import ( + OwnerAuthority, + SQLiteSessionOwnerStore, + SessionOwnershipConflictError, + SessionOwnershipConflictReason, + _datetime_from_sqlite_text, + _datetime_to_sqlite_text, +) +from coding_agent.server.stores.session_store import SQLiteSessionStore, SessionPayload + + +class SQLiteLocalDurableStore: + """SQLite-local protected mutation facade. + + This facade is intentionally local to the SQLite bundle path. It does not replace + the generic store protocols; it provides the transaction shape required by local + durable fencing: owner epoch check, target ownership check, and mutation in the + same SQLite transaction. + """ + + _SESSION_TAPES_SCHEMA_SQL: Final[str] = """ + CREATE TABLE IF NOT EXISTS session_tapes ( + session_id TEXT PRIMARY KEY, + tape_id TEXT UNIQUE NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS session_tapes_tape_id_idx + ON session_tapes (tape_id); + """ + + def __init__(self, path: Path) -> None: + self._path = path + self._lock = threading.Lock() + self._path.parent.mkdir(parents=True, exist_ok=True) + self._owner_store = SQLiteSessionOwnerStore(path) + with self._connect() as connection: + connection.executescript(SQLiteSessionStore._CREATE_TABLE_SQL) + connection.executescript(SQLiteTapeStore._CREATE_SCHEMA_SQL) + connection.executescript(SQLiteCheckpointStore._CREATE_SCHEMA_SQL) + connection.executescript(SQLiteRuntimeStore._CREATE_SCHEMA_SQL) + connection.executescript(self._SESSION_TAPES_SCHEMA_SQL) + + def _connect(self) -> sqlite3.Connection: + connection = sqlite3.connect(self._path, timeout=5.0) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout = 5000") + return connection + + async def acquire_owner( + self, + session_id: str, + owner_id: str, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + return await self._owner_store.acquire_authority( + session_id, + owner_id, + lease_seconds=lease_seconds, + ) + + async def renew_owner( + self, + authority: OwnerAuthority, + lease_seconds: float = 30.0, + ) -> OwnerAuthority: + return await self._owner_store.renew_authority( + authority, + lease_seconds=lease_seconds, + ) + + async def save_session( + self, + authority: OwnerAuthority, + payload: SessionPayload, + ) -> None: + _require_json_object("session payload", payload) + payload_id = payload.get("id") + if payload_id != authority.session_id: + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + payload_session_id = payload.get("session_id") + if ( + payload_session_id is not None + and payload_session_id != authority.session_id + ): + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + tape_id = payload.get("tape_id") + if tape_id is not None and not isinstance(tape_id, str): + raise TypeError("session payload tape_id must be a string") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + if tape_id: + self._bind_tape(connection, authority.session_id, tape_id) + connection.execute( + """ + INSERT INTO agent_http_sessions (session_id, payload, updated_at) + VALUES (?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(session_id) + DO UPDATE SET payload = excluded.payload, + updated_at = CURRENT_TIMESTAMP + """, + (authority.session_id, json.dumps(payload, sort_keys=True)), + ) + + async def delete_session(self, authority: OwnerAuthority) -> None: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + connection.execute( + "DELETE FROM agent_http_sessions WHERE session_id = ?", + (authority.session_id,), + ) + connection.execute( + "DELETE FROM session_tapes WHERE session_id = ?", + (authority.session_id,), + ) + + def load_session(self, session_id: str) -> SessionPayload | None: + _require_non_empty("session_id", session_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT payload FROM agent_http_sessions WHERE session_id = ?", + (session_id,), + ).fetchone() + if row is None: + return None + payload = json.loads(str(row["payload"])) + if not isinstance(payload, dict): + raise TypeError("sqlite session payload must be a JSON object") + return cast(SessionPayload, payload) + + async def append_tape_entries( + self, + authority: OwnerAuthority, + tape_id: str, + entries: list[dict[str, Any]], + ) -> None: + _require_non_empty("tape_id", tape_id) + if not entries: + return + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=tape_id, + session_id=authority.session_id, + ) + row = connection.execute( + "SELECT COALESCE(MAX(seq), -1) AS max_seq FROM tape_entries WHERE tape_id = ?", + (tape_id,), + ).fetchone() + max_seq = _row_required_int(row, "max_seq", context="tape max seq") + values: list[tuple[object, ...]] = [] + now = datetime.now(UTC).isoformat() + for offset, entry in enumerate(entries, start=1): + _require_json_object("tape entry", entry) + values.append( + ( + tape_id, + max_seq + offset, + json.dumps(entry, sort_keys=True), + _optional_entry_str(entry, "kind"), + _entry_nested_str(entry, "run_id"), + _entry_nested_str(entry, "tool_call_id"), + _entry_anchor_type(entry), + now, + ) + ) + connection.executemany( + """ + INSERT INTO tape_entries ( + tape_id, seq, entry_json, kind, run_id, tool_call_id, + anchor_type, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + values, + ) + + async def load_tape(self, tape_id: str) -> list[dict[str, Any]]: + _require_non_empty("tape_id", tape_id) + with self._lock, self._connect() as connection: + rows = connection.execute( + """ + SELECT entry_json + FROM tape_entries + WHERE tape_id = ? + ORDER BY seq + """, + (tape_id,), + ).fetchall() + return [ + _json_object_from_text(row["entry_json"], context="tape entry") + for row in rows + ] + + async def truncate_tape( + self, + authority: OwnerAuthority, + tape_id: str, + keep: int, + ) -> None: + _require_non_empty("tape_id", tape_id) + if keep < 0: + raise ValueError("keep must be >= 0") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=tape_id, + session_id=authority.session_id, + ) + connection.execute( + "DELETE FROM tape_entries WHERE tape_id = ? AND seq >= ?", + (tape_id, keep), + ) + + def session_id_for_tape(self, tape_id: str) -> str | None: + _require_non_empty("tape_id", tape_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM session_tapes WHERE tape_id = ?", + (tape_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("session_tapes session_id must be text") + return session_id + + async def create_agent_run( + self, + authority: OwnerAuthority, + record: AgentRunRecord, + ) -> AgentRunRecord: + if record.session_id != authority.session_id: + raise SessionOwnershipConflictError("agent run belongs to another session") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + existing = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (record.run_id,), + ).fetchone() + if existing is not None and existing["session_id"] != authority.session_id: + raise SessionOwnershipConflictError( + "agent run target belongs to another session" + ) + connection.execute( + """ + INSERT INTO agent_runs ( + run_id, session_id, tape_id, parent_run_id, agent_id, status, + started_at, ended_at, metadata, result, error + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(run_id) + DO UPDATE SET + tape_id = excluded.tape_id, + parent_run_id = excluded.parent_run_id, + agent_id = excluded.agent_id, + status = excluded.status, + started_at = excluded.started_at, + ended_at = excluded.ended_at, + metadata = excluded.metadata, + result = excluded.result, + error = excluded.error + """, + _agent_run_sqlite_values(record), + ) + return record + + async def update_agent_run( + self, + authority: OwnerAuthority, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: dict[str, Any], + result: dict[str, Any], + error: str | None, + ) -> AgentRunRecord: + _require_non_empty("run_id", run_id) + _require_non_empty("status", status) + _require_json_object("metadata", metadata) + _require_json_object("result", result) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=run_id, + session_id=authority.session_id, + ) + connection.execute( + """ + UPDATE agent_runs + SET status = ?, + ended_at = ?, + metadata = ?, + result = ?, + error = ? + WHERE run_id = ? + """, + ( + status, + None if ended_at is None else _datetime_to_json(ended_at), + _json_to_sql(cast(dict[str, Any], metadata)), + _json_to_sql(cast(dict[str, Any], result)), + error, + run_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found after update: {run_id}") + return _agent_run_from_sqlite_row(row) + + async def append_runtime_event( + self, + authority: OwnerAuthority, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT runtime_events.* + FROM runtime_events + JOIN agent_runs ON agent_runs.run_id = runtime_events.run_id + WHERE runtime_events.event_id = ? + """, + (record.event_id,), + ).fetchone() + if existing is not None: + if existing["run_id"] != record.run_id: + raise SessionOwnershipConflictError( + "runtime event target belongs to another run" + ) + return _runtime_event_from_sqlite_row(existing) + connection.execute( + """ + INSERT INTO runtime_events ( + event_id, run_id, event_kind, payload, created_at + ) + VALUES (?, ?, ?, ?, ?) + """, + ( + record.event_id, + record.run_id, + record.event_kind, + _json_to_sql(record.payload), + _datetime_to_json(record.created_at), + ), + ) + row = connection.execute( + "SELECT * FROM runtime_events WHERE event_id = ?", + (record.event_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite runtime event insert returned no row") + return _runtime_event_from_sqlite_row(row) + + async def replay_runtime_events( + self, + run_id: str, + *, + after_sequence: int = 0, + limit: int = 1000, + ) -> list[RuntimeEventRecord]: + _require_non_empty("run_id", run_id) + if after_sequence < 0: + raise ValueError("after_sequence must be >= 0") + if limit <= 0: + raise ValueError("limit must be positive") + with self._lock, self._connect() as connection: + rows = connection.execute( + """ + SELECT * FROM runtime_events + WHERE run_id = ? AND sequence > ? + ORDER BY sequence + LIMIT ? + """, + (run_id, after_sequence, limit), + ).fetchall() + return [_runtime_event_from_sqlite_row(row) for row in rows] + + async def save_message_snapshot( + self, + authority: OwnerAuthority, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT run_message_snapshots.run_id, agent_runs.session_id + FROM run_message_snapshots + JOIN agent_runs ON agent_runs.run_id = run_message_snapshots.run_id + WHERE run_message_snapshots.snapshot_id = ? + """, + (record.snapshot_id,), + ).fetchone() + if existing is not None and ( + existing["run_id"] != record.run_id + or existing["session_id"] != authority.session_id + ): + raise SessionOwnershipConflictError( + "message snapshot target belongs to another session" + ) + connection.execute( + """ + INSERT INTO run_message_snapshots ( + snapshot_id, run_id, messages, metadata, created_at + ) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(snapshot_id) + DO UPDATE SET + messages = excluded.messages, + metadata = excluded.metadata, + created_at = excluded.created_at + """, + ( + record.snapshot_id, + record.run_id, + _json_to_sql(cast(list[Any], record.messages)), + _json_to_sql(record.metadata), + _datetime_to_json(record.created_at), + ), + ) + return record + + async def create_agent_interaction( + self, + authority: OwnerAuthority, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_run_belongs_to_session( + connection, + run_id=record.run_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT agent_interactions.run_id, agent_runs.session_id + FROM agent_interactions + JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id + WHERE agent_interactions.interaction_id = ? + """, + (record.interaction_id,), + ).fetchone() + if existing is not None: + if ( + existing["run_id"] != record.run_id + or existing["session_id"] != authority.session_id + ): + raise SessionOwnershipConflictError( + "interaction target belongs to another session" + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (record.interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError( + "sqlite agent interaction lookup returned no row" + ) + return _interaction_from_sqlite_row(row) + connection.execute( + """ + INSERT INTO agent_interactions ( + interaction_id, run_id, interaction_kind, status, + request_payload, response_payload, metadata, + created_at, resolved_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + _interaction_sqlite_values(record), + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (record.interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite agent interaction insert returned no row") + return _interaction_from_sqlite_row(row) + + async def resolve_agent_interaction( + self, + authority: OwnerAuthority, + interaction_id: str, + *, + status: str, + response_payload: dict[str, Any], + resolved_at: datetime, + ) -> AgentInteractionRecord: + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + existing = connection.execute( + """ + SELECT agent_interactions.* + FROM agent_interactions + JOIN agent_runs ON agent_runs.run_id = agent_interactions.run_id + WHERE agent_interactions.interaction_id = ? + AND agent_runs.session_id = ? + """, + (interaction_id, authority.session_id), + ).fetchone() + if existing is None: + raise SessionOwnershipConflictError( + "interaction target belongs to another session" + ) + if existing["resolved_at"] is None: + connection.execute( + """ + UPDATE agent_interactions + SET status = ?, + response_payload = ?, + resolved_at = ? + WHERE interaction_id = ? + """, + ( + status, + _json_to_sql(response_payload), + _datetime_to_json(resolved_at), + interaction_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_interactions WHERE interaction_id = ?", + (interaction_id,), + ).fetchone() + if row is None: + raise RuntimeError("sqlite agent interaction resolve returned no row") + return _interaction_from_sqlite_row(row) + + async def claim_attached_executor_run( + self, + authorities: Mapping[str, OwnerAuthority], + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + _require_non_empty("executor_kind", executor_kind) + _require_json_object("claim_metadata", claim_metadata) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + filter_sql = "" + values: list[object] = [] + if session_id is not None: + filter_sql = "AND session_id = ?" + values.append(session_id) + rows = connection.execute( + f""" + SELECT * FROM agent_runs + WHERE status IN ('requested', 'expired') + {filter_sql} + ORDER BY started_at, run_id + """, + values, + ).fetchall() + candidates = [_agent_run_from_sqlite_row(row) for row in rows] + candidates = [ + run + for run in candidates + if run.metadata.get("executor_ref_kind") + in {"external_worker", "local_attached"} + and run.metadata.get("executor_kind") == executor_kind + ] + selected: AgentRunRecord | None = None + authority: OwnerAuthority | None = None + for candidate in candidates: + candidate_authority = authorities.get(candidate.session_id) + if candidate_authority is None: + continue + try: + self._assert_authority(connection, candidate_authority) + except SessionOwnershipConflictError: + continue + selected = candidate + authority = candidate_authority + break + if selected is None or authority is None: + return None + metadata = cast(dict[str, Any], {**selected.metadata, **claim_metadata}) + connection.execute( + """ + UPDATE agent_runs + SET status = 'claimed', + metadata = ? + WHERE run_id = ? + AND session_id = ? + """, + ( + _json_to_sql(metadata), + selected.run_id, + authority.session_id, + ), + ) + row = connection.execute( + "SELECT * FROM agent_runs WHERE run_id = ?", + (selected.run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found after claim: {selected.run_id}") + return _agent_run_from_sqlite_row(row) + + def session_id_for_run(self, run_id: str) -> str | None: + _require_non_empty("run_id", run_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("agent run session_id must be text") + return session_id + + async def save_checkpoint( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + ) -> None: + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError("checkpoint belongs to another session") + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=meta.tape_id, + session_id=authority.session_id, + ) + existing = connection.execute( + """ + SELECT session_id, tape_id + FROM checkpoints + WHERE checkpoint_id = ? + """, + (meta.checkpoint_id,), + ).fetchone() + if existing is not None and ( + existing["session_id"] != authority.session_id + or existing["tape_id"] != meta.tape_id + ): + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + payload = _checkpoint_snapshot_to_payload(snapshot) + connection.execute( + """ + INSERT INTO checkpoints ( + checkpoint_id, tape_id, session_id, entry_count, window_start, + label, snapshot_json, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(checkpoint_id) + DO UPDATE SET + entry_count = excluded.entry_count, + window_start = excluded.window_start, + label = excluded.label, + snapshot_json = excluded.snapshot_json, + created_at = excluded.created_at + """, + ( + meta.checkpoint_id, + meta.tape_id, + meta.session_id, + meta.entry_count, + meta.window_start, + meta.label, + json.dumps(payload, sort_keys=True), + meta.created_at.isoformat(), + ), + ) + + async def load_checkpoint(self, checkpoint_id: str) -> CheckpointSnapshot | None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT snapshot_json FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return None + return _checkpoint_snapshot_from_payload( + _json_object_from_text(row["snapshot_json"], context="checkpoint snapshot") + ) + + async def delete_checkpoint( + self, + authority: OwnerAuthority, + checkpoint_id: str, + ) -> None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + row = connection.execute( + "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return + if row["session_id"] != authority.session_id: + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + connection.execute( + "DELETE FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ) + + async def restore_checkpoint_state( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + session_payload: SessionPayload, + ) -> None: + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError("checkpoint belongs to another session") + if session_payload.get("id") != authority.session_id: + raise SessionOwnershipConflictError( + "session payload belongs to another owner" + ) + if session_payload.get("tape_id") != meta.tape_id: + raise SessionOwnershipConflictError( + "checkpoint restore session payload has mismatched tape id" + ) + with self._lock, self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + self._assert_authority(connection, authority) + self._assert_tape_belongs_to_session( + connection, + tape_id=meta.tape_id, + session_id=authority.session_id, + ) + checkpoint_row = connection.execute( + """ + SELECT session_id, tape_id + FROM checkpoints + WHERE checkpoint_id = ? + """, + (meta.checkpoint_id,), + ).fetchone() + if checkpoint_row is None: + raise KeyError(f"checkpoint not found: {meta.checkpoint_id}") + if ( + checkpoint_row["session_id"] != authority.session_id + or checkpoint_row["tape_id"] != meta.tape_id + ): + raise SessionOwnershipConflictError( + "checkpoint target belongs to another session" + ) + connection.execute( + "DELETE FROM tape_entries WHERE tape_id = ?", + (meta.tape_id,), + ) + now = datetime.now(UTC).isoformat() + values: list[tuple[object, ...]] = [] + for seq, entry in enumerate(snapshot.tape_entries): + _require_json_object("tape entry", entry) + values.append( + ( + meta.tape_id, + seq, + json.dumps(entry, sort_keys=True), + _optional_entry_str(entry, "kind"), + _entry_nested_str(entry, "run_id"), + _entry_nested_str(entry, "tool_call_id"), + _entry_anchor_type(entry), + now, + ) + ) + connection.executemany( + """ + INSERT INTO tape_entries ( + tape_id, seq, entry_json, kind, run_id, tool_call_id, + anchor_type, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + values, + ) + connection.execute( + """ + INSERT INTO agent_http_sessions (session_id, payload, updated_at) + VALUES (?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(session_id) + DO UPDATE SET payload = excluded.payload, + updated_at = CURRENT_TIMESTAMP + """, + (authority.session_id, json.dumps(session_payload, sort_keys=True)), + ) + connection.execute( + """ + DELETE FROM checkpoints + WHERE tape_id = ? + AND session_id = ? + AND entry_count > ? + """, + (meta.tape_id, authority.session_id, meta.entry_count), + ) + + def session_id_for_checkpoint(self, checkpoint_id: str) -> str | None: + _require_non_empty("checkpoint_id", checkpoint_id) + with self._lock, self._connect() as connection: + row = connection.execute( + "SELECT session_id FROM checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + if row is None: + return None + session_id = row["session_id"] + if not isinstance(session_id, str): + raise TypeError("checkpoint session_id must be text") + return session_id + + def _assert_authority( + self, + connection: sqlite3.Connection, + authority: OwnerAuthority, + ) -> None: + row = connection.execute( + """ + SELECT owner_id, lease_expires_at, fencing_token + FROM session_owners + WHERE session_id = ? + """, + (authority.session_id,), + ).fetchone() + if row is None: + raise SessionOwnershipConflictError( + "session owner lease is missing", + reason=SessionOwnershipConflictReason.MISSING_OWNER, + ) + lease_expires_at = _datetime_from_sqlite_text(row["lease_expires_at"]) + if lease_expires_at <= datetime.now(UTC): + raise SessionOwnershipConflictError( + "session owner lease has expired", + reason=SessionOwnershipConflictReason.EXPIRED_LEASE, + ) + if ( + row["owner_id"] != authority.owner_id + or row["fencing_token"] != authority.epoch + ): + raise SessionOwnershipConflictError("stale owner or fencing token rejected") + + def _bind_tape( + self, + connection: sqlite3.Connection, + session_id: str, + tape_id: str, + ) -> None: + now = _datetime_to_sqlite_text(datetime.now(UTC)) + session_row = connection.execute( + "SELECT tape_id FROM session_tapes WHERE session_id = ?", + (session_id,), + ).fetchone() + if session_row is not None and session_row["tape_id"] != tape_id: + raise SessionOwnershipConflictError("session tape target cannot be rebound") + tape_row = connection.execute( + "SELECT session_id FROM session_tapes WHERE tape_id = ?", + (tape_id,), + ).fetchone() + if tape_row is not None and tape_row["session_id"] != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + connection.execute( + """ + INSERT INTO session_tapes (session_id, tape_id, created_at, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(session_id) + DO UPDATE SET + tape_id = excluded.tape_id, + updated_at = excluded.updated_at + """, + (session_id, tape_id, now, now), + ) + + def _assert_tape_belongs_to_session( + self, + connection: sqlite3.Connection, + *, + tape_id: str, + session_id: str, + ) -> None: + row = connection.execute( + """ + SELECT session_id + FROM session_tapes + WHERE tape_id = ? + """, + (tape_id,), + ).fetchone() + if row is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + if row["session_id"] != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + + def _assert_run_belongs_to_session( + self, + connection: sqlite3.Connection, + *, + run_id: str, + session_id: str, + ) -> None: + row = connection.execute( + "SELECT session_id FROM agent_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + raise KeyError(f"agent run not found: {run_id}") + if row["session_id"] != session_id: + raise SessionOwnershipConflictError("agent run belongs to another session") + + +def _require_non_empty(field_name: str, value: str) -> None: + if not isinstance(value, str) or not value: + raise ValueError(f"{field_name} must be non-empty") + + +def _require_json_object(field_name: str, value: object) -> None: + if not isinstance(value, dict): + raise TypeError(f"{field_name} must be a JSON object") + + +def _row_required_int(row: sqlite3.Row | None, key: str, *, context: str) -> int: + if row is None: + raise TypeError(f"{context} row is missing") + value = row[key] + if not isinstance(value, int): + raise TypeError(f"{context} {key} must be an int") + return value + + +class FencedSQLiteTapeStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteTapeStore(path) + self._authority_for_session = authority_for_session + + async def save(self, tape_id: str, entries: list[dict[str, Any]]) -> None: + session_id = self._require_session_id_for_tape(tape_id) + await self._durable_store.append_tape_entries( + self._authority_for_session(session_id), + tape_id, + entries, + ) + + async def load(self, tape_id: str) -> list[dict[str, Any]]: + return await self._delegate.load(tape_id) + + async def list_ids(self) -> list[str]: + return await self._delegate.list_ids() + + async def truncate(self, tape_id: str, keep: int) -> None: + session_id = self._require_session_id_for_tape(tape_id) + await self._durable_store.truncate_tape( + self._authority_for_session(session_id), + tape_id, + keep, + ) + + async def info(self, tape_id: str) -> TapeInfo | None: + return await self._delegate.info(tape_id) + + async def search( + self, + *, + tape_id: str | None = None, + kind: str | None = None, + run_id: str | None = None, + tool_call_id: str | None = None, + anchor_type: str | None = None, + limit: int = 100, + ) -> list[TapeSearchResult]: + return await self._delegate.search( + tape_id=tape_id, + kind=kind, + run_id=run_id, + tool_call_id=tool_call_id, + anchor_type=anchor_type, + limit=limit, + ) + + def _require_session_id_for_tape(self, tape_id: str) -> str: + session_id = self._durable_store.session_id_for_tape(tape_id) + if session_id is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + return session_id + + +class FencedSQLiteCheckpointStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteCheckpointStore(path) + self._authority_for_session = authority_for_session + + async def save(self, snapshot: CheckpointSnapshot) -> None: + session_id = snapshot.meta.session_id + if session_id is None: + raise SessionOwnershipConflictError( + "checkpoint target is not bound to a session" + ) + await self._durable_store.save_checkpoint( + self._authority_for_session(session_id), + snapshot, + ) + + async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: + return await self._delegate.load(checkpoint_id) + + async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: + return await self._delegate.list_by_tape(tape_id) + + async def delete(self, checkpoint_id: str) -> None: + session_id = self._durable_store.session_id_for_checkpoint(checkpoint_id) + if session_id is None: + return + await self._durable_store.delete_checkpoint( + self._authority_for_session(session_id), + checkpoint_id, + ) + + +class FencedSQLiteRuntimeStore: + def __init__( + self, + *, + durable_store: SQLiteLocalDurableStore, + path: Path, + authority_for_session: Callable[[str], OwnerAuthority], + authorities: Callable[[], Mapping[str, OwnerAuthority]], + ) -> None: + self._durable_store = durable_store + self._delegate = SQLiteRuntimeStore(path) + self._authority_for_session = authority_for_session + self._authorities = authorities + + async def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: + return await self._durable_store.create_agent_run( + self._authority_for_session(record.session_id), + record, + ) + + async def update_agent_run( + self, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: dict[str, Any], + result: dict[str, Any], + error: str | None, + ) -> AgentRunRecord: + session_id = self._require_session_id_for_run(run_id) + return await self._durable_store.update_agent_run( + self._authority_for_session(session_id), + run_id, + status=status, + ended_at=ended_at, + metadata=metadata, + result=result, + error=error, + ) + + async def load_agent_run(self, run_id: str) -> AgentRunRecord | None: + return await self._delegate.load_agent_run(run_id) + + async def list_agent_runs(self, session_id: str) -> list[AgentRunRecord]: + return await self._delegate.list_agent_runs(session_id) + + async def claim_attached_executor_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + return await self._durable_store.claim_attached_executor_run( + self._authorities(), + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def claim_external_worker_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: dict[str, Any], + ) -> AgentRunRecord | None: + return await self._durable_store.claim_attached_executor_run( + self._authorities(), + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def append_runtime_event( + self, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.append_runtime_event( + self._authority_for_session(session_id), + record, + ) + + async def load_runtime_event( + self, + event_id: str, + ) -> RuntimeEventRecord | None: + return await self._delegate.load_runtime_event(event_id) + + async def replay_runtime_events( + self, + run_id: str, + *, + after_sequence: int = 0, + limit: int = 1000, + ) -> list[RuntimeEventRecord]: + return await self._delegate.replay_runtime_events( + run_id, + after_sequence=after_sequence, + limit=limit, + ) + + async def save_message_snapshot( + self, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.save_message_snapshot( + self._authority_for_session(session_id), + record, + ) + + async def load_message_snapshot( + self, + snapshot_id: str, + ) -> RunMessageSnapshotRecord | None: + return await self._delegate.load_message_snapshot(snapshot_id) + + async def list_message_snapshots( + self, + run_id: str, + ) -> list[RunMessageSnapshotRecord]: + return await self._delegate.list_message_snapshots(run_id) + + async def create_agent_interaction( + self, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + session_id = self._require_session_id_for_run(record.run_id) + return await self._durable_store.create_agent_interaction( + self._authority_for_session(session_id), + record, + ) + + async def load_agent_interaction( + self, + interaction_id: str, + ) -> AgentInteractionRecord | None: + return await self._delegate.load_agent_interaction(interaction_id) + + async def list_agent_interactions( + self, + run_id: str, + ) -> list[AgentInteractionRecord]: + return await self._delegate.list_agent_interactions(run_id) + + async def resolve_agent_interaction( + self, + interaction_id: str, + *, + status: str, + response_payload: dict[str, Any], + resolved_at: datetime, + ) -> AgentInteractionRecord: + interaction = await self._delegate.load_agent_interaction(interaction_id) + if interaction is None: + raise KeyError(f"agent interaction not found: {interaction_id}") + session_id = self._require_session_id_for_run(interaction.run_id) + return await self._durable_store.resolve_agent_interaction( + self._authority_for_session(session_id), + interaction_id, + status=status, + response_payload=response_payload, + resolved_at=resolved_at, + ) + + def _require_session_id_for_run(self, run_id: str) -> str: + session_id = self._durable_store.session_id_for_run(run_id) + if session_id is None: + raise KeyError(f"agent run not found: {run_id}") + return session_id diff --git a/coding-agent/src/coding_agent/stores/durable_pg.py b/coding-agent/src/coding_agent/stores/durable_pg.py new file mode 100644 index 000000000..07f7b83ef --- /dev/null +++ b/coding-agent/src/coding_agent/stores/durable_pg.py @@ -0,0 +1,1110 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from datetime import datetime +import json +from typing import Any, cast + +from agentkit.checkpoint.models import CheckpointMeta, CheckpointSnapshot +from agentkit.storage.pg import ( + PGCheckpointStore, + PGPool, + PGSessionOwnerStore, + PGTapeStore, +) +from agentkit.storage.protocols import TapeInfo, TapeSearchResult + +from coding_agent.runtime_store import ( + AgentInteractionRecord, + AgentRunRecord, + JSONObject, + PGRuntimeStore, + RunMessageSnapshotRecord, + RuntimeEventRecord, + _agent_run_from_row, + _interaction_from_row, + _message_snapshot_from_row, + _runtime_event_from_row, +) +from coding_agent.server.stores.session_owner_store import ( + OwnerAuthority, + SessionOwnershipConflictError, + SessionOwnershipConflictReason, +) +from coding_agent.server.stores.session_store import PGSessionMetadataStore + + +class PGDurableStore: + """PostgreSQL protected write facade for durable owner fencing.""" + + _UPSERT_SESSION_SQL = """ + INSERT INTO agent_http_sessions (session_id, payload) + VALUES ($1, $2::jsonb) + ON CONFLICT (session_id) + DO UPDATE SET payload = EXCLUDED.payload, updated_at = NOW() + """ + _DELETE_SESSION_SQL = "DELETE FROM agent_http_sessions WHERE session_id = $1" + _SELECT_SESSION_FOR_UPDATE_SQL = """ + SELECT payload + FROM agent_http_sessions + WHERE session_id = $1 + FOR UPDATE + """ + _SELECT_SESSION_BY_TAPE_SQL = """ + SELECT session_id + FROM session_tapes + WHERE tape_id = $1 + LIMIT 1 + """ + _CREATE_SESSION_TAPES_SQL = """ + CREATE TABLE IF NOT EXISTS session_tapes ( + session_id TEXT PRIMARY KEY, + tape_id TEXT UNIQUE NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ); + CREATE INDEX IF NOT EXISTS session_tapes_tape_id_idx + ON session_tapes (tape_id); + """ + _SELECT_SESSION_TAPE_FOR_UPDATE_SQL = """ + SELECT tape_id + FROM session_tapes + WHERE session_id = $1 + FOR UPDATE + """ + _SELECT_TAPE_SESSION_FOR_UPDATE_SQL = """ + SELECT session_id + FROM session_tapes + WHERE tape_id = $1 + FOR UPDATE + """ + _UPSERT_SESSION_TAPE_SQL = """ + INSERT INTO session_tapes (session_id, tape_id) + VALUES ($1, $2) + ON CONFLICT DO NOTHING + """ + _SELECT_OWNER_FOR_UPDATE_SQL = """ + SELECT owner_id, lease_expires_at, fencing_token + FROM session_owners + WHERE session_id = $1 + FOR UPDATE + """ + _SELECT_RUN_SESSION_FOR_UPDATE_SQL = """ + SELECT session_id + FROM agent_runs + WHERE run_id = $1 + FOR UPDATE + """ + _SELECT_INTERACTION_RUN_FOR_UPDATE_SQL = """ + SELECT run_id + FROM agent_interactions + WHERE interaction_id = $1 + FOR UPDATE + """ + _SELECT_CHECKPOINT_META_FOR_UPDATE_SQL = """ + SELECT meta + FROM agent_checkpoints + WHERE checkpoint_id = $1 + FOR UPDATE + """ + _UPSERT_OWNED_RUN_SQL = """ + INSERT INTO agent_runs ( + run_id, + session_id, + tape_id, + parent_run_id, + agent_id, + status, + started_at, + ended_at, + metadata, + result, + error + ) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::jsonb, $10::jsonb, $11) + ON CONFLICT (run_id) + DO UPDATE SET + tape_id = EXCLUDED.tape_id, + parent_run_id = EXCLUDED.parent_run_id, + agent_id = EXCLUDED.agent_id, + status = EXCLUDED.status, + started_at = EXCLUDED.started_at, + ended_at = EXCLUDED.ended_at, + metadata = EXCLUDED.metadata, + result = EXCLUDED.result, + error = EXCLUDED.error, + updated_at = NOW() + WHERE agent_runs.session_id = $2 + RETURNING * + """ + _INSERT_OWNED_EVENT_SQL = """ + WITH inserted AS ( + INSERT INTO runtime_events ( + event_id, + run_id, + event_kind, + payload, + created_at + ) + VALUES ($1, $2, $3, $4::jsonb, $5) + ON CONFLICT (event_id) DO NOTHING + RETURNING * + ) + SELECT * FROM inserted + UNION ALL + SELECT runtime_events.* + FROM runtime_events + WHERE runtime_events.event_id = $1 + AND runtime_events.run_id = $2 + AND NOT EXISTS (SELECT 1 FROM inserted) + """ + _UPSERT_OWNED_MESSAGE_SNAPSHOT_SQL = """ + INSERT INTO run_message_snapshots ( + snapshot_id, + run_id, + messages, + metadata, + created_at + ) + VALUES ($1, $2, $3::jsonb, $4::jsonb, $5) + ON CONFLICT (snapshot_id) + DO UPDATE SET + messages = EXCLUDED.messages, + metadata = EXCLUDED.metadata, + created_at = EXCLUDED.created_at, + updated_at = NOW() + WHERE run_message_snapshots.run_id = $2 + RETURNING * + """ + _INSERT_OWNED_INTERACTION_SQL = """ + WITH inserted AS ( + INSERT INTO agent_interactions ( + interaction_id, + run_id, + interaction_kind, + status, + request_payload, + response_payload, + metadata, + created_at, + resolved_at + ) + VALUES ($1, $2, $3, $4, $5::jsonb, $6::jsonb, $7::jsonb, $8, $9) + ON CONFLICT (interaction_id) DO NOTHING + RETURNING * + ) + SELECT * FROM inserted + UNION ALL + SELECT agent_interactions.* + FROM agent_interactions + WHERE agent_interactions.interaction_id = $1 + AND agent_interactions.run_id = $2 + AND NOT EXISTS (SELECT 1 FROM inserted) + """ + _UPSERT_OWNED_CHECKPOINT_SQL = """ + INSERT INTO agent_checkpoints ( + checkpoint_id, + tape_id, + meta, + entries, + plugin_states, + extra + ) + VALUES ($1, $2, $3::jsonb, $4::jsonb, $5::jsonb, $6::jsonb) + ON CONFLICT (checkpoint_id) DO UPDATE SET + tape_id = EXCLUDED.tape_id, + meta = EXCLUDED.meta, + entries = EXCLUDED.entries, + plugin_states = EXCLUDED.plugin_states, + extra = EXCLUDED.extra + WHERE agent_checkpoints.meta->>'session_id' = $7 + RETURNING checkpoint_id + """ + _INSERT_TAPE_SQL = PGTapeStore._INSERT_SQL + _TRUNCATE_TAPE_SQL = PGTapeStore._TRUNCATE_SQL + _DELETE_CHECKPOINT_SQL = PGCheckpointStore._DELETE_SQL + _DELETE_NEWER_CHECKPOINTS_SQL = """ + DELETE FROM agent_checkpoints + WHERE tape_id = $1 + AND meta->>'session_id' = $2 + AND (meta->>'entry_count')::int > $3 + """ + _DELETE_SESSION_TAPE_SQL = "DELETE FROM session_tapes WHERE session_id = $1" + + def __init__(self, *, pool: PGPool) -> None: + self._pool = pool + self._schema_ready = False + + async def _ensure_schema(self) -> None: + pool = await self._pool.get_pool() + if self._schema_ready: + return + _ = await pool.execute(PGSessionOwnerStore._CREATE_TABLE_SQL) + _ = await pool.execute(PGSessionMetadataStore._CREATE_TABLE_SQL) + _ = await pool.execute(self._CREATE_SESSION_TAPES_SQL) + _ = await pool.execute(PGTapeStore._CREATE_TABLE_SQL) + _ = await pool.execute(PGCheckpointStore._CREATE_TABLE_SQL) + _ = await pool.execute(PGRuntimeStore._CREATE_SCHEMA_SQL) + self._schema_ready = True + + async def _with_transaction(self, body: Callable[[Any], Any]) -> Any: + await self._ensure_schema() + connection = await self._pool.acquire() + try: + _ = await connection.execute("BEGIN") + result = await body(connection) + _ = await connection.execute("COMMIT") + return result + except BaseException: + try: + _ = await connection.execute("ROLLBACK") + except Exception: + pass + raise + finally: + await self._pool.release(connection) + + async def session_id_for_tape(self, tape_id: str) -> str | None: + await self._ensure_schema() + pool = await self._pool.get_pool() + row = await pool.fetchrow(self._SELECT_SESSION_BY_TAPE_SQL, tape_id) + if row is None: + return None + return _required_str(dict(row), "session_id") + + async def save_session( + self, + authority: OwnerAuthority, + payload: dict[str, Any], + ) -> None: + _require_payload_session(authority, payload) + tape_id = payload.get("tape_id") + if tape_id is not None and not isinstance(tape_id, str): + raise TypeError("session payload tape_id must be a string") + + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + if tape_id: + await self._bind_tape(connection, authority.session_id, tape_id) + _ = await connection.fetchrow( + self._SELECT_SESSION_FOR_UPDATE_SQL, + authority.session_id, + ) + await connection.execute( + self._UPSERT_SESSION_SQL, + authority.session_id, + payload, + ) + + await self._with_transaction(body) + + async def delete_session(self, authority: OwnerAuthority) -> None: + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + _ = await connection.fetchrow( + self._SELECT_SESSION_FOR_UPDATE_SQL, + authority.session_id, + ) + await connection.execute(self._DELETE_SESSION_SQL, authority.session_id) + await connection.execute( + self._DELETE_SESSION_TAPE_SQL, authority.session_id + ) + + await self._with_transaction(body) + + async def append_tape_entries( + self, + authority: OwnerAuthority, + tape_id: str, + entries: list[dict[str, Any]], + ) -> None: + if not entries: + return + + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + await self._require_stable_tape(connection, authority, tape_id) + payload_values = [json.dumps(entry) for entry in entries] + await connection.execute(self._INSERT_TAPE_SQL, tape_id, payload_values) + + await self._with_transaction(body) + + async def truncate_tape( + self, + authority: OwnerAuthority, + tape_id: str, + keep: int, + ) -> None: + if keep < 0: + raise ValueError("keep must be >= 0") + + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + await self._require_stable_tape(connection, authority, tape_id) + await connection.execute(self._TRUNCATE_TAPE_SQL, tape_id, keep) + + await self._with_transaction(body) + + async def create_agent_run( + self, + authority: OwnerAuthority, + record: AgentRunRecord, + ) -> AgentRunRecord: + if record.session_id != authority.session_id: + raise SessionOwnershipConflictError("run target belongs to another owner") + if record.tape_id is None: + raise SessionOwnershipConflictError("run target is not bound to a tape") + + async def body(connection: Any) -> AgentRunRecord: + await self._require_owner(connection, authority) + await self._require_stable_tape(connection, authority, record.tape_id) + row = await connection.fetchrow( + self._UPSERT_OWNED_RUN_SQL, + record.run_id, + record.session_id, + record.tape_id, + record.parent_run_id, + record.agent_id, + record.status, + record.started_at, + record.ended_at, + record.metadata, + record.result, + record.error, + ) + return _agent_run_from_row( + _required_owned_row(row, "run target belongs to another owner") + ) + + return cast(AgentRunRecord, await self._with_transaction(body)) + + async def update_agent_run( + self, + authority: OwnerAuthority, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: JSONObject, + result: JSONObject, + error: str | None, + ) -> AgentRunRecord: + async def body(connection: Any) -> AgentRunRecord: + await self._require_owner(connection, authority) + await self._require_run_owner(connection, authority, run_id) + row = await connection.fetchrow( + PGRuntimeStore._UPDATE_RUN_SQL, + run_id, + status, + ended_at, + metadata, + result, + error, + ) + return _agent_run_from_row(_required_row(row, "agent run update")) + + return cast(AgentRunRecord, await self._with_transaction(body)) + + async def append_runtime_event( + self, + authority: OwnerAuthority, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + async def body(connection: Any) -> RuntimeEventRecord: + await self._require_owner(connection, authority) + await self._require_run_owner(connection, authority, record.run_id) + row = await connection.fetchrow( + self._INSERT_OWNED_EVENT_SQL, + record.event_id, + record.run_id, + record.event_kind, + record.payload, + record.created_at, + ) + return _runtime_event_from_row( + _required_owned_row( + row, + "runtime event target belongs to another owner", + ) + ) + + return cast(RuntimeEventRecord, await self._with_transaction(body)) + + async def save_message_snapshot( + self, + authority: OwnerAuthority, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + async def body(connection: Any) -> RunMessageSnapshotRecord: + await self._require_owner(connection, authority) + await self._require_run_owner(connection, authority, record.run_id) + row = await connection.fetchrow( + self._UPSERT_OWNED_MESSAGE_SNAPSHOT_SQL, + record.snapshot_id, + record.run_id, + record.messages, + record.metadata, + record.created_at, + ) + return _message_snapshot_from_row( + _required_owned_row( + row, + "message snapshot target belongs to another owner", + ) + ) + + return cast(RunMessageSnapshotRecord, await self._with_transaction(body)) + + async def create_agent_interaction( + self, + authority: OwnerAuthority, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + async def body(connection: Any) -> AgentInteractionRecord: + await self._require_owner(connection, authority) + await self._require_run_owner(connection, authority, record.run_id) + row = await connection.fetchrow( + self._INSERT_OWNED_INTERACTION_SQL, + record.interaction_id, + record.run_id, + record.interaction_kind, + record.status, + record.request_payload, + record.response_payload, + record.metadata, + record.created_at, + record.resolved_at, + ) + return _interaction_from_row( + _required_owned_row(row, "interaction target belongs to another owner") + ) + + return cast(AgentInteractionRecord, await self._with_transaction(body)) + + async def resolve_agent_interaction( + self, + authority: OwnerAuthority, + interaction_id: str, + *, + status: str, + response_payload: JSONObject, + resolved_at: datetime, + ) -> AgentInteractionRecord: + async def body(connection: Any) -> AgentInteractionRecord: + await self._require_owner(connection, authority) + existing = await connection.fetchrow( + self._SELECT_INTERACTION_RUN_FOR_UPDATE_SQL, + interaction_id, + ) + if existing is None: + raise KeyError(f"agent interaction not found: {interaction_id}") + await self._require_run_owner( + connection, + authority, + _required_str(dict(existing), "run_id"), + ) + row = await connection.fetchrow( + PGRuntimeStore._RESOLVE_INTERACTION_SQL, + interaction_id, + status, + response_payload, + resolved_at, + ) + return _interaction_from_row( + _required_row(row, "agent interaction resolve") + ) + + return cast(AgentInteractionRecord, await self._with_transaction(body)) + + async def claim_attached_executor_run( + self, + authorities: Mapping[str, OwnerAuthority], + *, + session_id: str | None, + executor_kind: str, + claim_metadata: JSONObject, + ) -> AgentRunRecord | None: + if session_id is not None: + authority = authorities.get(session_id) + if authority is None: + return None + return await self._claim_attached_executor_run_for_authority( + authority, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + for authority in sorted(authorities.values(), key=lambda item: item.session_id): + claimed = await self._claim_attached_executor_run_for_authority( + authority, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + if claimed is not None: + return claimed + return None + + async def _claim_attached_executor_run_for_authority( + self, + authority: OwnerAuthority, + *, + executor_kind: str, + claim_metadata: JSONObject, + ) -> AgentRunRecord | None: + async def body(connection: Any) -> AgentRunRecord | None: + await self._require_owner(connection, authority) + row = await connection.fetchrow( + PGRuntimeStore._CLAIM_ATTACHED_EXECUTOR_RUN_SQL, + authority.session_id, + executor_kind, + claim_metadata, + ) + if row is None: + return None + return _agent_run_from_row(row) + + return cast(AgentRunRecord | None, await self._with_transaction(body)) + + async def save_checkpoint( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + ) -> None: + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError( + "checkpoint target belongs to another owner" + ) + + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + await self._require_stable_tape(connection, authority, meta.tape_id) + row = await connection.fetchrow( + self._UPSERT_OWNED_CHECKPOINT_SQL, + meta.checkpoint_id, + meta.tape_id, + _checkpoint_meta_payload(meta), + list(snapshot.tape_entries), + snapshot.plugin_states, + snapshot.extra, + authority.session_id, + ) + _required_owned_row(row, "checkpoint target belongs to another owner") + + await self._with_transaction(body) + + async def delete_checkpoint( + self, + authority: OwnerAuthority, + checkpoint_id: str, + ) -> None: + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + await self._require_checkpoint_owner(connection, authority, checkpoint_id) + await connection.execute(self._DELETE_CHECKPOINT_SQL, checkpoint_id) + + await self._with_transaction(body) + + async def restore_checkpoint_state( + self, + authority: OwnerAuthority, + snapshot: CheckpointSnapshot, + session_payload: dict[str, Any], + ) -> None: + _require_payload_session(authority, session_payload) + meta = snapshot.meta + if meta.session_id != authority.session_id: + raise SessionOwnershipConflictError( + "checkpoint target belongs to another owner" + ) + if session_payload.get("tape_id") != meta.tape_id: + raise SessionOwnershipConflictError( + "checkpoint restore session payload has mismatched tape id" + ) + + async def body(connection: Any) -> None: + await self._require_owner(connection, authority) + await self._require_stable_tape(connection, authority, meta.tape_id) + await self._require_checkpoint_owner( + connection, + authority, + meta.checkpoint_id, + ) + await connection.execute(self._TRUNCATE_TAPE_SQL, meta.tape_id, 0) + if snapshot.tape_entries: + payload_values = [json.dumps(entry) for entry in snapshot.tape_entries] + await connection.execute( + self._INSERT_TAPE_SQL, meta.tape_id, payload_values + ) + await connection.execute( + self._UPSERT_SESSION_SQL, + authority.session_id, + session_payload, + ) + await connection.execute( + self._DELETE_NEWER_CHECKPOINTS_SQL, + meta.tape_id, + authority.session_id, + meta.entry_count, + ) + + await self._with_transaction(body) + + async def _require_owner(self, connection: Any, authority: OwnerAuthority) -> None: + row = await connection.fetchrow( + self._SELECT_OWNER_FOR_UPDATE_SQL, + authority.session_id, + ) + if row is None: + raise SessionOwnershipConflictError( + "session owner lease is missing", + reason=SessionOwnershipConflictReason.MISSING_OWNER, + ) + row_dict = dict(row) + lease_expires_at = row_dict.get("lease_expires_at") + if not isinstance(lease_expires_at, datetime): + raise TypeError("postgres owner row missing datetime lease_expires_at") + now = datetime.now(lease_expires_at.tzinfo) + if lease_expires_at <= now: + raise SessionOwnershipConflictError( + "session owner lease has expired", + reason=SessionOwnershipConflictReason.EXPIRED_LEASE, + ) + if ( + row_dict.get("owner_id") != authority.owner_id + or row_dict.get("fencing_token") != authority.epoch + ): + raise SessionOwnershipConflictError("stale owner or fencing token rejected") + + async def _require_stable_tape( + self, + connection: Any, + authority: OwnerAuthority, + tape_id: str, + ) -> None: + row = await connection.fetchrow( + self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, + tape_id, + ) + if row is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + if dict(row).get("session_id") != authority.session_id: + raise SessionOwnershipConflictError("tape target belongs to another owner") + + async def _bind_tape( + self, + connection: Any, + session_id: str, + tape_id: str, + ) -> None: + await connection.execute(self._UPSERT_SESSION_TAPE_SQL, session_id, tape_id) + session_row = await connection.fetchrow( + self._SELECT_SESSION_TAPE_FOR_UPDATE_SQL, + session_id, + ) + if session_row is None: + tape_row = await connection.fetchrow( + self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, + tape_id, + ) + if tape_row is not None and dict(tape_row).get("session_id") != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + raise SessionOwnershipConflictError("session tape target is not bound") + if dict(session_row).get("tape_id") != tape_id: + raise SessionOwnershipConflictError("session tape target cannot be rebound") + tape_row = await connection.fetchrow( + self._SELECT_TAPE_SESSION_FOR_UPDATE_SQL, + tape_id, + ) + if tape_row is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + if dict(tape_row).get("session_id") != session_id: + raise SessionOwnershipConflictError( + "tape target belongs to another session" + ) + + async def _session_id_for_run( + self, + connection: Any, + run_id: str, + *, + for_update: bool, + ) -> str | None: + del for_update + row = await connection.fetchrow(self._SELECT_RUN_SESSION_FOR_UPDATE_SQL, run_id) + if row is None: + return None + return _required_str(dict(row), "session_id") + + async def _require_run_owner( + self, + connection: Any, + authority: OwnerAuthority, + run_id: str, + ) -> None: + session_id = await self._session_id_for_run( + connection, + run_id, + for_update=True, + ) + if session_id is None: + raise KeyError(f"agent run not found: {run_id}") + if session_id != authority.session_id: + raise SessionOwnershipConflictError("run target belongs to another owner") + + async def _require_checkpoint_owner( + self, + connection: Any, + authority: OwnerAuthority, + checkpoint_id: str, + ) -> None: + row = await connection.fetchrow( + self._SELECT_CHECKPOINT_META_FOR_UPDATE_SQL, + checkpoint_id, + ) + if row is None: + raise KeyError(f"checkpoint not found: {checkpoint_id}") + meta = _required_dict(dict(row), "meta") + if meta.get("session_id") != authority.session_id: + raise SessionOwnershipConflictError( + "checkpoint target belongs to another owner" + ) + + +class FencedPGTapeStore: + def __init__( + self, + *, + durable_store: PGDurableStore, + pool: PGPool, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = PGTapeStore(pool=pool) + self._authority_for_session = authority_for_session + + async def save(self, tape_id: str, entries: list[dict[str, Any]]) -> None: + session_id = await self._require_session_id_for_tape_async(tape_id) + await self._durable_store.append_tape_entries( + self._authority_for_session(session_id), + tape_id, + entries, + ) + + async def load(self, tape_id: str) -> list[dict[str, Any]]: + return await self._delegate.load(tape_id) + + async def list_ids(self) -> list[str]: + return await self._delegate.list_ids() + + async def truncate(self, tape_id: str, keep: int) -> None: + session_id = await self._require_session_id_for_tape_async(tape_id) + await self._durable_store.truncate_tape( + self._authority_for_session(session_id), + tape_id, + keep, + ) + + async def info(self, tape_id: str) -> TapeInfo | None: + return await self._delegate.info(tape_id) + + async def search( + self, + *, + tape_id: str | None = None, + kind: str | None = None, + run_id: str | None = None, + tool_call_id: str | None = None, + anchor_type: str | None = None, + limit: int = 100, + ) -> list[TapeSearchResult]: + return await self._delegate.search( + tape_id=tape_id, + kind=kind, + run_id=run_id, + tool_call_id=tool_call_id, + anchor_type=anchor_type, + limit=limit, + ) + + async def _require_session_id_for_tape_async(self, tape_id: str) -> str: + session_id = await self._durable_store.session_id_for_tape(tape_id) + if session_id is None: + raise SessionOwnershipConflictError("tape target is not bound to a session") + return session_id + + +class FencedPGCheckpointStore: + def __init__( + self, + *, + durable_store: PGDurableStore, + pool: PGPool, + authority_for_session: Callable[[str], OwnerAuthority], + ) -> None: + self._durable_store = durable_store + self._delegate = PGCheckpointStore(pool=pool) + self._authority_for_session = authority_for_session + + async def save(self, snapshot: CheckpointSnapshot) -> None: + session_id = snapshot.meta.session_id + if session_id is None: + raise SessionOwnershipConflictError( + "checkpoint target is not bound to a session" + ) + await self._durable_store.save_checkpoint( + self._authority_for_session(session_id), + snapshot, + ) + + async def load(self, checkpoint_id: str) -> CheckpointSnapshot | None: + return await self._delegate.load(checkpoint_id) + + async def list_by_tape(self, tape_id: str) -> list[CheckpointMeta]: + return await self._delegate.list_by_tape(tape_id) + + async def delete(self, checkpoint_id: str) -> None: + snapshot = await self._delegate.load(checkpoint_id) + if snapshot is None: + return + session_id = snapshot.meta.session_id + if session_id is None: + raise SessionOwnershipConflictError( + "checkpoint target is not bound to a session" + ) + await self._durable_store.delete_checkpoint( + self._authority_for_session(session_id), + checkpoint_id, + ) + + +class FencedPGRuntimeStore: + def __init__( + self, + *, + durable_store: PGDurableStore, + pool: PGPool, + authority_for_session: Callable[[str], OwnerAuthority], + authorities: Callable[[], Mapping[str, OwnerAuthority]], + ) -> None: + self._durable_store = durable_store + self._delegate = PGRuntimeStore(pool=pool) + self._authority_for_session = authority_for_session + self._authorities = authorities + + async def create_agent_run(self, record: AgentRunRecord) -> AgentRunRecord: + return await self._durable_store.create_agent_run( + self._authority_for_session(record.session_id), + record, + ) + + async def update_agent_run( + self, + run_id: str, + *, + status: str, + ended_at: datetime | None, + metadata: JSONObject, + result: JSONObject, + error: str | None, + ) -> AgentRunRecord: + existing = await self._delegate.load_agent_run(run_id) + if existing is None: + raise KeyError(f"agent run not found: {run_id}") + return await self._durable_store.update_agent_run( + self._authority_for_session(existing.session_id), + run_id, + status=status, + ended_at=ended_at, + metadata=metadata, + result=result, + error=error, + ) + + async def load_agent_run(self, run_id: str) -> AgentRunRecord | None: + return await self._delegate.load_agent_run(run_id) + + async def list_agent_runs(self, session_id: str) -> list[AgentRunRecord]: + return await self._delegate.list_agent_runs(session_id) + + async def claim_attached_executor_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: JSONObject, + ) -> AgentRunRecord | None: + return await self._durable_store.claim_attached_executor_run( + self._authorities(), + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def claim_external_worker_run( + self, + *, + session_id: str | None, + executor_kind: str, + claim_metadata: JSONObject, + ) -> AgentRunRecord | None: + return await self.claim_attached_executor_run( + session_id=session_id, + executor_kind=executor_kind, + claim_metadata=claim_metadata, + ) + + async def append_runtime_event( + self, + record: RuntimeEventRecord, + ) -> RuntimeEventRecord: + existing = await self._delegate.load_agent_run(record.run_id) + if existing is None: + raise KeyError(f"agent run not found: {record.run_id}") + return await self._durable_store.append_runtime_event( + self._authority_for_session(existing.session_id), + record, + ) + + async def replay_runtime_events( + self, + run_id: str, + *, + after_sequence: int = 0, + limit: int = 1000, + ) -> list[RuntimeEventRecord]: + return await self._delegate.replay_runtime_events( + run_id, + after_sequence=after_sequence, + limit=limit, + ) + + async def load_runtime_event( + self, + event_id: str, + ) -> RuntimeEventRecord | None: + return await self._delegate.load_runtime_event(event_id) + + async def save_message_snapshot( + self, + record: RunMessageSnapshotRecord, + ) -> RunMessageSnapshotRecord: + existing = await self._delegate.load_agent_run(record.run_id) + if existing is None: + raise KeyError(f"agent run not found: {record.run_id}") + return await self._durable_store.save_message_snapshot( + self._authority_for_session(existing.session_id), + record, + ) + + async def load_message_snapshot( + self, + snapshot_id: str, + ) -> RunMessageSnapshotRecord | None: + return await self._delegate.load_message_snapshot(snapshot_id) + + async def list_message_snapshots( + self, + run_id: str, + ) -> list[RunMessageSnapshotRecord]: + return await self._delegate.list_message_snapshots(run_id) + + async def create_agent_interaction( + self, + record: AgentInteractionRecord, + ) -> AgentInteractionRecord: + existing = await self._delegate.load_agent_run(record.run_id) + if existing is None: + raise KeyError(f"agent run not found: {record.run_id}") + return await self._durable_store.create_agent_interaction( + self._authority_for_session(existing.session_id), + record, + ) + + async def resolve_agent_interaction( + self, + interaction_id: str, + *, + status: str, + response_payload: JSONObject, + resolved_at: datetime, + ) -> AgentInteractionRecord: + existing = await self._delegate.load_agent_interaction(interaction_id) + if existing is None: + raise KeyError(f"agent interaction not found: {interaction_id}") + run = await self._delegate.load_agent_run(existing.run_id) + if run is None: + raise KeyError(f"agent run not found: {existing.run_id}") + return await self._durable_store.resolve_agent_interaction( + self._authority_for_session(run.session_id), + interaction_id, + status=status, + response_payload=response_payload, + resolved_at=resolved_at, + ) + + async def load_agent_interaction( + self, + interaction_id: str, + ) -> AgentInteractionRecord | None: + return await self._delegate.load_agent_interaction(interaction_id) + + async def list_agent_interactions( + self, + run_id: str, + ) -> list[AgentInteractionRecord]: + return await self._delegate.list_agent_interactions(run_id) + + +def _require_payload_session( + authority: OwnerAuthority, + payload: dict[str, Any], +) -> None: + payload_id = payload.get("id") + if payload_id != authority.session_id: + raise SessionOwnershipConflictError("session payload belongs to another owner") + payload_session_id = payload.get("session_id") + if payload_session_id is not None and payload_session_id != authority.session_id: + raise SessionOwnershipConflictError("session payload belongs to another owner") + + +def _checkpoint_meta_payload(meta: CheckpointMeta) -> dict[str, Any]: + return { + "checkpoint_id": meta.checkpoint_id, + "tape_id": meta.tape_id, + "session_id": meta.session_id, + "entry_count": meta.entry_count, + "window_start": meta.window_start, + "created_at": meta.created_at.isoformat(), + "label": meta.label, + } + + +def _required_row(row: dict[str, object] | None, context: str) -> dict[str, object]: + if row is None: + raise RuntimeError(f"postgres {context} returned no row") + return row + + +def _required_owned_row( + row: dict[str, object] | None, + conflict_message: str, +) -> dict[str, object]: + if row is None: + raise SessionOwnershipConflictError(conflict_message) + return row + + +def _required_str(row: dict[str, object], key: str) -> str: + value = row.get(key) + if not isinstance(value, str): + raise TypeError(f"postgres row must include string {key}") + return value + + +def _required_dict(row: dict[str, object], key: str) -> dict[str, object]: + value = row.get(key) + if not isinstance(value, dict): + raise TypeError(f"postgres row must include dict {key}") + return cast(dict[str, object], value) diff --git a/coding-agent/src/coding_agent/stores/migration.py b/coding-agent/src/coding_agent/stores/migration.py new file mode 100644 index 000000000..63ff3c5ac --- /dev/null +++ b/coding-agent/src/coding_agent/stores/migration.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast + +from agentkit.storage.checkpoint_fs import FSCheckpointStore +from agentkit.storage.sqlite import SQLiteCheckpointStore, SQLiteTapeStore +from coding_agent.local_storage import local_sqlite_path + + +@dataclass(frozen=True) +class StoreMigrationReport: + scanned: int = 0 + migrated: int = 0 + skipped: int = 0 + + +@dataclass(frozen=True) +class LegacySQLiteMigrationReport: + tapes: StoreMigrationReport + checkpoints: StoreMigrationReport + + +async def migrate_legacy_storage_to_sqlite( + data_dir: Path, + *, + tapes_dir: Path | None = None, + checkpoints_dir: Path | None = None, + tape_sqlite_path: Path | None = None, + checkpoint_sqlite_path: Path | None = None, + replace_tapes: bool = False, + dry_run: bool = False, +) -> LegacySQLiteMigrationReport: + resolved_data_dir = Path(data_dir) + return LegacySQLiteMigrationReport( + tapes=await migrate_jsonl_tapes_to_sqlite( + tapes_dir or resolved_data_dir / "tapes", + tape_sqlite_path or local_sqlite_path(resolved_data_dir), + replace=replace_tapes, + dry_run=dry_run, + ), + checkpoints=await migrate_fs_checkpoints_to_sqlite( + checkpoints_dir or resolved_data_dir / "checkpoints", + checkpoint_sqlite_path or local_sqlite_path(resolved_data_dir), + dry_run=dry_run, + ), + ) + + +async def migrate_jsonl_tapes_to_sqlite( + source_dir: Path, + sqlite_path: Path, + *, + replace: bool = False, + dry_run: bool = False, +) -> StoreMigrationReport: + source_dir = Path(source_dir) + sqlite_path = Path(sqlite_path) + if not source_dir.exists(): + return StoreMigrationReport() + if not source_dir.is_dir(): + raise NotADirectoryError(str(source_dir)) + + store = ( + None if dry_run and not sqlite_path.exists() else SQLiteTapeStore(sqlite_path) + ) + scanned = 0 + migrated = 0 + skipped = 0 + for path in sorted(source_dir.glob("*.jsonl")): + scanned += 1 + tape_id = path.stem + entries = _read_jsonl_objects(path) + existing = [] if store is None else await store.load(tape_id) + if not entries and not existing: + skipped += 1 + continue + if existing: + if existing == entries: + skipped += 1 + continue + if not replace: + raise ValueError( + f"tape {tape_id!r} already exists with different entries" + ) + if not dry_run: + if store is None: + store = SQLiteTapeStore(sqlite_path) + await store.truncate(tape_id, 0) + if not dry_run: + if store is None: + store = SQLiteTapeStore(sqlite_path) + await store.save(tape_id, entries) + migrated += 1 + return StoreMigrationReport(scanned=scanned, migrated=migrated, skipped=skipped) + + +async def migrate_fs_checkpoints_to_sqlite( + source_dir: Path, + sqlite_path: Path, + *, + dry_run: bool = False, +) -> StoreMigrationReport: + source_dir = Path(source_dir) + sqlite_path = Path(sqlite_path) + if not source_dir.exists(): + return StoreMigrationReport() + if not source_dir.is_dir(): + raise NotADirectoryError(str(source_dir)) + + source_store = FSCheckpointStore(source_dir) + target_store = ( + None + if dry_run and not sqlite_path.exists() + else SQLiteCheckpointStore(sqlite_path) + ) + scanned = 0 + migrated = 0 + skipped = 0 + for meta_path in sorted(source_dir.glob("*.meta.json")): + checkpoint_id = meta_path.name.removesuffix(".meta.json") + scanned += 1 + snapshot = await source_store.load(checkpoint_id) + if snapshot is None: + raise RuntimeError( + f"checkpoint metadata {meta_path} did not load checkpoint snapshot" + ) + existing = ( + None if target_store is None else await target_store.load(checkpoint_id) + ) + if existing == snapshot: + skipped += 1 + continue + if not dry_run: + if target_store is None: + target_store = SQLiteCheckpointStore(sqlite_path) + await target_store.save(snapshot) + migrated += 1 + return StoreMigrationReport(scanned=scanned, migrated=migrated, skipped=skipped) + + +def _read_jsonl_objects(path: Path) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + with path.open(encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + stripped = line.strip() + if not stripped: + continue + loaded = json.loads(stripped) + if not isinstance(loaded, dict): + raise TypeError( + f"{path}:{line_number} must contain a JSON object tape entry" + ) + entries.append(cast(dict[str, Any], loaded)) + return entries + + +__all__ = [ + "LegacySQLiteMigrationReport", + "StoreMigrationReport", + "migrate_fs_checkpoints_to_sqlite", + "migrate_jsonl_tapes_to_sqlite", + "migrate_legacy_storage_to_sqlite", +] diff --git a/coding-agent/src/coding_agent/workspace_archive.py b/coding-agent/src/coding_agent/workspace_archive.py index d44463d4e..0c098ca18 100644 --- a/coding-agent/src/coding_agent/workspace_archive.py +++ b/coding-agent/src/coding_agent/workspace_archive.py @@ -1,236 +1,6 @@ -from __future__ import annotations +"""Deprecated import location. Moved to coding_agent.environment.archive.""" +import sys as _sys -import base64 -import gzip -import io -import os -import shutil -import tarfile -import tempfile -import zlib -from pathlib import Path -from typing import Protocol, override +from coding_agent.environment import archive as _moved - -_PRESERVED_ROOT_NAMES = frozenset({".git"}) -_EXCLUDED_DIRECTORY_NAMES = frozenset({"__pycache__"}) -_EXCLUDED_FILE_SUFFIXES = (".pyc", ".pyo") -_MAX_WORKSPACE_ARCHIVE_BYTES = 8 * 1024 * 1024 -_MAX_WORKSPACE_ARCHIVE_BASE64_CHARS = 4 * ((_MAX_WORKSPACE_ARCHIVE_BYTES + 2) // 3) -_MAX_WORKSPACE_TAR_STREAM_BYTES = 16 * 1024 * 1024 -_MAX_WORKSPACE_ARCHIVE_MEMBERS = 4096 - - -def create_workspace_archive_base64(workspace_root: Path) -> str: - root = workspace_root.expanduser().resolve() - if not root.is_dir(): - raise ValueError(f"workspace root does not exist: {root}") - - buffer = io.BytesIO() - total_input_bytes = 0 - with tarfile.open(fileobj=buffer, mode="w:gz") as archive: - for path in sorted(root.rglob("*")): - relative = path.relative_to(root) - if should_exclude_workspace_archive_path(relative): - continue - if path.is_symlink(): - raise ValueError( - f"workspace archive does not support symlinks: {relative}" - ) - if not path.is_file(): - continue - stat_result = path.stat() - total_input_bytes += stat_result.st_size - _raise_if_archive_too_large(total_input_bytes) - data = path.read_bytes() - info = tarfile.TarInfo(name=relative.as_posix()) - info.size = len(data) - info.mode = stat_result.st_mode & 0o777 - info.mtime = int(stat_result.st_mtime) - archive.addfile(info, io.BytesIO(data)) - archive_bytes = buffer.getvalue() - _raise_if_archive_too_large(len(archive_bytes)) - return base64.b64encode(archive_bytes).decode("ascii") - - -def should_exclude_workspace_archive_path(relative: Path) -> bool: - if not relative.parts: - return True - if relative.parts[0] in _PRESERVED_ROOT_NAMES: - return True - if any(part in _EXCLUDED_DIRECTORY_NAMES for part in relative.parts): - return True - return relative.name.endswith(_EXCLUDED_FILE_SUFFIXES) - - -def extract_workspace_archive_base64(workspace_root: Path, archive_base64: str) -> None: - root = workspace_root.expanduser().resolve() - if not root.exists(): - root.mkdir(parents=True, exist_ok=True) - if not root.is_dir(): - raise ValueError(f"workspace root is not a directory: {root}") - - _raise_if_archive_base64_too_large(archive_base64) - try: - archive_bytes = base64.b64decode(archive_base64.encode("ascii"), validate=True) - except ValueError as exc: - raise ValueError("workspace archive must be valid base64") from exc - _raise_if_archive_too_large(len(archive_bytes)) - - with tempfile.TemporaryDirectory() as temp_dir: - temp_root = Path(temp_dir).resolve() - try: - _extract_workspace_archive_to_temp(temp_root, archive_bytes) - except (tarfile.TarError, gzip.BadGzipFile, EOFError, zlib.error) as exc: - raise ValueError("workspace archive is not a valid tar.gz") from exc - - _clear_extract_target(root) - _copy_extracted_workspace(temp_root, root) - - -def _extract_workspace_archive_to_temp(temp_root: Path, archive_bytes: bytes) -> None: - archive_buffer = io.BytesIO(archive_bytes) - gzip_stream = gzip.GzipFile(fileobj=archive_buffer, mode="rb") - counted_stream = _CountingReader( - gzip_stream, - max_bytes=_MAX_WORKSPACE_TAR_STREAM_BYTES, - ) - extracted_bytes = 0 - member_count = 0 - with tarfile.open(fileobj=counted_stream, mode="r|") as archive: - for member in archive: - member_count += 1 - if member_count > _MAX_WORKSPACE_ARCHIVE_MEMBERS: - raise ValueError("workspace archive contains too many members") - _reject_preserved_root_member(member.name) - target_path = _safe_archive_target(temp_root, member.name) - if member.isdir(): - target_path.mkdir(parents=True, exist_ok=True) - target_path.chmod(_member_mode(member)) - continue - if not member.isfile(): - raise ValueError( - f"workspace archive only supports regular files and directories: {member.name}" - ) - if member.size < 0: - raise ValueError("workspace archive is not a valid tar.gz") - extracted_bytes += member.size - _raise_if_archive_too_large(extracted_bytes) - file_obj = archive.extractfile(member) - if file_obj is None: - raise ValueError( - f"workspace archive member is unreadable: {member.name}" - ) - _copy_archive_file(file_obj, target_path, member) - - -def _copy_archive_file( - file_obj: _ReadableBinary, target_path: Path, member: tarfile.TarInfo -) -> None: - target_path.parent.mkdir(parents=True, exist_ok=True) - copied_bytes = 0 - with target_path.open("wb") as target_file: - while True: - chunk = file_obj.read(64 * 1024) - if not chunk: - break - copied_bytes += len(chunk) - if copied_bytes > member.size: - raise ValueError("workspace archive is not a valid tar.gz") - _ = target_file.write(chunk) - if copied_bytes != member.size: - raise ValueError("workspace archive is not a valid tar.gz") - target_path.chmod(_member_mode(member)) - os.utime(target_path, (member.mtime, member.mtime)) - - -def _copy_extracted_workspace(temp_root: Path, workspace_root: Path) -> None: - for source_path in sorted(temp_root.rglob("*")): - relative = source_path.relative_to(temp_root) - target_path = workspace_root / relative - if source_path.is_dir(): - target_path.mkdir(parents=True, exist_ok=True) - target_path.chmod(source_path.stat().st_mode & 0o777) - continue - target_path.parent.mkdir(parents=True, exist_ok=True) - _ = shutil.copy2(source_path, target_path) - - -def _clear_extract_target(workspace_root: Path) -> None: - for child in workspace_root.iterdir(): - if child.name in _PRESERVED_ROOT_NAMES: - continue - if child.is_dir() and not child.is_symlink(): - shutil.rmtree(child) - continue - child.unlink() - - -def _member_mode(member: tarfile.TarInfo) -> int: - mode = member.mode & 0o777 - if mode != 0: - return mode - if member.isdir(): - return 0o755 - return 0o644 - - -def _safe_archive_target(workspace_root: Path, member_name: str) -> Path: - member_path = Path(member_name) - if member_path.is_absolute() or ".." in member_path.parts: - raise ValueError( - f"workspace archive member escapes workspace root: {member_name}" - ) - if not member_path.parts: - raise ValueError("workspace archive member name must be non-empty") - candidate = (workspace_root / member_path).resolve() - try: - _ = candidate.relative_to(workspace_root) - except ValueError as exc: - raise ValueError( - f"workspace archive member escapes workspace root: {member_name}" - ) from exc - return candidate - - -def _reject_preserved_root_member(member_name: str) -> None: - member_path = Path(member_name) - if member_path.parts and member_path.parts[0] in _PRESERVED_ROOT_NAMES: - raise ValueError( - f"workspace archive must not contain preserved root entry: {member_name}" - ) - - -def _raise_if_archive_too_large(size_bytes: int) -> None: - if size_bytes > _MAX_WORKSPACE_ARCHIVE_BYTES: - raise ValueError("workspace archive exceeds 8 MiB limit") - - -def _raise_if_archive_base64_too_large(archive_base64: str) -> None: - if len(archive_base64) > _MAX_WORKSPACE_ARCHIVE_BASE64_CHARS: - raise ValueError("workspace archive exceeds 8 MiB limit") - - -class _ReadableBinary(Protocol): - def read(self, size: int = -1, /) -> bytes: ... - - -class _CountingReader(io.BufferedIOBase): - def __init__(self, stream: _ReadableBinary, *, max_bytes: int) -> None: - self._stream: _ReadableBinary = stream - self._max_bytes: int = max_bytes - self._bytes_read: int = 0 - - @override - def readable(self) -> bool: - return True - - @override - def read(self, size: int | None = -1) -> bytes: - if size is None: - size = -1 - chunk = self._stream.read(size) - self._bytes_read += len(chunk) - if self._bytes_read > self._max_bytes: - raise ValueError("workspace archive exceeds 8 MiB limit") - return chunk +_sys.modules[__name__] = _moved