create-windy 0.5.2 → 0.6.0

package/template/apps/document-server/src/southwind_document_server/service_recovery.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+
+from .contracts import (
+    ErrorResponse,
+    StartDocumentRunRequest,
+)
+from .grants import OutputGrantCell
+from .sink import OutputArtifactSinkClient
+from .store import DocumentRunStore, RunConflictError, RunRecord
+
+RECOVERY_ERROR = ErrorResponse(
+    code="RUNTIME_UNAVAILABLE",
+    message="Document Runtime 重启恢复期限已到",
+    request_id="document_recovery",
+    retryable=True,
+).model_dump(by_alias=True)
+type CleanupResult = tuple[str, int, bool]
+
+
+class RunRecoveryCoordinator:
+    def __init__(
+        self,
+        store: DocumentRunStore,
+        sink: OutputArtifactSinkClient,
+        grants: OutputGrantCell,
+        tasks: dict[str, asyncio.Task[None]],
+        requests: dict[str, StartDocumentRunRequest],
+        *,
+        clock: Callable[[], float],
+        timeout_seconds: float,
+        interval_seconds: float,
+    ) -> None:
+        if timeout_seconds <= 0 or interval_seconds <= 0:
+            raise ValueError("Document recovery interval 必须为正数")
+        self.store = store
+        self.sink = sink
+        self.grants = grants
+        self.tasks = tasks
+        self.requests = requests
+        self.clock = clock
+        self.timeout_seconds = timeout_seconds
+        self.interval_seconds = interval_seconds
+        self.task: asyncio.Task[None] | None = None
+        self.store.reconcile_startup(
+            self.clock(),
+            self.timeout_seconds,
+            RECOVERY_ERROR,
+        )
+
+    def new_deadline(self) -> float:
+        return self.clock() + self.timeout_seconds
+
+    def start(self) -> None:
+        if self.task is None or self.task.done():
+            self.task = asyncio.create_task(self._loop())
+
+    async def close(self) -> None:
+        task = self.task
+        self.task = None
+        if task is None:
+            return
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+    async def reconcile_once(self) -> None:
+        expired = self.store.expire_recoveries(self.clock(), RECOVERY_ERROR)
+        for record in expired:
+            task = self.tasks.get(record.run_id)
+            if task:
+                task.cancel()
+            self.requests.pop(record.run_id, None)
+        for record in self.store.pending_cleanups():
+            result = await self.cleanup_run(record.run_id)
+            if result and result[2]:
+                await self.finalize_cleanup(record.run_id, result)
+
+    async def cleanup_run(self, run_id: str) -> CleanupResult | None:
+        record = self.store.get(run_id)
+        if record is None or record.output_scope is None:
+            return None
+        return (
+            record.output_scope,
+            record.grant_generation,
+            await self._cleanup_record(record),
+        )
+
+    async def preserve_cancelled_ack(
+        self,
+        run_id: str,
+        scope: str,
+        generation: int,
+    ) -> None:
+        try:
+            self.store.begin_recovery(
+                run_id,
+                self.clock(),
+                self.timeout_seconds,
+            )
+        except Exception:
+            try:
+                await asyncio.shield(
+                    self.grants.discard_if_generation(scope, generation)
+                )
+            except BaseException:
+                pass
+
+    async def cleanup_terminal(self, run_id: str) -> None:
+        result = await self.cleanup_run(run_id)
+        if result:
+            await self.finalize_cleanup(run_id, result)
+
+    async def finalize_cleanup(
+        self,
+        run_id: str,
+        result: CleanupResult | None,
+    ) -> None:
+        if not result:
+            return
+        scope, generation, cleaned = result
+        persisted = self._mark_cleanup(run_id, generation, cleaned)
+        if cleaned and persisted:
+            await self.grants.discard_if_generation(scope, generation)
+
+    def _mark_cleanup(
+        self,
+        run_id: str,
+        generation: int,
+        cleaned: bool,
+    ) -> bool:
+        try:
+            self.store.mark_cleanup(
+                run_id,
+                generation,
+                "deleted" if cleaned else "pending-host-ttl",
+            )
+            return True
+        except (RunConflictError, RuntimeError):
+            current = self.store.get(run_id)
+            if (
+                current is not None
+                and current.terminal
+                and current.has_output_sink
+                and current.grant_generation != generation
+            ):
+                try:
+                    self.store.mark_cleanup(
+                        run_id,
+                        current.grant_generation,
+                        "pending-host-ttl",
+                    )
+                except (RunConflictError, RuntimeError):
+                    pass
+            return False
+
+    async def _loop(self) -> None:
+        while True:
+            await asyncio.sleep(self.interval_seconds)
+            try:
+                await self.reconcile_once()
+            except Exception:
+                pass
+
+    async def _cleanup_record(self, record: RunRecord) -> bool:
+        if record.output_scope is None:
+            return False
+        try:
+            aborted = await asyncio.shield(
+                self.sink.abort_scope(
+                    record.output_scope,
+                    record.grant_generation,
+                )
+            )
+        except Exception:
+            return False
+        if not aborted:
+            return False
+        return True

