wehub-resource-sync
115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from collections.abc import Iterable
|
|
from typing import Any, ClassVar, Literal, get_args, get_origin
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, field_validator, model_serializer
|
|
|
|
from ..manifest import Manifest
|
|
from ..snapshot import SnapshotBase
|
|
|
|
SessionStateClass = type["SandboxSessionState"]
|
|
|
|
|
|
class SandboxSessionState(BaseModel):
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
type: str
|
|
session_id: uuid.UUID = Field(default_factory=uuid.uuid4)
|
|
snapshot: SerializeAsAny[SnapshotBase]
|
|
manifest: Manifest
|
|
exposed_ports: tuple[int, ...] = Field(default_factory=tuple)
|
|
snapshot_fingerprint: str | None = None
|
|
snapshot_fingerprint_version: str | None = None
|
|
workspace_root_ready: bool = False
|
|
|
|
_subclass_registry: ClassVar[dict[str, SessionStateClass]] = {}
|
|
|
|
@classmethod
|
|
def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
|
|
"""Auto-register every subclass by its ``type`` field default."""
|
|
super().__pydantic_init_subclass__(**kwargs)
|
|
|
|
type_field = cls.model_fields.get("type")
|
|
if type_field is None:
|
|
return
|
|
|
|
annotation = type_field.annotation
|
|
if get_origin(annotation) is not Literal:
|
|
return
|
|
|
|
args = get_args(annotation)
|
|
if not args:
|
|
return
|
|
|
|
type_default = type_field.default
|
|
if not isinstance(type_default, str) or type_default == "":
|
|
return
|
|
|
|
SandboxSessionState._subclass_registry[type_default] = cls
|
|
|
|
@classmethod
|
|
def parse(cls, payload: object) -> SandboxSessionState:
|
|
"""Deserialize *payload* into the correct registered subclass.
|
|
|
|
Accepts a ``SandboxSessionState`` instance (returned as-is if already a
|
|
subclass, or upgraded via ``model_dump`` -> registry lookup if it is a
|
|
bare base instance) or a plain ``dict``.
|
|
"""
|
|
if isinstance(payload, SandboxSessionState):
|
|
if type(payload) is not SandboxSessionState:
|
|
return payload
|
|
payload = payload.model_dump()
|
|
|
|
if isinstance(payload, dict):
|
|
state_type = payload.get("type")
|
|
if not isinstance(state_type, str):
|
|
raise ValueError("sandbox session state payload must include a string `type`")
|
|
|
|
subclass = SandboxSessionState._subclass_registry.get(state_type)
|
|
if subclass is None:
|
|
raise ValueError(f"unknown sandbox session state type `{state_type}`")
|
|
|
|
return subclass.model_validate(payload)
|
|
|
|
raise TypeError("session state payload must be a SandboxSessionState or dict")
|
|
|
|
@model_serializer(mode="wrap")
|
|
def _serialize_always_include_defaults(self, handler: Any) -> dict[str, Any]:
|
|
data: dict[str, Any] = handler(self)
|
|
if self.type:
|
|
data["type"] = self.type
|
|
if self.session_id:
|
|
data["session_id"] = self.session_id
|
|
return data
|
|
|
|
@field_validator("snapshot", mode="before")
|
|
@classmethod
|
|
def _coerce_snapshot(cls, value: object) -> SnapshotBase:
|
|
return SnapshotBase.parse(value)
|
|
|
|
@field_validator("exposed_ports", mode="before")
|
|
@classmethod
|
|
def _coerce_exposed_ports(cls, value: object) -> tuple[int, ...]:
|
|
if value is None:
|
|
return ()
|
|
if isinstance(value, int):
|
|
ports: Iterable[object] = (value,)
|
|
elif isinstance(value, Iterable) and not isinstance(value, str | bytes | bytearray):
|
|
ports = value
|
|
else:
|
|
raise TypeError("exposed_ports must be an iterable of TCP port integers")
|
|
|
|
normalized: list[int] = []
|
|
seen: set[int] = set()
|
|
for port in ports:
|
|
if not isinstance(port, int):
|
|
raise TypeError("exposed_ports must contain integers")
|
|
if port < 1 or port > 65535:
|
|
raise ValueError("exposed_ports entries must be between 1 and 65535")
|
|
if port in seen:
|
|
continue
|
|
seen.add(port)
|
|
normalized.append(port)
|
|
return tuple(normalized)
|