agentengine-sdk-python 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

ksadk/resource_runtime/snapshots.py

@@ -0,0 +1,180 @@
+"""Credential-free resource snapshots consumed by Build activation and workers."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Literal
+from urllib.parse import urlsplit, urlunsplit
+
+from pydantic import Field, field_validator, model_validator
+
+from ksadk.plugins.bridges.dsh import DshProfileBuildSnapshot
+from ksadk.plugins.contracts import PluginContractModel
+from ksadk.resource_runtime.contracts import (
+    Digest,
+    Identifier,
+    ResourceConfig,
+    validate_resource_bindings,
+)
+
+
+class ConnectionTarget(PluginContractModel):
+    connection_ref: Identifier
+    tenant_ref: Identifier
+    principal_ref: Identifier
+    endpoint: str
+    auth_mode: Literal["signed", "sts", "token"]
+
+    @field_validator("endpoint")
+    @classmethod
+    def canonical_endpoint(cls, value: str) -> str:
+        if any(char.isspace() for char in value) or "\\" in value:
+            raise ValueError("Invalid resource endpoint")
+        url = urlsplit(value)
+        if (
+            url.scheme not in {"http", "https"}
+            or not url.hostname
+            or url.username is not None
+            or url.password is not None
+            or url.query
+            or url.fragment
+            or "%" in url.hostname
+        ):
+            raise ValueError("Resource endpoint must be an explicit HTTP service URL")
+        port = url.port  # validates port range and syntax
+        host = url.hostname.encode("idna").decode("ascii").lower()
+        if ":" in host:
+            host = f"[{host}]"
+        if port is not None and (url.scheme, port) not in {("http", 80), ("https", 443)}:
+            host += f":{port}"
+        return urlunsplit((url.scheme, host, url.path.rstrip("/"), "", ""))
+
+
+class SkillExecutionTarget(PluginContractModel):
+    backend: Literal["e2b"] = "e2b"
+    connection: ConnectionTarget
+    domain: str = Field(strict=True, pattern=r"^[a-zA-Z0-9]+(?:[.-][a-zA-Z0-9-]+)*$")
+    template_id: Identifier
+    timeout: int = Field(default=900, strict=True, ge=1, le=900)
+    allow_internet_access: bool = Field(default=False, strict=True)
+
+    @model_validator(mode="after")
+    def explicit_target(self):
+        url = urlsplit(self.connection.endpoint)
+        if self.connection.auth_mode != "token" or url.scheme != "https" or url.path:
+            raise ValueError("Skill execution requires an explicit HTTPS token connection")
+        return self
+
+
+class MemoryRecallPolicy(PluginContractModel):
+    """Frozen projection of Agent memory recall limits, never a second draft config."""
+
+    top_k: int = Field(default=8, strict=True, ge=1, le=64)
+    max_tokens: int = Field(default=1600, strict=True, ge=0, le=100000)
+    min_score: float = Field(default=0, strict=True, ge=0, le=1)
+
+    @model_validator(mode="after")
+    def supported_score(self):
+        if self.min_score != 0:
+            raise ValueError(
+                "MEMORY_SEARCH_POLICY_UNSUPPORTED: AICP score threshold is unavailable"
+            )
+        return self
+
+
+class FrozenResourceBinding(PluginContractModel):
+    config: ResourceConfig
+    connection: ConnectionTarget
+    # The connection target deliberately contains no credential references.
+    # Freeze the repository revision separately so rotating those references
+    # cannot make an old Build silently consume a different credential set.
+    connection_revision: int = Field(default=1, strict=True, ge=1)
+    skill_execution: SkillExecutionTarget | None = None
+    memory_recall: MemoryRecallPolicy | None = None
+
+    @model_validator(mode="after")
+    def matching_connection(self) -> FrozenResourceBinding:
+        if self.config.binding.connection_ref != self.connection.connection_ref:
+            raise ValueError("Resource binding and connection reference do not match")
+        if (
+            self.memory_recall is not None
+            and self.config.binding.resource.kind != "memory-instance"
+        ):
+            raise ValueError("Memory recall policy requires a memory binding")
+        if self.skill_execution is not None and (
+            self.config.binding.resource.kind != "skill-space"
+            or self.config.execution_mode != "isolated"
+            or self.skill_execution.connection.tenant_ref != self.connection.tenant_ref
+        ):
+            raise ValueError("Skill execution target requires an isolated Skill in the same tenant")
+        return self
+
+
+class ResourceSnapshot(PluginContractModel):
+    schema_version: Literal[1] = 1
+    plugin_lock_digest: Digest
+    bindings: tuple[FrozenResourceBinding, ...]
+    # Older standalone Worker fixtures may omit this. Studio Build/activation
+    # requires the actual host snapshot, bound into the resource digest below.
+    dsh_profile: DshProfileBuildSnapshot | None = None
+
+    @field_validator("schema_version", mode="before")
+    @classmethod
+    def strict_version(cls, value: object) -> object:
+        if type(value) is not int:
+            raise ValueError("Resource snapshot version must be an integer")
+        return value
+
+    @model_validator(mode="after")
+    def unique_bindings(self) -> ResourceSnapshot:
+        validate_resource_bindings(tuple(binding.config for binding in self.bindings))
+        return self
+
+    @property
+    def digest(self) -> str:
+        payload = {
+            "schemaVersion": self.schema_version,
+            "pluginLockDigest": self.plugin_lock_digest,
+            **({"dshProfileDigest": self.dsh_profile.digest} if self.dsh_profile else {}),
+            "bindings": [
+                {
+                    "config": json.loads(binding.config.canonical_bytes()),
+                    "connection": binding.connection.model_dump(by_alias=True, mode="json"),
+                    "connectionRevision": binding.connection_revision,
+                    **(
+                        {
+                            "memoryRecall": binding.memory_recall.model_dump(
+                                by_alias=True, mode="json"
+                            )
+                        }
+                        if binding.memory_recall is not None
+                        else {}
+                    ),
+                    **(
+                        {
+                            "skillExecution": binding.skill_execution.model_dump(
+                                by_alias=True, mode="json"
+                            )
+                        }
+                        if binding.skill_execution is not None
+                        else {}
+                    ),
+                }
+                for binding in sorted(self.bindings, key=lambda item: item.config.binding.id)
+            ],
+        }
+        raw = json.dumps(
+            payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False
+        ).encode("utf-8")
+        return "sha256:" + hashlib.sha256(raw).hexdigest()
+
+    def binding(self, binding_id: str) -> FrozenResourceBinding:
+        for binding in self.bindings:
+            if binding.config.binding.id == binding_id:
+                return binding
+        raise ValueError("RESOURCE_BINDING_NOT_IN_BUILD")
+
+    def verify_connection(self, binding_id: str, current: ConnectionTarget) -> None:
+        if self.binding(binding_id).connection != current:
+            raise ValueError("RESOURCE_CONNECTION_CHANGED")

ksadk/resource_runtime/socket_server.py

@@ -0,0 +1,103 @@
+"""Private Unix socket ingress for the existing host's resource broker."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import struct
+import tempfile
+from pathlib import Path
+
+from ksadk.resource_runtime.broker import ResourceBroker
+from ksadk.resource_runtime.ipc import MAX_FRAME_BYTES, FrameError, decode_body, encode_frame
+
+
+class ResourceSocketServer:
+    def __init__(self, broker: ResourceBroker, *, max_connections: int = 64):
+        if type(max_connections) is not int or not 1 <= max_connections <= 1024:
+            raise ValueError("Socket connection limit must be 1..1024")
+        self._broker = broker
+        self._max_connections = max_connections
+        self._directory: tempfile.TemporaryDirectory | None = None
+        self._server: asyncio.AbstractServer | None = None
+        self._clients: set[asyncio.Task] = set()
+        self._closing = False
+        self.path: Path | None = None
+
+    async def start(self) -> Path:
+        if os.name != "posix":
+            raise RuntimeError("RESOURCE_IPC_UNSUPPORTED")
+        if self._directory is not None or self._closing:
+            raise RuntimeError("Resource socket server cannot be restarted")
+        directory = tempfile.TemporaryDirectory(prefix="ksr-")
+        self._directory = directory
+        self.path = Path(directory.name) / "broker.sock"
+        try:
+            os.chmod(directory.name, 0o700)
+            # Private parent directory prevents access during the bind/chmod gap.
+            self._server = await asyncio.start_unix_server(self._accept, path=str(self.path))
+            os.chmod(self.path, 0o600)
+            return self.path
+        except BaseException:
+            await self.aclose()
+            raise
+
+    async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        task = asyncio.current_task()
+        assert task is not None
+        if self._closing or len(self._clients) >= self._max_connections:
+            writer.close()
+            await writer.wait_closed()
+            return
+        self._clients.add(task)
+        try:
+            while not self._closing:
+                payload = await asyncio.wait_for(self._read(reader), 15)
+                if payload is None:
+                    return
+                response = await self._broker.dispatch(payload)
+                writer.write(encode_frame(response))
+                await asyncio.wait_for(writer.drain(), 5)
+        except (
+            FrameError,
+            asyncio.TimeoutError,
+            asyncio.IncompleteReadError,
+            ConnectionError,
+            OSError,
+        ):
+            # Closing rejects malformed/slow ingress; never echo raw input.
+            pass
+        finally:
+            self._clients.discard(task)
+            writer.close()
+            try:
+                await writer.wait_closed()
+            except (ConnectionError, OSError):
+                pass
+
+    @staticmethod
+    async def _read(reader: asyncio.StreamReader) -> dict | None:
+        try:
+            prefix = await reader.readexactly(4)
+        except asyncio.IncompleteReadError as exc:
+            if not exc.partial:
+                return None
+            raise
+        (size,) = struct.unpack("!I", prefix)
+        if not 0 < size <= MAX_FRAME_BYTES:
+            raise FrameError("Invalid IPC frame length")
+        return decode_body(await reader.readexactly(size))
+
+    async def aclose(self) -> None:
+        self._closing = True
+        if self._server is not None:
+            self._server.close()
+            await self._server.wait_closed()
+        self._broker.revoke_generation()
+        clients = tuple(self._clients)
+        for task in clients:
+            task.cancel()
+        if clients:
+            await asyncio.gather(*clients, return_exceptions=True)
+        if self._directory is not None:
+            self._directory.cleanup()

