cortex-agent-sdk 0.0.1__py3-none-any.whl

cortex_agent_sdk/redis/store.py

@@ -0,0 +1,293 @@
+import asyncio
+import time
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from datetime import UTC, datetime
+from hashlib import sha256
+from typing import Self, cast
+from uuid import uuid4
+
+from redis.asyncio import Redis
+from redis.exceptions import RedisError
+
+from cortex_agent_sdk.errores import AppError, CodigoError
+from cortex_agent_sdk.redis import scripts
+from cortex_agent_sdk.sessions.codec import SessionCodec
+from cortex_agent_sdk.sessions.models import SessionRecord
+from cortex_agent_sdk.sessions.store import SessionLease
+
+type _RedisArgument = str | bytes | int | float
+
+
+class _RedisLease:
+    def __init__(
+        self,
+        store: "RedisSessionStore",
+        session_id: str,
+        owner: str,
+        fencing_token: int,
+    ) -> None:
+        self._store = store
+        self._session_id = session_id
+        self._owner = owner
+        self._fencing_token = fencing_token
+        self._active = True
+
+    @property
+    def session_id(self) -> str:
+        return self._session_id
+
+    @property
+    def fencing_token(self) -> int:
+        return self._fencing_token
+
+    @property
+    def renewal_interval_seconds(self) -> float:
+        return self._store.renewal_interval_seconds
+
+    async def verify_ownership(self, fencing_token: int) -> None:
+        self._ensure_active(fencing_token)
+        await self._store._verify(self._session_id, self._lease_value)
+
+    async def renew(self, fencing_token: int) -> None:
+        self._ensure_active(fencing_token)
+        await self._store._renew(self._session_id, self._lease_value)
+
+    async def load(self, fencing_token: int) -> SessionRecord | None:
+        self._ensure_active(fencing_token)
+        return await self._store._load(self._session_id, self._lease_value)
+
+    async def save(self, record: SessionRecord, fencing_token: int) -> SessionRecord:
+        self._ensure_active(fencing_token)
+        if record.session_id != self._session_id:
+            raise AppError(CodigoError.SESION_INVALIDA, "session_id no coincide con el lease")
+        return await self._store._save(record, self._lease_value)
+
+    def close(self) -> None:
+        self._active = False
+
+    @property
+    def _lease_value(self) -> str:
+        return f"{self._owner}:{self._fencing_token}"
+
+    def _ensure_active(self, fencing_token: int) -> None:
+        if not self._active or fencing_token != self._fencing_token:
+            raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+
+class RedisSessionStore:
+    """Store persistente con lease renovable y CAS atómico en Redis."""
+
+    def __init__(
+        self,
+        url: str | None = None,
+        *,
+        client: Redis | None = None,
+        key_prefix: str = "cortex:sessions",
+        ttl_seconds: float = 86_400,
+        lease_seconds: float = 30,
+        retry_interval_seconds: float = 0.05,
+        io_timeout_seconds: float = 10,
+    ) -> None:
+        if (url is None) == (client is None):
+            raise AppError(CodigoError.CONFIG_INVALIDA, "indica url o client de Redis")
+        durations = (ttl_seconds, lease_seconds, retry_interval_seconds, io_timeout_seconds)
+        if not key_prefix or min(durations) <= 0:
+            raise AppError(CodigoError.CONFIG_INVALIDA, "configuración de Redis inválida")
+
+        self._client = client or Redis.from_url(cast(str, url), decode_responses=False)
+        self._owns_client = client is None
+        self._key_prefix = key_prefix.rstrip(":")
+        self._ttl_milliseconds = _milliseconds(ttl_seconds)
+        self._lease_milliseconds = _milliseconds(lease_seconds)
+        self._retry_interval_seconds = retry_interval_seconds
+        self._io_timeout_seconds = io_timeout_seconds
+        self._codec = SessionCodec()
+        self._closed = False
+
+    @property
+    def renewal_interval_seconds(self) -> float:
+        return self._lease_milliseconds / 3_000
+
+    async def __aenter__(self) -> Self:
+        self._ensure_open()
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        await self.aclose()
+
+    @asynccontextmanager
+    async def acquire(
+        self,
+        session_id: str,
+        timeout_seconds: float,
+    ) -> AsyncIterator[SessionLease]:
+        self._ensure_open()
+        if not session_id or timeout_seconds <= 0:
+            raise AppError(CodigoError.CONFIG_INVALIDA, "adquisición de sesión inválida")
+
+        owner = uuid4().hex
+        fencing_token = await self._acquire(session_id, owner, timeout_seconds)
+        lease = _RedisLease(self, session_id, owner, fencing_token)
+        try:
+            yield lease
+        finally:
+            lease.close()
+            with suppress(AppError):
+                await self._release(session_id, f"{owner}:{fencing_token}")
+
+    async def reset(self, session_id: str, timeout_seconds: float = 5.0) -> None:
+        async with self.acquire(session_id, timeout_seconds) as lease:
+            result = await self._eval(
+                scripts.RESET,
+                self._keys(session_id)[:2],
+                self._lease_value(lease),
+            )
+            if int(cast(int, result)) != 1:
+                raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+    async def aclose(self) -> None:
+        if self._closed:
+            return
+        if self._owns_client:
+            try:
+                async with asyncio.timeout(self._io_timeout_seconds):
+                    await self._client.aclose()
+            except TimeoutError as error:
+                raise AppError(
+                    CodigoError.SESION_STORE_FALLO,
+                    "Redis no cerró su cliente dentro del timeout",
+                ) from error
+            except RedisError as error:
+                raise AppError(
+                    CodigoError.SESION_STORE_FALLO,
+                    f"Redis lanzó {type(error).__name__} al cerrar",
+                ) from error
+        self._closed = True
+
+    async def _acquire(self, session_id: str, owner: str, timeout_seconds: float) -> int:
+        lease_key, _, fence_key = self._keys(session_id)
+        deadline = time.monotonic() + timeout_seconds
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise AppError(CodigoError.SESION_OCUPADA, f"sesión ocupada: {session_id}")
+            result = await self._eval(
+                scripts.ACQUIRE,
+                (lease_key, fence_key),
+                owner,
+                self._lease_milliseconds,
+                timeout_seconds=remaining,
+            )
+            fencing_token = int(cast(int, result))
+            if fencing_token > 0:
+                return fencing_token
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise AppError(CodigoError.SESION_OCUPADA, f"sesión ocupada: {session_id}")
+            await asyncio.sleep(min(self._retry_interval_seconds, remaining))
+
+    async def _verify(self, session_id: str, lease_value: str) -> None:
+        lease_key, _, _ = self._keys(session_id)
+        result = await self._eval(scripts.VERIFY, (lease_key,), lease_value)
+        if int(cast(int, result)) != 1:
+            raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+    async def _renew(self, session_id: str, lease_value: str) -> None:
+        lease_key, _, _ = self._keys(session_id)
+        result = await self._eval(
+            scripts.RENEW,
+            (lease_key,),
+            lease_value,
+            self._lease_milliseconds,
+        )
+        if int(cast(int, result)) != 1:
+            raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+    async def _load(self, session_id: str, lease_value: str) -> SessionRecord | None:
+        lease_key, session_key, _ = self._keys(session_id)
+        raw = await self._eval(
+            scripts.LOAD,
+            (lease_key, session_key),
+            lease_value,
+            self._ttl_milliseconds,
+        )
+        result = cast(list[object], raw)
+        if not result or int(cast(int, result[0])) != 1:
+            raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+        if len(result) == 1:
+            return None
+        return self._codec.decode(cast(bytes, result[1]))
+
+    async def _save(self, record: SessionRecord, lease_value: str) -> SessionRecord:
+        saved = record.model_copy(
+            update={
+                "version": record.version + 1,
+                "updated_at": datetime.now(UTC),
+            }
+        )
+        lease_key, session_key, _ = self._keys(record.session_id)
+        result = await self._eval(
+            scripts.SAVE,
+            (lease_key, session_key),
+            lease_value,
+            record.version,
+            saved.version,
+            self._codec.encode(saved),
+            self._ttl_milliseconds,
+        )
+        status = int(cast(int, result))
+        if status == 1:
+            return saved
+        if status == -1:
+            raise AppError(CodigoError.SESION_CONFLICTO, "la versión de sesión cambió")
+        raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+    async def _release(self, session_id: str, lease_value: str) -> None:
+        lease_key, _, _ = self._keys(session_id)
+        await self._eval(scripts.RELEASE, (lease_key,), lease_value)
+
+    async def _eval(
+        self,
+        script: str,
+        keys: tuple[str, ...],
+        *arguments: _RedisArgument,
+        timeout_seconds: float | None = None,
+    ) -> object:
+        self._ensure_open()
+        timeout = timeout_seconds or self._io_timeout_seconds
+        try:
+            async with asyncio.timeout(min(timeout, self._io_timeout_seconds)):
+                return await self._client.eval(script, len(keys), *keys, *arguments)
+        except TimeoutError as error:
+            raise AppError(
+                CodigoError.SESION_STORE_FALLO,
+                "Redis agotó el timeout de I/O",
+            ) from error
+        except RedisError as error:
+            raise AppError(
+                CodigoError.SESION_STORE_FALLO,
+                f"Redis lanzó {type(error).__name__}",
+            ) from error
+
+    def _keys(self, session_id: str) -> tuple[str, str, str]:
+        digest = sha256(session_id.encode()).hexdigest()
+        base = f"{self._key_prefix}:{{{digest}}}"
+        return f"{base}:lease", f"{base}:record", f"{base}:fence"
+
+    def _lease_value(self, lease: SessionLease) -> str:
+        if not isinstance(lease, _RedisLease):
+            raise AppError(CodigoError.SESION_INVALIDA, "lease de Redis inválido")
+        return lease._lease_value
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise AppError(CodigoError.RECURSO_CERRADO, "store Redis cerrado")
+
+
+def _milliseconds(seconds: float) -> int:
+    value = int(seconds * 1_000)
+    if value < 1:
+        raise AppError(CodigoError.CONFIG_INVALIDA, "duración Redis demasiado pequeña")
+    return value
+

