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

ksadk/kernel/execution_grants_memory.py

@@ -9,28 +9,45 @@ from ksadk.kernel.execution_grants import (
     ExecutionGrantRecord,
     ExecutionGrantSpec,
     GrantState,
+    LocalExecutionGrantClock,
+    admission_operation_digest,
     command_grant_error,
     execution_grant_id,
+    grant_deadline_budget,
     grant_operation_digest,
+    grant_renewal_digest,
     make_grant_barrier,
+    renew_grant,
     require_grant_scope,
+    transition_admission,
     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:
+class MemoryExecutionGrantMixin(LocalExecutionGrantClock):
+    async def ensure_execution_grant(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        expires_at: str | None = None,
+        remaining_ttl_seconds: float | None = None,
+    ) -> ExecutionGrantRecord:
+        observed_deadline = grant_deadline_budget(remaining_ttl_seconds)
         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()
+                **spec.model_dump(),
+                expires_at=expires_at,
+                created_at=now_iso(),
+                updated_at=now_iso(),
             )
             self._execution_grants[spec.grant_id] = record
+            self._anchor_grant(record, observed_deadline=observed_deadline)
             return record
 
     def _memory_grant_barrier(self, record: ExecutionGrantRecord) -> ExecutionGrantBarrier:
@@ -56,6 +73,69 @@ class MemoryExecutionGrantMixin:
             require_grant_scope(record, spec)
             return self._memory_grant_barrier(record)
 
