agentengine-sdk-python 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

ksadk/sessions/resilience.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Iterator
+
+from ksadk.sessions.errors import SessionBackendUnavailable
+
+logger = logging.getLogger(__name__)
+
+
+def session_backend_timeout_seconds() -> float:
+    raw = (
+        os.getenv("KSADK_SESSION_CONNECT_TIMEOUT")
+        or os.getenv("KSADK_SESSION_PG_CONNECT_TIMEOUT")
+        or ""
+    ).strip()
+    if not raw:
+        return 5.0
+    try:
+        value = float(raw)
+    except ValueError:
+        logger.warning("Invalid KSADK_SESSION_CONNECT_TIMEOUT=%r; using 5 seconds", raw)
+        return 5.0
+    return max(0.1, value)
+
+
+def _exception_chain(exc: BaseException) -> Iterator[BaseException]:
+    seen: set[int] = set()
+    current: BaseException | None = exc
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        yield current
+        original = getattr(current, "orig", None)
+        if isinstance(original, BaseException) and id(original) not in seen:
+            current = original
+            continue
+        current = current.__cause__ or current.__context__
+
+
+def is_session_backend_failure(exc: BaseException) -> bool:
+    """Return whether an exception represents an unavailable database backend."""
+    for current in _exception_chain(exc):
+        if isinstance(
+            current,
+            (SessionBackendUnavailable, TimeoutError, OSError, ConnectionError),
+        ):
+            return True
+        module = type(current).__module__
+        if module.startswith("asyncpg.") or module.startswith("sqlalchemy."):
+            return True
+    return False