cortex_agent_sdk/results.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from enum import StrEnum
+
+from pydantic import BaseModel
+
+from cortex_agent_sdk.engine import Usage
+from cortex_agent_sdk.history.models import Turn
+
+
+class AgentExitReason(StrEnum):
+    COMPLETED = "completed"
+    FINAL_ANSWER = "final_answer"
+    MAX_STEPS = "max_steps"
+    TOOL_FAILURE_LIMIT = "tool_failure_limit"
+
+
+@dataclass(frozen=True, slots=True)
+class AgentResult:
+    text: str | None
+    reason: AgentExitReason
+    usage: Usage
+    steps: int
+    tool_calls: int
+    provider: str
+    model: str
+    history: tuple[Turn, ...]
+    raw_responses: tuple[object, ...]
+    structured: BaseModel | None = None
+

cortex_agent_sdk/runtime.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from cortex_agent_sdk.engine import Usage
+from cortex_agent_sdk.history.models import Turn
+from cortex_agent_sdk.sessions.lease import LeaseKeeper
+from cortex_agent_sdk.sessions.models import SessionRecord
+
+
+class AgentOptions(BaseModel):
+    """Límites del loop, hooks y sesiones."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    max_steps: int = Field(default=8, ge=1)
+    max_history_turns: int = Field(default=20, ge=1)
+    max_consecutive_tool_failures: int = Field(default=3, ge=1)
+    provider_timeout_seconds: float = Field(default=120.0, gt=0)
+    tool_timeout_seconds: float = Field(default=30.0, gt=0)
+    hook_timeout_seconds: float = Field(default=10.0, gt=0)
+    session_lock_timeout_seconds: float = Field(default=5.0, gt=0)
+    shutdown_timeout_seconds: float = Field(default=120.0, gt=0)
+
+
+@dataclass(slots=True)
+class RunState:
+    history: list[Turn]
+    record: SessionRecord | None
+    lease_keeper: LeaseKeeper | None
+    instructions: str
+    usage: Usage = field(default_factory=Usage)
+    raw_responses: list[object] = field(default_factory=list)
+    tool_calls: int = 0
+    consecutive_tool_failures: int = 0
+    last_text: str | None = None
+    structured: BaseModel | None = None
+    effective_model: str | None = None
+

cortex_agent_sdk/sessions/__init__.py

@@ -0,0 +1,4 @@
+from cortex_agent_sdk.sessions.codec import SessionCodec
+from cortex_agent_sdk.sessions.memory import MemorySessionStore
+from cortex_agent_sdk.sessions.models import SessionRecord
+from cortex_agent_sdk.sessions.store import SessionLease, SessionStore

cortex_agent_sdk/sessions/codec.py

@@ -0,0 +1,27 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+from cortex_agent_sdk.errores import AppError, CodigoError
+from cortex_agent_sdk.sessions.models import SessionRecord
+
+
+class _SessionEnvelope(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    format_version: Literal[1] = 1
+    record: SessionRecord
+
+
+class SessionCodec:
+    """Serializa sesiones a JSON validado, sin pickle."""
+
+    def encode(self, record: SessionRecord) -> bytes:
+        return _SessionEnvelope(record=record).model_dump_json().encode()
+
+    def decode(self, payload: bytes) -> SessionRecord:
+        try:
+            return _SessionEnvelope.model_validate_json(payload).record
+        except ValidationError as error:
+            raise AppError(CodigoError.SESION_INVALIDA, "payload de sesión inválido") from error
+

cortex_agent_sdk/sessions/lease.py

@@ -0,0 +1,85 @@
+import asyncio
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+
+from cortex_agent_sdk.errores import AppError, CodigoError
+from cortex_agent_sdk.sessions.models import SessionRecord
+from cortex_agent_sdk.sessions.store import SessionLease
+
+
+class LeaseKeeper:
+    def __init__(self, lease: SessionLease) -> None:
+        self._lease = lease
+        self._fencing_token = lease.fencing_token
+        self._heartbeat: asyncio.Task[None] | None = None
+
+    async def start(self) -> None:
+        await self.verify()
+        interval = self._lease.renewal_interval_seconds
+        if interval is not None and interval <= 0:
+            raise AppError(CodigoError.CONFIG_INVALIDA, "intervalo de renovación inválido")
+        if interval is not None:
+            self._heartbeat = asyncio.create_task(self._renew_loop())
+
+    async def load(self) -> SessionRecord | None:
+        return await self.run(lambda: self._lease.load(self._fencing_token))
+
+    async def save(self, record: SessionRecord) -> SessionRecord:
+        return await self.run(lambda: self._lease.save(record, self._fencing_token))
+
+    async def verify(self) -> None:
+        self._raise_heartbeat_failure()
+        await self._lease.verify_ownership(self._fencing_token)
+
+    async def run[T](self, operation: Callable[[], Awaitable[T]]) -> T:
+        await self.verify()
+        if self._heartbeat is None:
+            return await operation()
+
+        operation_task = asyncio.ensure_future(operation())
+        try:
+            done, _ = await asyncio.wait(
+                (operation_task, self._heartbeat),
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            if self._heartbeat in done:
+                operation_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await operation_task
+                self._raise_heartbeat_failure()
+            return await operation_task
+        finally:
+            if not operation_task.done():
+                operation_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await operation_task
+
+    async def aclose(self) -> None:
+        if self._heartbeat is None:
+            return
+        self._heartbeat.cancel()
+        with suppress(asyncio.CancelledError, AppError):
+            await self._heartbeat
+
+    async def _renew_loop(self) -> None:
+        interval = self._lease.renewal_interval_seconds
+        if interval is None:
+            return
+        try:
+            while True:
+                await asyncio.sleep(interval)
+                await self._lease.renew(self._fencing_token)
+        except asyncio.CancelledError:
+            raise
+        except AppError:
+            raise
+        except Exception as error:
+            raise AppError(
+                CodigoError.SESION_LEASE_PERDIDO,
+                f"renovación lanzó {type(error).__name__}",
+            ) from error
+
+    def _raise_heartbeat_failure(self) -> None:
+        if self._heartbeat is not None and self._heartbeat.done():
+            self._heartbeat.result()
+

cortex_agent_sdk/sessions/memory.py

@@ -0,0 +1,186 @@
+import asyncio
+import time
+from collections import OrderedDict
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Self
+
+from cortex_agent_sdk.errores import AppError, CodigoError
+from cortex_agent_sdk.sessions.models import SessionRecord
+from cortex_agent_sdk.sessions.store import SessionLease
+
+
+@dataclass(slots=True)
+class _Entry:
+    record: SessionRecord
+    expires_at: float
+
+
+@dataclass(slots=True)
+class _SessionLock:
+    lock: asyncio.Lock
+    users: int = 0
+
+
+class _MemoryLease:
+    def __init__(self, store: "MemorySessionStore", session_id: str, fencing_token: int) -> None:
+        self._store = store
+        self._session_id = session_id
+        self._fencing_token = fencing_token
+        self._active = True
+
+    @property
+    def session_id(self) -> str:
+        return self._session_id
+
+    @property
+    def fencing_token(self) -> int:
+        return self._fencing_token
+
+    @property
+    def renewal_interval_seconds(self) -> float | None:
+        return None
+
+    async def verify_ownership(self, fencing_token: int) -> None:
+        self._ensure_active(fencing_token)
+        self._store._ensure_open()
+
+    async def renew(self, fencing_token: int) -> None:
+        await self.verify_ownership(fencing_token)
+
+    async def load(self, fencing_token: int) -> SessionRecord | None:
+        await self.verify_ownership(fencing_token)
+        return await self._store._load(self._session_id)
+
+    async def save(self, record: SessionRecord, fencing_token: int) -> SessionRecord:
+        await self.verify_ownership(fencing_token)
+        if record.session_id != self._session_id:
+            raise AppError(CodigoError.SESION_INVALIDA, "session_id no coincide con el lease")
+        return await self._store._save(record)
+
+    def close(self) -> None:
+        self._active = False
+
+    def _ensure_active(self, fencing_token: int) -> None:
+        if not self._active or fencing_token != self._fencing_token:
+            raise AppError(CodigoError.SESION_LEASE_PERDIDO, "ownership de sesión perdido")
+
+
+class MemorySessionStore:
+    """Store local con exclusión por sesión, CAS, TTL y LRU."""
+
+    def __init__(self, max_sessions: int = 1_000, ttl_seconds: float = 3_600) -> None:
+        if max_sessions < 1 or ttl_seconds <= 0:
+            raise AppError(CodigoError.CONFIG_INVALIDA, "límites de sesión inválidos")
+        self._max_sessions = max_sessions
+        self._ttl_seconds = ttl_seconds
+        self._entries: OrderedDict[str, _Entry] = OrderedDict()
+        self._locks: dict[str, _SessionLock] = {}
+        self._guard = asyncio.Lock()
+        self._closed = False
+        self._next_fencing_token = 1
+
+    async def __aenter__(self) -> Self:
+        self._ensure_open()
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        await self.aclose()
+
+    @asynccontextmanager
+    async def acquire(
+        self,
+        session_id: str,
+        timeout_seconds: float,
+    ) -> AsyncIterator[SessionLease]:
+        self._ensure_open()
+        if not session_id or timeout_seconds <= 0:
+            raise AppError(CodigoError.CONFIG_INVALIDA, "adquisición de sesión inválida")
+
+        async with self._guard:
+            lock_entry = self._locks.setdefault(session_id, _SessionLock(asyncio.Lock()))
+            lock_entry.users += 1
+            fencing_token = self._next_fencing_token
+            self._next_fencing_token += 1
+
+        try:
+            await asyncio.wait_for(lock_entry.lock.acquire(), timeout_seconds)
+        except TimeoutError as error:
+            await self._release_lock_reference(session_id, lock_entry)
+            raise AppError(CodigoError.SESION_OCUPADA, f"sesión ocupada: {session_id}") from error
+        except BaseException:
+            await self._release_lock_reference(session_id, lock_entry)
+            raise
+
+        lease = _MemoryLease(self, session_id, fencing_token)
+        try:
+            yield lease
+        finally:
+            lease.close()
+            lock_entry.lock.release()
+            await self._release_lock_reference(session_id, lock_entry)
+
+    async def reset(self, session_id: str, timeout_seconds: float = 5.0) -> None:
+        async with self.acquire(session_id, timeout_seconds), self._guard:
+            self._entries.pop(session_id, None)
+
+    async def aclose(self) -> None:
+        async with self._guard:
+            self._closed = True
+            self._entries.clear()
+            self._locks.clear()
+
+    async def _load(self, session_id: str) -> SessionRecord | None:
+        async with self._guard:
+            self._ensure_open()
+            entry = self._entries.get(session_id)
+            if entry is None:
+                return None
+            if entry.expires_at <= time.monotonic():
+                del self._entries[session_id]
+                return None
+            entry.expires_at = time.monotonic() + self._ttl_seconds
+            self._entries.move_to_end(session_id)
+            return entry.record
+
+    async def _save(self, record: SessionRecord) -> SessionRecord:
+        async with self._guard:
+            self._ensure_open()
+            current = self._entries.get(record.session_id)
+            current_version = current.record.version if current else 0
+            if current_version != record.version:
+                raise AppError(
+                    CodigoError.SESION_CONFLICTO,
+                    f"versión esperada {record.version}, vigente {current_version}",
+                )
+
+            saved = record.model_copy(
+                update={
+                    "version": record.version + 1,
+                    "updated_at": datetime.now(UTC),
+                }
+            )
+            expires_at = time.monotonic() + self._ttl_seconds
+            self._entries[record.session_id] = _Entry(saved, expires_at)
+            self._entries.move_to_end(record.session_id)
+            while len(self._entries) > self._max_sessions:
+                self._entries.popitem(last=False)
+            return saved
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise AppError(CodigoError.RECURSO_CERRADO, "store de sesiones cerrado")
+
+    async def _release_lock_reference(
+        self,
+        session_id: str,
+        lock_entry: _SessionLock,
+    ) -> None:
+        async with self._guard:
+            lock_entry.users -= 1
+            current = self._locks.get(session_id)
+            if current is lock_entry and lock_entry.users == 0:
+                del self._locks[session_id]
+