agentengine-sdk-python 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

ksadk/kernel/execution_grants.py

@@ -0,0 +1,221 @@
+"""Revocable execution admission, independent of any scheduling business domain.
+
+The host authorizes grant management. A grant is scoped to an exact tenant,
+AgentInstance and session; its owner is a trusted host reference. Grant state
+and Inbox claim MUST be serialized by the same control-store transaction.
+Claim is the start-qualification boundary, not evidence of a running model.
+Revocation never revives a command and does not cancel already qualified work.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any, Literal
+from uuid import NAMESPACE_URL, uuid5
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ksadk.kernel.contracts import AgentControlCommand
+from ksadk.kernel.errors import InvalidCommandError
+
+GrantState = Literal["active", "suspended", "revoked"]
+
+
+class ExecutionGrantSpec(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid", str_strip_whitespace=True)
+
+    grant_id: str = Field(min_length=1, max_length=512)
+    tenant_id: str = Field(min_length=1)
+    agent_instance_id: str = Field(min_length=1)
+    session_id: str = Field(min_length=1)
+    owner_ref: str = Field(min_length=1)
+
+
+class ExecutionGrantRecord(ExecutionGrantSpec):
+    state: GrantState = "active"
+    revision: int = Field(default=1, ge=1)
+    created_at: str
+    updated_at: str
+
+
+class ExecutionGrantCommand(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    message_id: str
+    command_id: str
+    idempotency_key: str
+    inbox_state: str
+    run_id: str | None = None
+    run_state: str | None = None
+
+
+class ExecutionGrantBarrier(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    grant: ExecutionGrantRecord
+    queued_message_ids: tuple[str, ...] = ()
+    in_flight_message_ids: tuple[str, ...] = ()
+    discarded_message_ids: tuple[str, ...] = ()
+    settled_message_ids: tuple[str, ...] = ()
+    commands: tuple[ExecutionGrantCommand, ...] = ()
+
+
+class ExecutionGrantBlocked(InvalidCommandError):
+    """A listed command lost eligibility before the transactional claim."""
+
+    def __init__(self, state: str, *, message_id: str = "") -> None:
+        self.grant_state = state
+        self.message_id = message_id
+        super().__init__(
+            f"execution grant is {state}",
+            details={"reason": f"execution_grant_{state}", "message_id": message_id},
+        )
+
+
+def execution_grant_id(command: AgentControlCommand) -> str | None:
+    value = command.payload.get("execution_grant_id")
+    if value is None:
+        return None
+    if command.command_type != "enqueue" or not isinstance(value, str) or not value.strip():
+        raise InvalidCommandError("execution_grant_id requires enqueue and a nonempty string")
+    if value != value.strip() or len(value) > 512:
+        raise InvalidCommandError("invalid execution_grant_id")
+    return value
+
+
+def execution_grant_run_id(command: AgentControlCommand) -> str:
+    """Stable lookup identity; never use a new Run to replay uncertain starts."""
+    return str(
+        uuid5(
+            NAMESPACE_URL,
+            "ksadk:execution-grant:"
+            + json.dumps(
+                [
+                    command.tenant_id,
+                    command.agent_instance_id,
+                    command.session_id,
+                    str(command.command_id),
+                ],
+                separators=(",", ":"),
+            ),
+        )
+    )
+
+
+def require_grant_scope(record: ExecutionGrantRecord, spec: ExecutionGrantSpec) -> None:
+    if any(getattr(record, key) != value for key, value in spec.model_dump().items()):
+        raise InvalidCommandError("execution grant scope or owner mismatch")
+
+
+def command_grant_error(
+    command: AgentControlCommand,
+    record: ExecutionGrantRecord | None,
+) -> str | None:
+    if record is None:
+        return "execution_grant_not_found"
+    if any(
+        getattr(record, key) != getattr(command, key)
+        for key in (
+            "tenant_id",
+            "agent_instance_id",
+            "session_id",
+        )
+    ):
+        return "execution_grant_scope_mismatch"
+    if record.state == "revoked":
+        return "execution_grant_revoked"
+    return None
+
+
+def grant_operation_digest(
+    spec: ExecutionGrantSpec, state: GrantState, expected_revision: int
+) -> str:
+    if state not in {"active", "suspended", "revoked"}:
+        raise InvalidCommandError("unknown execution grant state")
+    if (
+        isinstance(expected_revision, bool)
+        or not isinstance(expected_revision, int)
+        or expected_revision < 1
+    ):
+        raise InvalidCommandError("expected_revision must be a positive integer")
+    return hashlib.sha256(
+        json.dumps(
+            {
+                "spec": spec.model_dump(),
+                "state": state,
+                "expected_revision": expected_revision,
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        ).encode()
+    ).hexdigest()
+
+
+def transition_grant(
+    record: ExecutionGrantRecord,
+    state: GrantState,
+    expected_revision: int,
+    now: str,
+) -> ExecutionGrantRecord:
+    if record.revision != expected_revision:
+        raise InvalidCommandError(
+            "execution grant revision mismatch",
+            details={
+                "expected_revision": expected_revision,
+                "actual_revision": record.revision,
+            },
+        )
+    if record.state == "revoked" and state != "revoked":
+        raise InvalidCommandError("revoked execution grant cannot be reactivated")
+    return record.model_copy(
+        update={
+            "state": state,
+            "revision": record.revision + 1,
+            "updated_at": now,
+        }
+    )
+
+
+def make_grant_barrier(
+    record: ExecutionGrantRecord,
+    messages: list[Any],
+    runs: list[Any],
+) -> ExecutionGrantBarrier:
+    from ksadk.kernel.state import TERMINAL_RUN_STATES
+
+    runs_by_command = {str(run.metadata.get("command_id")): run for run in runs}
+    queued, inflight, discarded, settled, commands = [], [], [], [], []
+    for message in messages:
+        command = message.command
+        if command is None or execution_grant_id(command) != record.grant_id:
+            continue
+        run = runs_by_command.get(str(command.command_id))
+        status = str(message.status)
+        commands.append(
+            ExecutionGrantCommand(
+                message_id=message.message_id,
+                command_id=str(command.command_id),
+                idempotency_key=command.idempotency_key,
+                inbox_state=status,
+                run_id=run.run_id if run else None,
+                run_state=str(run.state) if run else None,
+            )
+        )
+        if status == "accepted":
+            queued.append(message.message_id)
+        elif status == "discarded":
+            discarded.append(message.message_id)
+        elif run is not None and run.state in TERMINAL_RUN_STATES:
+            settled.append(message.message_id)
+        else:
+            # Completed Inbox means adapter.start returned, not model completion.
+            inflight.append(message.message_id)
+    return ExecutionGrantBarrier(
+        grant=record,
+        queued_message_ids=tuple(queued),
+        in_flight_message_ids=tuple(inflight),
+        discarded_message_ids=tuple(discarded),
+        settled_message_ids=tuple(settled),
+        commands=tuple(commands),
+    )

ksadk/kernel/execution_grants_memory.py

@@ -0,0 +1,133 @@
+"""In-memory execution-grant transactions, using the Inbox session lock."""
+
+from __future__ import annotations
+
+from ksadk.kernel.errors import InvalidCommandError
+from ksadk.kernel.execution_grants import (
+    ExecutionGrantBarrier,
+    ExecutionGrantBlocked,
+    ExecutionGrantRecord,
+    ExecutionGrantSpec,
+    GrantState,
+    command_grant_error,
+    execution_grant_id,
+    grant_operation_digest,
+    make_grant_barrier,
+    require_grant_scope,
+    transition_grant,
+)
+from ksadk.kernel.state import InboxState
+from ksadk.kernel.store import now_iso
+
+
+class MemoryExecutionGrantMixin:
+    async def ensure_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantRecord:
+        async with self._lock(spec.agent_instance_id, spec.session_id):
+            existing = self._execution_grants.get(spec.grant_id)
+            if existing is not None:
+                require_grant_scope(existing, spec)
+                return existing
+            record = ExecutionGrantRecord(
+                **spec.model_dump(), created_at=now_iso(), updated_at=now_iso()
+            )
+            self._execution_grants[spec.grant_id] = record
+            return record
+
+    def _memory_grant_barrier(self, record: ExecutionGrantRecord) -> ExecutionGrantBarrier:
+        return make_grant_barrier(
+            record,
+            [
+                self._to_message(row)
+                for row in self._session_rows(record.agent_instance_id, record.session_id)
+            ],
+            [
+                run
+                for run in self._runs.values()
+                if run.agent_instance_id == record.agent_instance_id
+                and run.session_id == record.session_id
+            ],
+        )
+
+    async def get_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantBarrier | None:
+        async with self._lock(spec.agent_instance_id, spec.session_id):
+            record = self._execution_grants.get(spec.grant_id)
+            if record is None:
+                return None
+            require_grant_scope(record, spec)
+            return self._memory_grant_barrier(record)
+
+    async def set_execution_grant_state(
+        self,
+        spec: ExecutionGrantSpec,
+        state: GrantState,
+        *,
+        expected_revision: int,
+        idempotency_key: str,
+    ) -> ExecutionGrantBarrier:
+        if not idempotency_key or not idempotency_key.strip():
+            raise InvalidCommandError("grant mutation requires idempotency_key")
+        digest = grant_operation_digest(spec, state, expected_revision)
+        async with self._lock(spec.agent_instance_id, spec.session_id):
+            record = self._execution_grants.get(spec.grant_id)
+            if record is None:
+                raise InvalidCommandError("execution grant not found")
+            require_grant_scope(record, spec)
+            operation_key = (spec.grant_id, idempotency_key)
+            previous = self._execution_grant_operations.get(operation_key)
+            if previous is not None:
+                if previous[0] != digest:
+                    raise InvalidCommandError("execution grant idempotency conflict")
+                return previous[1]
+            updated = transition_grant(record, state, expected_revision, now_iso())
+            # No await between state change, queue invalidation and receipt.
+            # claim_message uses exactly this same session lock.
+            self._execution_grants[spec.grant_id] = updated
+            if state == "revoked":
+                for row in self._session_rows(spec.agent_instance_id, spec.session_id):
+                    if (
+                        row["status"] == InboxState.ACCEPTED
+                        and execution_grant_id(row["command"]) == spec.grant_id
+                    ):
+                        row["status"] = InboxState.DISCARDED
+            barrier = self._memory_grant_barrier(updated)
+            self._execution_grant_operations[operation_key] = (digest, barrier)
+            return barrier
+
+    def _memory_admission_grant_error(self, command) -> str | None:
+        grant_id = execution_grant_id(command)
+        return (
+            command_grant_error(command, self._execution_grants.get(grant_id)) if grant_id else None
+        )
+
+    def _memory_require_claim_grant(self, row) -> None:
+        command = row["command"]
+        grant_id = execution_grant_id(command)
+        if not grant_id:
+            return
+        record = self._execution_grants.get(grant_id)
+        error = command_grant_error(command, record)
+        # A previously qualified command remains in flight across a barrier.
+        if row["status"] == InboxState.CLAIMED and error in (None, "execution_grant_revoked"):
+            return
+        if error:
+            raise ExecutionGrantBlocked(
+                error.removeprefix("execution_grant_"), message_id=row["message_id"]
+            )
+        if record.state == "suspended":
+            raise ExecutionGrantBlocked("suspended", message_id=row["message_id"])
+        earlier = any(
+            other["accepted_seq"] < row["accepted_seq"]
+            and other["status"] in (InboxState.ACCEPTED, InboxState.CLAIMED)
+            and other["command"].command_type == "enqueue"
+            for other in self._session_rows(row["agent_instance_id"], row["session_id"])
+        )
+        busy = any(
+            run.agent_instance_id == row["agent_instance_id"]
+            and run.session_id == row["session_id"]
+            and str(run.state) in {"pending", "running", "paused", "waiting"}
+            for run in self._runs.values()
+        )
+        if earlier or busy:
+            raise ExecutionGrantBlocked(
+                "queued" if earlier else "busy", message_id=row["message_id"]
+            )

ksadk/kernel/execution_grants_postgres.py

@@ -0,0 +1,217 @@
+"""PostgreSQL grants: lock order is grant -> Inbox for every grant writer.
+
+The grant row is locked FOR UPDATE before claim/revoke touch Inbox rows. Thus
+concurrent writers cannot qualify work after a committed revocation barrier.
+"""
+
+from __future__ import annotations
+
+import json
+
+from ksadk.kernel.errors import InvalidCommandError
+from ksadk.kernel.execution_grants import (
+    ExecutionGrantBarrier,
+    ExecutionGrantBlocked,
+    ExecutionGrantRecord,
+    ExecutionGrantSpec,
+    GrantState,
+    command_grant_error,
+    execution_grant_id,
+    grant_operation_digest,
+    make_grant_barrier,
+    require_grant_scope,
+    transition_grant,
+)
+from ksadk.kernel.store import RunRecord, now_iso
+
+
+class PostgresExecutionGrantMixin:
+    def _require_grant_tenant(self, tenant_id: str) -> None:
+        if tenant_id != self.tenant_id:
+            raise InvalidCommandError("execution grant tenant does not match kernel store")
+
+    async def _postgres_grant_record(
+        self, connection, grant_id: str
+    ) -> ExecutionGrantRecord | None:
+        row = await connection.fetchrow(
+            "SELECT * FROM kernel_execution_grants WHERE grant_id=$1 FOR UPDATE",
+            grant_id,
+        )
+        return ExecutionGrantRecord.model_validate(dict(row)) if row is not None else None
+
+    async def ensure_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantRecord:
+        self._require_grant_tenant(spec.tenant_id)
+        now = now_iso()
+        async with self._connection() as connection:
+            async with connection.transaction():
+                await connection.execute(
+                    "INSERT INTO kernel_execution_grants (grant_id, tenant_id, "
+                    "agent_instance_id, session_id, "
+                    "owner_ref, state, revision, created_at, updated_at) VALUES "
+                    "($1,$2,$3,$4,$5,'active',1,$6,$6) "
+                    "ON CONFLICT (grant_id) DO NOTHING",
+                    spec.grant_id,
+                    spec.tenant_id,
+                    spec.agent_instance_id,
+                    spec.session_id,
+                    spec.owner_ref,
+                    now,
+                )
+                record = await self._postgres_grant_record(connection, spec.grant_id)
+                assert record is not None
+                require_grant_scope(record, spec)
+                return record
+
+    async def _postgres_grant_barrier(self, connection, record) -> ExecutionGrantBarrier:
+        rows = await connection.fetch(
+            "SELECT * FROM kernel_inbox WHERE tenant_id=$1 AND agent_instance_id=$2 AND "
+            "session_id=$3 "
+            "ORDER BY accepted_seq",
+            record.tenant_id,
+            record.agent_instance_id,
+            record.session_id,
+        )
+        messages = [self._row_to_message(row) for row in rows]
+        rows = await connection.fetch(
+            "SELECT * FROM kernel_runs WHERE tenant_id=$1 AND agent_instance_id=$2 AND "
+            "session_id=$3",
+            record.tenant_id,
+            record.agent_instance_id,
+            record.session_id,
+        )
+        runs = [
+            RunRecord(
+                run_id=row["run_id"],
+                agent_instance_id=row["agent_instance_id"],
+                session_id=row["session_id"],
+                state=row["state"],
+                metadata=json.loads(row["metadata"])
+                if isinstance(row["metadata"], str)
+                else dict(row["metadata"]),
+            )
+            for row in rows
+        ]
+        return make_grant_barrier(record, [m for m in messages if m], runs)
+
+    async def get_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantBarrier | None:
+        self._require_grant_tenant(spec.tenant_id)
+        async with self._connection() as connection:
+            async with connection.transaction():
+                record = await self._postgres_grant_record(connection, spec.grant_id)
+                if record is None:
+                    return None
+                require_grant_scope(record, spec)
+                return await self._postgres_grant_barrier(connection, record)
+
+    async def set_execution_grant_state(
+        self,
+        spec: ExecutionGrantSpec,
+        state: GrantState,
+        *,
+        expected_revision: int,
+        idempotency_key: str,
+    ) -> ExecutionGrantBarrier:
+        self._require_grant_tenant(spec.tenant_id)
+        if not idempotency_key or not idempotency_key.strip():
+            raise InvalidCommandError("grant mutation requires idempotency_key")
+        digest = grant_operation_digest(spec, state, expected_revision)
+        async with self._connection() as connection:
+            async with connection.transaction():
+                record = await self._postgres_grant_record(connection, spec.grant_id)
+                if record is None:
+                    raise InvalidCommandError("execution grant not found")
+                require_grant_scope(record, spec)
+                previous = await connection.fetchrow(
+                    "SELECT * FROM kernel_execution_grant_operations WHERE grant_id=$1 AND "
+                    "idempotency_key=$2",
+                    spec.grant_id,
+                    idempotency_key,
+                )
+                if previous is not None:
+                    if previous["request_digest"] != digest:
+                        raise InvalidCommandError("execution grant idempotency conflict")
+                    return ExecutionGrantBarrier.model_validate_json(previous["receipt_json"])
+                updated = transition_grant(record, state, expected_revision, now_iso())
+                await connection.execute(
+                    "UPDATE kernel_execution_grants SET state=$1, revision=$2, updated_at=$3 "
+                    "WHERE grant_id=$4",
+                    updated.state,
+                    updated.revision,
+                    updated.updated_at,
+                    spec.grant_id,
+                )
+                if state == "revoked":
+                    await connection.execute(
+                        "UPDATE kernel_inbox SET status='discarded' WHERE tenant_id=$1 AND "
+                        "agent_instance_id=$2 "
+                        "AND session_id=$3 AND status='accepted' AND "
+                        "payload->>'command_type'='enqueue' "
+                        "AND payload->'payload'->>'execution_grant_id'=$4",
+                        spec.tenant_id,
+                        spec.agent_instance_id,
+                        spec.session_id,
+                        spec.grant_id,
+                    )
+                barrier = await self._postgres_grant_barrier(connection, updated)
+                await connection.execute(
+                    "INSERT INTO kernel_execution_grant_operations (grant_id, idempotency_key, "
+                    "request_digest, receipt_json) "
+                    "VALUES ($1,$2,$3,$4)",
+                    spec.grant_id,
+                    idempotency_key,
+                    digest,
+                    barrier.model_dump_json(),
+                )
+                return barrier
+
+    async def _postgres_lock_command_grant(self, connection, command):
+        grant_id = execution_grant_id(command)
+        if not grant_id:
+            return None
+        if command.tenant_id != self.tenant_id:
+            return None
+        return await self._postgres_grant_record(connection, grant_id)
+
+    async def _postgres_admission_grant_error(self, connection, command) -> str | None:
+        if not execution_grant_id(command):
+            return None
+        if command.tenant_id != self.tenant_id:
+            return "execution_grant_scope_mismatch"
+        return command_grant_error(
+            command, await self._postgres_lock_command_grant(connection, command)
+        )
+
+    async def _postgres_require_claim_grant(self, connection, row, command, record) -> None:
+        if not execution_grant_id(command):
+            return
+        error = command_grant_error(command, record)
+        if row["status"] == "claimed" and error in (None, "execution_grant_revoked"):
+            return
+        if error:
+            raise ExecutionGrantBlocked(
+                error.removeprefix("execution_grant_"), message_id=str(row["message_id"])
+            )
+        if record.state == "suspended":
+            raise ExecutionGrantBlocked("suspended", message_id=str(row["message_id"]))
+        earlier = await connection.fetchval(
+            "SELECT EXISTS (SELECT 1 FROM kernel_inbox WHERE tenant_id=$1 "
+            "AND agent_instance_id=$2 AND session_id=$3 AND accepted_seq < $4 "
+            "AND status IN ('accepted','claimed') AND payload->>'command_type'='enqueue')",
+            record.tenant_id,
+            row["agent_instance_id"],
+            row["session_id"],
+            row["accepted_seq"],
+        )
+        busy = await connection.fetchval(
+            "SELECT EXISTS (SELECT 1 FROM kernel_runs WHERE tenant_id=$1 "
+            "AND agent_instance_id=$2 AND session_id=$3 "
+            "AND state IN ('pending','running','paused','waiting'))",
+            record.tenant_id,
+            row["agent_instance_id"],
+            row["session_id"],
+        )
+        if earlier or busy:
+            raise ExecutionGrantBlocked(
+                "queued" if earlier else "busy",
+                message_id=str(row["message_id"]),
+            )