create-windy 0.3.2 → 0.3.5

package/template/apps/agent-server/src/southwind_agent_server/postgres_store.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+import uuid
+
+from .contracts import AgentError, JsonValue
+from .postgres_store_backend import PostgresStoreBackend
+from .sqlite_store import RunStoreCapacityError
+from .store import (
+    CanonicalEvent,
+    IdempotencyConflictError,
+    RunRecord,
+    status_for,
+)
+from .store_cipher import RunStoreCipher
+
+
+class PostgresRunStore:
+    """Multi-replica Run store with encrypted payloads and renewable leases."""
+
+    def __init__(
+        self,
+        database_url: str,
+        cipher: RunStoreCipher,
+        *,
+        max_runs: int = 100_000,
+        max_events_per_run: int = 2_000,
+        max_output_bytes: int = 16 * 1024 * 1024,
+        terminal_ttl_seconds: int = 7 * 24 * 60 * 60,
+        lease_seconds: int = 30,
+    ) -> None:
+        self._cipher = cipher
+        self._max_runs = max_runs
+        self._max_events = max_events_per_run
+        self._max_output_bytes = max_output_bytes
+        self._terminal_ttl = terminal_ttl_seconds
+        self._lease_seconds = lease_seconds
+        self._owner = f"agent_{uuid.uuid4().hex}"
+        self._backend = PostgresStoreBackend(
+            database_url,
+            cipher,
+            self._owner,
+            terminal_ttl_seconds,
+        )
+        self._pool = self._backend.pool
+
+    async def create_or_get(
+        self,
+        record: RunRecord,
+        invocation: JsonValue = None,
+    ) -> tuple[RunRecord, bool]:
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            async with connection.transaction():
+                await connection.execute("select pg_advisory_xact_lock(74123001)")
+                await self._backend.prune(connection)
+                current = await self._backend.by_idempotency(
+                    connection,
+                    record.idempotency_key,
+                )
+                if current:
+                    if current.fingerprint != record.fingerprint:
+                        raise IdempotencyConflictError
+                    return current, False
+                count = await connection.execute(
+                    "select count(*) as value from agent_runtime_runs"
+                )
+                row = await count.fetchone()
+                if not row:
+                    raise RuntimeError("Agent Run Store 容量查询失败")
+                if int(row["value"]) >= self._max_runs:
+                    raise RunStoreCapacityError("Agent Run Store 容量已满")
+                release_id = await self._backend.ensure_release(
+                    connection,
+                    record,
+                )
+                now = time.time()
+                inserted = await connection.execute(
+                    """
+                    insert into agent_runtime_runs (
+                      run_id, operation_key, idempotency_key, fingerprint,
+                      status, release_id, release_json, trace_id, invocation_cipher,
+                      lease_owner, lease_until,
+                      created_at, updated_at, expires_at
+                    ) values (
+                      %s, %s, %s, %s, 'accepted', %s, %s, %s, %s, %s,
+                      extract(epoch from clock_timestamp()) + %s, %s, %s, %s
+                    )
+                    on conflict (idempotency_key) do nothing
+                    returning run_id
+                    """,
+                    (
+                        record.run_id,
+                        record.operation_key,
+                        record.idempotency_key,
+                        record.fingerprint,
+                        release_id,
+                        json.dumps(record.release, ensure_ascii=False),
+                        record.trace_id,
+                        self._cipher.seal(invocation)
+                        if invocation is not None
+                        else None,
+                        self._owner,
+                        self._lease_seconds,
+                        now,
+                        now,
+                        now + self._terminal_ttl,
+                    ),
+                )
+                if await inserted.fetchone() is None:
+                    current = await self._backend.by_idempotency(
+                        connection,
+                        record.idempotency_key,
+                    )
+                    if current is None or current.fingerprint != record.fingerprint:
+                        raise IdempotencyConflictError
+                    return current, False
+                return record, True
+
+    async def get(self, run_id: str) -> RunRecord | None:
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            return await self._backend.read(connection, run_id)
+
+    async def append(
+        self,
+        run_id: str,
+        event_type: str,
+        fields: dict[str, object] | None = None,
+    ) -> CanonicalEvent | None:
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            async with connection.transaction():
+                record = await self._backend.locked(connection, run_id)
+                if record.terminal:
+                    return None
+                sequence = len(record.events) + 1
+                if sequence >= self._max_events and not event_type.startswith("run."):
+                    event_type = "run.failed"
+                    fields = failure_fields(
+                        sequence,
+                        "Agent 事件超过安全上限",
+                    )
+                event: CanonicalEvent = {
+                    "type": event_type,
+                    "runId": run_id,
+                    "sequence": sequence,
+                    **(fields or {}),
+                }
+                await self._backend.insert_event(
+                    connection,
+                    run_id,
+                    event,
+                    status_for(event_type, record.status),
+                )
+                return event
+
+    async def complete(self, run_id: str, output: JsonValue) -> None:
+        encoded = json.dumps(output, ensure_ascii=False, separators=(",", ":"))
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            async with connection.transaction():
+                record = await self._backend.locked(connection, run_id)
+                if record.terminal:
+                    return
+                sequence = len(record.events) + 1
+                if len(encoded.encode()) > self._max_output_bytes:
+                    event: CanonicalEvent = {
+                        "type": "run.failed",
+                        "runId": run_id,
+                        "sequence": sequence,
+                        **failure_fields(sequence, "Agent 输出超过安全上限"),
+                    }
+                    await self._backend.insert_event(
+                        connection,
+                        run_id,
+                        event,
+                        "failed",
+                    )
+                    return
+                event = {
+                    "type": "run.completed",
+                    "runId": run_id,
+                    "sequence": sequence,
+                }
+                await self._backend.insert_event(
+                    connection,
+                    run_id,
+                    event,
+                    "completed",
+                    self._cipher.seal(output),
+                )
+
+    async def events_after(
+        self,
+        run_id: str,
+        after: int,
+        wait_seconds: float,
+    ) -> tuple[list[CanonicalEvent], bool]:
+        deadline = time.monotonic() + wait_seconds
+        while True:
+            record = await self.get(run_id)
+            if record is None:
+                raise KeyError(run_id)
+            events = [event for event in record.events if sequence_of(event) > after]
+            if events or record.terminal or time.monotonic() >= deadline:
+                return events, record.terminal
+            await asyncio.sleep(min(0.5, max(0, deadline - time.monotonic())))
+
+    async def save_invocation(self, run_id: str, value: JsonValue) -> None:
+        await self._backend.update_run(
+            run_id,
+            "invocation_cipher = %s",
+            (self._cipher.seal(value),),
+        )
+
+    async def recoverable(self) -> list[tuple[RunRecord, JsonValue]]:
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            async with connection.transaction():
+                cursor = await connection.execute(
+                    """
+                    select run_id from agent_runtime_runs
+                    where status not in ('completed', 'failed', 'cancelled')
+                      and invocation_cipher is not null
+                      and (
+                        lease_until < extract(epoch from clock_timestamp())
+                        or lease_owner = %s
+                      )
+                    order by created_at for update skip locked limit 20
+                    """,
+                    (self._owner,),
+                )
+                rows = await cursor.fetchall()
+                result: list[tuple[RunRecord, JsonValue]] = []
+                for row in rows:
+                    run_id = str(row["run_id"])
+                    await connection.execute(
+                        """
+                        update agent_runtime_runs
+                        set lease_owner = %s,
+                          lease_until = extract(epoch from clock_timestamp()) + %s
+                        where run_id = %s
+                        """,
+                        (self._owner, self._lease_seconds, run_id),
+                    )
+                    record = await self._backend.read(connection, run_id)
+                    cipher = await connection.execute(
+                        """
+                        select invocation_cipher from agent_runtime_runs
+                        where run_id = %s
+                        """,
+                        (run_id,),
+                    )
+                    value = await cipher.fetchone()
+                    if record and value:
+                        result.append(
+                            (
+                                record,
+                                self._cipher.open(str(value["invocation_cipher"])),
+                            )
+                        )
+                return result
+
+    async def clear_invocation(self, run_id: str) -> None:
+        await self._backend.update_run(
+            run_id,
+            "invocation_cipher = null",
+            (),
+        )
+
+    async def renew_lease(self, run_id: str) -> None:
+        await self._backend.update_run(
+            run_id,
+            "lease_until = extract(epoch from clock_timestamp()) + %s",
+            (self._lease_seconds,),
+            owner_required=True,
+            active_required=True,
+        )
+
+    async def release_lease(self, run_id: str) -> None:
+        await self._backend.update_run(
+            run_id,
+            "lease_owner = null, lease_until = 0",
+            (),
+            owner_required=True,
+        )
+
+    async def record_evaluation(
+        self,
+        release_id: str,
+        evaluation_set_key: str,
+        evaluation_set_version: str,
+        metrics: dict[str, float],
+    ) -> str:
+        return await self._backend.record_evaluation(
+            release_id,
+            evaluation_set_key,
+            evaluation_set_version,
+            metrics,
+        )
+
+    async def close(self) -> None:
+        await self._backend.close()
+
+    async def probe(self) -> None:
+        await self._backend.ready()
+        async with self._pool.connection() as connection:
+            await connection.execute("select 1")
+
+
+def sequence_of(event: CanonicalEvent) -> int:
+    value = event.get("sequence")
+    if not isinstance(value, int):
+        raise RuntimeError("Agent Event sequence 无效")
+    return value
+
+
+def failure_fields(sequence: int, message: str) -> dict[str, object]:
+    return {
+        "error": AgentError(
+            code="BUDGET_EXCEEDED",
+            message=message,
+            request_id=f"request_store_{sequence}",
+            retryable=False,
+        ).model_dump(by_alias=True)
+    }