ksadk/sessions/resilient.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Optional, cast
+
+from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState
+from ksadk.sessions.in_memory import InMemorySessionService
+from ksadk.sessions.resilience import is_session_backend_failure
+
+logger = logging.getLogger(__name__)
+
+
+class ResilientSessionService(BaseSessionService):
+    """Keep live agent sessions available when durable persistence is unavailable.
+
+    The in-memory service is authoritative for the lifetime of this process. The
+    configured durable service is used as a read-through source and a best-effort
+    write-through sink. After its first failure it stays disabled until a
+    background probe confirms the durable backend is reachable again, at which
+    point it is re-enabled and an INFO log is emitted.
+    """
+
+    _probe_interval_seconds: float = 30.0
+
+    def __init__(
+        self,
+        primary: BaseSessionService,
+        fallback: InMemorySessionService | None = None,
+    ) -> None:
+        self.primary = primary
+        self.fallback = fallback or InMemorySessionService()
+        self._primary_enabled = True
+        self._hydrate_lock = asyncio.Lock()
+        self._primary_session_lock = asyncio.Lock()
+        self._primary_session_ids: set[str] = set()
+        self._probe_task: asyncio.Task[None] | None = None
+
+    @property
+    def degraded(self) -> bool:
+        return not self._primary_enabled
+
+    async def _call_primary(self, method_name: str, *args: Any, **kwargs: Any) -> tuple[bool, Any]:
+        if not self._primary_enabled:
+            return False, None
+        try:
+            method = getattr(self.primary, method_name)
+            return True, await method(*args, **kwargs)
+        except Exception as exc:
+            if not is_session_backend_failure(exc):
+                raise
+            self._disable_primary(exc)
+            return False, None
+
+    def _disable_primary(self, exc: Exception) -> None:
+        if not self._primary_enabled:
+            return
+        self._primary_enabled = False
+        logger.error(
+            "KSADK session persistence degraded; using in-memory live session: %s",
+            exc,
+            extra={
+                "session_backend_state": "degraded",
+                "session_backend": type(self.primary).__name__,
+            },
+        )
+        self._start_probe()
+
+    def _start_probe(self) -> None:
+        if self._probe_task is not None and not self._probe_task.done():
+            return
+        self._probe_task = asyncio.create_task(self._probe_loop())
+
+    async def _probe_loop(self) -> None:
+        while not self._primary_enabled:
+            await asyncio.sleep(self._probe_interval_seconds)
+            if self._primary_enabled:
+                break
+            try:
+                await self.primary.get_session("__ksadk_probe__")
+            except Exception:
+                continue
+            self._primary_enabled = True
+            logger.info(
+                "KSADK session persistence recovered; durable backend re-enabled",
+                extra={
+                    "session_backend_state": "recovered",
+                    "session_backend": type(self.primary).__name__,
+                },
+            )
+
+    async def _hydrate(self, session: Session) -> Session:
+        async with self._hydrate_lock:
+            self._primary_session_ids.add(session.id)
+            existing = await self.fallback.get_session(session.id)
+            if existing is None:
+                await self.fallback.create_session(
+                    session.agent_id,
+                    session.user_id,
+                    session_id=session.id,
+                )
+            existing_events = await self.fallback.get_events(session.id)
+            existing_ids = {event.id for event in existing_events}
+            for event in sorted(session.events, key=lambda item: item.seq_id):
+                if event.id not in existing_ids:
+                    await self.fallback.append_event(session.id, event)
+            await self.fallback.update_session_metadata(
+                session.id,
+                title=session.title,
+                title_source=session.title_source,
+                summary=session.summary,
+                first_prompt=session.first_prompt,
+                last_prompt=session.last_prompt,
+            )
+            current = await self.fallback.get_session(session.id)
+            if current is not None and session.state != current.state:
+                await self.fallback.update_state(
+                    agent_id=session.agent_id,
+                    user_id=session.user_id,
+                    session_id=session.id,
+                    scope="session",
+                    state_delta=session.state,
+                )
+            hydrated = await self.fallback.get_session(session.id)
+            if hydrated is None:
+                raise RuntimeError(f"Failed to hydrate live session {session.id}")
+            return hydrated
+
+    async def create_session(
+        self,
+        agent_id: str,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        if session_id:
+            ok, durable = await self._call_primary("get_session", session_id)
+            if ok and durable is not None:
+                return await self._hydrate(durable)
+            existing = await self.fallback.get_session(session_id)
+            if existing is not None:
+                return existing
+
+        live = await self.fallback.create_session(agent_id, user_id, session_id=session_id)
+        ok, durable = await self._call_primary(
+            "create_session",
+            agent_id,
+            user_id,
+            session_id=live.id,
+        )
+        if ok and durable is not None:
+            self._primary_session_ids.add(durable.id)
+            return await self._hydrate(durable)
+        return live
+
+    async def get_session(self, session_id: str) -> Optional[Session]:
+        live = await self.fallback.get_session(session_id)
+        ok, durable = await self._call_primary("get_session", session_id)
+        if ok and durable is not None:
+            return await self._hydrate(durable)
+        return live
+
+    async def list_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> list[Session]:
+        ok, durable_sessions = await self._call_primary(
+            "list_sessions",
+            agent_id,
+            user_id,
+            offset,
+            limit,
+        )
+        if ok:
+            for session in durable_sessions or []:
+                await self._hydrate(session)
+        return cast(
+            list[Session],
+            await self.fallback.list_sessions(agent_id, user_id, offset, limit),
+        )
+
+    async def count_sessions(self, agent_id: str, user_id: Optional[str] = None) -> int:
+        sessions = await self.list_sessions(agent_id, user_id)
+        return len(sessions)
+
+    async def delete_session(self, session_id: str) -> bool:
+        deleted = await self.fallback.delete_session(session_id)
+        ok, durable_deleted = await self._call_primary("delete_session", session_id)
+        if ok:
+            self._primary_session_ids.discard(session_id)
+        return deleted or bool(durable_deleted) if ok else deleted
+
+    async def update_session_metadata(
+        self,
+        session_id: str,
+        *,
+        title: Optional[str] = None,
+        title_source: Optional[str] = None,
+        summary: Optional[str] = None,
+        first_prompt: Optional[str] = None,
+        last_prompt: Optional[str] = None,
+    ) -> Session:
+        if await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.update_session_metadata(
+            session_id,
+            title=title,
+            title_source=title_source,
+            summary=summary,
+            first_prompt=first_prompt,
+            last_prompt=last_prompt,
+        )
+        await self._ensure_primary_session(session_id)
+        await self._call_primary(
+            "update_session_metadata",
+            session_id,
+            title=title,
+            title_source=title_source,
+            summary=summary,
+            first_prompt=first_prompt,
+            last_prompt=last_prompt,
+        )
+        return live
+
+    async def _ensure_primary_session(self, session_id: str) -> None:
+        """Create the session in PG if it only exists in memory (degraded-era)."""
+        if not self._primary_enabled or session_id in self._primary_session_ids:
+            return
+        async with self._primary_session_lock:
+            if not self._primary_enabled or session_id in self._primary_session_ids:
+                return
+            ok, durable = await self._call_primary("get_session", session_id)
+            if not ok:
+                return
+            if durable is None:
+                live = await self.fallback.get_session(session_id)
+                if live is None:
+                    return
+                ok, durable = await self._call_primary(
+                    "create_session",
+                    live.agent_id,
+                    live.user_id,
+                    session_id=live.id,
+                )
+                if not ok or durable is None:
+                    return
+            self._primary_session_ids.add(session_id)
+
+    async def append_event(self, session_id: str, event: SessionEvent) -> SessionEvent:
+        if await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.append_event(session_id, event)
+        await self._ensure_primary_session(session_id)
+        await self._call_primary("append_event", session_id, event)
+        return live
+
+    async def get_events(
+        self,
+        session_id: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        await self.get_session(session_id)
+        return cast(
+            list[SessionEvent],
+            await self.fallback.get_events(
+                session_id,
+                offset,
+                limit,
+                after_seq_id,
+                before_seq_id,
+            ),
+        )
+
+    async def count_events(
+        self,
+        session_id: str,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> int:
+        await self.get_session(session_id)
+        return cast(
+            int,
+            await self.fallback.count_events(session_id, after_seq_id, before_seq_id),
+        )
+
+    async def get_state(
+        self,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str = "session",
+    ) -> Optional[SessionState]:
+        if session_id:
+            await self.get_session(session_id)
+        live = await self.fallback.get_state(agent_id, user_id, session_id, scope)
+        if live is not None or not self._primary_enabled:
+            return live
+        ok, durable = await self._call_primary(
+            "get_state",
+            agent_id,
+            user_id,
+            session_id,
+            scope,
+        )
+        if ok and durable is not None:
+            return await self.fallback.update_state(
+                agent_id=agent_id,
+                user_id=user_id,
+                session_id=session_id,
+                scope=scope,
+                state_delta=durable.state,
+            )
+        return live
+
+    async def update_state(
+        self,
+        *,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str,
+        state_delta: dict[str, Any],
+    ) -> SessionState:
+        if session_id and await self.fallback.get_session(session_id) is None:
+            await self.get_session(session_id)
+        live = await self.fallback.update_state(
+            agent_id=agent_id,
+            user_id=user_id,
+            session_id=session_id,
+            scope=scope,
+            state_delta=state_delta,
+        )
+        if session_id:
+            await self._ensure_primary_session(session_id)
+        await self._call_primary(
+            "update_state",
+            agent_id=agent_id,
+            user_id=user_id,
+            session_id=session_id,
+            scope=scope,
+            state_delta=state_delta,
+        )
+        return live
+
+    async def aclose(self) -> None:
+        if self._probe_task is not None and not self._probe_task.done():
+            self._probe_task.cancel()
+            try:
+                await self._probe_task
+            except asyncio.CancelledError:
+                pass
+        for service in (self.primary, self.fallback):
+            close = getattr(service, "aclose", None)
+            if close is not None:
+                await close()

ksadk/sessions/sqlite_service.py

@@ -0,0 +1,13 @@
+"""Compatibility wrapper for the legacy sqlite-backed local session service path."""
+
+from ksadk.sessions.local_service import (
+    LocalSessionService,
+    create_local_session_service,
+    resolve_local_session_dir,
+)
+
+__all__ = [
+    "LocalSessionService",
+    "create_local_session_service",
+    "resolve_local_session_dir",
+]

ksadk/skills/__init__.py

@@ -0,0 +1,9 @@
+"""Skill Center runtime consumption helpers."""
+
+from ksadk.skills.models import ContentHash, SkillListResponse, SkillRef
+
+__all__ = [
+    "ContentHash",
+    "SkillListResponse",
+    "SkillRef",
+]

ksadk/skills/loader.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class LocalSkill:
+    name: str
+    description: str
+    root_dir: Path
+    skill_md: Path
+    body: str
+
+
+def load_local_skill(root_dir: str | Path) -> LocalSkill:
+    root = Path(root_dir)
+    skill_md = root / "SKILL.md"
+    if not skill_md.exists():
+        raise FileNotFoundError(f"SKILL.md not found: {skill_md}")
+    content = skill_md.read_text(encoding="utf-8")
+    metadata, body = _parse_frontmatter(content)
+    return LocalSkill(
+        name=metadata.get("name") or root.name,
+        description=metadata.get("description") or "",
+        root_dir=root,
+        skill_md=skill_md,
+        body=body,
+    )
+
+
+def _parse_frontmatter(content: str) -> tuple[dict[str, str], str]:
+    if not content.startswith("---\n"):
+        return {}, content
+    end = content.find("\n---\n", 4)
+    if end == -1:
+        return {}, content
+    raw_meta = content[4:end]
+    body = content[end + 5 :]
+    metadata: dict[str, str] = {}
+    for line in raw_meta.splitlines():
+        if ":" not in line:
+            continue
+        key, value = line.split(":", 1)
+        metadata[key.strip()] = value.strip().strip("\"'")
+    return metadata, body

ksadk/skills/models.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ContentHash:
+    algorithm: str
+    value: str
+
+    @classmethod
+    def parse(cls, raw: str | None) -> "ContentHash | None":
+        normalized = str(raw or "").strip()
+        if not normalized:
+            return None
+        if ":" not in normalized:
+            return cls("sha256", normalized.lower())
+        algorithm, value = normalized.split(":", 1)
+        return cls(algorithm.strip().lower(), value.strip().lower())
+
+    def render(self) -> str:
+        return f"{self.algorithm}:{self.value}"
+
+
+@dataclass(frozen=True)
+class SkillRef:
+    skill_id: str
+    version_id: str
+    version: str
+    name: str
+    description: str = ""
+    status: str = ""
+    content_hash: ContentHash | None = None
+    archive_uri: str = ""
+    aliases: tuple[str, ...] = ()
+    tags: tuple[str, ...] = ()
+    examples: tuple[str, ...] = ()
+    input_schema: dict[str, Any] | None = None
+    runtime_requirements: dict[str, Any] | None = None
+
+    @classmethod
+    def from_payload(cls, payload: dict[str, Any]) -> "SkillRef":
+        return cls(
+            skill_id=str(payload.get("SkillId") or payload.get("skill_id") or ""),
+            version_id=str(payload.get("VersionId") or payload.get("version_id") or ""),
+            version=str(payload.get("Version") or payload.get("version") or ""),
+            name=str(payload.get("Name") or payload.get("name") or ""),
+            description=str(payload.get("Description") or payload.get("description") or ""),
+            status=str(payload.get("Status") or payload.get("status") or ""),
+            content_hash=ContentHash.parse(payload.get("ContentHash") or payload.get("content_hash")),
+            archive_uri=str(payload.get("ArchiveUri") or payload.get("archive_uri") or ""),
+            aliases=_string_tuple(payload.get("Aliases") or payload.get("aliases")),
+            tags=_string_tuple(payload.get("Tags") or payload.get("tags")),
+            examples=_string_tuple(payload.get("Examples") or payload.get("examples")),
+            input_schema=_dict_or_none(payload.get("InputSchema") or payload.get("input_schema")),
+            runtime_requirements=_dict_or_none(
+                payload.get("RuntimeRequirements") or payload.get("runtime_requirements")
+            ),
+        )
+
+    @property
+    def cache_key(self) -> str:
+        identity = self.skill_id or self.name
+        version_key = self.version_id or self.version
+        if not version_key and self.content_hash:
+            version_key = self.content_hash.value
+        return "__".join(part.replace("/", "_").replace(":", "_") for part in (identity, version_key) if part)
+
+    @property
+    def is_active(self) -> bool:
+        return not self.status or self.status.lower() in {"active", "available", "enabled"}
+
+
+@dataclass(frozen=True)
+class SkillListResponse:
+    request_id: str
+    space_id: str
+    space_name: str
+    skills: list[SkillRef]
+    code: int = 0
+    message: str = ""
+
+    @classmethod
+    def from_payload(
+        cls,
+        payload: dict[str, Any],
+        *,
+        space_id: str = "",
+        space_name: str = "",
+    ) -> "SkillListResponse":
+        data = payload.get("Data") or payload.get("data") or {}
+        skills_payload = data.get("Skills") or data.get("skills") or data.get("Items") or data.get("items") or []
+        return cls(
+            code=int(payload.get("Code") or payload.get("code") or 0),
+            message=str(payload.get("Message") or payload.get("message") or ""),
+            request_id=str(payload.get("RequestId") or payload.get("request_id") or ""),
+            space_id=str(data.get("SkillSpaceId") or data.get("skill_space_id") or space_id),
+            space_name=str(data.get("SkillSpaceName") or data.get("skill_space_name") or space_name),
+            skills=[SkillRef.from_payload(item) for item in skills_payload if isinstance(item, dict)],
+        )
+
+    def active_skills(self) -> list[SkillRef]:
+        return [skill for skill in self.skills if skill.is_active]
+
+
+def _string_tuple(raw: Any) -> tuple[str, ...]:
+    if raw is None:
+        return ()
+    if isinstance(raw, str):
+        values = raw.split(",")
+    elif isinstance(raw, (list, tuple, set)):
+        values = raw
+    else:
+        return ()
+    return tuple(value for value in (str(item).strip() for item in values) if value)
+
+
+def _dict_or_none(raw: Any) -> dict[str, Any] | None:
+    return dict(raw) if isinstance(raw, dict) else None