package/template/apps/document-server/src/southwind_document_server/sink.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+from collections.abc import Sequence
+from typing import Literal
+from urllib.parse import quote
+
+import httpx
+from pydantic import Field
+
+from .artifacts import ProducedArtifact
+from .config import DocumentRuntimeConfig
+from .contracts import (
+    CanonicalModel,
+    OutputArtifactReference,
+    OutputArtifactSinkDescriptor,
+)
+from .grants import GrantStateError, OutputGrantCell
+
+PART_MAX_BYTES = 8 * 1024 * 1024
+RESPONSE_MAX_BYTES = 64 * 1024
+
+
+class SinkClientError(Exception):
+    def __init__(self, *, ambiguous: bool = False) -> None:
+        super().__init__("Output Artifact Sink request failed")
+        self.ambiguous = ambiguous
+
+
+class _BeginReceipt(CanonicalModel):
+    contract_version: Literal[1]
+    artifact_id: str = Field(pattern=r"^outartifact_[a-f0-9]{32}$")
+    status: Literal["uploading", "completed"]
+
+
+class _Completion(CanonicalModel):
+    contract_version: Literal[1]
+    artifact_id: str = Field(pattern=r"^outartifact_[a-f0-9]{32}$")
+    output_key: str
+    status: Literal["completed"]
+    media_type: str
+    byte_size: int
+    sha256: str
+
+
+class _SealResponse(CanonicalModel):
+    contract_version: Literal[1]
+    artifacts: list[OutputArtifactReference]
+
+
+class OutputArtifactSinkClient:
+    def __init__(
+        self,
+        config: DocumentRuntimeConfig,
+        grants: OutputGrantCell,
+        transport: httpx.AsyncBaseTransport | None = None,
+    ) -> None:
+        self._base_url = config.output_artifact_url.rstrip("/")
+        self._internal_token = config.internal_token
+        self._timeout = httpx.Timeout(
+            config.request_timeout_seconds,
+            connect=config.connect_timeout_seconds,
+        )
+        self._grants = grants
+        self._transport = transport
+
+    async def accept(self, sink: OutputArtifactSinkDescriptor) -> None:
+        body = await self._send(
+            "POST",
+            (
+                f"/scopes/{quote(sink.scope, safe='')}/grants/"
+                f"{sink.grant_generation}/accept"
+            ),
+            sink.scope,
+            expected_status=204,
+            exact_generation=sink.grant_generation,
+        )
+        if body:
+            raise SinkClientError
+
+    async def publish(
+        self,
+        sink: OutputArtifactSinkDescriptor,
+        artifacts: Sequence[ProducedArtifact],
+    ) -> list[OutputArtifactReference]:
+        validate_publish_set(sink, artifacts)
+        for artifact in artifacts:
+            await self._publish_one(sink, artifact)
+        sealed = _SealResponse.model_validate_json(
+            await self._send(
+                "POST",
+                f"/scopes/{quote(sink.scope, safe='')}/seal",
+                sink.scope,
+                expected_status=200,
+            )
+        )
+        return ordered_references(artifacts, sealed.artifacts)
+
+    async def abort(self, sink: OutputArtifactSinkDescriptor) -> bool:
+        return await self.abort_scope(sink.scope, sink.grant_generation)
+
+    async def abort_scope(self, scope: str, generation: int) -> bool:
+        try:
+            await self._send(
+                "DELETE",
+                f"/scopes/{quote(scope, safe='')}",
+                scope,
+                expected_status=204,
+                exact_generation=generation,
+            )
+        except SinkClientError:
+            return False
+        return True
+
+    async def _publish_one(
+        self,
+        sink: OutputArtifactSinkDescriptor,
+        artifact: ProducedArtifact,
+    ) -> None:
+        declaration = artifact.declaration
+        manifest = {
+            "contractVersion": 1,
+            **declaration.model_dump(by_alias=True),
+            "outputKind": sink.output_kind,
+        }
+        receipt = _BeginReceipt.model_validate_json(
+            await self._send(
+                "POST",
+                f"/scopes/{quote(sink.scope, safe='')}/artifacts",
+                sink.scope,
+                expected_status=201,
+                headers={"content-type": "application/json"},
+                content=json.dumps(
+                    manifest,
+                    ensure_ascii=False,
+                    separators=(",", ":"),
+                ).encode(),
+            )
+        )
+        if receipt.status != "completed":
+            await self._upload_parts(sink, receipt.artifact_id, artifact)
+        completed = _Completion.model_validate_json(
+            await self._send(
+                "POST",
+                (
+                    f"/scopes/{quote(sink.scope, safe='')}/artifacts/"
+                    f"{quote(receipt.artifact_id, safe='')}/complete"
+                ),
+                sink.scope,
+                expected_status=200,
+            )
+        )
+        if (
+            completed.artifact_id != receipt.artifact_id
+            or completed.output_key != declaration.output_key
+            or completed.media_type != declaration.media_type
+            or completed.byte_size != declaration.byte_size
+            or completed.sha256 != declaration.sha256
+        ):
+            raise SinkClientError
+
+    async def _upload_parts(
+        self,
+        sink: OutputArtifactSinkDescriptor,
+        artifact_id: str,
+        artifact: ProducedArtifact,
+    ) -> None:
+        total = 0
+        digest = hashlib.sha256()
+        part_number = 0
+        with artifact.path.open("rb") as source:
+            while True:
+                chunk = await asyncio.to_thread(source.read, PART_MAX_BYTES)
+                if not chunk:
+                    break
+                part_digest = hashlib.sha256(chunk).hexdigest()
+                await self._send(
+                    "PUT",
+                    (
+                        f"/scopes/{quote(sink.scope, safe='')}/artifacts/"
+                        f"{quote(artifact_id, safe='')}/parts/{part_number}"
+                    ),
+                    sink.scope,
+                    expected_status=200,
+                    headers={
+                        "content-type": "application/octet-stream",
+                        "content-length": str(len(chunk)),
+                        "x-upload-offset": str(total),
+                        "x-content-sha256": part_digest,
+                    },
+                    content=chunk,
+                )
+                total += len(chunk)
+                digest.update(chunk)
+                part_number += 1
+        if (
+            total != artifact.declaration.byte_size
+            or digest.hexdigest() != artifact.declaration.sha256
+        ):
+            raise SinkClientError
+
+    async def _send(
+        self,
+        method: str,
+        path: str,
+        credential_scope: str,
+        *,
+        expected_status: int,
+        headers: dict[str, str] | None = None,
+        content: bytes | None = None,
+        exact_generation: int | None = None,
+    ) -> bytes:
+        for attempt in range(2):
+            try:
+                token = (
+                    await self._grants.token_for(
+                        credential_scope,
+                        exact_generation,
+                    )
+                    if exact_generation is not None
+                    else await self._grants.token(credential_scope)
+                )
+                return await self._send_once(
+                    method,
+                    path,
+                    expected_status,
+                    token,
+                    headers,
+                    content,
+                )
+            except SinkClientError:
+                raise
+            except httpx.TransportError:
+                if attempt == 0:
+                    continue
+                raise SinkClientError(ambiguous=True) from None
+            except (GrantStateError, httpx.HTTPError, ValueError):
+                raise SinkClientError from None
+        raise SinkClientError
+
+    async def _send_once(
+        self,
+        method: str,
+        path: str,
+        expected_status: int,
+        token: str,
+        headers: dict[str, str] | None,
+        content: bytes | None,
+    ) -> bytes:
+        request_headers = {
+            "authorization": f"Bearer {self._internal_token}",
+            "x-ai-output-artifact-grant": token,
+            **(headers or {}),
+        }
+        async with httpx.AsyncClient(
+            timeout=self._timeout,
+            transport=self._transport,
+            follow_redirects=False,
+        ) as client:
+            async with client.stream(
+                method,
+                f"{self._base_url}{path}",
+                headers=request_headers,
+                content=content,
+            ) as response:
+                if response.status_code != expected_status:
+                    raise SinkClientError
+                body = bytearray()
+                async for chunk in response.aiter_bytes():
+                    body.extend(chunk)
+                    if len(body) > RESPONSE_MAX_BYTES:
+                        raise SinkClientError
+                return bytes(body)
+
+
+def validate_publish_set(
+    sink: OutputArtifactSinkDescriptor,
+    artifacts: Sequence[ProducedArtifact],
+) -> None:
+    declarations = [artifact.declaration for artifact in artifacts]
+    if (
+        len(declarations) > sink.limits.max_artifacts
+        or len({item.output_key for item in declarations}) != len(declarations)
+        or sum(item.byte_size for item in declarations) > sink.limits.max_total_bytes
+        or any(
+            item.byte_size > sink.limits.max_artifact_bytes
+            or item.media_type not in sink.allowed_media_types
+            for item in declarations
+        )
+    ):
+        raise SinkClientError
+
+
+def ordered_references(
+    artifacts: Sequence[ProducedArtifact],
+    references: Sequence[OutputArtifactReference],
+) -> list[OutputArtifactReference]:
+    by_key = {reference.output_key: reference for reference in references}
+    if len(by_key) != len(references) or set(by_key) != {
+        artifact.declaration.output_key for artifact in artifacts
+    }:
+        raise SinkClientError
+    ordered: list[OutputArtifactReference] = []
+    for artifact in artifacts:
+        reference = by_key[artifact.declaration.output_key]
+        declaration = artifact.declaration
+        if (
+            reference.media_type != declaration.media_type
+            or reference.byte_size != declaration.byte_size
+            or reference.sha256 != declaration.sha256
+        ):
+            raise SinkClientError
+        ordered.append(reference)
+    return ordered