package/template/apps/agent-server/src/southwind_agent_server/postgres_store_backend.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+import uuid
+from typing import Any, cast
+
+from psycopg import AsyncConnection
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
+
+from .contracts import JsonValue, RunStatus
+from .store import (
+    CanonicalEvent,
+    IdempotencyConflictError,
+    ReleaseNotFoundError,
+    RunRecord,
+)
+from .store_cipher import RunStoreCipher
+
+type DatabaseConnection = AsyncConnection[dict[str, Any]]
+
+
+class PostgresStoreBackend:
+    def __init__(
+        self,
+        database_url: str,
+        cipher: RunStoreCipher,
+        owner: str,
+        terminal_ttl_seconds: int,
+    ) -> None:
+        self.pool = cast(
+            AsyncConnectionPool[DatabaseConnection],
+            AsyncConnectionPool(
+                database_url,
+                min_size=1,
+                max_size=10,
+                open=False,
+                kwargs={"row_factory": dict_row},
+            ),
+        )
+        self.cipher = cipher
+        self.owner = owner
+        self.terminal_ttl = terminal_ttl_seconds
+        self._open_lock = asyncio.Lock()
+        self._opened = False
+
+    async def ready(self) -> None:
+        if self._opened:
+            return
+        async with self._open_lock:
+            if not self._opened:
+                await self.pool.open()
+                self._opened = True
+
+    async def close(self) -> None:
+        if self._opened:
+            await self.pool.close()
+
+    async def read(
+        self,
+        connection: DatabaseConnection,
+        run_id: str,
+    ) -> RunRecord | None:
+        cursor = await connection.execute(
+            "select * from agent_runtime_runs where run_id = %s",
+            (run_id,),
+        )
+        row = await cursor.fetchone()
+        if not row:
+            return None
+        events_cursor = await connection.execute(
+            """
+            select event_cipher from agent_runtime_events
+            where run_id = %s order by sequence
+            """,
+            (run_id,),
+        )
+        events = [
+            cast(CanonicalEvent, self.cipher.open(str(item["event_cipher"])))
+            for item in await events_cursor.fetchall()
+        ]
+        output = (
+            self.cipher.open(str(row["output_cipher"]))
+            if row["output_cipher"]
+            else None
+        )
+        return RunRecord(
+            run_id=str(row["run_id"]),
+            operation_key=str(row["operation_key"]),
+            input=None,
+            idempotency_key=str(row["idempotency_key"]),
+            fingerprint=str(row["fingerprint"]),
+            release=cast(JsonValue, row["release_json"]),
+            trace_id=str(row["trace_id"]),
+            status=cast(RunStatus, str(row["status"])),
+            events=events,
+            output=output,
+        )
+
+    async def locked(
+        self,
+        connection: DatabaseConnection,
+        run_id: str,
+    ) -> RunRecord:
+        await connection.execute(
+            "select run_id from agent_runtime_runs where run_id = %s for update",
+            (run_id,),
+        )
+        record = await self.read(connection, run_id)
+        if record is None:
+            raise KeyError(run_id)
+        return record
+
+    async def insert_event(
+        self,
+        connection: DatabaseConnection,
+        run_id: str,
+        event: CanonicalEvent,
+        status: RunStatus,
+        output_cipher: str | None = None,
+    ) -> None:
+        await connection.execute(
+            """
+            insert into agent_runtime_events (run_id, sequence, event_cipher)
+            values (%s, %s, %s)
+            """,
+            (
+                run_id,
+                event["sequence"],
+                self.cipher.seal(cast(JsonValue, event)),
+            ),
+        )
+        now = time.time()
+        await connection.execute(
+            """
+            update agent_runtime_runs
+            set status = %s, output_cipher = coalesce(%s, output_cipher),
+              updated_at = %s, expires_at = %s
+            where run_id = %s
+            """,
+            (
+                status,
+                output_cipher,
+                now,
+                now + self.terminal_ttl,
+                run_id,
+            ),
+        )
+
+    async def update_run(
+        self,
+        run_id: str,
+        assignment: str,
+        values: tuple[object, ...],
+        *,
+        owner_required: bool = False,
+        active_required: bool = False,
+    ) -> None:
+        await self.ready()
+        owner_clause = " and lease_owner = %s" if owner_required else ""
+        active_clause = (
+            " and status not in ('completed', 'failed', 'cancelled')"
+            if active_required
+            else ""
+        )
+        parameters = (
+            *values,
+            run_id,
+            *((self.owner,) if owner_required else ()),
+        )
+        async with self.pool.connection() as connection:
+            cursor = await connection.execute(
+                f"update agent_runtime_runs set {assignment} "
+                f"where run_id = %s{owner_clause}{active_clause}",
+                parameters,
+            )
+            if cursor.rowcount != 1:
+                raise KeyError(run_id)
+
+    async def by_idempotency(
+        self,
+        connection: DatabaseConnection,
+        idempotency_key: str,
+    ) -> RunRecord | None:
+        cursor = await connection.execute(
+            """
+            select run_id from agent_runtime_runs
+            where idempotency_key = %s
+            """,
+            (idempotency_key,),
+        )
+        row = await cursor.fetchone()
+        return await self.read(connection, str(row["run_id"])) if row else None
+
+    async def ensure_release(
+        self,
+        connection: DatabaseConnection,
+        record: RunRecord,
+    ) -> str:
+        if not isinstance(record.release, dict):
+            raise IdempotencyConflictError
+        release_id = record.release.get("releaseId")
+        if not isinstance(release_id, str):
+            raise IdempotencyConflictError
+        encoded = json.dumps(
+            record.release,
+            sort_keys=True,
+            separators=(",", ":"),
+        )
+        await connection.execute(
+            """
+            insert into agent_operation_releases (
+              release_id, operation_key, snapshot_json, created_at
+            ) values (%s, %s, %s::jsonb, %s)
+            on conflict (release_id) do nothing
+            """,
+            (release_id, record.operation_key, encoded, time.time()),
+        )
+        cursor = await connection.execute(
+            """
+            select operation_key, snapshot_json from agent_operation_releases
+            where release_id = %s
+            """,
+            (release_id,),
+        )
+        existing = await cursor.fetchone()
+        if not existing:
+            raise IdempotencyConflictError
+        actual = json.dumps(
+            existing["snapshot_json"],
+            sort_keys=True,
+            separators=(",", ":"),
+        )
+        if str(existing["operation_key"]) != record.operation_key or actual != encoded:
+            raise IdempotencyConflictError
+        return release_id
+
+    async def prune(self, connection: DatabaseConnection) -> None:
+        await connection.execute(
+            """
+            delete from agent_runtime_runs
+            where status in ('completed', 'failed', 'cancelled')
+              and expires_at < %s
+            """,
+            (time.time(),),
+        )
+
+    async def record_evaluation(
+        self,
+        release_id: str,
+        evaluation_set_key: str,
+        evaluation_set_version: str,
+        metrics: dict[str, float],
+    ) -> str:
+        await self.ready()
+        encoded = json.dumps(metrics, sort_keys=True, separators=(",", ":"))
+        async with self.pool.connection() as connection:
+            async with connection.transaction():
+                release = await connection.execute(
+                    """
+                    select release_id from agent_operation_releases
+                    where release_id = %s
+                    """,
+                    (release_id,),
+                )
+                if await release.fetchone() is None:
+                    raise ReleaseNotFoundError
+                cursor = await connection.execute(
+                    """
+                    select evaluation_id, metrics_json
+                    from agent_evaluation_results
+                    where release_id = %s and evaluation_set_key = %s
+                      and evaluation_set_version = %s
+                    """,
+                    (release_id, evaluation_set_key, evaluation_set_version),
+                )
+                existing = await cursor.fetchone()
+                if existing:
+                    actual = json.dumps(
+                        existing["metrics_json"],
+                        sort_keys=True,
+                        separators=(",", ":"),
+                    )
+                    if actual != encoded:
+                        raise IdempotencyConflictError
+                    return str(existing["evaluation_id"])
+                evaluation_id = f"eval_{uuid.uuid4().hex}"
+                inserted = await connection.execute(
+                    """
+                    insert into agent_evaluation_results (
+                      evaluation_id, release_id, evaluation_set_key,
+                      evaluation_set_version, metrics_json, created_at
+                    ) values (%s, %s, %s, %s, %s::jsonb, %s)
+                    on conflict (
+                      release_id, evaluation_set_key, evaluation_set_version
+                    ) do nothing
+                    returning evaluation_id
+                    """,
+                    (
+                        evaluation_id,
+                        release_id,
+                        evaluation_set_key,
+                        evaluation_set_version,
+                        encoded,
+                        time.time(),
+                    ),
+                )
+                row = await inserted.fetchone()
+                if row:
+                    return str(row["evaluation_id"])
+                repeated = await connection.execute(
+                    """
+                    select evaluation_id, metrics_json
+                    from agent_evaluation_results
+                    where release_id = %s and evaluation_set_key = %s
+                      and evaluation_set_version = %s
+                    """,
+                    (release_id, evaluation_set_key, evaluation_set_version),
+                )
+                raced = await repeated.fetchone()
+                if not raced:
+                    raise IdempotencyConflictError
+                actual = json.dumps(
+                    raced["metrics_json"],
+                    sort_keys=True,
+                    separators=(",", ":"),
+                )
+                if actual != encoded:
+                    raise IdempotencyConflictError
+                return str(raced["evaluation_id"])