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

ksadk/sandbox/local_launcher.py

@@ -0,0 +1,124 @@
+"""Isolated Python entrypoint that applies rlimits before exec.
+
+It is started with ``python -I -S`` so untrusted workspace modules cannot replace
+the imported launcher or the standard-library ``resource`` module.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import signal
+import sys
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--limits", required=True)
+    parser.add_argument("--report-fd", required=True, type=int)
+    parser.add_argument("--owner-fd", type=int)
+    parser.add_argument("command", nargs=argparse.REMAINDER)
+    arguments = parser.parse_args(argv)
+    command = list(arguments.command)
+    if command[:1] == ["--"]:
+        command = command[1:]
+    report = _apply_limits(json.loads(arguments.limits))
+    os.write(arguments.report_fd, json.dumps(report, sort_keys=True).encode("utf-8"))
+    os.close(arguments.report_fd)
+    status = report["status"]
+    if status != "applied" and not (status == "partial" and sys.platform == "darwin"):
+        return 125
+    if not command:
+        return 125
+    if arguments.owner_fd is not None:
+        try:
+            _watch_owner(arguments.owner_fd)
+        except OSError:
+            return 125
+    try:
+        os.execvpe(command[0], command, os.environ)
+    except OSError as exc:
+        print(f"local-process launcher exec failed: {type(exc).__name__}: {exc}", file=sys.stderr)
+        return 126
+
+
+def _watch_owner(owner_fd: int) -> None:
+    """Bind this exec group's lifetime to the caller's private pipe writer.
+
+    A thread cannot survive exec. The small watcher stays in the command's
+    owned group, reads no workspace input and holds no command output pipes.
+    Normal cleanup stops the whole group; caller death closes the writer and
+    lets the watcher stop it even if the caller cannot execute a finally block.
+    """
+    group = os.getpgrp()
+    if group != os.getpid():
+        raise OSError("local command owner watcher requires its own process group")
+    if os.fork() != 0:
+        os.close(owner_fd)
+        return
+    try:
+        for descriptor in (0, 1, 2):
+            try:
+                os.close(descriptor)
+            except OSError:
+                pass
+        while os.read(owner_fd, 1):
+            pass
+        os.killpg(group, signal.SIGKILL)
+    finally:
+        os._exit(0)
+
+
+def _apply_limits(configured: dict[str, object]) -> dict[str, object]:
+    try:
+        import resource
+    except ImportError:
+        return {"status": "unsupported", "applied": {}, "unsupported": ["resource_module"]}
+
+    names = {
+        "cpu_seconds": "RLIMIT_CPU",
+        "address_space_bytes": "RLIMIT_AS",
+        "max_processes": "RLIMIT_NPROC",
+        "max_open_files": "RLIMIT_NOFILE",
+        "max_file_bytes": "RLIMIT_FSIZE",
+    }
+    applied: dict[str, int] = {}
+    unsupported: list[str] = []
+    failed: list[str] = []
+    for name, constant_name in names.items():
+        limit = configured.get(name)
+        if type(limit) is not int or limit < 1:
+            failed.append(f"{name}:invalid")
+            continue
+        if name == "address_space_bytes" and sys.platform == "darwin":
+            # Darwin exposes RLIMIT_AS but rejects useful values on supported
+            # Python/macOS combinations. Linux is the deployment target.
+            unsupported.append(name)
+            continue
+        if name == "max_processes" and sys.platform == "darwin":
+            # RLIMIT_NPROC is per UID. Developer desktops commonly already have
+            # more processes than the requested Skill budget.
+            unsupported.append(name)
+            continue
+        if not hasattr(resource, constant_name):
+            unsupported.append(name)
+            continue
+        resource_id = getattr(resource, constant_name)
+        try:
+            _, hard = resource.getrlimit(resource_id)
+            effective = limit if hard == resource.RLIM_INFINITY else min(limit, hard)
+            resource.setrlimit(resource_id, (effective, effective))
+            soft_after, hard_after = resource.getrlimit(resource_id)
+            if soft_after != effective or hard_after != effective:
+                failed.append(f"{name}:verification")
+            else:
+                applied[name] = effective
+        except (OSError, ValueError) as exc:
+            failed.append(f"{name}:{type(exc).__name__}")
+    status = "failed" if failed else ("partial" if unsupported else "applied")
+    return {"status": status, "applied": applied, "unsupported": unsupported, "failed": failed}
+
+
+if __name__ == "__main__":  # pragma: no cover - exercised through subprocesses
+    raise SystemExit(main())

ksadk/scheduler/dispatcher.py

@@ -30,7 +30,7 @@ class AgentControlSchedulerDispatcher:
     """Translate one occurrence into a single durable AgentControl enqueue.
 
     This dispatcher is intentionally local-runtime only.  A hosted scheduler
