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

ksadk/sessions/postgres_service.py

@@ -10,7 +10,14 @@ from typing import Any, Optional
 from urllib.parse import urlsplit, urlunsplit
 
 from ksadk.ids import new_session_id
+from ksadk.sessions._postgres_schema import _PostgresSchemaMixin
+from ksadk.sessions._postgres_tables import (
+    KSADK_PG_EVENTS_TABLE,
+    KSADK_PG_SESSIONS_TABLE,
+    KSADK_PG_STATES_TABLE,
+)
 from ksadk.sessions.base import (
+    CANONICAL_EVENT_STORAGE_CAPABILITIES,
     BaseSessionService,
     Session,
     SessionEvent,
@@ -19,15 +26,12 @@ from ksadk.sessions.base import (
 )
 from ksadk.sessions.errors import SessionBackendUnavailable
 
-KSADK_PG_SESSIONS_TABLE = "ksadk_sessions"
-KSADK_PG_EVENTS_TABLE = "ksadk_events"
-KSADK_PG_STATES_TABLE = "ksadk_states"
-PG_READABLE_EVENTS_VIEW = "ksadk_session_events_readable"
-
 logger = logging.getLogger(__name__)
 
 
-class PostgresSessionService(BaseSessionService):
+class PostgresSessionService(_PostgresSchemaMixin, BaseSessionService):
+    storage_capabilities = CANONICAL_EVENT_STORAGE_CAPABILITIES
+
     def __init__(
         self,
         *,
@@ -304,7 +308,9 @@ class PostgresSessionService(BaseSessionService):
                     seq_id=int(next_seq or 1),
                     invocation_id=event.invocation_id,
                     metadata=dict(event.metadata),
+                    seq_binding=event.seq_binding,
                 )
+                stored.bind_seq_id(int(next_seq or 1))
                 await connection.execute(
                     f"""
                     INSERT INTO {KSADK_PG_EVENTS_TABLE} (
@@ -349,6 +355,52 @@ class PostgresSessionService(BaseSessionService):
                 )
                 return stored
 
+    async def get_event_by_id(self, session_id: str, event_id: str) -> Optional[SessionEvent]:
+        await self._ensure_schema()
+        async with self._pool.acquire() as connection:
+            row = await connection.fetchrow(
+                f"""
+                SELECT id, session_id, author, event_type, content_json, timestamp,
+                       state_delta_json, seq_id, invocation_id, metadata_json
+                FROM {KSADK_PG_EVENTS_TABLE}
+                WHERE namespace = $1 AND session_id = $2 AND id = $3
+                """,
+                self.namespace,
+                session_id,
+                event_id,
+            )
+            return self._event_from_row(row) if row is not None else None
+
+    async def get_events_by_invocation_id(
+        self,
+        session_id: str,
+        invocation_id: str,
+        *,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        await self._ensure_schema()
+        conditions = ["namespace = $1", "session_id = $2", "invocation_id = $3"]
+        params: list[Any] = [self.namespace, session_id, invocation_id]
+        if after_seq_id is not None:
+            params.append(after_seq_id)
+            conditions.append(f"seq_id > ${len(params)}")
+        if before_seq_id is not None:
+            params.append(before_seq_id)
+            conditions.append(f"seq_id < ${len(params)}")
+        async with self._pool.acquire() as connection:
+            rows = await connection.fetch(
+                f"""
+                SELECT id, session_id, author, event_type, content_json, timestamp,
+                       state_delta_json, seq_id, invocation_id, metadata_json
+                FROM {KSADK_PG_EVENTS_TABLE}
+                WHERE {" AND ".join(conditions)}
+                ORDER BY seq_id ASC
+                """,
+                *params,
+            )
+            return [self._event_from_row(row) for row in rows]
+
     async def get_events(
         self,
         session_id: str,
@@ -687,148 +739,6 @@ class PostgresSessionService(BaseSessionService):
                     f"could not connect to {mask_postgres_session_dsn(self.dsn)}"
                 ) from exc
 
-    async def _ensure_schema(self) -> None:
-        if self._schema_ready:
-            return
-        async with self._schema_lock:
-            if self._schema_ready:
-                return
-            await self._ensure_pool()
-            async with self._pool.acquire() as connection:
-                await connection.execute(f"""
-                    CREATE TABLE IF NOT EXISTS {KSADK_PG_SESSIONS_TABLE} (
-                        namespace TEXT NOT NULL,
-                        tenant_id TEXT NOT NULL DEFAULT 'default',
-                        workspace_id TEXT NOT NULL DEFAULT 'default',
-                        id TEXT NOT NULL,
-                        agent_id TEXT NOT NULL,
-                        user_id TEXT NOT NULL,
-                        title TEXT NOT NULL DEFAULT '',
-                        title_source TEXT NOT NULL DEFAULT '',
-                        summary TEXT NOT NULL DEFAULT '',
-                        first_prompt TEXT NOT NULL DEFAULT '',
-                        last_prompt TEXT NOT NULL DEFAULT '',
-                        state_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
-                        created_at DOUBLE PRECISION NOT NULL,
-                        updated_at DOUBLE PRECISION NOT NULL,
-                        version INTEGER NOT NULL DEFAULT 0,
-                        PRIMARY KEY (namespace, id)
-                    );
-
-                    CREATE TABLE IF NOT EXISTS {KSADK_PG_EVENTS_TABLE} (
-                        namespace TEXT NOT NULL,
-                        tenant_id TEXT NOT NULL DEFAULT 'default',
-                        workspace_id TEXT NOT NULL DEFAULT 'default',
-                        id TEXT NOT NULL,
-                        session_id TEXT NOT NULL,
-                        author TEXT NOT NULL,
-                        event_type TEXT NOT NULL,
-                        content_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
-                        timestamp DOUBLE PRECISION NOT NULL,
-                        state_delta_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
-                        seq_id INTEGER NOT NULL,
-                        invocation_id TEXT,
-                        metadata_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
-                        PRIMARY KEY (namespace, id),
-                        UNIQUE (namespace, session_id, seq_id),
-                        FOREIGN KEY (namespace, session_id)
-                            REFERENCES {KSADK_PG_SESSIONS_TABLE}(namespace, id)
-                            ON DELETE CASCADE
-                    );
-
-                    CREATE INDEX IF NOT EXISTS idx_ksadk_pg_events_session_seq
-                    ON {KSADK_PG_EVENTS_TABLE} (namespace, session_id, seq_id);
-
-                    -- 跨会话事件查询（get_events_for_agent）需 JOIN sessions 按
-                    -- s.agent_id 过滤并按 e.timestamp 排序；events 表无 agent_id 列，
-                    -- 覆盖索引 (namespace, session_id, timestamp, id) 服务 JOIN 键
-                    -- s.id=e.session_id + ORDER BY e.timestamp DESC。
-                    CREATE INDEX IF NOT EXISTS idx_ksadk_pg_events_session_ts
-                    ON {KSADK_PG_EVENTS_TABLE} (namespace, session_id, timestamp, id);
-
-                    -- ListSessions 归并按 agent_id 过滤 + updated_at DESC 排序；
-                    -- sessions 表 PK 是 (namespace, id)，缺 agent_id 前缀索引。
-                    CREATE INDEX IF NOT EXISTS idx_ksadk_pg_sessions_agent_updated
-                    ON {KSADK_PG_SESSIONS_TABLE} (namespace, agent_id, updated_at DESC, id);
-
-                    CREATE TABLE IF NOT EXISTS {KSADK_PG_STATES_TABLE} (
-                        namespace TEXT NOT NULL,
-                        tenant_id TEXT NOT NULL DEFAULT 'default',
-                        workspace_id TEXT NOT NULL DEFAULT 'default',
-                        scope TEXT NOT NULL,
-                        agent_id TEXT NOT NULL,
-                        user_id TEXT NOT NULL DEFAULT '',
-                        session_id TEXT NOT NULL DEFAULT '',
-                        state_json JSONB NOT NULL DEFAULT '{{}}'::jsonb,
-                        version INTEGER NOT NULL DEFAULT 0,
-                        updated_at DOUBLE PRECISION NOT NULL,
-                        PRIMARY KEY (namespace, scope, agent_id, user_id, session_id)
-                    );
-
-                    ALTER TABLE {KSADK_PG_SESSIONS_TABLE}
-                    ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
-                    ALTER TABLE {KSADK_PG_SESSIONS_TABLE}
-                    ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
-                    ALTER TABLE {KSADK_PG_EVENTS_TABLE}
-                    ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
-                    ALTER TABLE {KSADK_PG_EVENTS_TABLE}
-                    ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
-                    ALTER TABLE {KSADK_PG_STATES_TABLE}
-                    ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
-                    ALTER TABLE {KSADK_PG_STATES_TABLE}
-                    ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT 'default';
-                    """)
-                try:
-                    await connection.execute(f"""
-                    CREATE OR REPLACE VIEW {PG_READABLE_EVENTS_VIEW} AS
-                    SELECT
-                        event_row.namespace,
-                        event_row.tenant_id,
-                        event_row.workspace_id,
-                        session_row.agent_id,
-                        session_row.user_id,
-                        session_row.title AS session_title,
-                        event_row.session_id,
-                        event_row.seq_id,
-                        event_row.id AS event_id,
-                        event_row.invocation_id,
-                        event_row.author,
-                        event_row.event_type,
-                        CASE
-                            WHEN event_row.event_type = 'user_message' THEN 'user'
-                            WHEN event_row.event_type IN (
-                                'assistant_message', 'reasoning', 'tool_call'
-                            ) THEN 'assistant'
-                            WHEN event_row.event_type = 'tool_result' THEN 'tool'
-                            ELSE NULL
-                        END AS message_role,
-                        COALESCE(
-                            NULLIF(event_row.content_json #>> '{{parts,0,text}}', ''),
-                            NULLIF(event_row.content_json ->> 'text', ''),
-                            NULLIF(event_row.metadata_json ->> 'reasoning', ''),
-                            NULLIF(event_row.metadata_json ->> 'tool_output', '')
-                        ) AS message_text,
-                        event_row.metadata_json ->> 'tool_name' AS tool_name,
-                        CASE
-                            WHEN event_row.event_type = 'run_status' THEN COALESCE(
-                                event_row.content_json ->> 'status',
-                                event_row.metadata_json ->> 'status'
-                            )
-                            ELSE NULL
-                        END AS lifecycle_status,
-                        to_timestamp(event_row.timestamp) AS created_at,
-                        event_row.content_json,
-                        event_row.state_delta_json,
-                        event_row.metadata_json
-                    FROM {KSADK_PG_EVENTS_TABLE} AS event_row
-                    JOIN {KSADK_PG_SESSIONS_TABLE} AS session_row
-                      ON session_row.namespace = event_row.namespace
-                     AND session_row.id = event_row.session_id;
-                        """)
-                except Exception as exc:
-                    logger.warning("Postgres readable session view unavailable: %s", exc)
-            self._schema_ready = True
-
     async def _get_session_with_connection(
         self,
         connection: Any,

ksadk/sessions/resilient.py

@@ -4,7 +4,12 @@ import asyncio
 import logging
 from typing import Any, Optional, cast
 
-from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState
+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
 
@@ -40,6 +45,12 @@ class ResilientSessionService(BaseSessionService):
     def degraded(self) -> bool:
         return not self._primary_enabled
 
+    # This service is intentionally live-first and writes two independently
+    # sequenced stores.  Even when both children can atomically bind a local
+    # seq, the wrapper cannot guarantee one shared physical seq/fact across
+    # both writes, so it inherits BaseSessionService's empty canonical storage
+    # capabilities and RuntimeEventStore fails closed before either write.
+
     async def _call_primary(self, method_name: str, *args: Any, **kwargs: Any) -> tuple[bool, Any]:
         if not self._primary_enabled:
             return False, None
@@ -260,6 +271,29 @@ class ResilientSessionService(BaseSessionService):
         await self._call_primary("append_event", session_id, event)
         return live
 
+    async def get_event_by_id(self, session_id: str, event_id: str) -> Optional[SessionEvent]:
+        if await self.fallback.get_session_metadata(session_id) is None:
+            await self.get_session(session_id)
+        return await self.fallback.get_event_by_id(session_id, event_id)
+
+    async def get_events_by_invocation_id(
+        self,
+        session_id: str,
+        invocation_id: str,
+        *,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        # ResilientSessionService is explicitly live-first: hydrate any durable
+        # prefix, then read the indexed in-memory authority used by get_events.
+        await self.get_session(session_id)
+        return await self.fallback.get_events_by_invocation_id(
+            session_id,
+            invocation_id,
+            after_seq_id=after_seq_id,
+            before_seq_id=before_seq_id,
+        )
+
     async def get_events(
         self,
         session_id: str,

ksadk/skills/service_client.py

@@ -119,6 +119,7 @@ class SkillServiceClient:
         download_url = self.get_skill_download_url(skill)
         if not download_url:
             raise ValueError(f"Skill Service did not return DownloadUrl for {skill.skill_id}")
+        download_url = _rewrite_ks3_to_internal(download_url)
         with httpx.Client(**self._client_kwargs()) as client:
             response = client.get(download_url)
             response.raise_for_status()
@@ -268,3 +269,50 @@ def _normalize_base_url(base_url: str) -> str:
     elif path.endswith("/docs"):
         path = path[: -len("/docs")] + "/api/v1"
     return urlunsplit((parsed.scheme, parsed.netloc, path, "", "")).rstrip("/")
+
+
+def _rewrite_ks3_to_internal(url: str) -> str:
+    """Rewrite a KS3 public endpoint to internal only when public is unreachable.
+
+    AICP Skill Service returns pre-signed download URLs with public KS3 domains
+    (e.g. skill.ks3-cn-beijing.ksyuncs.com -> 60.x public IP). On private_only
+    compute pods those public IPs are unreachable. This probes the public domain
+    first; if reachable keeps it, otherwise rewrites to the internal domain
+    (ks3-cn-beijing-internal.ksyuncs.com -> 198.18.96.x).
+    """
+    if not url:
+        return url
+    try:
+        from urllib.parse import urlsplit
+
+        from ksadk.common.constants import get_ks3_endpoints
+
+        region = os.environ.get(
+            "KSADK_SKILL_SERVICE_REGION", "KSYUN_REGION"
+        ) or "cn-beijing-6"
+        public_ep, internal_ep = get_ks3_endpoints(region)
+        if not public_ep or not internal_ep:
+            return url
+        if public_ep not in url:
+            return url
+
+        # Public reachable => keep it.
+        import socket
+
+        parsed = urlsplit(url)
+        host = (parsed.hostname or "").lower()
+        port = parsed.port or (443 if parsed.scheme == "https" else 80)
+        try:
+            s = socket.socket()
+            s.settimeout(1.5)
+            s.connect((host, port))
+            s.close()
+            return url
+        except OSError:
+            pass
+
+        # Public unreachable => rewrite to internal.
+        return url.replace(public_ep, internal_ep)
+    except Exception:
+        pass
+    return url

ksadk/studio/agent_lifecycle.py

@@ -24,28 +24,34 @@ def delete_framework_agent(studio: Any, agent_id: str, *, purge: bool) -> None:
             status_code=409,
             details={"agentId": agent_id, "runIds": running},
         )
+    builds = studio.builds.list_for_agent(agent_id)
+    removed_bindings = studio.plugin_compositions.unbind_builds(builds)
     trash_directory = None
-    if not purge:
-        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
-        trash_directory = studio.workspace.resolve(
-            Path(".agentkit/trash/agents") / f"{agent_id}-{timestamp}"
+    try:
+        if not purge:
+            timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
+            trash_directory = studio.workspace.resolve(
+                Path(".agentkit/trash/agents") / f"{agent_id}-{timestamp}"
+            )
+            trash_directory.mkdir(parents=True, exist_ok=False)
+        studio.event_store.delete_agent(
+            agent_id,
+            purge=purge,
+            trash_directory=trash_directory,
         )
-        trash_directory.mkdir(parents=True, exist_ok=False)
-    studio.event_store.delete_agent(
-        agent_id,
-        purge=purge,
-        trash_directory=trash_directory,
-    )
-    studio.builds.delete_for_agent(
-        agent_id,
-        purge=purge,
-        trash_directory=trash_directory,
-    )
-    studio.drafts.delete(
-        agent_id,
-        purge=purge,
-        trash_directory=trash_directory,
-    )
+        studio.builds.delete_for_agent(
+            agent_id,
+            purge=purge,
+            trash_directory=trash_directory,
+        )
+        studio.drafts.delete(
+            agent_id,
+            purge=purge,
+            trash_directory=trash_directory,
+        )
+    except Exception:
+        studio.plugin_compositions.restore_bindings(removed_bindings)
+        raise
 
 
 __all__ = ["delete_framework_agent"]