+    async def require_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantRecord:
+        async with self._lock(spec.agent_instance_id, spec.session_id):
+            record = self._execution_grants.get(spec.grant_id)
+            if record is None:
+                raise ExecutionGrantBlocked("not_found")
+            require_grant_scope(record, spec)
+            self._require_local_grant(record)
+            return record
+
+    async def lookup_execution_grant_operation(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        idempotency_key: str,
+    ) -> 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)
+            receipt = self._execution_grant_operations.get((spec.grant_id, idempotency_key))
+            return receipt[1] if receipt else None
+
+    async def renew_execution_grant(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        expected_revision: int,
+        expires_at: str,
+        renewal_id: str,
+        remaining_ttl_seconds: float | None = None,
+    ) -> ExecutionGrantBarrier:
+        observed_deadline = grant_deadline_budget(remaining_ttl_seconds)
+        digest = grant_renewal_digest(
+            spec,
+            expected_revision=expected_revision,
+            expires_at=expires_at,
+            renewal_id=renewal_id,
+        )
+        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)
+            key = (spec.grant_id, renewal_id)
+            previous = self._execution_grant_operations.get(key)
+            if previous:
+                if previous[0] != digest:
+                    raise InvalidCommandError("execution grant idempotency conflict")
+                return previous[1]
+            self._require_local_renewal(record)
+            updated = renew_grant(
+                record,
+                expected_revision=expected_revision,
+                expires_at=expires_at,
+                now=now_iso(),
+            )
+            self._execution_grants[spec.grant_id] = updated
+            barrier = self._memory_grant_barrier(updated)
+            self._execution_grant_operations[key] = (digest, barrier)
+            self._anchor_grant(updated, observed_deadline=observed_deadline)
+            return barrier
+
     async def set_execution_grant_state(
         self,
         spec: ExecutionGrantSpec,
@@ -93,11 +173,41 @@ class MemoryExecutionGrantMixin:
             self._execution_grant_operations[operation_key] = (digest, barrier)
             return barrier
 
+    async def set_execution_admission(
+        self,
+        spec: ExecutionGrantSpec,
+        allowed: bool,
+        *,
+        expected_revision: int,
+        idempotency_key: str,
+    ) -> ExecutionGrantBarrier:
+        if not idempotency_key or not idempotency_key.strip():
+            raise InvalidCommandError("admission mutation requires idempotency_key")
+        digest = admission_operation_digest(spec, allowed, 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)
+            key = (spec.grant_id, idempotency_key)
+            previous = self._execution_grant_operations.get(key)
+            if previous:
+                if previous[0] != digest:
+                    raise InvalidCommandError("execution grant idempotency conflict")
+                return previous[1]
+            self._require_local_grant(record)
+            updated = transition_admission(record, allowed, expected_revision, now_iso())
+            self._execution_grants[spec.grant_id] = updated
+            barrier = self._memory_grant_barrier(updated)
+            self._execution_grant_operations[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
-        )
+        if not grant_id:
+            return None
+        record = self._execution_grants.get(grant_id)
+        return command_grant_error(command, record) or self._local_grant_error(record)
 
     def _memory_require_claim_grant(self, row) -> None:
         command = row["command"]
@@ -105,9 +215,14 @@ class MemoryExecutionGrantMixin:
         if not grant_id:
             return
         record = self._execution_grants.get(grant_id)
-        error = command_grant_error(command, record)
+        error = command_grant_error(command, record) or self._local_grant_error(record)
         # A previously qualified command remains in flight across a barrier.
-        if row["status"] == InboxState.CLAIMED and error in (None, "execution_grant_revoked"):
+        if (
+            row["status"] == InboxState.CLAIMED
+            and record is not None
+            and record.expires_at is None
+            and error in (None, "execution_grant_revoked")
+        ):
             return
         if error:
             raise ExecutionGrantBlocked(
@@ -115,6 +230,10 @@ class MemoryExecutionGrantMixin:
             )
         if record.state == "suspended":
             raise ExecutionGrantBlocked("suspended", message_id=row["message_id"])
+        if row["status"] == InboxState.CLAIMED:
+            return
+        if not record.admission_allowed:
+            raise ExecutionGrantBlocked("admission_paused", message_id=row["message_id"])
         earlier = any(
             other["accepted_seq"] < row["accepted_seq"]
             and other["status"] in (InboxState.ACCEPTED, InboxState.CLAIMED)

ksadk/kernel/execution_grants_postgres.py

@@ -15,17 +15,27 @@ from ksadk.kernel.execution_grants import (
     ExecutionGrantRecord,
     ExecutionGrantSpec,
     GrantState,
+    admission_operation_digest,
     command_grant_error,
     execution_grant_id,
+    grant_deadline_budget,
     grant_operation_digest,
+    grant_renewal_digest,
     make_grant_barrier,
+    renew_grant,
+    require_active_grant,
     require_grant_scope,
+    transition_admission,
     transition_grant,
 )
 from ksadk.kernel.store import RunRecord, now_iso
 
 
 class PostgresExecutionGrantMixin:
+    async def _postgres_grant_now(self, connection) -> str:
+        # clock_timestamp(), not transaction_timestamp(): lock waits consume TTL.
+        return (await connection.fetchval("SELECT clock_timestamp()")).isoformat()
+
     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")
@@ -39,16 +49,35 @@ class PostgresExecutionGrantMixin:
         )
         return ExecutionGrantRecord.model_validate(dict(row)) if row is not None else None
 
-    async def ensure_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantRecord:
+    async def ensure_execution_grant(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        expires_at: str | None = None,
+        remaining_ttl_seconds: float | None = None,
+    ) -> ExecutionGrantRecord:
+        # PostgreSQL uses its authoritative clock, not a local transport anchor.
+        grant_deadline_budget(remaining_ttl_seconds)
         self._require_grant_tenant(spec.tenant_id)
         now = now_iso()
-        async with self._connection() as connection:
+        async with self._connection(write=True) as connection:
             async with connection.transaction():
+                existing = await self._postgres_grant_record(connection, spec.grant_id)
+                if existing is not None:
+                    require_grant_scope(existing, spec)
+                    return existing
+                candidate = ExecutionGrantRecord(
+                    **spec.model_dump(),
+                    expires_at=expires_at,
+                    created_at=now,
+                    updated_at=now,
+                )
                 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) "
+                    "owner_ref, state, revision, created_at, updated_at, "
+                    "attempt_epoch, expires_at) VALUES "
+                    "($1,$2,$3,$4,$5,'active',1,$6,$6,$7,$8) "
                     "ON CONFLICT (grant_id) DO NOTHING",
                     spec.grant_id,
                     spec.tenant_id,
@@ -56,12 +85,103 @@ class PostgresExecutionGrantMixin:
                     spec.session_id,
                     spec.owner_ref,
                     now,
+                    candidate.attempt_epoch,
+                    candidate.expires_at,
                 )
                 record = await self._postgres_grant_record(connection, spec.grant_id)
                 assert record is not None
                 require_grant_scope(record, spec)
                 return record
 
