cortex-suite-sdk 1.0.2__py3-none-any.whl

cortex_sdk/__init__.py

@@ -0,0 +1,4 @@
+from .client import CortexClient
+from .errors import CortexError
+
+__all__ = ["CortexClient", "CortexError"]

cortex_sdk/_generated_constants.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+# Generated from shared/constants.json. Do not edit manually.
+
+DEFAULT_AUTH_URL: str = 'https://auth.cortexsuite.app'
+AUTH_TOKEN_PATH: str = '/auth/token'
+AUTH_REFRESH_PATH: str = '/auth/refresh'
+WS_SUBPROTOCOL: str = 'cortex-sdk.v1'
+WS_SUBPROTOCOL_JWT_PREFIX: str = 'cortex-sdk.jwt.'
+SCHEMA_VERSION: str = '1.0'
+
+DEFAULT_CONNECT_TIMEOUT_MS: int = 10000
+DEFAULT_SEND_TIMEOUT_MS: int = 10000
+DEFAULT_RESYNC_TIMEOUT_MS: int = 15000
+DEFAULT_PING_INTERVAL_MS: int = 15000
+DEFAULT_PONG_TIMEOUT_MS: int = 5000
+DEFAULT_STALE_THRESHOLD_MS: int = 45000
+TOKEN_REFRESH_BUFFER_MS: int = 60000
+
+RECONNECT_BACKOFF_MS: tuple[int, ...] = (1000, 2000, 5000, 10000, 20000, 30000,)
+
+# Deprecated compatibility aliases. Prefer DEFAULT_AUTH_URL/AUTH_TOKEN_PATH,
+# DEFAULT_AUTH_URL/AUTH_REFRESH_PATH, and WS_SUBPROTOCOL directly.
+CORTEX_AUTH_URL: str = DEFAULT_AUTH_URL + AUTH_TOKEN_PATH
+CORTEX_REFRESH_URL: str = DEFAULT_AUTH_URL + AUTH_REFRESH_PATH
+WS_SUBPROTOCOL_BASE: str = WS_SUBPROTOCOL

cortex_sdk/_generated_errors.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Generated from sdk/shared/errors.json. Do not edit manually.
+
+@dataclass(frozen=True)
+class GeneratedErrorEntry:
+    code: str
+    retryable: bool
+    fatal: bool
+
+GENERATED_ERROR_CATALOG: tuple[GeneratedErrorEntry, ...] = (
+    GeneratedErrorEntry('auth_invalid', retryable=False, fatal=True),
+    GeneratedErrorEntry('auth_expired', retryable=True, fatal=False),
+    GeneratedErrorEntry('auth_refresh_failed', retryable=False, fatal=True),
+    GeneratedErrorEntry('transport_connect_timeout', retryable=True, fatal=False),
+    GeneratedErrorEntry('transport_send_timeout', retryable=True, fatal=False),
+    GeneratedErrorEntry('transport_protocol_violation', retryable=False, fatal=True),
+    GeneratedErrorEntry('session_not_found', retryable=False, fatal=True),
+    GeneratedErrorEntry('session_terminal', retryable=False, fatal=True),
+    GeneratedErrorEntry('resync_timeout', retryable=True, fatal=False),
+    GeneratedErrorEntry('replay_unavailable', retryable=True, fatal=False),
+    GeneratedErrorEntry('upload_failed', retryable=True, fatal=False),
+    GeneratedErrorEntry('upload_too_large', retryable=False, fatal=False),
+    GeneratedErrorEntry('upload_type_rejected', retryable=False, fatal=False),
+)

cortex_sdk/auth.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import base64
+import json
+import time
+
+import httpx
+
+from .constants import (
+    DEFAULT_AUTH_URL,
+    AUTH_TOKEN_PATH,
+    AUTH_REFRESH_PATH,
+    TOKEN_REFRESH_BUFFER,
+)
+from .errors import make_error
+from .types import AuthTokenResponse
+
+
+def _parse_jwt_exp(token: str) -> float | None:
+    """Extract exp (as Unix timestamp in seconds) from a JWT without verifying signature."""
+    try:
+        parts = token.split(".")
+        if len(parts) != 3:
+            return None
+        # base64url → base64 standard (add padding)
+        payload_b64 = parts[1]
+        # Add padding
+        padding = 4 - len(payload_b64) % 4
+        if padding != 4:
+            payload_b64 += "=" * padding
+        payload_json = base64.urlsafe_b64decode(payload_b64)
+        payload = json.loads(payload_json)
+        exp = payload.get("exp")
+        return float(exp) if isinstance(exp, (int, float)) else None
+    except Exception:
+        return None
+
+
+def is_token_expiring_soon(access_token: str) -> bool:
+    """Return True if the token expires within TOKEN_REFRESH_BUFFER seconds."""
+    exp = _parse_jwt_exp(access_token)
+    if exp is None:
+        return False
+    return time.time() > exp - TOKEN_REFRESH_BUFFER
+
+
+async def exchange_api_key(
+    api_key: str,
+    *,
+    auth_base_url: str = DEFAULT_AUTH_URL,
+) -> AuthTokenResponse:
+    """Exchange an API key for tokens and WS URL."""
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(
+            _build_auth_endpoint(auth_base_url, AUTH_TOKEN_PATH),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"ApiKey {api_key}",
+            },
+        )
+
+    if not resp.is_success:
+        try:
+            body = resp.json()
+            code = body.get("error", "auth_invalid")
+            message = body.get("message", "API key rejected")
+        except Exception:
+            code, message = "auth_invalid", "API key rejected"
+        raise make_error(str(code), str(message))
+
+    return resp.json()  # type: ignore[no-any-return]
+
+
+async def refresh_access_token(
+    refresh_token: str,
+    *,
+    auth_base_url: str = DEFAULT_AUTH_URL,
+) -> str:
+    """Refresh the access token using the refresh token."""
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(
+            _build_auth_endpoint(auth_base_url, AUTH_REFRESH_PATH),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {refresh_token}",
+            },
+        )
+
+    if not resp.is_success:
+        raise make_error("auth_refresh_failed", "Refresh token expired or invalid")
+
+    body = resp.json()
+    return str(body["access_token"])
+
+
+def normalize_auth_base_url(auth_url: str) -> str:
+    import warnings
+    normalized = auth_url.rstrip("/")
+    # Guard: consumer may have passed a full endpoint instead of the base URL.
+    # Strip known auth paths and warn so developers catch the misconfiguration early.
+    for known_path in [AUTH_TOKEN_PATH, AUTH_REFRESH_PATH]:
+        if normalized.endswith(known_path):
+            base = normalized[: -len(known_path)]
+            warnings.warn(
+                f"[CortexSDK] auth_url must be a base URL (origin only), not a full endpoint. "
+                f'Received "{auth_url}" — "{known_path}" has been stripped automatically. '
+                f'Pass the base URL only, e.g., "{base}".',
+                UserWarning,
+                stacklevel=3,
+            )
+            normalized = base
+            break
+    return normalized or DEFAULT_AUTH_URL
+
+
+def _build_auth_endpoint(auth_base_url: str, path: str) -> str:
+    return f"{normalize_auth_base_url(auth_base_url)}{path}"

