gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/store/gateway.py

@@ -0,0 +1,983 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""GatewayTracker and GatewayLightningStore - Gateway tracking integration.
+
+This module provides Gateway tracking integration:
+
+1. **GatewayLightningStore** (recommended): Creates a real GatewayRun for experiment
+   tracking with automatic session management, heartbeats, and crash detection.
+   ```python
+   from gatewaysdk.store import GatewayLightningStore
+
+   store = GatewayLightningStore(
+       api_key="pk-lf-abc123",
+       secret_key="sk-lf-xyz789",
+       gateway_url="https://app.gateway.dev",
+       project_id="proj-123",
+       run_name="gepa-experiment-v1",
+       config={"lr": 0.01, "epochs": 100},
+   )
+   # Store creates a real GatewayRun record in the database
+   # Rollouts create real TraceSession records
+   # All traces are linked via run_id
+   ```
+
+2. **GatewayTracker** (DEPRECATED): Legacy event-based tracking with virtual job IDs.
+   Use GatewayLightningStore instead, which uses real GatewayRun records.
+
+Features:
+- Real GatewayRun records with lifecycle management (heartbeats, crash detection)
+- Session creation for rollouts (TraceSession records linked to runs)
+- OTLP trace export support
+- Automatic reward-to-score conversion
+"""
+
+from __future__ import annotations
+
+import atexit
+import base64
+import logging
+import os
+import queue
+import threading
+import time
+import uuid
+import warnings
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from .gateway_listener import GatewayRunListener
+from .memory import InMemoryLightningStore  # noqa: F401
+
+if TYPE_CHECKING:
+    from gatewaysdk.run import Run
+    from .client_server import LightningStoreClient  # noqa: F401
+    from .redis_stream import RedisStreamLightningStore  # noqa: F401
+
+import httpx
+
+from gatewaysdk.emitter.reward import get_rewards_from_span, is_reward_span
+from gatewaysdk.types import (
+    Attempt,
+    ResourcesUpdate,
+    Rollout,
+    Span,
+)
+
+from .listener import StorageListener
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GatewayEvent:
+    """An event to be sent to the Gateway API."""
+
+    id: str
+    type: str
+    timestamp: str  # ISO 8601 datetime string
+    data: Dict[str, Any]
+
+
+class GatewayTracker:
+    """DEPRECATED: A StorageListener that streams storage events to Gateway.
+
+    .. deprecated::
+        GatewayTracker is deprecated. Use GatewayLightningStore instead, which
+        creates real GatewayRun records with proper lifecycle management, heartbeats,
+        and crash detection.
+
+    This tracker implements the StorageListener protocol and can be attached
+    to any LightningStore to enable Gateway tracking. However, it uses virtual
+    job IDs that are not persisted to the database.
+
+    For new code, use GatewayLightningStore which:
+    - Creates real GatewayRun records in PostgreSQL
+    - Has automatic heartbeats and crash detection
+    - Creates real TraceSession records for rollouts
+    - Links all traces via run_id
+
+    Args:
+        api_key: Gateway public API key for authentication.
+        secret_key: Gateway secret key for authentication.
+        gateway_url: Gateway server URL (e.g., "http://localhost:3000").
+        project_id: Project ID to associate events with.
+        flush_interval: Seconds between automatic flushes (default: 1.0).
+        max_buffer_size: Maximum events before forcing a flush (default: 1000).
+        request_timeout: HTTP request timeout in seconds (default: 10.0).
+        max_retries: Maximum retry attempts for failed requests (default: 3).
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        secret_key: str,
+        gateway_url: str,
+        project_id: str,
+        flush_interval: float = 1.0,
+        max_buffer_size: int = 1000,
+        request_timeout: float = 10.0,
+        max_retries: int = 3,
+    ) -> None:
+        warnings.warn(
+            "GatewayTracker is deprecated. Use GatewayLightningStore instead, which "
+            "creates real GatewayRun records with proper lifecycle management.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        # Store configuration
+        self._api_key = api_key
+        self._secret_key = secret_key
+        self._gateway_url = gateway_url.rstrip("/")
+        self._project_id = project_id
+        self._flush_interval = flush_interval
+        self._max_buffer_size = max_buffer_size
+        self._request_timeout = request_timeout
+        self._max_retries = max_retries
+
+        # Generate a unique job ID for this tracker instance
+        self._job_id = f"job-{uuid.uuid4().hex[:12]}"
+
+        # Track if job has been completed to prevent double-complete
+        self._completed = False
+
+        # Event buffer (thread-safe queue)
+        self._event_buffer: queue.Queue[GatewayEvent] = queue.Queue()
+
+        # Background sender thread control
+        self._stop_event = threading.Event()
+        self._sender_thread: Optional[threading.Thread] = None
+
+        # Start the background sender thread
+        self._start_sender_thread()
+
+        # Emit job.created event immediately
+        self._emit_job_created()
+
+        # Register cleanup on exit
+        atexit.register(self._cleanup)
+
+        logger.info(
+            f"GatewayTracker initialized - streaming to {self._gateway_url} "
+            f"(project={self._project_id}, job={self._job_id})"
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # StorageListener Protocol Implementation
+    # ─────────────────────────────────────────────────────────────
+
+    @property
+    def capabilities(self) -> Dict[str, bool]:
+        """Return the capabilities of the listener."""
+        return {
+            "otlp_traces": True,  # Enable OTLP trace export to Gateway
+        }
+
+    @property
+    def job_id(self) -> str:
+        """Return the job ID for this tracker instance."""
+        return self._job_id
+
+    def otlp_traces_endpoint(self) -> Optional[str]:
+        """Return the OTLP/HTTP traces endpoint."""
+        endpoint = f"{self._gateway_url}/api/public/otel/v1/traces"
+        logger.debug(f"OTLP traces endpoint: {endpoint}")
+        return endpoint
+
+    def get_otlp_headers(self) -> Dict[str, str]:
+        """Return the authentication headers for OTLP export.
+
+        Gateway's OTLP endpoint uses Basic Auth with format: public_key:secret_key
+        """
+        credentials = f"{self._api_key}:{self._secret_key}"
+        encoded = base64.b64encode(credentials.encode()).decode()
+        return {
+            "Authorization": f"Basic {encoded}",
+        }
+
+    async def on_job_created(self, job_id: str, project_id: Optional[str] = None) -> None:
+        """Called when a job is created. (Usually handled internally)"""
+        pass  # We emit job.created in __init__
+
+    async def on_rollout_created(self, rollout: Rollout) -> None:
+        """Called when a rollout is created."""
+        self._emit(
+            "rollout.created",
+            {
+                "id": rollout.rollout_id,
+                "input": rollout.input,
+                "status": rollout.status,
+                "resource_id": rollout.resources_id,
+                "mode": rollout.mode,
+                "start_time": rollout.start_time,
+                "config": {
+                    "max_attempts": rollout.config.max_attempts,
+                    "retry_condition": rollout.config.retry_condition,
+                    "timeout_seconds": rollout.config.timeout_seconds,
+                    "unresponsive_seconds": rollout.config.unresponsive_seconds,
+                },
+                "metadata": rollout.metadata,
+            },
+        )
+
+    async def on_rollout_updated(self, rollout: Rollout) -> None:
+        """Called when a rollout is updated."""
+        data: Dict[str, Any] = {
+            "id": rollout.rollout_id,
+            "status": rollout.status,
+        }
+        if rollout.end_time is not None:
+            data["end_time"] = rollout.end_time
+        self._emit("rollout.status_changed", data)
+
+    async def on_attempt_created(self, attempt: Attempt) -> None:
+        """Called when an attempt is created."""
+        self._emit(
+            "attempt.created",
+            {
+                "id": attempt.attempt_id,
+                "rollout_id": attempt.rollout_id,
+                "sequence_id": attempt.sequence_id,
+                "status": attempt.status,
+                "start_time": attempt.start_time,
+                "worker_id": attempt.worker_id,
+            },
+        )
+
+    async def on_attempt_updated(self, attempt: Attempt, rollout_id: str) -> None:
+        """Called when an attempt is updated."""
+        data: Dict[str, Any] = {
+            "id": attempt.attempt_id,
+            "rollout_id": rollout_id,
+            "status": attempt.status,
+        }
+        if attempt.end_time is not None:
+            data["end_time"] = attempt.end_time
+        self._emit("attempt.status_changed", data)
+
+    async def on_span_created(self, span: Span) -> None:
+        """Called when a span is added.
+
+        If the span is a reward span, also sends the reward as an Gateway score.
+        """
+        # Emit span event to the event buffer
+        self._emit(
+            "span.emitted",
+            {
+                "trace_id": span.trace_id,
+                "span_id": span.span_id,
+                "parent_id": span.parent_id,
+                "attempt_id": span.attempt_id,
+                "rollout_id": span.rollout_id,
+                "sequence_id": span.sequence_id,
+                "name": span.name,
+                "status": {
+                    "status_code": span.status.status_code,
+                    "description": span.status.description,
+                },
+                "attributes": span.attributes,
+                "start_time": span.start_time,
+                "end_time": span.end_time,
+            },
+        )
+
+        # Check if this is a reward span and send as Gateway score
+        if is_reward_span(span):
+            rewards = get_rewards_from_span(span)
+            for reward in rewards:
+                self._send_score(
+                    name=reward.name,
+                    value=reward.value,
+                    trace_id=span.trace_id,
+                    observation_id=span.span_id,
+                    rollout_id=span.rollout_id,
+                    attempt_id=span.attempt_id,
+                )
+
+    async def on_resource_registered(self, resource: ResourcesUpdate) -> None:
+        """Called when a resource is registered/updated.
+
+        Sends FULL resource content for complete experiment tracking.
+        """
+        # Serialize resources to JSON-compatible dicts
+        # Resources are Pydantic models (PromptTemplate, LLM, etc.) which need model_dump()
+        serialized_resources: Dict[str, Any] = {}
+        for name, res in resource.resources.items():
+            if hasattr(res, "model_dump"):
+                serialized_resources[name] = res.model_dump()
+            elif hasattr(res, "dict"):
+                # Fallback for older Pydantic v1 models
+                serialized_resources[name] = res.dict()
+            else:
+                # Already a dict or primitive
+                serialized_resources[name] = res
+
+        self._emit(
+            "resource.registered",
+            {
+                "id": resource.resources_id,
+                "version": resource.version,
+                "create_time": resource.create_time,
+                "update_time": resource.update_time,
+                "resources": serialized_resources,
+            },
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # Event Emission Helpers
+    # ─────────────────────────────────────────────────────────────
+
+    def _emit(self, event_type: str, data: Dict[str, Any]) -> None:
+        """Add an event to the buffer for later sending."""
+        # Format: 2026-01-17T19:51:40.123Z (Zod datetime expects 'Z' suffix for UTC)
+        now = datetime.now(timezone.utc)
+        timestamp = now.strftime("%Y-%m-%dT%H:%M:%S.") + f"{now.microsecond // 1000:03d}Z"
+        event = GatewayEvent(
+            id=f"evt-{uuid.uuid4().hex[:8]}",
+            type=event_type,
+            timestamp=timestamp,
+            data=data,
+        )
+        try:
+            self._event_buffer.put_nowait(event)
+        except queue.Full:
+            logger.warning(f"Event buffer full, dropping event: {event_type}")
+            return
+
+        # Force flush if buffer is at capacity
+        if self._event_buffer.qsize() >= self._max_buffer_size:
+            self._trigger_flush()
+
+    def _emit_job_created(self) -> None:
+        """Emit the job.created event when the tracker is initialized."""
+        self._emit(
+            "job.created",
+            {
+                "project_id": self._project_id,
+                "type": "agent",
+            },
+        )
+        # Immediately flush to ensure job is created before other events
+        self._flush_events()
+
+    def _send_score(
+        self,
+        name: str,
+        value: float,
+        trace_id: str,
+        observation_id: Optional[str] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> None:
+        """Send a score to Gateway's /api/public/scores endpoint.
+
+        Args:
+            name: Score name (e.g., "primary", "task_completion").
+            value: Numeric score value.
+            trace_id: OTEL trace ID to link the score to.
+            observation_id: Optional span ID to link to specific observation.
+            rollout_id: Optional rollout ID for metadata.
+            attempt_id: Optional attempt ID for metadata.
+        """
+        score_payload: Dict[str, Any] = {
+            "name": name,
+            "value": value,
+            "dataType": "NUMERIC",
+            "traceId": trace_id,
+        }
+
+        if observation_id:
+            score_payload["observationId"] = observation_id
+
+        # Add rollout/attempt context as metadata
+        metadata: Dict[str, Any] = {
+            "source": "agent_lightning",
+            "job_id": self._job_id,
+        }
+        if rollout_id:
+            metadata["rollout_id"] = rollout_id
+        if attempt_id:
+            metadata["attempt_id"] = attempt_id
+        score_payload["metadata"] = metadata
+
+        # Create Basic auth header
+        auth_string = f"{self._api_key}:{self._secret_key}"
+        auth_bytes = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
+
+        # Send score in a non-blocking way (fire and forget with retries)
+        def send_score_async() -> None:
+            for attempt in range(self._max_retries):
+                try:
+                    with httpx.Client(timeout=self._request_timeout) as client:
+                        response = client.post(
+                            f"{self._gateway_url}/api/public/scores",
+                            headers={
+                                "Authorization": f"Basic {auth_bytes}",
+                                "Content-Type": "application/json",
+                            },
+                            json=score_payload,
+                        )
+                        response.raise_for_status()
+                        logger.debug(f"Successfully sent score '{name}={value}' to Gateway")
+                        return
+                except httpx.HTTPStatusError as e:
+                    if e.response.status_code == 401:
+                        logger.error("Unauthorized (401) sending score to Gateway. Check API credentials.")
+                        return
+                    else:
+                        try:
+                            error_body = e.response.json()
+                        except Exception:
+                            error_body = e.response.text
+                        logger.warning(
+                            f"HTTP error sending score (attempt {attempt + 1}/{self._max_retries}): "
+                            f"{e.response.status_code} - {error_body}"
+                        )
+                except httpx.HTTPError as e:
+                    logger.warning(f"Failed to send score (attempt {attempt + 1}/{self._max_retries}): {e}")
+
+                # Exponential backoff
+                if attempt < self._max_retries - 1:
+                    backoff_time = 2**attempt
+                    time.sleep(backoff_time)
+
+            logger.error(f"Failed to send score '{name}' after {self._max_retries} retries")
+
+        # Run in a thread to avoid blocking
+        threading.Thread(target=send_score_async, daemon=True, name="gateway-score-sender").start()
+
+    # ─────────────────────────────────────────────────────────────
+    # Background Sender Thread
+    # ─────────────────────────────────────────────────────────────
+
+    def _start_sender_thread(self) -> None:
+        """Start the background sender thread."""
+
+        def sender_loop() -> None:
+            while not self._stop_event.is_set():
+                # Wait for the flush interval or until stopped
+                self._stop_event.wait(timeout=self._flush_interval)
+                if not self._stop_event.is_set():
+                    self._flush_events()
+
+        self._sender_thread = threading.Thread(
+            target=sender_loop,
+            daemon=True,
+            name="gateway-sender",
+        )
+        self._sender_thread.start()
+
+    def _trigger_flush(self) -> None:
+        """Trigger an immediate flush by interrupting the wait."""
+        # The flush will happen on the next iteration since we're using a timeout
+        pass
+
+    def _flush_events(self) -> None:
+        """Flush all buffered events to the Gateway API."""
+        events: List[GatewayEvent] = []
+
+        # Drain the queue
+        while True:
+            try:
+                event = self._event_buffer.get_nowait()
+                events.append(event)
+            except queue.Empty:
+                break
+
+        if not events:
+            return
+
+        self._send_events(events)
+
+    def _send_events(self, events: List[GatewayEvent]) -> None:
+        """Send events to the Gateway API with retry logic."""
+        if not events:
+            return
+
+        payload = {
+            "job_id": self._job_id,
+            # Note: project_id is not sent - it's derived from the API key auth
+            "events": [
+                {
+                    "id": e.id,
+                    "type": e.type,
+                    "timestamp": e.timestamp,
+                    "data": e.data,
+                }
+                for e in events
+            ],
+        }
+
+        # Create Basic auth header from api_key:secret_key
+        auth_string = f"{self._api_key}:{self._secret_key}"
+        auth_bytes = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
+
+        for attempt in range(self._max_retries):
+            try:
+                with httpx.Client(timeout=self._request_timeout) as client:
+                    response = client.post(
+                        f"{self._gateway_url}/api/public/v1/agent/ingest",
+                        headers={
+                            "Authorization": f"Basic {auth_bytes}",
+                            "Content-Type": "application/json",
+                        },
+                        json=payload,
+                    )
+                    response.raise_for_status()
+                    logger.debug(f"Successfully sent {len(events)} events to Gateway")
+                    return
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 401:
+                    # Unauthorized - bad API key, won't fix itself
+                    logger.error("Unauthorized (401) sending events to Gateway. Check API credentials.")
+                    return
+                elif e.response.status_code == 429:
+                    # Rate limited - back off
+                    logger.warning("Rate limited (429) sending events, retrying with backoff...")
+                else:
+                    # Log the full response for debugging
+                    try:
+                        error_body = e.response.json()
+                    except Exception:
+                        error_body = e.response.text
+                    logger.warning(
+                        f"HTTP error sending events (attempt {attempt + 1}/{self._max_retries}): "
+                        f"{e.response.status_code} - {error_body}"
+                    )
+            except httpx.HTTPError as e:
+                logger.warning(f"Failed to send events (attempt {attempt + 1}/{self._max_retries}): {e}")
+
+            # Exponential backoff
+            if attempt < self._max_retries - 1:
+                backoff_time = 2**attempt
+                time.sleep(backoff_time)
+
+        logger.error(f"Failed to send {len(events)} events after {self._max_retries} retries - events dropped")
+
+    # ─────────────────────────────────────────────────────────────
+    # Lifecycle Methods
+    # ─────────────────────────────────────────────────────────────
+
+    def complete(self, summary: Optional[Dict[str, Any]] = None) -> None:
+        """Mark the job as complete and flush all remaining events."""
+        if self._completed:
+            return
+        self._completed = True
+        self._emit("job.completed", {"summary": summary or {}})
+        self._stop_event.set()
+        self._flush_events()
+        logger.info(f"GatewayTracker job {self._job_id} completed")
+
+    def fail(self, error: str) -> None:
+        """Mark the job as failed and flush all remaining events."""
+        self._emit("job.failed", {"error": error})
+        self._stop_event.set()
+        self._flush_events()
+        logger.error(f"GatewayTracker job {self._job_id} failed: {error}")
+
+    def _cleanup(self) -> None:
+        """Cleanup resources on exit."""
+        if not self._stop_event.is_set():
+            # Flush any remaining events before exiting
+            self._flush_events()
+            self._stop_event.set()
+
+    def __enter__(self) -> "GatewayTracker":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> None:
+        if exc_type:
+            self.fail(str(exc_val) if exc_val else "Unknown error")
+        elif not self._completed:
+            self.complete()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GatewayLightningStore - Uses GatewayRunListener for consolidated tracking
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def GatewayLightningStore(
+    *,
+    api_key: str,
+    secret_key: str,
+    gateway_url: str,
+    project_id: str,
+    # Run configuration
+    run_name: Optional[str] = None,
+    run_id: Optional[str] = None,
+    config: Optional[Dict[str, Any]] = None,
+    labels: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    source: str = "gateway_store",
+    heartbeat_interval: int = 30,
+    # Store configuration
+    thread_safe: bool = False,
+    init_tracing: bool = False,
+    # Legacy parameters (passed to listener)
+    flush_interval: float = 1.0,
+    max_buffer_size: int = 1000,
+    request_timeout: float = 10.0,
+    max_retries: int = 3,
+    **kwargs: Any,
+) -> "InMemoryLightningStore":
+    """Create an InMemoryLightningStore with Gateway tracking.
+
+    This function creates a real GatewayRun record in the database for experiment
+    tracking. All Gateway communication is consolidated in GatewayRunListener:
+    - Run creation and lifecycle management
+    - Heartbeats for crash detection
+    - Metrics logging
+    - Session creation for rollouts
+    - Score submission
+
+    Args:
+        api_key: Gateway public API key for authentication.
+        secret_key: Gateway secret key for authentication.
+        gateway_url: Gateway server URL (e.g., "http://localhost:3000").
+        project_id: Project ID to associate the run with.
+        run_name: Optional name for the run (auto-generated if not provided).
+        config: Hyperparameters/config dict (logged for reproducibility).
+        labels: Arbitrary key-value labels for filtering runs.
+        tags: Simple string tags for categorization.
+        source: Source identifier (default: "gateway_store").
+        heartbeat_interval: Seconds between heartbeats (default: 30).
+        thread_safe: Whether the underlying store is thread-safe (default: False).
+        init_tracing: If True, automatically initialize gatewaysdk.tracing with:
+            - Gateway OTLP exporter configured from the provided credentials
+            - TracingContextSpanProcessor for automatic context injection
+            - OpenInference instrumentors (litellm, openai, anthropic, etc.)
+
+    Returns:
+        An InMemoryLightningStore with GatewayRunListener attached.
+
+    Example:
+        ```python
+        # Basic usage
+        store = GatewayLightningStore(
+            api_key="pk-lf-abc123",
+            secret_key="sk-lf-xyz789",
+            gateway_url="https://app.gateway.dev",
+            project_id="proj-123",
+            run_name="gepa-experiment-v1",
+            config={"lr": 0.01, "epochs": 100},
+        )
+
+        # Log metrics directly on the store
+        store.log({"loss": 0.5, "accuracy": 0.92}, step=100)
+
+        # With automatic tracing SDK setup (recommended)
+        store = GatewayLightningStore(
+            api_key="pk-lf-abc123",
+            secret_key="sk-lf-xyz789",
+            gateway_url="https://app.gateway.dev",
+            project_id="proj-123",
+            init_tracing=True,
+        )
+
+        trainer = Trainer(algorithm=GEPA(...), store=store)
+        # Run is tracked with heartbeats, metrics, sessions, etc.
+
+        # When done, call complete() or fail()
+        store.complete()
+        ```
+    """
+    import os
+
+    # Set env vars for tracing SDK and other components
+    os.environ["GATEWAY_HOST"] = gateway_url
+    os.environ["GATEWAY_PUBLIC_KEY"] = api_key
+    os.environ["GATEWAY_SECRET_KEY"] = secret_key
+    os.environ["GATEWAY_PROJECT_ID"] = project_id
+
+    logger.debug(
+        f"GatewayLightningStore: Set env vars - "
+        f"GATEWAY_HOST={gateway_url}, "
+        f"GATEWAY_PUBLIC_KEY={api_key[:20]}..., "
+        f"GATEWAY_PROJECT_ID={project_id}"
+    )
+
+    # Initialize tracing SDK if requested
+    if init_tracing:
+        import gatewaysdk.tracing as tracing
+
+        tracing.init(instrument_default=True)
+
+    # Create consolidated listener - handles ALL Gateway communication
+    listener = GatewayRunListener(
+        api_key=api_key,
+        secret_key=secret_key,
+        gateway_url=gateway_url,
+        project_id=project_id,
+        run_name=run_name,
+        run_id=run_id,  # Reuse existing run if provided (e.g., platform runs)
+        config=config,
+        labels=labels,
+        tags=tags,
+        source=source,
+        heartbeat_interval=heartbeat_interval,
+        request_timeout=request_timeout,
+        max_retries=max_retries,
+    )
+
+    # Start the run (creates DB record, starts heartbeat, registers handlers)
+    listener.start()
+
+    # Create store with listener attached
+    store = InMemoryLightningStore(thread_safe=thread_safe, listeners=[listener], **kwargs)
+
+    # Attach listener to store for access
+    store._gateway_listener = listener  # type: ignore[attr-defined]
+
+    # Expose clean API methods on store instance
+    def log(
+        metrics,
+        value=None,
+        *,
+        step=None,
+        labels=None,
+        unit="",
+        metric_type="gauge",
+    ):
+        """Log metrics for this run."""
+        listener.log(metrics, value, step=step, labels=labels, unit=unit, metric_type=metric_type)
+
+    def complete(summary: Optional[Dict[str, Any]] = None) -> None:
+        """Mark the run as completed."""
+        listener.complete()
+
+    def fail(error: Optional[str] = None) -> None:
+        """Mark the run as failed."""
+        listener.fail(error)
+
+    def log_state(state: Dict[str, Any]) -> None:
+        """Log rich structured state for live visualization."""
+        listener.log_state(state)
+
+    def flush() -> None:
+        """Manually flush buffered metrics."""
+        listener.flush()
+
+    # Monkey-patch the API methods onto the store instance
+    store.log = log  # type: ignore[attr-defined]
+    store.complete = complete  # type: ignore[attr-defined]
+    store.fail = fail  # type: ignore[attr-defined]
+    store.log_state = log_state  # type: ignore[attr-defined]
+    store.flush = flush  # type: ignore[attr-defined]
+
+    # Add run_id property for convenience
+    store.run_id = listener.run_id  # type: ignore[attr-defined]
+    store.job_id = listener.run_id  # type: ignore[attr-defined]  # Backward compat
+
+    # Deprecated: _run attribute for backward compatibility
+    # Creates a simple object that forwards to the listener
+    class _RunCompat:
+        """Backward compatibility wrapper for store._run access."""
+
+        def __init__(self, listener):
+            self._listener = listener
+
+        @property
+        def id(self):
+            return self._listener.run_id
+
+        def log(self, *args, **kwargs):
+            return self._listener.log(*args, **kwargs)
+
+        def log_state(self, state):
+            return self._listener.log_state(state)
+
+        def flush(self):
+            return self._listener.flush()
+
+        def finish(self, state="completed"):
+            return self._listener.finish(state)
+
+    store._run = _RunCompat(listener)  # type: ignore[attr-defined]
+
+    logger.info(
+        f"GatewayLightningStore created with GatewayRun {listener.run_id} "
+        f"(project={project_id}, name={listener.run_name})"
+    )
+
+    return store
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# RedisStreamGatewayLightningStore - For platform runs with distributed coord
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def RedisStreamGatewayLightningStore(
+    *,
+    redis_url: str,
+    run_id: str,
+    api_key: str,
+    secret_key: str,
+    gateway_url: str,
+    project_id: str,
+    run_name: Optional[str] = None,
+    config: Optional[Dict[str, Any]] = None,
+    labels: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    source: str = "platform",
+    heartbeat_interval: int = 30,
+    consumer_group: str = "gateway-runners",
+    thread_safe: bool = False,
+    init_tracing: bool = False,
+) -> "RedisStreamLightningStore":
+    """Create a RedisStreamLightningStore with GatewayRunListener attached.
+
+    Used by the Python BullMQ worker for platform-managed runs. Combines
+    Redis Streams for distributed rollout coordination with GatewayRunListener
+    for run lifecycle management (heartbeats, metrics, sessions).
+
+    Args:
+        redis_url: Redis connection string.
+        run_id: Existing GatewayRun ID (created by the platform start endpoint).
+        api_key: Gateway public API key.
+        secret_key: Gateway secret key.
+        gateway_url: Gateway server URL.
+        project_id: Project ID.
+        run_name: Optional run name.
+        config: Run config dict.
+        labels: Run labels.
+        tags: Run tags.
+        source: Run source (default: "platform").
+        heartbeat_interval: Heartbeat interval in seconds.
+        consumer_group: Redis consumer group name.
+        thread_safe: Whether the store is thread-safe.
+        init_tracing: If True, auto-initialize gatewaysdk.tracing.
+
+    Returns:
+        RedisStreamLightningStore with GatewayRunListener attached.
+    """
+    import os
+
+    from .redis_stream import RedisStreamLightningStore as RedisStore
+
+    # Set non-secret env vars for SDK components that read from env
+    os.environ["GATEWAY_HOST"] = gateway_url
+    os.environ["GATEWAY_PUBLIC_KEY"] = api_key
+    os.environ["GATEWAY_PROJECT_ID"] = project_id
+
+    if init_tracing:
+        # Temporarily set secret key for tracing init, then remove (F6)
+        os.environ["GATEWAY_SECRET_KEY"] = secret_key
+        import gatewaysdk.tracing as tracing
+
+        tracing.init(instrument_default=True)
+        os.environ.pop("GATEWAY_SECRET_KEY", None)
+
+    listener = GatewayRunListener(
+        api_key=api_key,
+        secret_key=secret_key,
+        gateway_url=gateway_url,
+        project_id=project_id,
+        run_name=run_name,
+        config=config,
+        labels=labels,
+        tags=tags,
+        source=source,
+        heartbeat_interval=heartbeat_interval,
+    )
+    listener.start()
+
+    store = RedisStore(
+        redis_url=redis_url,
+        run_id=run_id,
+        consumer_group=consumer_group,
+        thread_safe=thread_safe,
+        listeners=[listener],
+    )
+
+    store.log = listener.log  # type: ignore[attr-defined]
+    store.complete = listener.complete  # type: ignore[attr-defined]
+    store.fail = listener.fail  # type: ignore[attr-defined]
+    store.log_state = listener.log_state  # type: ignore[attr-defined]
+    store.flush = listener.flush  # type: ignore[attr-defined]
+
+    from gatewaysdk.utils.redact import redact_url
+
+    logger.info(
+        f"RedisStreamGatewayLightningStore created for run {run_id} "
+        f"(redis={redact_url(redis_url)}, project={project_id})"
+    )
+
+    return store
+
+
+def create_platform_store_client_from_env() -> "LightningStoreClient":
+    """Create a store client for agent containers running on the Gateway platform.
+
+    Only requires ``GATEWAY_STORE_URL`` (set by the orchestrator). All Gateway
+    integration (OTLP export, sessions, scores, heartbeats) is handled
+    server-side by the worker's GatewayRunListener.
+
+    Returns:
+        A :class:`LightningStoreClient` connected to the worker's store server.
+
+    Raises:
+        EnvironmentError: If GATEWAY_STORE_URL is not set.
+    """
+    import os
+
+    from .client_server import LightningStoreClient
+
+    store_url = os.environ.get("GATEWAY_STORE_URL")
+    if not store_url:
+        raise EnvironmentError(
+            "GATEWAY_STORE_URL is not set. This variable is injected automatically "
+            "by the Gateway platform orchestrator when starting agent containers."
+        )
+
+    logger.info(f"Creating store client for {store_url}")
+    return LightningStoreClient(server_address=store_url)
+
+
+def create_platform_store_from_env() -> "RedisStreamLightningStore":
+    """DEPRECATED: Use create_platform_store_client_from_env() instead.
+
+    This function creates a RedisStreamLightningStore which has been
+    replaced by LightningStoreClient for platform runs.
+    """
+    import os
+    import warnings
+
+    warnings.warn(
+        "create_platform_store_from_env() is deprecated. "
+        "Use create_platform_store_client_from_env() instead, which uses "
+        "LightningStoreClient (HTTP) rather than Redis Streams.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+    required = ["GATEWAY_RUN_ID", "GATEWAY_HOST", "GATEWAY_PUBLIC_KEY", "GATEWAY_SECRET_KEY", "GATEWAY_PROJECT_ID"]
+    missing = [k for k in required if not os.environ.get(k)]
+    if missing:
+        raise EnvironmentError(
+            f"Missing required environment variables: {', '.join(missing)}. "
+            "These are set automatically by the Gateway platform orchestrator."
+        )
+
+    run_id = os.environ["GATEWAY_RUN_ID"]
+    host = os.environ.get("REDIS_HOST", "localhost")
+    port = os.environ.get("REDIS_PORT", "6379")
+    auth = os.environ.get("REDIS_AUTH", "")
+    redis_url = f"redis://:{auth}@{host}:{port}" if auth else f"redis://{host}:{port}"
+
+    return RedisStreamGatewayLightningStore(
+        redis_url=redis_url,
+        run_id=run_id,
+        api_key=os.environ["GATEWAY_PUBLIC_KEY"],
+        secret_key=os.environ["GATEWAY_SECRET_KEY"],
+        gateway_url=os.environ["GATEWAY_HOST"],
+        project_id=os.environ["GATEWAY_PROJECT_ID"],
+    )