+    async def require_execution_grant(self, spec: ExecutionGrantSpec) -> ExecutionGrantRecord:
+        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:
+                    raise ExecutionGrantBlocked("not_found")
+                require_grant_scope(record, spec)
+                require_active_grant(record, now=await self._postgres_grant_now(connection))
+                return record
+
+    async def lookup_execution_grant_operation(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        idempotency_key: str,
+    ) -> 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)
+                row = await connection.fetchrow(
+                    "SELECT receipt_json FROM kernel_execution_grant_operations "
+                    "WHERE grant_id=$1 AND idempotency_key=$2",
+                    spec.grant_id,
+                    idempotency_key,
+                )
+                return (
+                    ExecutionGrantBarrier.model_validate_json(row["receipt_json"]) if row else None
+                )
+
+    async def renew_execution_grant(
+        self,
+        spec: ExecutionGrantSpec,
+        *,
+        expected_revision: int,
+        expires_at: str,
+        renewal_id: str,
+        remaining_ttl_seconds: float | None = None,
+    ) -> ExecutionGrantBarrier:
+        grant_deadline_budget(remaining_ttl_seconds)
+        self._require_grant_tenant(spec.tenant_id)
+        digest = grant_renewal_digest(
+            spec, expected_revision=expected_revision, expires_at=expires_at, renewal_id=renewal_id
+        )
+        async with self._connection(write=True) 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,
+                    renewal_id,
+                )
+                if previous:
+                    if previous["request_digest"] != digest:
+                        raise InvalidCommandError("execution grant idempotency conflict")
+                    return ExecutionGrantBarrier.model_validate_json(previous["receipt_json"])
+                updated = renew_grant(
+                    record,
+                    expected_revision=expected_revision,
+                    expires_at=expires_at,
+                    now=await self._postgres_grant_now(connection),
+                )
+                await connection.execute(
+                    "UPDATE kernel_execution_grants SET expires_at=$1, revision=$2, updated_at=$3 "
+                    "WHERE grant_id=$4",
+                    updated.expires_at,
+                    updated.revision,
+                    updated.updated_at,
+                    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,
+                    renewal_id,
+                    digest,
+                    barrier.model_dump_json(),
+                )
+                return barrier
+
     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 "
@@ -115,7 +235,7 @@ class PostgresExecutionGrantMixin:
         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 self._connection(write=True) as connection:
             async with connection.transaction():
                 record = await self._postgres_grant_record(connection, spec.grant_id)
                 if record is None:
@@ -164,6 +284,56 @@ class PostgresExecutionGrantMixin:
                 )
                 return barrier
 
+    async def set_execution_admission(
+        self,
+        spec: ExecutionGrantSpec,
+        allowed: bool,
+        *,
+        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("admission mutation requires idempotency_key")
+        digest = admission_operation_digest(spec, allowed, expected_revision)
+        async with self._connection(write=True) 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:
+                    if previous["request_digest"] != digest:
+                        raise InvalidCommandError("execution grant idempotency conflict")
+                    return ExecutionGrantBarrier.model_validate_json(previous["receipt_json"])
+                updated = transition_admission(
+                    record, allowed, expected_revision, await self._postgres_grant_now(connection)
+                )
+                await connection.execute(
+                    "UPDATE kernel_execution_grants SET admission_allowed=$1, "
+                    "admission_revision=$2, updated_at=$3 WHERE grant_id=$4",
+                    updated.admission_allowed,
+                    updated.admission_revision,
+                    updated.updated_at,
+                    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:
@@ -177,15 +347,19 @@ class PostgresExecutionGrantMixin:
             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)
-        )
+        record = await self._postgres_lock_command_grant(connection, command)
+        return command_grant_error(command, record, now=await self._postgres_grant_now(connection))
 
     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"):
+        error = command_grant_error(command, record, now=await self._postgres_grant_now(connection))
+        if (
+            row["status"] == "claimed"
+            and record is not None
+            and record.expires_at is None
+            and error in (None, "execution_grant_revoked")
+        ):
             return
         if error:
             raise ExecutionGrantBlocked(
@@ -193,6 +367,10 @@ class PostgresExecutionGrantMixin:
             )
         if record.state == "suspended":
             raise ExecutionGrantBlocked("suspended", message_id=str(row["message_id"]))
+        if row["status"] == "claimed":
+            return
+        if not record.admission_allowed:
+            raise ExecutionGrantBlocked("admission_paused", 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 "