cortex_sdk/client.py

@@ -0,0 +1,331 @@
+from __future__ import annotations
+
+import asyncio
+import secrets
+from typing import BinaryIO
+from urllib.parse import urlsplit, urlunsplit
+
+from .auth import (
+    exchange_api_key,
+    refresh_access_token,
+    is_token_expiring_soon,
+    normalize_auth_base_url,
+)
+from .constants import (
+    DEFAULT_AUTH_URL,
+    DEFAULT_CONNECT_TIMEOUT,
+    DEFAULT_SEND_TIMEOUT,
+    DEFAULT_RESYNC_TIMEOUT,
+    DEFAULT_PING_INTERVAL,
+    DEFAULT_PONG_TIMEOUT,
+    DEFAULT_STALE_THRESHOLD,
+    TOKEN_REFRESH_BUFFER,
+    RECONNECT_BACKOFF,
+)
+from .errors import CortexError, make_error
+from .liveness import LivenessMonitor
+from .session import SessionManager
+from .transport import Transport
+from .types import ChannelState, CortexMessage, MessageCallback, SessionState
+from .upload import upload_file
+
+
+class CortexClient:
+    """Cortex SDK client — one instance per session.
+
+    Public API:
+        connect() / disconnect()
+        send_message(content, attachments)
+        upload_attachment(file)
+        stop()
+        session_state / channel_state / session_id  (read-only properties)
+
+    Time options are in **seconds** (not milliseconds).
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        on_message: MessageCallback,
+        auth_url: str | None = None,
+        connect_timeout: float = DEFAULT_CONNECT_TIMEOUT,
+        send_timeout: float = DEFAULT_SEND_TIMEOUT,
+        resync_timeout: float = DEFAULT_RESYNC_TIMEOUT,
+        ping_interval: float = DEFAULT_PING_INTERVAL,
+        pong_timeout: float = DEFAULT_PONG_TIMEOUT,
+        stale_threshold: float = DEFAULT_STALE_THRESHOLD,
+        # Private test override — not part of the public API
+        _upload_url: str | None = None,
+    ) -> None:
+        self._api_key = api_key
+        self._on_message_cb = on_message
+        self._auth_url = normalize_auth_base_url(auth_url or DEFAULT_AUTH_URL)
+        self._connect_timeout = connect_timeout
+        self._send_timeout = send_timeout
+        self._resync_timeout = resync_timeout
+        self._ping_interval = ping_interval
+        self._pong_timeout = pong_timeout
+        self._stale_threshold = stale_threshold
+
+        self._upload_url = _upload_url  # None → runtime-side upload URL heuristic
+
+        # Internal state
+        self._channel_state: ChannelState = "CLOSED"
+        self._access_token: str | None = None
+        self._refresh_token: str | None = None
+        self._ws_url: str | None = None
+        self._channel_id: str = f"ch_{secrets.token_hex(4)}"
+        self._reconnect_attempt: int = 0
+        self._disconnect_requested: bool = False
+
+        # Components
+        self._transport = Transport(connect_timeout, send_timeout)
+        self._session = SessionManager(
+            on_message=self._dispatch_message,
+            on_fatal_error=self._handle_fatal_error,
+        )
+        self._liveness: LivenessMonitor | None = None
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._token_refresh_task: asyncio.Task[None] | None = None
+
+        # Wire transport callbacks
+        self._transport.on_message = self._session.handle_incoming
+        self._transport.on_close = self._handle_close
+        self._transport.on_error = None  # handled via on_close
+
+    # ── public properties ─────────────────────────────────────────────────────
+
+    @property
+    def session_state(self) -> SessionState:
+        return self._session.session_state
+
+    @property
+    def channel_state(self) -> ChannelState:
+        return self._channel_state
+
+    @property
+    def session_id(self) -> str | None:
+        return self._session.session_id
+
+    # ── public methods ────────────────────────────────────────────────────────
+
+    async def connect(self) -> None:
+        """Full bootstrap: auth exchange → WS open → system::init → liveness start."""
+        self._disconnect_requested = False
+        self._reconnect_attempt = 0
+
+        auth = await exchange_api_key(self._api_key, auth_base_url=self._auth_url)
+        self._access_token = auth["access_token"]
+        self._refresh_token = auth["refresh_token"]
+        self._ws_url = auth["ws_url"]
+        if self._upload_url is None:
+            self._upload_url = _derive_upload_url_from_ws_url(self._ws_url)
+
+        await self._open_channel()
+
+        self._session.set_transport(self._transport, self._send_timeout)
+        await self._session.send_init(auth["runtime_bootstrap"])
+
+        self._start_liveness()
+        self._schedule_token_refresh()
+
+    async def disconnect(self) -> None:
+        """Close the session and WebSocket connection cleanly."""
+        self._disconnect_requested = True
+        self._stop_liveness()
+        self._stop_token_refresh()
+
+        if self._reconnect_task and not self._reconnect_task.done():
+            self._reconnect_task.cancel()
+            try:
+                await self._reconnect_task
+            except (asyncio.CancelledError, Exception):
+                pass
+            self._reconnect_task = None
+
+        self._channel_state = "CLOSED"
+        await self._transport.aclose()
+
+    async def send_message(
+        self, content: str, attachments: list[str] | None = None
+    ) -> None:
+        await self._session.send_chat_message(content, attachments)
+
+    async def upload_attachment(self, file: str | bytes | BinaryIO) -> str:
+        if not self._access_token:
+            raise make_error("auth_invalid", "Not connected")
+        url = self._upload_url or "/upload"
+        return await upload_file(file, self._access_token, upload_url=url)
+
+    async def stop(self) -> None:
+        await self._session.send_stop()
+
+    # ── internal helpers ──────────────────────────────────────────────────────
+
+    async def _open_channel(self) -> None:
+        if not self._ws_url or not self._access_token:
+            raise RuntimeError("Auth not completed before _open_channel")
+        self._channel_state = "CONNECTING"
+        await self._transport.open(self._ws_url, self._access_token)
+        self._channel_state = "OPEN"
+        self._reconnect_attempt = 0
+
+    def _start_liveness(self) -> None:
+        self._stop_liveness()
+        self._liveness = LivenessMonitor(
+            transport=self._transport,
+            ping_interval=self._ping_interval,
+            pong_timeout=self._pong_timeout,
+            stale_threshold=self._stale_threshold,
+            on_stale=self._handle_stale,
+            get_session_id=lambda: self._session.session_id,
+            get_channel_id=lambda: self._channel_id,
+        )
+        self._liveness.start()
+
+    def _stop_liveness(self) -> None:
+        if self._liveness:
+            self._liveness.stop()
+            self._liveness = None
+
+    def _schedule_token_refresh(self) -> None:
+        self._stop_token_refresh()
+
+        async def _refresh_loop() -> None:
+            try:
+                while not self._disconnect_requested:
+                    await asyncio.sleep(TOKEN_REFRESH_BUFFER / 2)
+                    if not self._access_token or not self._refresh_token:
+                        continue
+                    if is_token_expiring_soon(self._access_token):
+                        try:
+                            self._access_token = await refresh_access_token(
+                                self._refresh_token,
+                                auth_base_url=self._auth_url,
+                            )
+                        except Exception:
+                            pass  # surfaced at next reconnect
+            except asyncio.CancelledError:
+                pass
+
+        self._token_refresh_task = asyncio.create_task(_refresh_loop())
+
+    def _stop_token_refresh(self) -> None:
+        if self._token_refresh_task and not self._token_refresh_task.done():
+            self._token_refresh_task.cancel()
+        self._token_refresh_task = None
+
+    # ── callbacks from transport / session ────────────────────────────────────
+
+    def _dispatch_message(self, msg: CortexMessage) -> None:
+        # system::pong is internal — route to liveness, not to user callback
+        if msg.get("type") == "system::pong":
+            hb_id = msg["payload"].get("heartbeat_id")
+            if isinstance(hb_id, str) and self._liveness:
+                self._liveness.record_pong(hb_id)
+            return
+        self._on_message_cb(msg)
+
+    def _handle_fatal_error(self, err: CortexError) -> None:
+        self._channel_state = "AUTH_FAILED"
+        self._stop_liveness()
+        self._stop_token_refresh()
+        self._transport.close()
+        # Surface error as a system::error message
+        error_msg: CortexMessage = {
+            "type": "system::error",
+            "schema": "1.0",
+            "session_id": self._session.session_id or "",
+            "payload": {"code": err.code, "message": str(err)},
+            "ts": _now_iso(),
+        }
+        self._on_message_cb(error_msg)
+
+    def _handle_stale(self) -> None:
+        if self._channel_state in ("STALE", "RECONNECTING"):
+            return
+        self._channel_state = "STALE"
+        self._stop_liveness()
+        self._transport.close(1001, "stale")
+        # on_close will trigger reconnect
+
+    def _handle_close(self, code: int, reason: str) -> None:
+        if self._disconnect_requested:
+            return
+        if self._channel_state == "AUTH_FAILED":
+            return
+        if code == 4001:
+            self._channel_state = "AUTH_FAILED"
+            return
+        self._channel_state = "RECONNECTING"
+        self._reconnect_task = asyncio.ensure_future(self._reconnect_loop())
+
+    async def _reconnect_loop(self) -> None:
+        try:
+            while (
+                not self._disconnect_requested
+                and self._channel_state != "AUTH_FAILED"
+            ):
+                idx = min(self._reconnect_attempt, len(RECONNECT_BACKOFF) - 1)
+                backoff = RECONNECT_BACKOFF[idx]
+                self._reconnect_attempt += 1
+
+                await asyncio.sleep(backoff)
+                if self._disconnect_requested:
+                    break
+
+                # Proactively refresh token if needed
+                try:
+                    await self._maybe_refresh_token()
+                except CortexError:
+                    self._channel_state = "AUTH_FAILED"
+                    return
+
+                try:
+                    await self._open_channel()
+                except Exception:
+                    continue  # retry after next backoff
+
+                # Re-attach session to the new transport
+                self._session.set_transport(self._transport, self._send_timeout)
+
+                # Resync with timeout
+                try:
+                    await asyncio.wait_for(
+                        self._session.send_resync(),
+                        timeout=self._resync_timeout,
+                    )
+                except Exception:
+                    self._transport.close()
+                    continue
+
+                # Reconnect successful — restart liveness + token refresh
+                self._start_liveness()
+                self._schedule_token_refresh()
+                return
+        except asyncio.CancelledError:
+            pass
+
+    async def _maybe_refresh_token(self) -> None:
+        if not self._refresh_token:
+            raise make_error("auth_refresh_failed", "No refresh token")
+        if not self._access_token or is_token_expiring_soon(self._access_token):
+            self._access_token = await refresh_access_token(
+                self._refresh_token,
+                auth_base_url=self._auth_url,
+            )
+
+
+def _now_iso() -> str:
+    from datetime import datetime, timezone
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _derive_upload_url_from_ws_url(ws_url: str | None) -> str:
+    if not ws_url:
+        return "/upload"
+
+    parsed = urlsplit(ws_url)
+    scheme = parsed.scheme.lower()
+    http_scheme = "https" if scheme == "wss" else "http" if scheme == "ws" else scheme
+    return urlunsplit((http_scheme, parsed.netloc, "/upload", "", ""))

cortex_sdk/constants.py

@@ -0,0 +1,50 @@
+from ._generated_constants import (
+    DEFAULT_AUTH_URL,
+    AUTH_TOKEN_PATH,
+    AUTH_REFRESH_PATH,
+    WS_SUBPROTOCOL,
+    WS_SUBPROTOCOL_JWT_PREFIX,
+    CORTEX_AUTH_URL,
+    CORTEX_REFRESH_URL,
+    WS_SUBPROTOCOL_BASE,
+    SCHEMA_VERSION,
+    DEFAULT_CONNECT_TIMEOUT_MS,
+    DEFAULT_SEND_TIMEOUT_MS,
+    DEFAULT_RESYNC_TIMEOUT_MS,
+    DEFAULT_PING_INTERVAL_MS,
+    DEFAULT_PONG_TIMEOUT_MS,
+    DEFAULT_STALE_THRESHOLD_MS,
+    TOKEN_REFRESH_BUFFER_MS,
+    RECONNECT_BACKOFF_MS,
+)
+
+__all__ = [
+    "DEFAULT_AUTH_URL",
+    "AUTH_TOKEN_PATH",
+    "AUTH_REFRESH_PATH",
+    "WS_SUBPROTOCOL",
+    "WS_SUBPROTOCOL_JWT_PREFIX",
+    "CORTEX_AUTH_URL",
+    "CORTEX_REFRESH_URL",
+    "WS_SUBPROTOCOL_BASE",
+    "SCHEMA_VERSION",
+    "DEFAULT_CONNECT_TIMEOUT",
+    "DEFAULT_SEND_TIMEOUT",
+    "DEFAULT_RESYNC_TIMEOUT",
+    "DEFAULT_PING_INTERVAL",
+    "DEFAULT_PONG_TIMEOUT",
+    "DEFAULT_STALE_THRESHOLD",
+    "TOKEN_REFRESH_BUFFER",
+    "RECONNECT_BACKOFF",
+]
+
+
+DEFAULT_CONNECT_TIMEOUT: float = DEFAULT_CONNECT_TIMEOUT_MS / 1000
+DEFAULT_SEND_TIMEOUT: float = DEFAULT_SEND_TIMEOUT_MS / 1000
+DEFAULT_RESYNC_TIMEOUT: float = DEFAULT_RESYNC_TIMEOUT_MS / 1000
+DEFAULT_PING_INTERVAL: float = DEFAULT_PING_INTERVAL_MS / 1000
+DEFAULT_PONG_TIMEOUT: float = DEFAULT_PONG_TIMEOUT_MS / 1000
+DEFAULT_STALE_THRESHOLD: float = DEFAULT_STALE_THRESHOLD_MS / 1000
+TOKEN_REFRESH_BUFFER: float = TOKEN_REFRESH_BUFFER_MS / 1000
+
+RECONNECT_BACKOFF: list[float] = [value / 1000 for value in RECONNECT_BACKOFF_MS]

cortex_sdk/errors.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ._generated_errors import GENERATED_ERROR_CATALOG
+
+
+@dataclass(frozen=True)
+class _ErrorEntry:
+    code: str
+    retryable: bool
+    fatal: bool
+
+
+_CATALOG: list[_ErrorEntry] = [
+    _ErrorEntry(entry.code, retryable=entry.retryable, fatal=entry.fatal)
+    for entry in GENERATED_ERROR_CATALOG
+]
+
+_CATALOG_MAP: dict[str, _ErrorEntry] = {e.code: e for e in _CATALOG}
+
+
+class CortexError(Exception):
+    """SDK error with canonical code, retryable, and fatal flags."""
+
+    def __init__(self, code: str, message: str, retryable: bool, fatal: bool) -> None:
+        super().__init__(message)
+        self.code = code
+        self.retryable = retryable
+        self.fatal = fatal
+
+    def __repr__(self) -> str:
+        return (
+            f"CortexError(code={self.code!r}, message={str(self)!r}, "
+            f"retryable={self.retryable}, fatal={self.fatal})"
+        )
+
+
+def make_error(code: str, message: str) -> CortexError:
+    """Build a CortexError, filling retryable/fatal from the catalog."""
+    entry = _CATALOG_MAP.get(code)
+    if entry is not None:
+        return CortexError(entry.code, message, entry.retryable, entry.fatal)
+    # Unknown code — conservative defaults
+    return CortexError(code, message, retryable=False, fatal=False)
+
+
+def lookup_error(code: str) -> _ErrorEntry | None:
+    return _CATALOG_MAP.get(code)

cortex_sdk/liveness.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import asyncio
+import secrets
+import time
+from typing import Callable
+
+from .constants import SCHEMA_VERSION
+from .transport import Transport
+
+
+class _LivenessCallbacks:
+    on_stale: Callable[[], None]
+    get_session_id: Callable[[], str | None]
+    get_channel_id: Callable[[], str]
+
+
+class LivenessMonitor:
+    """Sends system::ping on a fixed interval and monitors pong responses."""
+
+    def __init__(
+        self,
+        transport: Transport,
+        ping_interval: float,
+        pong_timeout: float,
+        stale_threshold: float,
+        on_stale: Callable[[], None],
+        get_session_id: Callable[[], str | None],
+        get_channel_id: Callable[[], str],
+    ) -> None:
+        self._transport = transport
+        self._ping_interval = ping_interval
+        self._pong_timeout = pong_timeout
+        self._stale_threshold = stale_threshold
+        self._on_stale = on_stale
+        self._get_session_id = get_session_id
+        self._get_channel_id = get_channel_id
+
+        self._running = False
+        self._last_pong_time = time.monotonic()
+        self._pending_heartbeat_id: str | None = None
+        self._task: asyncio.Task[None] | None = None
+
+    def start(self) -> None:
+        self._running = True
+        self._last_pong_time = time.monotonic()
+        self._task = asyncio.create_task(self._ping_loop())
+
+    def stop(self) -> None:
+        self._running = False
+        if self._task and not self._task.done():
+            self._task.cancel()
+        self._task = None
+        self._pending_heartbeat_id = None
+
+    def record_pong(self, heartbeat_id: str) -> None:
+        """Called when a system::pong is received."""
+        if heartbeat_id == self._pending_heartbeat_id:
+            self._pending_heartbeat_id = None
+            self._last_pong_time = time.monotonic()
+
+    async def _ping_loop(self) -> None:
+        try:
+            while self._running:
+                await asyncio.sleep(self._ping_interval)
+                if not self._running:
+                    break
+
+                # Check stale threshold before sending ping
+                if time.monotonic() - self._last_pong_time > self._stale_threshold:
+                    self._on_stale()
+                    return
+
+                heartbeat_id = f"hb_{secrets.token_hex(4)}"
+                self._pending_heartbeat_id = heartbeat_id
+
+                envelope: dict[str, object] = {
+                    "type": "system::ping",
+                    "schema": SCHEMA_VERSION,
+                    "payload": {
+                        "heartbeat_id": heartbeat_id,
+                        "channel_id": self._get_channel_id(),
+                    },
+                    "ts": _now_iso(),
+                }
+                session_id = self._get_session_id()
+                if session_id:
+                    envelope["session_id"] = session_id
+
+                try:
+                    await self._transport.send(envelope)
+                except Exception:
+                    pass  # send errors are handled via transport on_close
+
+                # Schedule pong-timeout check as a separate task
+                asyncio.create_task(self._check_pong_timeout(heartbeat_id))
+        except asyncio.CancelledError:
+            pass
+
+    async def _check_pong_timeout(self, heartbeat_id: str) -> None:
+        await asyncio.sleep(self._pong_timeout)
+        if not self._running:
+            return
+        if self._pending_heartbeat_id == heartbeat_id:
+            if time.monotonic() - self._last_pong_time > self._stale_threshold:
+                self._on_stale()
+
+
+def _now_iso() -> str:
+    from datetime import datetime, timezone
+    return datetime.now(timezone.utc).isoformat()

cortex_sdk/py.typed

@@ -0,0 +1 @@
+

cortex_sdk/session.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import json
+from typing import Callable
+
+from .constants import SCHEMA_VERSION
+from .errors import make_error, lookup_error, CortexError
+from .transport import Transport
+from .types import CortexMessage, RuntimeBootstrap, SessionState
+
+
+_TERMINAL_STATES: set[SessionState] = {
+    "COMPLETED",
+    "FAILED",
+    "STOPPED",
+    "TIMEOUT",
+    "CANCELLED",
+}
+_LIFECYCLE_STATE_MAP: dict[str, SessionState] = {
+    "active": "ACTIVE",
+    "waiting": "WAITING",
+    "completed": "COMPLETED",
+    "failed": "FAILED",
+    "stopped": "STOPPED",
+    "timeout": "TIMEOUT",
+    "cancelled": "CANCELLED",
+}
+
+
+class SessionManager:
+    """Manages session state and envelope construction/parsing."""
+
+    def __init__(
+        self,
+        on_message: Callable[[CortexMessage], None],
+        on_fatal_error: Callable[[CortexError], None],
+    ) -> None:
+        self._on_message = on_message
+        self._on_fatal_error = on_fatal_error
+
+        self._session_id: str | None = None
+        self._session_state: SessionState = "CREATED"
+        self._last_seq: int = 0
+        self._transport: Transport | None = None
+        self._send_timeout: float = 10.0
+
+    def _set_session_state(self, next_state: SessionState) -> None:
+        if self._session_state in _TERMINAL_STATES and self._session_state != next_state:
+            return
+        self._session_state = next_state
+
+    # ── accessors ─────────────────────────────────────────────────────────────
+
+    @property
+    def session_id(self) -> str | None:
+        return self._session_id
+
+    @property
+    def session_state(self) -> SessionState:
+        return self._session_state
+
+    @property
+    def last_seq(self) -> int:
+        return self._last_seq
+
+    # ── configuration ─────────────────────────────────────────────────────────
+
+    def set_transport(self, transport: Transport, send_timeout: float) -> None:
+        self._transport = transport
+        self._send_timeout = send_timeout
+
+    # ── outbound helpers ──────────────────────────────────────────────────────
+
+    def _build_envelope(
+        self, type_: str, payload: dict[str, object]
+    ) -> dict[str, object]:
+        env: dict[str, object] = {
+            "type": type_,
+            "schema": SCHEMA_VERSION,
+            "payload": payload,
+            "ts": _now_iso(),
+        }
+        if self._session_id:
+            env["session_id"] = self._session_id
+        return env
+
+    async def _send(self, envelope: dict[str, object]) -> None:
+        if self._transport is None:
+            raise make_error("transport_send_timeout", "No transport configured")
+        await self._transport.send(envelope)
+
+    async def send_init(self, bootstrap: RuntimeBootstrap) -> None:
+        """Send system::init — intentionally omits session_id."""
+        self._session_state = "INITIALIZING"
+        envelope: dict[str, object] = {
+            "type": "system::init",
+            "schema": SCHEMA_VERSION,
+            "payload": dict(bootstrap),
+            "ts": _now_iso(),
+        }
+        await self._send(envelope)
+
+    async def send_resync(self) -> None:
+        await self._send(
+            self._build_envelope("system::resync", {"last_seq": self._last_seq})
+        )
+
+    async def send_stop(self) -> None:
+        await self._send(self._build_envelope("sandbox::stop", {}))
+
+    async def send_chat_message(
+        self, content: str, attachments: list[str] | None
+    ) -> None:
+        payload: dict[str, object] = {"content": content, "role": "user"}
+        if attachments:
+            payload["attachments"] = attachments
+        await self._send(self._build_envelope("chat::message", payload))
+
+    # ── inbound handling ──────────────────────────────────────────────────────
+
+    def handle_incoming(self, raw: str) -> None:
+        """Parse a raw JSON frame and dispatch callbacks."""
+        try:
+            parsed: dict[str, object] = json.loads(raw)
+            msg: CortexMessage = parsed  # type: ignore[assignment]
+        except Exception:
+            return  # malformed frame — ignore
+
+        # Learn session_id from first server message
+        if not self._session_id and msg.get("session_id"):
+            self._session_id = msg["session_id"]
+            if self._session_state not in _TERMINAL_STATES:
+                self._session_state = "ACTIVE"
+
+        # Track highest sequence number
+        seq = msg.get("seq")
+        if isinstance(seq, int) and seq > self._last_seq:
+            self._last_seq = seq
+
+        self._update_session_state(msg)
+
+        # Fatal error detection
+        if msg.get("type") == "system::error":
+            self._handle_system_error(msg["payload"])
+
+        self._on_message(msg)
+
+    def _handle_system_error(self, payload: dict[str, object]) -> None:
+        code = payload.get("code", "session_terminal")
+        message = payload.get("message", "Runtime error")
+        entry = lookup_error(str(code))
+        if entry and entry.fatal:
+            self._set_session_state("FAILED")
+            err = make_error(str(code), str(message))
+            self._on_fatal_error(err)
+
+    def _update_session_state(self, msg: CortexMessage) -> None:
+        msg_type = msg.get("type")
+
+        if msg_type == "sandbox::snapshot":
+            state = msg["payload"].get("state")
+            if isinstance(state, str) and state.lower() == "waiting":
+                self._set_session_state("WAITING")
+            return
+
+        if msg_type == "sandbox::lifecycle":
+            status = msg["payload"].get("status")
+            if isinstance(status, str):
+                next_state = _LIFECYCLE_STATE_MAP.get(status.lower())
+                if next_state is not None:
+                    self._set_session_state(next_state)
+            return
+
+        if msg_type == "system::error":
+            code = msg["payload"].get("code")
+            if isinstance(code, str):
+                entry = lookup_error(code)
+                if entry and entry.fatal:
+                    self._set_session_state("FAILED")
+
+
+def _now_iso() -> str:
+    from datetime import datetime, timezone
+    return datetime.now(timezone.utc).isoformat()

cortex_sdk/transport.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Callable
+
+import websockets
+import websockets.asyncio.client
+import websockets.exceptions
+
+from .constants import WS_SUBPROTOCOL, WS_SUBPROTOCOL_JWT_PREFIX
+from .errors import make_error
+
+
+class Transport:
+    """Thin WebSocket abstraction. Platform-agnostic — no business logic."""
+
+    def __init__(self, connect_timeout: float, send_timeout: float) -> None:
+        self._connect_timeout = connect_timeout
+        self._send_timeout = send_timeout
+        self._ws: websockets.asyncio.client.ClientConnection | None = None
+        self._reader_task: asyncio.Task[None] | None = None
+
+        # Callbacks set by the consumer (client.py)
+        self.on_message: Callable[[str], None] | None = None
+        self.on_close: Callable[[int, str], None] | None = None
+        self.on_error: Callable[[Exception], None] | None = None
+
+    async def open(self, ws_url: str, access_token: str) -> None:
+        """Open the WebSocket connection and start the reader loop."""
+        subprotocols: list[str] = [
+            WS_SUBPROTOCOL,
+            f"{WS_SUBPROTOCOL_JWT_PREFIX}{access_token}",
+        ]
+        try:
+            from websockets.typing import Subprotocol
+            typed_protocols = [Subprotocol(p) for p in subprotocols]
+            ws = await asyncio.wait_for(
+                websockets.connect(
+                    ws_url,
+                    subprotocols=typed_protocols,
+                    ping_interval=None,  # SDK manages its own liveness
+                ),
+                timeout=self._connect_timeout,
+            )
+        except asyncio.TimeoutError as exc:
+            raise make_error(
+                "transport_connect_timeout", "WebSocket connect timed out"
+            ) from exc
+        except Exception as exc:
+            raise make_error(
+                "transport_connect_timeout", f"WebSocket connect failed: {exc}"
+            ) from exc
+
+        self._ws = ws
+        self._reader_task = asyncio.create_task(self._reader_loop())
+
+    async def send(self, message: dict[str, object]) -> None:
+        """Serialize and send a message. Raises transport_send_timeout on timeout."""
+        ws = self._ws
+        if ws is None:
+            raise make_error("transport_send_timeout", "No open connection")
+        data = json.dumps(message)
+        try:
+            await asyncio.wait_for(ws.send(data), timeout=self._send_timeout)
+        except asyncio.TimeoutError as exc:
+            raise make_error("transport_send_timeout", "Send timed out") from exc
+        except Exception as exc:
+            raise make_error("transport_send_timeout", str(exc)) from exc
+
+    def close(self, code: int = 1000, reason: str = "disconnect") -> None:
+        """Synchronous close — cancels reader task and schedules WS close."""
+        task = self._reader_task
+        ws = self._ws
+        self._reader_task = None
+        self._ws = None
+
+        if task and not task.done():
+            task.cancel()
+        if ws is not None:
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(ws.close(code, reason))
+            except RuntimeError:
+                pass  # no running loop — connection already gone
+
+    async def aclose(self, code: int = 1000, reason: str = "disconnect") -> None:
+        """Async close — waits for reader task to finish."""
+        task = self._reader_task
+        ws = self._ws
+        self._reader_task = None
+        self._ws = None
+
+        if task and not task.done():
+            task.cancel()
+            try:
+                await task
+            except (asyncio.CancelledError, Exception):
+                pass
+        if ws is not None:
+            try:
+                await ws.close(code, reason)
+            except Exception:
+                pass
+
+    async def _reader_loop(self) -> None:
+        ws = self._ws
+        if ws is None:
+            return
+        try:
+            async for raw in ws:
+                data = raw if isinstance(raw, str) else raw.decode()
+                if self.on_message:
+                    self.on_message(data)
+        except websockets.exceptions.ConnectionClosed as exc:
+            # Only fire on_close if we weren't already cleaned up
+            if self._ws is None:
+                return  # transport.close() was called — suppress callback
+            # .rcvd is the received close frame (None if connection was aborted)
+            rcvd = exc.rcvd
+            code = rcvd.code if rcvd is not None else 1006
+            reason = rcvd.reason if rcvd is not None else ""
+            if self.on_close:
+                self.on_close(code, reason)
+        except asyncio.CancelledError:
+            pass  # clean shutdown from close() / aclose()
+        except Exception as exc:
+            if self.on_error:
+                self.on_error(exc)

cortex_sdk/types.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import Callable, Literal
+
+from typing_extensions import TypedDict, NotRequired
+
+
+class RuntimeBootstrap(TypedDict):
+    execution_mode: str
+    bundle_url: str
+    checksum: str
+    artifact_id: NotRequired[str | None]
+    artifact_kind: NotRequired[str | None]
+    run_mode: NotRequired[str | None]
+
+
+class AuthTokenResponse(TypedDict):
+    ws_url: str
+    access_token: str
+    refresh_token: str
+    runtime_bootstrap: RuntimeBootstrap
+
+
+class CortexMessage(TypedDict):
+    type: str
+    schema: str
+    session_id: str
+    seq: NotRequired[int | None]
+    payload: dict[str, object]
+    meta: NotRequired[dict[str, object] | None]
+    ts: str
+
+
+SessionState = Literal[
+    "CREATED",
+    "INITIALIZING",
+    "ACTIVE",
+    "WAITING",
+    "COMPLETED",
+    "FAILED",
+    "STOPPED",
+    "TIMEOUT",
+    "CANCELLED",
+]
+
+ChannelState = Literal[
+    "CONNECTING",
+    "OPEN",
+    "STALE",
+    "RECONNECTING",
+    "CLOSED",
+    "AUTH_FAILED",
+]
+
+MessageCallback = Callable[["CortexMessage"], None]

cortex_sdk/upload.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import BinaryIO
+
+import httpx
+
+from .errors import make_error
+
+# Default upload endpoint — overridden via client's _upload_url for tests
+_DEFAULT_UPLOAD_URL = "/upload"
+
+
+async def upload_file(
+    file: str | bytes | BinaryIO,
+    access_token: str,
+    upload_url: str = _DEFAULT_UPLOAD_URL,
+) -> str:
+    """Upload a file and return the attachment_id.
+
+    Args:
+        file: file path string, raw bytes, or a file-like (BinaryIO) object.
+        access_token: Bearer token for Authorization header.
+        upload_url: full URL of the upload endpoint.
+    """
+    if isinstance(file, str):
+        with open(file, "rb") as fh:
+            data = fh.read()
+    elif isinstance(file, bytes):
+        data = file
+    else:
+        data = file.read()  # BinaryIO
+
+    files = {"file": ("upload", data, "application/octet-stream")}
+
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(
+            upload_url,
+            headers={"Authorization": f"Bearer {access_token}"},
+            files=files,
+        )
+
+    if not resp.is_success:
+        if resp.status_code == 413:
+            raise make_error("upload_too_large", "File exceeds the allowed size limit")
+        if resp.status_code == 415:
+            raise make_error("upload_type_rejected", "File type not accepted by the runtime")
+        raise make_error("upload_failed", f"Upload failed with status {resp.status_code}")
+
+    body = resp.json()
+    return str(body["attachment_id"])

cortex_suite_sdk-1.0.2.dist-info/METADATA

@@ -0,0 +1,77 @@
+Metadata-Version: 2.4
+Name: cortex-suite-sdk
+Version: 1.0.2
+Summary: Cortex SDK — transport client for Python
+License: MIT
+Requires-Python: >=3.10
+Requires-Dist: httpx>=0.27
+Requires-Dist: websockets>=12.0
+Provides-Extra: dev
+Requires-Dist: anyio[trio]; extra == 'dev'
+Requires-Dist: build>=1.2; extra == 'dev'
+Requires-Dist: jsonschema>=4.23; extra == 'dev'
+Requires-Dist: mypy>=1.9; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# Cortex SDK for Python
+
+## Install
+
+`pip install cortex-suite-sdk`
+
+## Quick start
+
+```python
+import asyncio
+
+from cortex_sdk import CortexClient
+
+
+def on_message(msg):
+    print(msg["type"], msg["payload"])
+
+
+async def main():
+    client = CortexClient(
+        api_key="your-api-key",
+        # auth_url="https://auth.cortexsuite.app",  # optional override
+        on_message=on_message,
+    )
+    await client.connect()
+    await client.send_message(content="Hello")
+
+
+asyncio.run(main())
+```
+
+## API
+
+See the [full API reference](../docs/04_api_reference.md).
+`auth_url` is optional; if omitted, the SDK uses its default auth base URL.
+
+## Error handling
+
+```python
+from cortex_sdk import CortexClient, CortexError
+
+
+def on_message(msg):
+    if msg["type"] == "system::error":
+        print("Runtime error:", msg["payload"]["code"])
+        if msg["payload"].get("fatal"):
+            pass  # handle unrecoverable session
+
+
+client = CortexClient(
+    api_key="your-api-key",
+    on_message=on_message,
+)
+
+try:
+    await client.connect()
+except CortexError as err:
+    if err.code in ("auth_invalid", "auth_refresh_failed"):
+        pass  # prompt re-authentication
+```

cortex_suite_sdk-1.0.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+cortex_sdk/__init__.py,sha256=Uyo3yc5EOYIL-BKuBvxJryUB0QSdcp8cz8G1CmQu_Aw,108
+cortex_sdk/_generated_constants.py,sha256=MFCjxFMCRETv58DJMrf2jQ_pJutFnXHcsa2xVdAAO3E,1033
+cortex_sdk/_generated_errors.py,sha256=vLvglZq0IUwpM-FOxlGr_LdD_eeOCc2yjim07SD1ZMw,1315
+cortex_sdk/auth.py,sha256=oKnHtQUB5156SQVL_doUEQ6xdTFt74RqQS2_YiFmHTw,3758
+cortex_sdk/client.py,sha256=WkAC47psjKP648YI7p9pSLgTjx0LYiRgUUHY_cxfQHc,12421
+cortex_sdk/constants.py,sha256=yUcfAazvIcQWOOIRpT0UYy9JLPHoEdkPBGRNrQcAhJY,1483
+cortex_sdk/errors.py,sha256=D1Na6ZBIt8fWt6JsMsy3Gq-vYSVHRxDE9rPCjljRZYk,1415
+cortex_sdk/liveness.py,sha256=kfHvNWIViwAyKko6PALs4p0uf1LBj2woPipb8DaOEdg,3725
+cortex_sdk/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+cortex_sdk/session.py,sha256=Pjtxmr4t7VYU5IwkGEiucwWHqB4ryty6jtXfhax3xnA,6622
+cortex_sdk/transport.py,sha256=NxS72EGKL77G81hrooMFT-p_ip8ejd8eOPF7U7U8Zq8,4890
+cortex_sdk/types.py,sha256=98pTdLj-NIi5fPttpvfDaCeYwwoHgHJF_wrl65yD5DQ,1041
+cortex_sdk/upload.py,sha256=vC72ogwmmunyEjY-cQOMcjDD2CQArsCpONGVy_gWMdw,1523
+cortex_suite_sdk-1.0.2.dist-info/METADATA,sha256=5wYvE55J7uNHr7DqROQWyGL27a-MRQAf-XPxDL7BEjw,1712
+cortex_suite_sdk-1.0.2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+cortex_suite_sdk-1.0.2.dist-info/RECORD,,

cortex_suite_sdk-1.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any