ksadk/resource_runtime/supervisor.py

@@ -0,0 +1,380 @@
+"""Generation-owned resource lifecycle for trusted Studio/runtime adapters."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+from ksadk.resource_runtime.broker import ResourceBroker, ResourceWriteAuthorizer
+from ksadk.resource_runtime.discovery_receipts import (
+    DiscoverySelectionReceipts,
+    discovery_scope_key,
+)
+from ksadk.resource_runtime.ipc import ResourceRequest
+from ksadk.resource_runtime.leases import (
+    LeaseRejected,
+    ResourceLease,
+    ResourceLeaseRegistry,
+    ResourceScope,
+)
+from ksadk.resource_runtime.operation_ledger import OperationLedger
+from ksadk.resource_runtime.process import ResourceWorkerProcess
+from ksadk.resource_runtime.socket_server import ResourceSocketServer
+from ksadk.resource_runtime.worker import DiscoverySkillWorkerBinding, WorkerInitialization
+
+ResourceRevalidator = Callable[[tuple[ResourceScope, ...]], Awaitable[bool]]
+
+
+@dataclass(frozen=True)
+class ActiveResources:
+    activation_id: str
+    socket_path: Path
+    leases: tuple[ResourceLease, ...] = field(repr=False)
+    worker_pid: int
+
+
+@dataclass
+class _Running:
+    worker: ResourceWorkerProcess
+    resources: ActiveResources
+    lease_changed: asyncio.Event = field(default_factory=asyncio.Event)
+
+
+class _ActivationApprovals:
+    """Dispatch only to the approval owner admitted for this exact activation."""
+
+    def __init__(self):
+        self.owners: dict[str, tuple[tuple[ResourceScope, ...], ResourceWriteAuthorizer]] = {}
+
+    def register(self, scopes: tuple[ResourceScope, ...], owner: ResourceWriteAuthorizer) -> None:
+        activation_id = scopes[0].activation_id
+        if activation_id in self.owners:
+            raise ValueError("RESOURCE_APPROVAL_OWNER_IMMUTABLE")
+        self.owners[activation_id] = (scopes, owner)
+
+    def revoke(self, activation_id: str) -> None:
+        self.owners.pop(activation_id, None)
+
+    async def authorize(self, request: ResourceRequest, scope: ResourceScope) -> str | None:
+        admitted = self.owners.get(scope.activation_id)
+        if admitted is None or scope not in admitted[0]:
+            return None
+        event_id = await admitted[1].authorize(request, scope)
+        # An approval may finish after the run was stopped. Do not revive it.
+        return event_id if self.owners.get(scope.activation_id) is admitted else None
+
+
+class ResourceSupervisor:
+    """One owner per Core generation, with one worker per activation.
+
+    Admission and real credential/principal verification precede activate().
+    This class never derives user identity from tool arguments, reads credential
+    environment variables, restarts failed activations or replays requests.
+    """
+
+    def __init__(
+        self,
+        generation_id: str,
+        *,
+        max_workers: int = 32,
+        write_authorizer: ResourceWriteAuthorizer | None = None,
+        operation_ledger: OperationLedger | None = None,
+    ):
+        if type(max_workers) is not int or not 1 <= max_workers <= 128:
+            raise ValueError("Resource worker limit must be 1..128")
+        self.generation_id = generation_id
+        self._max_workers = max_workers
+        if write_authorizer is not None and operation_ledger is None:
+            raise ValueError("Write approval requires a durable operation ledger")
+        self._default_write_authorizer = write_authorizer
+        self._has_operation_ledger = operation_ledger is not None
+        self._discovery_receipts = (
+            DiscoverySelectionReceipts(operation_ledger) if operation_ledger is not None else None
+        )
+        self._approvals = _ActivationApprovals()
+        self._registry = ResourceLeaseRegistry()
+        self._broker = ResourceBroker(
+            self._registry,
+            generation_id,
+            write_authorizer=self._approvals if operation_ledger is not None else None,
+            operation_ledger=operation_ledger,
+        )
+        self._socket = ResourceSocketServer(self._broker)
+        self._lock = asyncio.Lock()
+        self._running: dict[str, _Running] = {}
+        self._used: set[str] = set()
+        self._monitors: set[asyncio.Task] = set()
+        self._closed = False
+
+    async def start_broker(self) -> Path:
+        """Prepare private ingress before Core startup, without workers or grants."""
+        async with self._lock:
+            return await self._start_broker_locked()
+
+    async def _start_broker_locked(self) -> Path:
+        if self._closed:
+            raise RuntimeError("RESOURCE_SUPERVISOR_CLOSED")
+        if self._socket.path is None:
+            try:
+                await self._socket.start()
+            except BaseException:
+                self._closed = True
+                raise
+        assert self._socket.path is not None
+        return self._socket.path
+
+    async def activate(
+        self,
+        initialization: WorkerInitialization,
+        *,
+        write_authorizer: ResourceWriteAuthorizer | None = None,
+    ) -> ActiveResources:
+        # Revalidate even models assembled via unchecked model_copy/model_construct.
+        initialization = WorkerInitialization.model_validate(initialization.pipe_payload())
+        discovery_keys = {}
+        restored_bindings = list(initialization.bindings)
+        for index, binding in enumerate(initialization.bindings):
+            if isinstance(binding, DiscoverySkillWorkerBinding):
+                if self._discovery_receipts is None:
+                    raise ValueError("RESOURCE_DISCOVERY_RECEIPTS_REQUIRED")
+                binding_scope = next(
+                    item for item in initialization.scopes
+                    if item.binding_id == binding.config.binding.id
+                )
+                key = discovery_scope_key(binding_scope, binding.selection_run_ref)
+                discovery_keys[binding_scope.binding_id] = key
+                restored = await asyncio.to_thread(self._discovery_receipts.read, key)
+                # Only host ledger contents may restore selections. Ignore caller
+                # supplied preloads at this public lifecycle entry.
+                restored_bindings[index] = binding.model_copy(
+                    update={"restored_selections": restored}
+                )
+        initialization = initialization.model_copy(update={"bindings": tuple(restored_bindings)})
+        scope = initialization.scopes[0]
+        owner = write_authorizer if write_authorizer is not None else self._default_write_authorizer
+        if owner is not None and not self._has_operation_ledger:
+            raise ValueError("Write approval requires a durable operation ledger")
+        if scope.generation_id != self.generation_id:
+            raise ValueError("Resource activation belongs to another generation")
+        async with self._lock:
+            if self._closed:
+                raise RuntimeError("RESOURCE_SUPERVISOR_CLOSED")
+            if scope.activation_id in self._used:
+                raise ValueError("RESOURCE_ACTIVATION_ALREADY_USED")
+            if len(self._used) >= 4096:
+                raise RuntimeError("RESOURCE_GENERATION_CAPACITY")
+            if len(self._running) >= self._max_workers:
+                raise RuntimeError("RESOURCE_WORKER_CAPACITY")
+            await self._start_broker_locked()
+            self._used.add(scope.activation_id)
+            worker = None
+            try:
+                if owner is not None:
+                    self._approvals.register(initialization.scopes, owner)
+                worker = await ResourceWorkerProcess.start(initialization)
+                leases = []
+                for binding_scope in initialization.scopes:
+                    binding = next(
+                        item
+                        for item in initialization.bindings
+                        if item.config.binding.id == binding_scope.binding_id
+                    )
+                    self._broker.register(
+                        binding_scope,
+                        worker,
+                        artifact_root=getattr(binding, "artifact_directory", None),
+                        discovery_key=discovery_keys.get(binding_scope.binding_id),
+                    )
+                    leases.append(self._registry.issue(binding_scope))
+                assert self._socket.path is not None
+                resources = ActiveResources(
+                    scope.activation_id, self._socket.path, tuple(leases), worker.pid
+                )
+                self._running[scope.activation_id] = _Running(worker, resources)
+                monitor = asyncio.create_task(self._monitor(scope.activation_id, worker))
+                self._monitors.add(monitor)
+                monitor.add_done_callback(self._monitors.discard)
+                return resources
+            except BaseException:
+                self._approvals.revoke(scope.activation_id)
+                self._broker.revoke_activation(scope.activation_id)
+                if worker is not None:
+                    await worker.aclose()
+                await self._broker.retire_activation(scope.activation_id)
+                raise
+
+    async def deactivate(self, activation_id: str) -> None:
+        async with self._lock:
+            await self._deactivate_locked(activation_id)
+
+    async def owns(self, current: ActiveResources) -> bool:
+        """Confirm an opaque activation handle still names this live generation."""
+
+        async with self._lock:
+            running = self._running.get(current.activation_id)
+            return (
+                not self._closed
+                and running is not None
+                and running.resources is current
+                and running.worker.is_running
+            )
+
+    async def renew(
+        self,
+        current: ActiveResources,
+        *,
+        revalidate: ResourceRevalidator,
+        lifetime: int = 600,
+    ) -> ActiveResources:
+        """Host-only renewal after fresh upstream permission checks.
+
+        The caller replaces its connector with the returned handles at a safe
+        turn boundary. No model operation, automatic timer, or default grant is
+        installed here. The verifier must check every supplied binding scope.
+        """
+        if type(lifetime) is not int or not 1 <= lifetime <= 600:
+            raise ValueError("Resource lease lifetime must be 1..600 seconds")
+        async with self._lock:
+            running = self._running.get(current.activation_id)
+            if (
+                self._closed
+                or running is None
+                or running.resources is not current
+                or not running.worker.is_running
+            ):
+                raise ValueError("RESOURCE_RENEWAL_STALE")
+            scopes = tuple(
+                ResourceScope.model_validate(lease.scope.model_dump()) for lease in current.leases
+            )
+        # Do not block immediate deactivation while checking the platform.
+        try:
+            if await asyncio.wait_for(revalidate(scopes), timeout=15) is not True:
+                raise PermissionError("RESOURCE_RENEWAL_DENIED")
+        except BaseException:
+            async with self._lock:
+                latest = self._running.get(current.activation_id)
+                if latest is running and latest.resources is current:
+                    await self._deactivate_locked(current.activation_id)
+            raise
+        async with self._lock:
+            if (
+                self._closed
+                or self._running.get(current.activation_id) is not running
+                or running.resources is not current
+                or not running.worker.is_running
+            ):
+                raise ValueError("RESOURCE_RENEWAL_STALE")
+            try:
+                leases = self._registry.renew_many(
+                    tuple(lease.handle for lease in current.leases),
+                    lifetime=lifetime,
+                )
+            except BaseException:
+                await self._deactivate_locked(current.activation_id)
+                raise
+            renewed = ActiveResources(
+                current.activation_id, current.socket_path, leases, current.worker_pid
+            )
+            running.resources = renewed
+            running.lease_changed.set()
+            return renewed
+
+    async def runtime_status(self, *, agent_id: str, activation_id: str) -> dict | None:
+        """Observe this live activation without issuing leases or probing upstream.
+
+        Unknown and differently owned activations are indistinguishable. This is
+        transport availability, not proof that current upstream grants are valid.
+        """
+        async with self._lock:
+            running = self._running.get(activation_id)
+            if running is None:
+                return None
+            scopes = [lease.scope for lease in running.resources.leases]
+            if any(scope.identity.agent_id != agent_id for scope in scopes):
+                return None
+            items = []
+            for lease in running.resources.leases:
+                available = running.worker.is_running and not self._closed
+                try:
+                    self._registry.authorize(
+                        lease.handle,
+                        operation=lease.scope.allowed_operations[0],
+                        generation_id=self.generation_id,
+                        activation_id=activation_id,
+                    )
+                except LeaseRejected:
+                    available = False
+                items.append(
+                    {
+                        "bindingId": lease.scope.binding_id,
+                        "runtimeState": "available" if available else "unavailable",
+                        "allowedOperations": list(lease.scope.allowed_operations)
+                        if available
+                        else [],
+                    }
+                )
+            first = scopes[0]
+            return {
+                "agentId": agent_id,
+                "activationId": activation_id,
+                "generationId": self.generation_id,
+                "buildDigest": first.build_digest,
+                "bindingSnapshotDigest": first.binding_snapshot_digest,
+                "observedAt": datetime.now(timezone.utc).isoformat(),
+                "upstreamVerification": "not-observed",
+                "items": items,
+            }
+
+    async def _deactivate_locked(self, activation_id: str) -> None:
+        self._approvals.revoke(activation_id)
+        self._broker.revoke_activation(activation_id)
+        running = self._running.pop(activation_id, None)
+        if running is not None:
+            await running.worker.aclose()
+        await self._broker.retire_activation(activation_id)
+
+    async def _monitor(self, activation_id: str, worker: ResourceWorkerProcess) -> None:
+        stopped = asyncio.create_task(worker.wait_stopped())
+        try:
+            while True:
+                async with self._lock:
+                    running = self._running.get(activation_id)
+                    if running is None or running.worker is not worker:
+                        return
+                    remaining = self._registry.remaining_lifetime(
+                        tuple(lease.handle for lease in running.resources.leases)
+                    )
+                    if stopped.done() or remaining <= 0:
+                        await self._deactivate_locked(activation_id)
+                        return
+                    # Read current leases and clear under the same lock as renew.
+                    # An old expiry wakeup never revokes a successfully renewed run.
+                    running.lease_changed.clear()
+                changed = asyncio.create_task(running.lease_changed.wait())
+                try:
+                    # Also notice complete binding revocation without another model
+                    # request. This timer cannot refresh credentials or issue leases.
+                    await asyncio.wait(
+                        (stopped, changed), timeout=min(remaining, 30.0),
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                finally:
+                    changed.cancel()
+                    await asyncio.gather(changed, return_exceptions=True)
+        finally:
+            if not stopped.done():
+                stopped.cancel()
+            await asyncio.gather(stopped, return_exceptions=True)
+
+    async def aclose(self) -> None:
+        async with self._lock:
+            self._closed = True
+            await self._socket.aclose()
+            for activation_id in tuple(self._running):
+                await self._deactivate_locked(activation_id)
+        # Monitors may be waiting on the supervisor lock, so join after releasing it.
+        if self._monitors:
+            await asyncio.gather(*tuple(self._monitors), return_exceptions=True)