-    must obtain a Server-issued permit at its admission boundary in Phase 3;
+    must obtain a Server-issued permit at its admission boundary in admission boundary;
     this class never makes local permits look valid against hosted JWKS.
     """
 

ksadk/server/factory.py

@@ -25,6 +25,7 @@ from __future__ import annotations
 import asyncio
 import contextvars
 import logging
+import sqlite3
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
 
@@ -200,6 +201,7 @@ class RuntimeAppConfig:
         kernel_adapter_provider: Callable[[], RuntimeAdapter] | None = None,
         session_service_provider: Callable[[], Any] | None = None,
         session_backend_provider: Callable[[], dict[str, Any]] | None = None,
+        teams_host_factory: Any | None = None,
     ) -> None:
         self.runtime_type = runtime_type
         self.route_groups: set[str] = (
@@ -217,6 +219,7 @@ class RuntimeAppConfig:
         self.kernel_adapter_provider = kernel_adapter_provider
         self.session_service_provider = session_service_provider
         self.session_backend_provider = session_backend_provider
+        self.teams_host_factory = teams_host_factory
 
 
 # ---------------------------------------------------------------------------
@@ -359,6 +362,9 @@ def create_runtime_app(
         session_service_provider=config.session_service_provider,
         session_backend_provider=config.session_backend_provider,
     )
+    from ksadk.kernel.teams_runtime_bootstrap import CloudTeamsRuntimeHost
+
+    teams_host = config.teams_host_factory or CloudTeamsRuntimeHost.from_env()
 
     @asynccontextmanager
     async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
@@ -372,19 +378,36 @@ def create_runtime_app(
         from ksadk.runtime.factory import kernel_start_request_defaults
 
         adapter_provider = _kernel_adapter_provider(config)
+        if teams_host is not None:
+            adapter_provider = teams_host.wrap_adapter_provider(adapter_provider)
         request_defaults = (
             kernel_start_request_defaults(config.launch_context)
             if config.launch_context is not None
             else {}
         )
-        kernel_runtime = await bootstrap_agent_kernel_runtime_from_env(
-            adapter_provider=adapter_provider,
-            runtime_executor=config.runtime_executor,
-            launch_context=config.launch_context,
-            start_request_defaults=request_defaults,
-            session_service=state.resolve_session_service(),
-        )
+        try:
+            kernel_runtime = await bootstrap_agent_kernel_runtime_from_env(
+                adapter_provider=adapter_provider,
+                runtime_executor=config.runtime_executor,
+                launch_context=config.launch_context,
+                start_request_defaults=request_defaults,
+                session_service=state.resolve_session_service(),
+                **({"before_start": teams_host.before_start} if teams_host is not None else {}),
+            )
+        except BaseException:
+            if teams_host is not None:
+                await teams_host.close()
+                _kernel_ingress.set_teams_ingress_gate(None)
+            await shutdown_runtime_resources(state)
+            if config.a2a_runtime_adapter is not None:
+                close_app_resources = getattr(config.a2a_runtime_adapter, "close_app_resources", None)
+                if callable(close_app_resources):
+                    await close_app_resources()
+            raise
         app.state.agent_kernel_runtime = kernel_runtime
+        if teams_host is not None and kernel_runtime is None:
+            await teams_host.close()
+            raise RuntimeError("Teams Host requires an enabled durable Kernel")
         if kernel_runtime is not None:
             # Kernel canonical events are the replay authority for the normal
             # Session APIs as well.  In memory/ephemeral mode bootstrap reuses
@@ -405,9 +428,16 @@ def create_runtime_app(
                 await state.a2a_bootstrap.stop()
             if kernel_runtime is not None:
                 await kernel_runtime.close()
+            if teams_host is not None:
+                await teams_host.close()
+                _kernel_ingress.set_teams_ingress_gate(None)
             clear_agent_kernel_runtime()
             _kernel_ingress.clear_agent_kernel()
             await shutdown_runtime_resources(state)
+            if config.a2a_runtime_adapter is not None:
+                close_app_resources = getattr(config.a2a_runtime_adapter, "close_app_resources", None)
+                if callable(close_app_resources):
+                    await close_app_resources()
 
     app = FastAPI(
         title="ADK Core API",
@@ -417,6 +447,11 @@ def create_runtime_app(
     )
     state.app = app
     app.state.runtime = state
+    if teams_host is not None:
+        from ksadk.kernel.teams_host_http import create_teams_host_router
+
+        app.include_router(create_teams_host_router(teams_host.resolve_host))
+        app.state.teams_host = teams_host
 
     # 中间件:把当前 app 的 state 写入请求级 contextvar(handler 经 get_state() 取)。
     @app.middleware("http")
@@ -499,6 +534,10 @@ def create_runtime_app(
                 response.headers["Expires"] = "0"
             return response
 
+    from ksadk.server.history_deadline import HistoryDeadlineMiddleware
+
+    app.add_middleware(HistoryDeadlineMiddleware)
+
     # CORS(与旧 app 一致,默认对 ADK 工具放开)。
     app.add_middleware(
         CORSMiddleware,
@@ -517,13 +556,84 @@ def create_runtime_app(
     async def session_backend_unavailable_handler(_request, exc: SessionBackendUnavailable):
         return Response(
             content=_json.dumps(
-                {"detail": {"code": "session_backend_unavailable", "message": str(exc)}},
+                {"Code": 503, "Message": str(exc), "Data": None,
+                 "ReasonCode": "SESSION_BACKEND_UNAVAILABLE",
+                 "detail": {"code": "session_backend_unavailable", "message": str(exc)}},
                 ensure_ascii=False,
             ),
             status_code=503,
             media_type="application/json",
         )
 
+    from ksadk.sessions.errors import (
+        CheckpointLookupTimeout,
+        HistoryProjectionBudgetExceeded,
+        HistoryReloadRequired,
+    )
+
+    @app.exception_handler(CheckpointLookupTimeout)
+    async def checkpoint_lookup_timeout_handler(_request, exc: CheckpointLookupTimeout):
+        return Response(
+            content=_json.dumps({"Code": 504, "Message": str(exc), "Data": None,
+                                 "ReasonCode": "CHECKPOINT_LOOKUP_TIMEOUT"}),
+            status_code=504, media_type="application/json",
+        )
+
+    @app.exception_handler(HistoryProjectionBudgetExceeded)
+    async def history_projection_budget_handler(_request, exc: HistoryProjectionBudgetExceeded):
+        return Response(
+            content=_json.dumps({"Code": 422, "Message": str(exc), "Data": None,
+                                 "ReasonCode": "HISTORY_PROJECTION_BUDGET_EXCEEDED"}),
+            status_code=422, media_type="application/json",
+        )
+
+    @app.exception_handler(HistoryReloadRequired)
+    async def history_reload_handler(_request, exc: HistoryReloadRequired):
+        return Response(
+            content=_json.dumps({"Code": 409, "Message": str(exc), "Data": None,
+                                 "ReasonCode": "HISTORY_RELOAD_REQUIRED"}),
+            status_code=409, media_type="application/json",
+        )
+
+    # PostgreSQL replies to these cases, so resilience must not turn them into
+    # independent live writes. Expose explicit HTTP classifications instead.
+    try:
+        import asyncpg
+    except ImportError:  # local-only SDK installation
+        pass
+    else:
+        @app.exception_handler(asyncpg.QueryCanceledError)
+        async def postgres_statement_timeout_handler(_request, _exc):
+            return Response(
+                content=_json.dumps({"Code": 504, "Message": "Database statement timed out",
+                                     "Data": None, "ReasonCode": "DATABASE_STATEMENT_TIMEOUT"}),
+                status_code=504, media_type="application/json",
+            )
+
+        @app.exception_handler(asyncpg.LockNotAvailableError)
+        async def postgres_lock_timeout_handler(_request, _exc):
+            return Response(
+                content=_json.dumps({"Code": 503, "Message": "Database lock unavailable",
+                                     "Data": None, "ReasonCode": "DATABASE_LOCK_TIMEOUT"}),
+                status_code=503, media_type="application/json",
+            )
+
+        @app.exception_handler(asyncpg.UniqueViolationError)
+        async def postgres_unique_handler(_request, _exc):
+            return Response(
+                content=_json.dumps({"Code": 409, "Message": "Conflicting database record",
+                                     "Data": None, "ReasonCode": "DATABASE_CONFLICT"}),
+                status_code=409, media_type="application/json",
+            )
+
+        @app.exception_handler(asyncpg.CheckViolationError)
+        async def postgres_check_handler(_request, _exc):
+            return Response(
+                content=_json.dumps({"Code": 422, "Message": "Database constraint rejected request",
+                                     "Data": None, "ReasonCode": "DATABASE_CONSTRAINT"}),
+                status_code=422, media_type="application/json",
+            )
+
     # A2A 端点须在 configure 之前装配:health_meta 含 GET catch-all ``/{requested_path:path}``,
     # 若 A2A 后挂,A2A 的 GET 路由(agent-card / tasks/{id})会被 catch-all 遮蔽成 404。
     _wire_a2a_if_enabled(app, state, config)
@@ -548,6 +658,9 @@ def _kernel_adapter_provider(
     if config.kernel_adapter_provider is not None:
         return config.kernel_adapter_provider
     if config.a2a_runtime_adapter is not None:
+        adopt_app_lifetime = getattr(config.a2a_runtime_adapter, "adopt_app_lifetime", None)
+        if callable(adopt_app_lifetime):
+            adopt_app_lifetime()
         return lambda: config.a2a_runtime_adapter
     if config.runtime_executor is not None and config.launch_context is not None:
         return lambda: config.runtime_executor.create_adapter(config.launch_context)
@@ -583,6 +696,11 @@ def _wire_a2a_if_enabled(app: FastAPI, state: RuntimeAppState, config: RuntimeAp
     runtime_type = adapter.runtime.runtime_type or config.runtime_type
 
     if isinstance(a2a_cfg, AgentEngineA2ABootstrap):
+        if config.runtime_executor is not None:
+            config.runtime_executor.configure_managed_a2a(a2a_cfg)
+        configure_adapter = getattr(adapter, "configure_managed_a2a", None)
+        if callable(configure_adapter):
+            configure_adapter(a2a_cfg)
         server = a2a_cfg.mount(
             app,
             runtime_adapter=adapter,
@@ -635,27 +753,62 @@ def _wire_agui_if_enabled(app: FastAPI, state: RuntimeAppState, config: RuntimeA
             self._service = service
             self._store = RuntimeEventStore(service)
 
-        async def append_one(self, session_id: str, event: Any) -> Any:
-            existing = await self._service.get_session(session_id)
+        async def _ensure_session(self, session_id: str, event: Any) -> None:
+            meta = event.source.metadata
+            asserted_agent = str(meta.get("agent_id") or "").strip()
+            asserted_user = str(meta.get("user_id") or "").strip()
+            agent_id = asserted_agent or "agent"
+            user_id = asserted_user or "user"
+            authority = getattr(self._service, "canonical_event_service", self._service)
+            header = getattr(self._service, "get_session_header", None)
+            if not callable(header):
+                header = self._service.get_session_metadata
+            existing = await header(session_id)
             if existing is None:
-                metadata = getattr(event, "source", None)
-                meta = metadata.metadata if metadata is not None else {}
-                await self._service.create_session(
-                    str(meta.get("agent_id") or "agent"),
-                    str(meta.get("user_id") or "user"),
-                    session_id,
-                )
+                try:
+                    ensure = getattr(self._service, "ensure_canonical_session_header", None)
+                    if not callable(ensure):
+                        ensure = getattr(authority, "ensure_session_header", None)
+                    existing = (await ensure(agent_id, user_id, session_id)
+                                if callable(ensure) else
+                                await authority.create_session(agent_id, user_id, session_id))
+                except Exception as exc:
+                    if not (getattr(exc, "sqlite_errorcode", None) in
+                            {sqlite3.SQLITE_CONSTRAINT_UNIQUE, sqlite3.SQLITE_CONSTRAINT_PRIMARYKEY}
+                            or getattr(exc, "sqlstate", None) == "23505"):
+                        raise
+                    existing = await header(session_id)
+                    if existing is None:
+                        raise
+            if ((asserted_agent and existing.agent_id != asserted_agent)
+                    or (asserted_user and existing.user_id != asserted_user)):
+                raise ValueError("AG-UI session id belongs to a different agent or user")
+            # Canonical events are persisted by the durable authority even when
+            # the session service has a dirty process-local fallback session.
+            degraded = getattr(self._service, "persistence_degraded", None)
+            if authority is not self._service and callable(degraded) and degraded(session_id):
+                authority_header = getattr(authority, "get_session_header", None)
+                if not callable(authority_header):
+                    authority_header = authority.get_session_metadata
+                durable = await authority_header(session_id)
+                if durable is None:
+                    from ksadk.sessions.errors import SessionBackendUnavailable
+
+                    raise SessionBackendUnavailable(
+                        "Canonical AG-UI session is unavailable in durable storage"
+                    )
+                if ((asserted_agent and durable.agent_id != asserted_agent)
+                        or (asserted_user and durable.user_id != asserted_user)):
+                    raise ValueError("AG-UI session id belongs to a different agent or user")
+
+        async def append_one(self, session_id: str, event: Any) -> Any:
+            await self._ensure_session(session_id, event)
             return await self._store.append_one(session_id, event)
 
         async def reserve_once(self, event: Any) -> Any:
-            existing = await self._service.get_session(event.session_id)
-            if existing is None:
-                await self._service.create_session(
-                    event.agent_id,
-                    event.user_id,
-                    event.session_id,
-                )
-            return await self._store.reserve_once(event)
+            session_id = str(event.source.metadata.get("session_id") or "")
+            await self._ensure_session(session_id, event)
+            return await self._store.persist_one(session_id, event)
 
         async def list(self, session_id: str, **kwargs: Any) -> Any:
             return await self._store.list(session_id, **kwargs)

ksadk/server/history_deadline.py

@@ -0,0 +1,158 @@
+"""Total ASGI deadline for bounded metadata/history actions.
+
+The response is buffered until completion so a timeout during serialization
+cannot produce a partial HTTP 200. This middleware never wraps Agent or SSE
+requests, which have their own lifetimes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from collections import Counter
+from typing import Any
+
+from starlette.datastructures import Headers
+
+from ksadk.server.invocation_identity import trusted_invocation_identity_from_headers
+from ksadk.sessions.history_budget import CURRENT_HISTORY_READ_BUDGET, HistoryReadBudget
+
+logger = logging.getLogger(__name__)
+INVALID_REMAINING_HEADER: Counter[str] = Counter()
+_ACTIONS = frozenset({
+    "GetSession", "ListSessions", "ListSessionEvents", "ListSessionMessages",
+    "ListSessionCheckpoints", "ListToolReceipts", "GetAgentUiBootstrap",
+    "GetCheckpointResumePreview",
+})
+_MAX_BUFFER_BYTES = 32 * 1024 * 1024
+_ERROR_SEND_SECONDS = 0.5
+_HISTORY_ACTIONS = frozenset({
+    "ListSessionEvents", "ListSessionMessages", "ListSessionCheckpoints",
+    "ListToolReceipts", "GetCheckpointResumePreview",
+})
+
+
+class _ResponseBufferLimitExceeded(RuntimeError):
+    pass
+
+
+def _remaining_ms(raw: str) -> int | None:
+    if raw.startswith("-") and raw[1:] and raw[1:].isascii() and raw[1:].isdecimal():
+        return 0
+    if not raw or not raw.isascii() or not raw.isdecimal():
+        return None
+    # This carrier only shortens our 2.6-second cap. Avoid Python's maximum
+    # integer-string conversion limit and safely saturate oversized values.
+    significant = raw.lstrip("0")
+    if not significant:
+        return 0
+    if len(significant) > 18:
+        return 2600
+    try:
+        return int(significant)
+    except ValueError:  # pragma: no cover - bounded ASCII digits above
+        return None
+
+
+class HistoryDeadlineMiddleware:
+    def __init__(self, app: Any) -> None:
+        self.app = app
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+        path = scope.get("path", "")
+        if scope.get("method") != "POST" or path.rsplit("/", 1)[-1] not in _ACTIONS:
+            await self.app(scope, receive, send)
+            return
+
+        started = time.monotonic()
+        budget = 2.6
+        headers = Headers(scope=scope)
+        try:
+            trusted = not trusted_invocation_identity_from_headers(headers).is_empty
+        except Exception:
+            trusted = False
+        if trusted:
+            raw = headers.get("x-agentengine-remaining-ms")
+            if raw is not None:
+                parsed = _remaining_ms(raw)
+                if parsed is None:
+                    INVALID_REMAINING_HEADER["malformed"] += 1
+                else:
+                    budget = min(budget, parsed / 1000)
+        budget -= time.monotonic() - started
+        if budget <= 0:
+            await self._timeout(send)
+            return
+
+        messages: list[dict[str, Any]] = []
+        buffered_bytes = 0
+
+        async def buffered_send(message: dict[str, Any]) -> None:
+            nonlocal buffered_bytes
+            buffered_bytes += len(message.get("body", b""))
+            if buffered_bytes > _MAX_BUFFER_BYTES:
+                raise _ResponseBufferLimitExceeded()
+            messages.append(message)
+
+        token = (CURRENT_HISTORY_READ_BUDGET.set(HistoryReadBudget())
+                 if path.rsplit("/", 1)[-1] in _HISTORY_ACTIONS else None)
+        try:
+            async with asyncio.timeout(budget):
+                await self.app(scope, receive, buffered_send)
+        except _ResponseBufferLimitExceeded:
+            await self._response_too_large(send)
+            return
+        except TimeoutError:
+            await self._timeout(send)
+            return
+        finally:
+            if token is not None:
+                CURRENT_HISTORY_READ_BUDGET.reset(token)
+        remaining = budget - (time.monotonic() - started)
+        if remaining <= 0:
+            await self._timeout(send)
+            return
+        emitted = False
+        try:
+            async with asyncio.timeout(remaining):
+                for message in messages:
+                    await send(message)
+                    emitted = True
+        except TimeoutError:
+            if emitted:
+                raise  # close a started response rather than append a second status
+            await self._timeout(send)
+
+    @staticmethod
+    async def _timeout(send: Any) -> None:
+        body = json.dumps({
+            "Code": 504, "Message": "History request deadline exceeded",
+            "Data": None, "ReasonCode": "RUNTIME_HISTORY_TIMEOUT",
+        }).encode()
+        await HistoryDeadlineMiddleware._send_error(send, 504, body)
+
+    @staticmethod
+    async def _response_too_large(send: Any) -> None:
+        body = json.dumps({
+            "Code": 422, "Message": "History response exceeds the bounded response size",
+            "Data": None, "ReasonCode": "HISTORY_PROJECTION_BUDGET_EXCEEDED",
+        }).encode()
+        await HistoryDeadlineMiddleware._send_error(send, 422, body)
+
+    @staticmethod
+    async def _send_error(send: Any, status: int, body: bytes) -> None:
+        # A slow/broken ASGI transport cannot extend the error path forever.
+        # Both messages share one cap; once the header was sent, a blocked
+        # body closes this response rather than trying another status.
+        async with asyncio.timeout(_ERROR_SEND_SECONDS):
+            await send({
+                "type": "http.response.start", "status": status,
+                "headers": [(b"content-type", b"application/json"),
+                            (b"content-length", str(len(body)).encode())],
+            })
+            await send({"type": "http.response.body", "body": body})