update dashboard test

.github/workflows/dashboard.yml

@@ -8,11 +8,25 @@ on:
 
   workflow_dispatch:
 
-  push:
-    branches: [ main, stable/**/* ]
+  repository_dispatch:
+    types: [ci-dashboard, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'PR #{0} - Label {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Dashboard - {0}', github.event_name) }}
 
 jobs:
   dashboard:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-dashboard' ||
+      github.event.action == 'ci-all'
     name: Chromatic
     runs-on: ubuntu-latest
     timeout-minutes: 15
@@ -26,7 +40,7 @@ jobs:
       - name: Install JavaScript dependencies
         run: cd dashboard && npm ci
       - name: Run Chromatic
-        uses: chromaui/action@v13
+        uses: chromaui/action@latest
         with:
           projectToken: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
           workingDir: dashboard

.github/workflows/pypi-nightly.yml

@@ -26,14 +26,6 @@ jobs:
       - name: Sync dependencies
         run: uv sync --frozen --no-default-groups --group dev
 
-      - uses: actions/setup-node@v6
-        with:
-          node-version: '22'
-      - name: Install JavaScript dependencies
-        run: cd dashboard && npm ci
-      - name: Build dashboard
-        run: cd dashboard && npm run build
-
       - name: Get current version
         id: get_version
         run: |

.github/workflows/pypi-release.yml

@@ -60,14 +60,6 @@ jobs:
       - name: Sync dependencies
         run: uv sync --frozen --no-default-groups --group dev
 
-      - uses: actions/setup-node@v6
-        with:
-          node-version: '22'
-      - name: Install JavaScript dependencies
-        run: cd dashboard && npm ci
-      - name: Build dashboard
-        run: cd dashboard && npm run build
-
       - name: Build package
         run: |
           uv build

.github/workflows/tests-full.yml

@@ -73,14 +73,6 @@ jobs:
           path: requirements-freeze.txt
           compression-level: 0
 
-      - uses: actions/setup-node@v6
-        with:
-          node-version: '22'
-      - name: Install JavaScript dependencies
-        run: cd dashboard && npm ci
-      - name: Build dashboard
-        run: cd dashboard && npm run build
-
       - name: Launch LiteLLM Proxy
         run: |
           ./scripts/litellm_run.sh

.github/workflows/tests.yml

@@ -155,15 +155,6 @@ jobs:
           name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
           compression-level: 0
-
-      - uses: actions/setup-node@v6
-        with:
-          node-version: '22'
-      - name: Install JavaScript dependencies
-        run: cd dashboard && npm ci
-      - name: Build dashboard
-        run: cd dashboard && npm run build
-
       - name: Run tests
         run: |
           uv run pytest -v --durations=0 tests

.gitignore

@@ -207,14 +207,3 @@ cython_debug/
 
 # Claude
 .claude/*.local.json
-
-# Temporary and backup files
-*.tmp
-*.bak
-*.backup
-
-# Dashboard generated files
-agentlightning/dashboard/**/*.css
-agentlightning/dashboard/**/*.js
-agentlightning/dashboard/**/*.html
-agentlightning/dashboard/**/*.svg

agentlightning/llm_proxy.py

@@ -7,26 +7,15 @@ import asyncio
 import logging
 import os
 import re
+import socket
 import tempfile
 import threading
-from contextlib import asynccontextmanager
-from typing import (
-    Any,
-    AsyncGenerator,
-    Awaitable,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    TypedDict,
-    Union,
-    cast,
-)
+import time
+from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional, Sequence, TypedDict, Union, cast
 
 import litellm
 import opentelemetry.trace as trace_api
+import uvicorn
 import yaml
 from fastapi import Request, Response
 from litellm.integrations.custom_logger import CustomLogger
@@ -37,12 +26,6 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from agentlightning.types import LLM, ProxyLLM
-from agentlightning.utils.server_launcher import (
-    LaunchMode,
-    PythonServerLauncher,
-    PythonServerLauncherArgs,
-    noop_context,
-)
 
 from .store.base import LightningStore
 
@@ -515,10 +498,9 @@ class LLMProxy:
     * [`stop()`][agentlightning.LLMProxy.stop] tears down the server and removes the temp config file.
     * [`restart()`][agentlightning.LLMProxy.restart] convenience wrapper to stop then start.
 
-    !!! note
-
-        As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
-        process from the main runner (i.e., tracer from agents). See `launch_mode` for how to change that.
+    Usage Note:
+    As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
+    process from the main runner (i.e., tracer from agents).
 
     !!! warning
 
@@ -530,64 +512,38 @@ class LLMProxy:
         with tracers like [`AgentOpsTracer`][agentlightning.AgentOpsTracer].
 
     Args:
-        port: TCP port to bind. Will bind to a random port if not provided.
+        port: TCP port to bind.
         model_list: LiteLLM `model_list` entries.
         store: LightningStore used for span sequence and persistence.
-        host: Publicly reachable host used in resource endpoints. See `host` of `launcher_args` for more details.
+        host: Publicly reachable host used in resource endpoints. Defaults to best-guess IPv4.
         litellm_config: Extra LiteLLM proxy config merged with `model_list`.
         num_retries: Default LiteLLM retry count injected into `litellm_settings`.
-        num_workers: Number of workers to run in the server. Only applicable for "mp" launch mode. Ignored if launcher_args is provided.
-            When `num_workers > 1`, the server will be run using [gunicorn](https://gunicorn.org/).
-        launch_mode: Launch mode for the server. Defaults to "mp". Cannot be used together with launcher_args. Ignored if launcher_args is provided.
-            It's recommended to use `launch_mode="mp"` to launch the proxy, which will launch the server in a separate process.
-            `launch_mode="thread"` can also be used if used in caution. It will launch the server in a separate thread.
-            `launch_mode="asyncio"` launches the server in the current thread as an asyncio task.
-            It is NOT recommended because it often causes hanging requests. Only use it if you know what you are doing.
-        launcher_args: Arguments for the server launcher. If this is provided, host, port, and launch_mode will be ignored. Cannot be used together with port, host, and launch_mode.
     """
 
     def __init__(
         self,
-        port: int | None = None,
+        port: int,
         model_list: List[ModelConfig] | None = None,
         store: Optional[LightningStore] = None,
         host: str | None = None,
         litellm_config: Dict[str, Any] | None = None,
         num_retries: int = 0,
-        num_workers: int = 1,
-        launch_mode: LaunchMode = "mp",
-        launcher_args: PythonServerLauncherArgs | None = None,
         _add_return_token_ids: bool = True,
     ):
         self.store = store
-
-        if launcher_args is not None and (
-            port is not None or host is not None or launch_mode != "mp" or num_workers != 1
-        ):
-            raise ValueError("port, host, launch_mode, and num_workers cannot be set when launcher_args is provided.")
-
-        self.server_launcher_args = launcher_args or PythonServerLauncherArgs(
-            port=port,
-            host=host,
-            launch_mode=launch_mode,
-            n_workers=num_workers,
-            # NOTE: This /health endpoint can be slow sometimes because it actually probes the backend LLM service.
-            healthcheck_url="/health",
-            startup_timeout=60.0,
-        )
-
-        if self.server_launcher_args.healthcheck_url is None:
-            logger.warning("healthcheck_url is not set. LLM Proxy will not be checked for healthiness after starting.")
-
+        self.host = host or _get_default_ipv4_address()
+        self.port = port
         self.model_list = model_list or []
         self.litellm_config = litellm_config or {}
 
         # Ensure num_retries is present inside the litellm_settings block.
         self.litellm_config.setdefault("litellm_settings", {})
         self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
-        self.server_launcher = PythonServerLauncher(app, self.server_launcher_args, noop_context())
 
+        self._server_thread = None
         self._config_file = None
+        self._uvicorn_server = None
+        self._ready_event = threading.Event()
 
         self._add_return_token_ids = _add_return_token_ids
 
@@ -608,15 +564,44 @@ class LLMProxy:
         self.store = store
 
     def update_model_list(self, model_list: List[ModelConfig]) -> None:
-        """Replace the in-memory model list.
+        """Replace the in-memory model list and hot-restart if running.
 
         Args:
             model_list: New list of model entries.
         """
         self.model_list = model_list
         logger.info(f"Updating LLMProxy model list to: {model_list}")
+        if self.is_running():
+            self.restart()
         # Do nothing if the server is not running.
 
+    def update_port(self, port: int) -> None:
+        """Update the port for the proxy.
+
+        Args:
+            port: The new port to use for the proxy.
+        """
+        self.port = port
+
+    def _wait_until_started(self, startup_timeout: float = 20.0):
+        """Block until the uvicorn server reports started or timeout.
+
+        Args:
+            startup_timeout: Maximum seconds to wait.
+        """
+        start = time.time()
+        while True:
+            if self._uvicorn_server is None:
+                break
+            if self._uvicorn_server.started:
+                self._ready_event.set()
+                break
+            if self._uvicorn_server.should_exit:
+                break
+            if time.time() - start > startup_timeout:
+                break
+            time.sleep(0.01)
+
     def initialize(self):
         """Initialize global middleware and LiteLLM callbacks.
 
@@ -660,12 +645,19 @@ class LLMProxy:
             # reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
             _reset_litellm_logging_worker()
 
-    @asynccontextmanager
-    async def _serve_context(self) -> AsyncGenerator[None, None]:
-        """Context manager to serve the proxy server.
+    def start(self):
+        """Start the proxy server thread and initialize global wiring.
 
-        See [`start`][agentlightning.LLMProxy.start] and [`stop`][agentlightning.LLMProxy.stop] for more details.
+        Side effects:
+
+        * Sets the module-level global store for middleware/exporter access.
+        * Calls `initialize()` once to register middleware and callbacks.
+        * Writes a temporary YAML config consumed by LiteLLM worker.
+        * Launches uvicorn in a daemon thread and waits for readiness.
         """
+        if self.is_running():
+            # Trigger restart
+            self.stop()
 
         if not self.store:
             raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
@@ -686,59 +678,24 @@ class LLMProxy:
 
         save_worker_config(config=self._config_file)
 
-        # NOTE: When running the _serve_context in current process, you might encounter the following problems:
+        # Bind to all interfaces to allow other hosts to reach it if needed.
+        self._uvicorn_server = uvicorn.Server(uvicorn.Config(app, host="0.0.0.0", port=self.port))
+
+        def run_server():
+            # Serve uvicorn in this background thread with its own event loop.
+            assert self._uvicorn_server is not None
+            asyncio.run(self._uvicorn_server.serve())
+
+        logger.info("Starting LLMProxy server thread...")
+        self._ready_event.clear()
+        # FIXME: This thread should either be reused or the whole proxy should live in another process.
         # Problem 1: in litellm worker, <Queue at 0x70f1d028cd90 maxsize=50000> is bound to a different event loop
         # Problem 2: Proxy has conflicted opentelemetry setup with the main process.
+        self._server_thread = threading.Thread(target=run_server, daemon=True)
+        self._server_thread.start()
+        self._wait_until_started()
 
-        # Ready
-        logger.info("LLMProxy preparation is done. Will start the server.")
-        yield
-
-        # Clean up
-
-        logger.info("LLMProxy server is cleaning up.")
-
-        # Remove worker config to avoid stale references.
-        if self._config_file and os.path.exists(self._config_file):
-            os.unlink(self._config_file)
-
-        logger.info("LLMProxy server finishes.")
-
-    async def start(self):
-        """Start the proxy server thread and initialize global wiring.
-
-        Side effects:
-
-        * Sets the module-level global store for middleware/exporter access.
-        * Calls `initialize()` once to register middleware and callbacks.
-        * Writes a temporary YAML config consumed by LiteLLM worker.
-        * Launches uvicorn in a daemon thread and waits for readiness.
-        """
-        # Refresh the serve context
-        self.server_launcher.serve_context = self._serve_context()
-
-        if self.store is None:
-            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
-
-        store_capabilities = self.store.capabilities()
-        if self.server_launcher.args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
-            raise RuntimeError(
-                "The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
-            )
-        elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
-            raise RuntimeError(
-                "The store is not thread-safe. Please use another store, or use asyncio mode to launch the server."
-            )
-        elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities["async_safe"]:
-            raise RuntimeError("The store is not async-safe. Please use another store.")
-
-        logger.info(
-            f"Starting LLMProxy server in {self.server_launcher.args.launch_mode} mode with store capabilities: {store_capabilities}"
-        )
-
-        await self.server_launcher.start()
-
-    async def stop(self):
+    def stop(self):
         """Stop the proxy server and clean up temporary artifacts.
 
         This is a best-effort graceful shutdown with a bounded join timeout.
@@ -747,19 +704,43 @@ class LLMProxy:
             logger.warning("LLMProxy is not running. Nothing to stop.")
             return
 
-        await self.server_launcher.stop()
+        # Remove worker config to avoid stale references.
+        if self._config_file and os.path.exists(self._config_file):
+            os.unlink(self._config_file)
 
-    async def restart(self, *, _port: int | None = None) -> None:
+        logger.info("Stopping LLMProxy server thread...")
+        stop_success = True
+        if self._server_thread is not None and self._uvicorn_server is not None and self._uvicorn_server.started:
+            self._uvicorn_server.should_exit = True
+            self._server_thread.join(timeout=10.0)  # Allow time for graceful shutdown.
+            if self._server_thread.is_alive():
+                logger.error(
+                    "LLMProxy server thread is still alive after 10 seconds. Cannot kill it because it's a thread."
+                )
+                stop_success = False
+            self._server_thread = None
+            self._uvicorn_server = None
+            self._config_file = None
+            self._ready_event.clear()
+            if not _check_port(self.host, self.port):
+                logger.error(f"Port {self.port} is still in use. Stopping LLMProxy is not successful.")
+                stop_success = False
+        if stop_success:
+            logger.info("LLMProxy server thread stopped.")
+        else:
+            logger.error("LLMProxy server is not stopped successfully.")
+
+    def restart(self, *, _port: int | None = None) -> None:
         """Restart the proxy if running, else start it.
 
         Convenience wrapper calling `stop()` followed by `start()`.
         """
         logger.info("Restarting LLMProxy server...")
         if self.is_running():
-            await self.stop()
+            self.stop()
         if _port is not None:
-            self.server_launcher_args.port = _port
-        await self.start()
+            self.port = _port
+        self.start()
 
     def is_running(self) -> bool:
         """Return whether the uvicorn server is active.
@@ -767,7 +748,7 @@ class LLMProxy:
         Returns:
             bool: True if server was started and did not signal exit.
         """
-        return self.server_launcher.is_running()
+        return self._uvicorn_server is not None and self._uvicorn_server.started
 
     def as_resource(
         self,
@@ -810,13 +791,13 @@ class LLMProxy:
 
         if rollout_id is None and attempt_id is None:
             return ProxyLLM(
-                endpoint=self.server_launcher.access_endpoint,
+                endpoint=f"http://{self.host}:{self.port}",
                 model=model,
                 sampling_parameters=dict(sampling_parameters or {}),
             )
         elif rollout_id is not None and attempt_id is not None:
             return LLM(
-                endpoint=f"{self.server_launcher.access_endpoint}/rollout/{rollout_id}/attempt/{attempt_id}",
+                endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
                 model=model,
                 sampling_parameters=dict(sampling_parameters or {}),
             )
@@ -897,6 +878,35 @@ def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
     return False
 
 
+def _get_default_ipv4_address() -> str:
+    """Determine the default outbound IPv4 address for this machine.
+
+    Implementation:
+        Opens a UDP socket and "connects" to a public address to force route
+        selection, then inspects the socket's local address. No packets are sent.
+
+    Returns:
+        str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except Exception:
+        return "127.0.0.1"
+    finally:
+        s.close()
+
+
+def _check_port(host: str, port: int) -> bool:
+    """Check if a port is available."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.settimeout(1)
+        result = s.connect_ex((host, port))
+        return result != 0  # True if unavailable
+
+
 def _check_tracer_provider() -> bool:
     """Check if the global tracer provider is properly initialized.
 

agentlightning/store/__init__.py

@@ -1,17 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import LightningStore, LightningStoreCapabilities
+from .base import LightningStore
 from .client_server import LightningStoreClient, LightningStoreServer
-from .database import SqlLightningStore
 from .memory import InMemoryLightningStore
 from .threading import LightningStoreThreaded
 
 __all__ = [
     "LightningStore",
-    "LightningStoreCapabilities",
     "LightningStoreClient",
     "LightningStoreServer",
     "InMemoryLightningStore",
     "LightningStoreThreaded",
-    "SqlLightningStore",
 ]

agentlightning/store/base.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal, Optional, Sequence, Union, TypedDict
+from typing import Any, Dict, List, Literal, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -52,17 +52,6 @@ UNSET = _UnsetType()
 Unset = _UnsetType  # Alias for convenience
 
 
-class LightningStoreCapabilities(TypedDict):
-    """Capability of a LightningStore implementation."""
-
-    thread_safe: bool
-    """Whether the store is thread-safe."""
-    async_safe: bool
-    """Whether the store is async-safe."""
-    zero_copy: bool
-    """Whether the store has only one copy across all threads/processes."""
-
-
 class LightningStore:
     """Contract for the persistent control-plane that coordinates training rollouts.
 
@@ -85,14 +74,6 @@ class LightningStore:
     Unless stated otherwise, missing identifiers should result in a `ValueError`.
     """
 
-    def capabilities(self) -> LightningStoreCapabilities:
-        """Return the capabilities of the store."""
-        return LightningStoreCapabilities(
-            thread_safe=False,
-            async_safe=False,
-            zero_copy=False,
-        )
-
     async def start_rollout(
         self,
         input: TaskInput,
@@ -267,7 +248,7 @@ class LightningStore:
 
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Union[Rollout, AttemptedRollout]]:
+    ) -> List[Rollout]:
         """Retrieve rollouts filtered by status and/or explicit identifiers.
 
         Args:
@@ -297,7 +278,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
         """Fetch a rollout by identifier without mutating its state.
 
         Args:
@@ -457,8 +438,6 @@ class LightningStore:
         This API is typically used by algorithms that maintain mutable resources (e.g., model
         checkpoints) under a stable identifier.
 
-        If `resources_id` does not exist, implementations should add it as a new snapshot.
-
         Args:
             resources_id: Identifier of the snapshot to replace.
             resources: Updated mapping of resource names to payloads.
@@ -468,6 +447,7 @@ class LightningStore:
 
         Raises:
             NotImplementedError: Subclasses must implement resource persistence.
+            ValueError: Implementations must raise when `resources_id` does not exist.
         """
         raise NotImplementedError()
 

agentlightning/store/client_server.py

@@ -9,17 +9,15 @@ import threading
 import time
 import traceback
 from contextlib import suppress
-from pathlib import Path
-from typing import Any, Awaitable, Callable, Dict, Generic, List, Literal, Optional, Sequence, TypeVar, Union
+from typing import Any, Awaitable, Callable, Dict, Generic, List, Literal, Optional, Sequence, TypeVar
 
 import aiohttp
 import uvicorn
-from fastapi import APIRouter, Body, Depends, FastAPI, HTTPException
+from fastapi import Body, Depends, FastAPI, HTTPException
 from fastapi import Query as FastAPIQuery
 from fastapi import Request, Response
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse, JSONResponse
-from fastapi.staticfiles import StaticFiles
+from fastapi.responses import JSONResponse
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel, Field, TypeAdapter
 
@@ -36,13 +34,11 @@ from agentlightning.types import (
     TaskInput,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
+from .base import UNSET, LightningStore, Unset
 
 logger = logging.getLogger(__name__)
 
-API_V1_PREFIX = "/v1"
-API_AGL_PREFIX = "/agl"
-API_V1_AGL_PREFIX = API_V1_PREFIX + API_AGL_PREFIX
+AGL_API_V1_PREFIX = "/v1/agl"
 
 T = TypeVar("T")
 
@@ -240,14 +236,6 @@ def _apply_filters_sort_paginate(
     return PaginatedResponse(items=paginated_items, limit=limit, offset=offset, total=total)
 
 
-class CachedStaticFiles(StaticFiles):
-    def file_response(self, *args: Any, **kwargs: Any) -> Response:
-        resp = super().file_response(*args, **kwargs)
-        # hashed filenames are safe to cache "forever"
-        resp.headers.setdefault("Cache-Control", "public, max-age=31536000, immutable")
-        return resp
-
-
 class LightningStoreServer(LightningStore):
     """
     Server wrapper that exposes a LightningStore via HTTP API.
@@ -298,14 +286,6 @@ class LightningStoreServer(LightningStore):
         self._owner_pid = os.getpid()
         self._client: Optional[LightningStoreClient] = None
 
-    def capabilities(self) -> LightningStoreCapabilities:
-        """Return the capabilities of the store."""
-        capabilities = self.store.capabilities().copy()
-        capabilities["async_safe"] = True
-        capabilities["thread_safe"] = True
-        capabilities["zero_copy"] = True
-        return capabilities
-
     def __getstate__(self):
         """
         Control pickling to prevent server state from being sent to subprocesses.
@@ -436,7 +416,7 @@ class LightningStoreServer(LightningStore):
         while time.time() - current_time < 10:
             async with aiohttp.ClientSession() as session:
                 with suppress(Exception):
-                    async with session.get(f"{self.endpoint}{API_V1_AGL_PREFIX}/health") as response:
+                    async with session.get(f"{self.endpoint}{AGL_API_V1_PREFIX}/health") as response:
                         if response.status == 200:
                             return True
             await asyncio.sleep(0.1)
@@ -548,7 +528,7 @@ class LightningStoreServer(LightningStore):
                 return await call_next(request)
             except Exception as exc:
                 # decide whether to convert this into your 400 JSONResponse
-                if request.url.path.startswith(API_V1_AGL_PREFIX):
+                if request.url.path.startswith(AGL_API_V1_PREFIX):
                     logger.exception("Unhandled application error", exc_info=exc)
                     payload = {
                         "detail": "Internal server error",
@@ -564,8 +544,7 @@ class LightningStoreServer(LightningStore):
         async def _log_time(  # pyright: ignore[reportUnusedFunction]
             request: Request, call_next: Callable[[Request], Awaitable[Response]]
         ):
-            # If not API request, just pass through
-            if not request.url.path.startswith(API_V1_AGL_PREFIX):
+            if not request.url.path.startswith("/v1/agl/"):
                 return await call_next(request)
 
             start = time.perf_counter()
@@ -583,13 +562,11 @@ class LightningStoreServer(LightningStore):
             )
             return response
 
-        api = APIRouter(prefix=API_V1_PREFIX)
-
-        @api.get(API_AGL_PREFIX + "/health")
+        @self.app.get(AGL_API_V1_PREFIX + "/health")
         async def health():  # pyright: ignore[reportUnusedFunction]
             return {"status": "ok"}
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=Rollout)
+        @self.app.post(AGL_API_V1_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=Rollout)
         async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.enqueue_rollout(
                 input=request.input,
@@ -599,11 +576,11 @@ class LightningStoreServer(LightningStore):
                 metadata=request.metadata,
             )
 
-        @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
+        @self.app.post(AGL_API_V1_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
         async def dequeue_rollout():  # pyright: ignore[reportUnusedFunction]
             return await self.dequeue_rollout()
 
-        @api.post(API_AGL_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
+        @self.app.post(AGL_API_V1_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.start_rollout(
                 input=request.input,
@@ -613,7 +590,7 @@ class LightningStoreServer(LightningStore):
                 metadata=request.metadata,
             )
 
-        @api.get(API_AGL_PREFIX + "/rollouts", response_model=PaginatedResponse[Union[AttemptedRollout, Rollout]])
+        @self.app.get(AGL_API_V1_PREFIX + "/rollouts", response_model=PaginatedResponse[Rollout])
         async def query_rollouts(params: QueryRolloutsRequest = Depends()):  # pyright: ignore[reportUnusedFunction]
             # Get all rollouts from the underlying store
             all_rollouts = await self.query_rollouts()
@@ -637,7 +614,7 @@ class LightningStoreServer(LightningStore):
                 params.offset,
             )
 
-        @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Union[AttemptedRollout, Rollout])
+        @self.app.get(AGL_API_V1_PREFIX + "/rollouts/{rollout_id}", response_model=Rollout)
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_rollout_by_id(rollout_id)
 
@@ -652,7 +629,7 @@ class LightningStoreServer(LightningStore):
             else:
                 return UNSET
 
-        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Rollout)
+        @self.app.post(AGL_API_V1_PREFIX + "/rollouts/{rollout_id}", response_model=Rollout)
         async def update_rollout(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, request: UpdateRolloutRequest = Body(...)
         ):
@@ -666,11 +643,13 @@ class LightningStoreServer(LightningStore):
                 metadata=request.metadata if "metadata" in request.model_fields_set else UNSET,
             )
 
-        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", status_code=201, response_model=AttemptedRollout)
+        @self.app.post(
+            AGL_API_V1_PREFIX + "/rollouts/{rollout_id}/attempts", status_code=201, response_model=AttemptedRollout
+        )
         async def start_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.start_attempt(rollout_id)
 
-        @api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
+        @self.app.post(AGL_API_V1_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
         async def update_attempt(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, attempt_id: str, request: UpdateAttemptRequest = Body(...)
         ):
@@ -683,7 +662,7 @@ class LightningStoreServer(LightningStore):
                 metadata=_get_mandatory_field_or_unset(request, "metadata"),
             )
 
-        @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", response_model=PaginatedResponse[Attempt])
+        @self.app.get(AGL_API_V1_PREFIX + "/rollouts/{rollout_id}/attempts", response_model=PaginatedResponse[Attempt])
         async def query_attempts(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, params: QueryAttemptsRequest = Depends()
         ):
@@ -700,11 +679,11 @@ class LightningStoreServer(LightningStore):
                 params.offset,
             )
 
-        @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/latest", response_model=Optional[Attempt])
+        @self.app.get(AGL_API_V1_PREFIX + "/rollouts/{rollout_id}/attempts/latest", response_model=Optional[Attempt])
         async def get_latest_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_latest_attempt(rollout_id)
 
-        @api.get(API_AGL_PREFIX + "/resources", response_model=PaginatedResponse[ResourcesUpdate])
+        @self.app.get(AGL_API_V1_PREFIX + "/resources", response_model=PaginatedResponse[ResourcesUpdate])
         async def query_resources(params: QueryResourcesRequest = Depends()):  # pyright: ignore[reportUnusedFunction]
             # Get all resources
             all_resources = await self.query_resources()
@@ -726,29 +705,29 @@ class LightningStoreServer(LightningStore):
                 params.offset,
             )
 
-        @api.post(API_AGL_PREFIX + "/resources", status_code=201, response_model=ResourcesUpdate)
+        @self.app.post(AGL_API_V1_PREFIX + "/resources", status_code=201, response_model=ResourcesUpdate)
         async def add_resources(resources: NamedResources):  # pyright: ignore[reportUnusedFunction]
             return await self.add_resources(resources)
 
-        @api.get(API_AGL_PREFIX + "/resources/latest", response_model=Optional[ResourcesUpdate])
+        @self.app.get(AGL_API_V1_PREFIX + "/resources/latest", response_model=Optional[ResourcesUpdate])
         async def get_latest_resources():  # pyright: ignore[reportUnusedFunction]
             return await self.get_latest_resources()
 
-        @api.post(API_AGL_PREFIX + "/resources/{resources_id}", response_model=ResourcesUpdate)
+        @self.app.post(AGL_API_V1_PREFIX + "/resources/{resources_id}", response_model=ResourcesUpdate)
         async def update_resources(  # pyright: ignore[reportUnusedFunction]
             resources_id: str, resources: NamedResources
         ):
             return await self.update_resources(resources_id, resources)
 
-        @api.get(API_AGL_PREFIX + "/resources/{resources_id}", response_model=Optional[ResourcesUpdate])
+        @self.app.get(AGL_API_V1_PREFIX + "/resources/{resources_id}", response_model=Optional[ResourcesUpdate])
         async def get_resources_by_id(resources_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_resources_by_id(resources_id)
 
-        @api.post(API_AGL_PREFIX + "/spans", status_code=201, response_model=Span)
+        @self.app.post(AGL_API_V1_PREFIX + "/spans", status_code=201, response_model=Span)
         async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
             return await self.add_span(span)
 
-        @api.get(API_AGL_PREFIX + "/spans", response_model=PaginatedResponse[Span])
+        @self.app.get(AGL_API_V1_PREFIX + "/spans", response_model=PaginatedResponse[Span])
         async def query_spans(params: QuerySpansRequest = Depends()):  # pyright: ignore[reportUnusedFunction]
             # Get all spans for the rollout/attempt
             all_spans = await self.query_spans(params.rollout_id, params.attempt_id)
@@ -776,76 +755,33 @@ class LightningStoreServer(LightningStore):
                 all_spans, filters, params.filter_logic, params.sort_by, params.sort_order, params.limit, params.offset
             )
 
-        @api.post(API_AGL_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
+        @self.app.post(AGL_API_V1_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
         async def get_next_span_sequence_id(request: NextSequenceIdRequest):  # pyright: ignore[reportUnusedFunction]
             sequence_id = await self.get_next_span_sequence_id(request.rollout_id, request.attempt_id)
             return NextSequenceIdResponse(sequence_id=sequence_id)
 
-        @api.post(API_AGL_PREFIX + "/waits/rollouts", response_model=List[Rollout])
+        @self.app.post(AGL_API_V1_PREFIX + "/waits/rollouts", response_model=List[Rollout])
         async def wait_for_rollouts(request: WaitForRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
             return await self.wait_for_rollouts(rollout_ids=request.rollout_ids, timeout=request.timeout)
 
         # Reserved methods for OTEL traces
         # https://opentelemetry.io/docs/specs/otlp/#otlphttp-request
-        @api.post("/traces")
+        @self.app.post("/v1/traces")
         async def otlp_traces():  # pyright: ignore[reportUnusedFunction]
             return Response(status_code=501)
 
-        @api.post("/metrics")
+        @self.app.post("/v1/metrics")
         async def otlp_metrics():  # pyright: ignore[reportUnusedFunction]
             return Response(status_code=501)
 
-        @api.post("/logs")
+        @self.app.post("/v1/logs")
         async def otlp_logs():  # pyright: ignore[reportUnusedFunction]
             return Response(status_code=501)
 
-        @api.post("/development/profiles")
+        @self.app.post("/v1/development/profiles")
         async def otlp_development_profiles():  # pyright: ignore[reportUnusedFunction]
             return Response(status_code=501)
 
-        # Mount the API router of /v1/...
-        self.app.include_router(api)
-
-        # Finally, mount the dashboard assets
-        self._setup_dashboard()
-
-    def _setup_dashboard(self):
-        """Setup the dashboard static files and SPA."""
-        assert self.app is not None
-
-        dashboard_dir = (Path(__file__).parent.parent / "dashboard").resolve()
-        if not dashboard_dir.exists():
-            logger.error("Dashboard directory not found at %s. Please build the dashboard first.", dashboard_dir)
-            return
-
-        dashboard_assets_dir = dashboard_dir / "assets"
-        if not dashboard_assets_dir.exists():
-            logger.error(
-                "Dashboard assets directory not found at %s. Please build the dashboard first.", dashboard_assets_dir
-            )
-            return
-
-        index_file = dashboard_dir / "index.html"
-        if not index_file.exists():
-            logger.error("Dashboard index file not found at %s. Please build the dashboard first.", index_file)
-            return
-
-        # Mount the static files in dashboard/assets
-        self.app.mount("/assets", CachedStaticFiles(directory=dashboard_assets_dir), name="assets")
-
-        # SPA fallback (client-side routing)
-        # Anything that's not /v1/* or a real file in /assets will serve index.html
-        @self.app.get("/", include_in_schema=False)
-        def root():  # pyright: ignore[reportUnusedFunction]
-            return FileResponse(index_file)
-
-        @self.app.get("/{full_path:path}", include_in_schema=False)
-        def spa_fallback(full_path: str):  # pyright: ignore[reportUnusedFunction]
-            # Let the frontend router handle it
-            return FileResponse(index_file)
-
-        logger.info("Agent-lightning dashboard will be available at %s", self.endpoint)
-
     # Delegate methods
     async def _call_store_method(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
         backend = self._backend()
@@ -1019,7 +955,7 @@ class LightningStoreClient(LightningStore):
         retry_delays: Sequence[float] = (1.0, 2.0, 5.0),
         health_retry_delays: Sequence[float] = (0.1, 0.2, 0.5),
     ):
-        self.server_address = server_address.rstrip("/") + API_V1_AGL_PREFIX
+        self.server_address = server_address.rstrip("/") + AGL_API_V1_PREFIX
         self._sessions: Dict[int, aiohttp.ClientSession] = {}  # id(loop) -> ClientSession
         self._lock = threading.RLock()
 
@@ -1031,14 +967,6 @@ class LightningStoreClient(LightningStore):
         self._dequeue_was_successful: bool = False
         self._dequeue_first_unsuccessful: bool = True
 
-    def capabilities(self) -> LightningStoreCapabilities:
-        """Return the capabilities of the store."""
-        return LightningStoreCapabilities(
-            thread_safe=True,
-            async_safe=True,
-            zero_copy=True,
-        )
-
     def __getstate__(self):
         """
         When LightningStoreClient is pickled (e.g., passed to a subprocess), we only
@@ -1288,14 +1216,7 @@ class LightningStoreClient(LightningStore):
 
         data = await self._request_json("get", "/rollouts", params=params if params else None)
         # Extract items from PaginatedResponse
-        return [
-            (
-                AttemptedRollout.model_validate(item)
-                if isinstance(item, dict) and "attempt" in item
-                else Rollout.model_validate(item)
-            )
-            for item in data["items"]
-        ]
+        return [Rollout.model_validate(item) for item in data["items"]]
 
     async def query_attempts(self, rollout_id: str) -> List[Attempt]:
         data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts")
@@ -1339,10 +1260,7 @@ class LightningStoreClient(LightningStore):
         """
         try:
             data = await self._request_json("get", f"/rollouts/{rollout_id}")
-            if isinstance(data, dict) and "attempt" in data:
-                return AttemptedRollout.model_validate(data)
-            else:
-                return Rollout.model_validate(data)
+            return Rollout.model_validate(data) if data else None
         except Exception as e:
             logger.error(f"get_rollout_by_id failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
             return None

agentlightning/store/database/__init__.py

@@ -1,7 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from .sqlite import SqlLightningStore
-
-__all__ = [
-    "SqlLightningStore",
-]

agentlightning/store/database/orm/__init__.py

@@ -1,20 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from .attempt import AttemptInDB, SpanSeqIdInDB
-from .base import (
-    AttemptStatusUpdateMessage,
-    SqlAlchemyBase,
-)
-from .resources import ResourcesUpdateInDB
-from .rollout import RolloutInDB
-from .span import SpanInDB
-
-__all__ = [
-    "SqlAlchemyBase",
-    "AttemptStatusUpdateMessage",
-    "RolloutInDB",
-    "AttemptInDB",
-    "ResourcesUpdateInDB",
-    "SpanSeqIdInDB",
-    "SpanInDB",
-]

agentlightning/store/database/orm/attempt.py

@@ -1,251 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import hashlib
-import logging
-import time
-import uuid
-from dataclasses import InitVar
-from typing import Any, Dict, List, Optional
-
-from sqlalchemy import JSON, Float, Integer, String, select
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
-from sqlalchemy.orm import Mapped, mapped_column
-
-from agentlightning.types import Attempt
-
-from .base import AttemptStatusUpdateMessage, SqlAlchemyBase
-
-logger = logging.getLogger(__name__)
-
-
-def _generate_attempt_id() -> str:
-    """We don't need that long because attempts are limited to rollouts."""
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
-    return "at-" + short_id
-
-
-class AttemptInDB(SqlAlchemyBase):
-    __tablename__ = "attempts"
-
-    rollout_id: Mapped[str] = mapped_column(String, nullable=False)
-    attempt_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_attempt_id)
-    sequence_id: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
-    start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
-    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
-    status: Mapped[str] = mapped_column(String, default="preparing", nullable=False)
-    worker_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
-    last_heartbeat_time: Mapped[Optional[float]] = mapped_column(Float, nullable=False, default_factory=time.time)
-    attempt_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
-
-    # addition columns for processing
-    max_duration: Mapped[Optional[float]] = mapped_column(
-        Float, nullable=True, default=None
-    )  # maximum duration allowed for this attempt in seconds
-    max_heartbeat_interval: Mapped[Optional[float]] = mapped_column(
-        Float, nullable=True, default=None
-    )  # maximum allowed heartbeat interval in seconds
-
-    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
-    __mapper_args__ = {
-        "version_id_col": version_id,
-    }
-
-    def is_unresponsive(self, current_time: float) -> bool:
-        """Check if the attempt is unresponsive based on the last heartbeat time and max_heartbeat_interval."""
-        if self.max_heartbeat_interval is None:
-            return False
-        if self.last_heartbeat_time is None:
-            return False
-        return (current_time - self.last_heartbeat_time) > self.max_heartbeat_interval
-
-    def is_timed_out(self, current_time: float) -> bool:
-        """Check if the attempt has timed out based on the start time and max_duration."""
-        if self.max_duration is None:
-            return False
-        return (current_time - self.start_time) > self.max_duration
-
-    def as_attempt(self) -> Attempt:
-        return Attempt(
-            **self.model_dump(
-                exclude={"max_duration", "max_heartbeat_interval", "version_id"},
-                mapper={"metadata": lambda obj: obj.attempt_metadata},  # type: ignore
-            )
-        )
-
-    def _validate_status_message(self, msg: Dict[str, Any]) -> None:
-        """This function validates the status update message from caller.
-        Raises ValueError if the message is invalid.
-        """
-        if "event" not in msg:
-            raise ValueError("Status update message must contain 'event' field.")
-        if "timestamp" not in msg:
-            msg["timestamp"] = time.time()
-        if msg["event"] not in [
-            "user_update",  # user update attempt status via dbstore.update_attempt()
-            "span_received",  # new span received
-            "single_step_timeout",  # single step timeout detected (from last span heartbeat)
-            "overall_timeout",  # overall timeout detected
-        ]:
-            raise ValueError(f"Unsupported event type: {msg['event']}")
-        if msg["event"] == "user_update" and "new_status" not in msg:
-            raise ValueError("User update event must contain 'new_status' field.")
-
-    def get_finished_statuses(self) -> List[str]:
-        """This function returns the list of statuses that are considered finished."""
-        return [
-            "succeeded",
-            "failed",
-            "timeout",
-        ]
-
-    def update_status(self, msg: Dict[str, Any]) -> Optional[AttemptStatusUpdateMessage]:
-        """This function updates the status of the attempt based on the event.
-        Args:
-            msg: A dictionary containing the status update message. It must contain an "event" field, and optionally a "new_status" field.
-                More details about the message format can be found in the `_validate_status_message`() method.
-            current_time: The current time to use for updating timestamps. If None, uses time.time().
-        Returns:
-            A dictionary containing the status update message: {"event": "attempt_status_updated", "old_status": old_status, "new_status": new_status}.
-                IF no meaningful status update is performed, returns None.
-        Raises:
-            ValueError: If the event is not recognized or the status transition is invalid.
-            NotImplementedError: If the event handling is not implemented for the current status.
-            RuntimeError: If the new status is not set after processing the event.
-        """
-        self._validate_status_message(msg)
-        event = msg["event"]
-        current_time = msg.get("timestamp", time.time())
-        old_status = self.status
-        new_status = msg.get("new_status", None)
-
-        # Step 1: Determine the new status based on the event and current status
-        if event == "user_update":
-            if not new_status:
-                raise ValueError("new_status must be provided for user_update event.")
-        elif event == "span_received":
-            self.last_heartbeat_time = current_time
-            if old_status in ["preparing", "unresponsive", "running"]:
-                new_status = "running"
-            elif old_status in self.get_finished_statuses():
-                logger.warning(
-                    f"Span received after attempt is already in status {self.status}. No status update performed."
-                )
-                return  # no further status update needed
-            else:
-                raise NotImplementedError(f"Event {event} is not implemented for status {old_status}.")
-        elif event == "single_step_timeout":
-            if old_status in [
-                "preparing",
-                "running",
-            ]:
-                new_status = "unresponsive"
-            else:
-                logger.warning(
-                    f"Single step timeout detected but attempt is in status {self.status}. No status update performed."
-                )
-                return  # no further status update needed
-        elif event == "overall_timeout":
-            if old_status not in self.get_finished_statuses():
-                new_status = "timeout"
-            else:
-                logger.warning(
-                    f"Overall timeout detected but attempt is in status {self.status}. No status update performed."
-                )
-                return  # no further status update needed
-        else:
-            raise NotImplementedError(f"Event {event} is not implemented for status update.")
-
-        # Step 2: Update the status
-        if not new_status:
-            raise RuntimeError(
-                f"new_status should not be {new_status} after processing event for {event} on status {old_status}."
-            )
-        if new_status == old_status:
-            return  # no status change
-        if new_status in self.get_finished_statuses():
-            # when attempt is finished, set end_time
-            self.end_time = current_time
-        self.status = new_status
-
-        # Step 3: Return the status update info for further processing
-        return AttemptStatusUpdateMessage(
-            attempt_id=self.attempt_id,
-            rollout_id=self.rollout_id,
-            timestamp=current_time,
-            old_status=old_status,
-            new_status=new_status,
-        )
-
-    @classmethod
-    async def get_latest_attempt_for_rollout(
-        cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
-    ) -> Optional[Attempt]:
-        async with session_factory() as session:
-            async with session.begin():
-                result = await session.scalars(
-                    select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.desc()).limit(1)
-                )
-                attempt_obj = result.one_or_none()
-                if attempt_obj is None:
-                    return None
-                return attempt_obj.as_attempt()
-
-    @classmethod
-    async def get_attempts_for_rollout(
-        cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
-    ) -> List[Attempt]:
-        async with session_factory() as session:
-            async with session.begin():
-                result = await session.scalars(
-                    select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.asc())
-                )
-                return [attempt.as_attempt() for attempt in result.all()]
-
-
-class SpanSeqIdInDB(SqlAlchemyBase):
-    __tablename__ = "span_sequence"
-
-    rollout_id: Mapped[str] = mapped_column(nullable=False, primary_key=True)
-
-    # FIXME InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
-    # attempt_id: Mapped[str] = mapped_column(nullable=False)
-    attempt_id: InitVar[str]  # not mapped column, just for type hinting
-
-    current_sequence: Mapped[int] = mapped_column(default=1, nullable=False)
-
-    # Versioning for optimistic concurrency control
-    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
-    __mapper_args__ = {
-        "version_id_col": version_id,
-        # "primary_key": [rollout_id, attempt_id],
-        # "primary_key": [rollout_id],
-    }
-
-    @classmethod
-    async def get_next_sequence_id(
-        cls: type[SpanSeqIdInDB],
-        session_factory: async_sessionmaker[AsyncSession],
-        rollout_id: str,
-        attempt_id: str,
-        external_seq_id: Optional[int] = None,
-    ) -> int:
-        """Get the next sequence ID with retries to handle race conditions.
-        IF external_seq_id is provided and is greater than current_sequence, set current_sequence to external_seq_id.
-        """
-        async with session_factory() as session:
-            async with session.begin():
-                seq_obj = await session.get(cls, rollout_id)
-                # seq_obj = await session.get(cls, [rollout_id, attempt_id])
-                if seq_obj is None:
-                    raise ValueError(f"Rollout {rollout_id} not found")
-                else:
-                    current_seq = (
-                        external_seq_id
-                        if external_seq_id is not None and external_seq_id > seq_obj.current_sequence
-                        else seq_obj.current_sequence
-                    )
-                    seq_obj.current_sequence = current_seq + 1
-                    await session.flush()
-                    return current_seq

agentlightning/store/database/orm/base.py

@@ -1,186 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import json
-import time
-from typing import Any, Callable, Dict, List, Optional
-
-from pydantic import BaseModel, Field, TypeAdapter, computed_field
-
-# from dataclasses import asdict
-from sqlalchemy import JSON, TypeDecorator
-from sqlalchemy.ext.asyncio import AsyncAttrs
-from sqlalchemy.orm import DeclarativeBase, MappedAsDataclass
-
-
-class SqlAlchemyBase(AsyncAttrs, MappedAsDataclass, DeclarativeBase):
-    pass
-
-    def model_dump(
-        self,
-        exclude: set[str] | None = None,
-        mapper: Dict[str, Callable[["SqlAlchemyBase"], Any]] | None = None,
-    ) -> Dict[str, Any]:
-        """Dump the SQLAlchemy model to a dictionary.
-        Args:
-            exclude: set[str]
-                The set of field names to exclude.
-            mapper: Dict[str, Callable[[SqlAlchemyBase], Any]]
-                A mapping from field names to functions that take the model instance and return the value to be used for that field.
-                If the key is "*", the function should return a dictionary of additional fields to be added to the output.
-        Returns:
-            Dict[str, Any]: The dumped model as a dictionary.
-        """
-        exclude = exclude or set()
-        mapper = mapper or {}
-        dic = {k: getattr(self, k) for k in self.__table__.columns.keys() if k not in exclude}
-        for k, func in mapper.items():
-            if k == "*":
-                dic.update(func(self))
-            else:
-                dic[k] = func(self)
-        return dic
-
-
-class PydanticInDB(TypeDecorator[BaseModel]):
-    """Custom SQLAlchemy type to store pydantic.BaseModel as JSON in the database.
-    Attributes:
-        target_type: type[BaseModel], the type of the pydantic model to be stored.
-    """
-
-    impl = JSON
-    target_type: type[BaseModel] | None = None
-
-    def process_bind_param(self, value: BaseModel | None, dialect: Any) -> Optional[str]:
-        if value is None:
-            return None
-        if self.target_type is not None:
-            return TypeAdapter(self.target_type).validate_python(value).model_dump_json()  # type: ignore
-        return json.dumps(value)
-
-    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[BaseModel]:
-        if value is None:
-            return None
-        if self.target_type is not None:
-            return TypeAdapter(self.target_type).validate_json(value)  # type: ignore
-        dic = json.loads(value)
-        return dic  # type: ignore
-
-
-class PydanticListInDB(TypeDecorator[list[BaseModel]]):
-    """Custom SQLAlchemy type to store List[pydantic.BaseModel] as JSON in the database.
-    Attributes:
-        value_type: type[BaseModel], the type of the pydantic model to be stored in the list.
-    """
-
-    impl = JSON
-    value_type: type[BaseModel] | None = None
-
-    def process_bind_param(self, value: List[BaseModel] | None, dialect: Any) -> Optional[str]:
-        if value is None:
-            return None
-        if self.value_type is not None:
-            lst = [TypeAdapter(self.value_type).validate_python(v).model_dump() for v in value]
-            return json.dumps(lst)
-        raise ValueError("target_type must be set for PydanticListInDB")
-
-    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[List[BaseModel]]:
-        if value is None:
-            return None
-        if self.value_type is not None:
-            dic = json.loads(value)
-            return [TypeAdapter(self.value_type).validate_python(v) for v in dic]  # type: ignore
-        raise ValueError("target_type must be set for PydanticListInDB")
-
-
-class NamedDictBase(TypeDecorator[Dict[str, Any]]):
-    """Custom SQLAlchemy type to store Dict[str, pydantic.BaseModel] as JSON in the database.
-    Attributes:
-        target_alias: type[Dict[str, BaseModel]], the alias type of the dict.
-        value_type: type[BaseModel], the type of the values in the dict.
-
-    For example, given NamedResources = Dict[str, ResourceUnion],
-    we can define NamedDictBase with target_alias=NamedResources and target_type=ResourceUnion.
-    """
-
-    impl = JSON
-    target_alias: type | None = None
-    value_type: type[BaseModel] | Any = None
-
-    def process_bind_param(self, value: Dict[str, Any] | None, dialect: Any) -> Optional[str]:
-        if value is None:
-            return None
-
-        # ignore target_alias for when dumping because Dict is not a pydantic model
-        if self.value_type is not None:
-            dic = {
-                k: TypeAdapter(self.value_type).validate_python(v).model_dump() if isinstance(v, BaseModel) else v
-                for k, v in value.items()
-            }
-            return json.dumps(dic)
-        dic = {k: v.model_dump() if isinstance(v, BaseModel) else v for k, v in value.items()}
-        return json.dumps(dic)
-
-    def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[Dict[str, Any]]:
-        if value is None:
-            return None
-        if self.target_alias is not None:
-            return TypeAdapter(self.target_alias).validate_json(value)  # type: ignore
-        if self.value_type is not None:
-            dic = json.loads(value)
-            return {k: TypeAdapter(self.value_type).validate_python(v) for k, v in dic.items()}  # type: ignore
-        return json.loads(value)
-
-
-class DatabaseRuntimeError(Exception):
-    """Raised when a runtime error occurs during database operations.
-    Particularly used when the execution of a query fails.
-    """
-
-    pass
-
-
-class RaceConditionError(Exception):
-    """Raised when a race condition is detected during database operations."""
-
-    pass
-
-
-class NoRolloutToDequeueError(Exception):
-    """Raised when there is no rollout available to dequeue."""
-
-    pass
-
-
-class AttemptStatusUpdateMessage(BaseModel):
-    attempt_id: str
-    rollout_id: str
-    timestamp: float = Field(default_factory=time.time)
-    old_status: Optional[str] = None
-    new_status: str
-
-    @computed_field
-    @property
-    def event(self) -> str:
-        return "attempt_status_update"
-
-    @computed_field
-    @property
-    def is_failed(self) -> bool:
-        return self.new_status in ["failed", "timeout", "unresponsive"]
-
-    @computed_field
-    @property
-    def is_succeeded(self) -> bool:
-        return self.new_status == "succeeded"
-
-    @computed_field
-    @property
-    def is_finished(self) -> bool:
-        return self.is_failed or self.is_succeeded
-
-    @computed_field
-    @property
-    def is_running(self) -> bool:
-        return self.new_status in ["running", "preparing"]

agentlightning/store/database/orm/resources.py

@@ -1,55 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import hashlib
-import time
-import uuid
-from typing import Optional
-
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
-from sqlalchemy.orm import Mapped, mapped_column
-
-from agentlightning.types import NamedResources, ResourcesUpdate
-
-from .base import NamedDictBase, SqlAlchemyBase
-
-
-def _generate_resources_id() -> str:
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
-    return "rs-" + short_id
-
-
-class NamedResourcesInDB(NamedDictBase):
-    """Custom SQLAlchemy type to store NamedResources as JSON in the database."""
-
-    target_alias = NamedResources
-
-
-class ResourcesUpdateInDB(SqlAlchemyBase):
-    __tablename__ = "resources"
-    resources: Mapped[NamedResources] = mapped_column(
-        NamedResourcesInDB, nullable=False
-    )  # JSON serialized, convert to NamedResources when needed
-    resources_id: Mapped[str] = mapped_column(primary_key=True, default_factory=_generate_resources_id)
-    create_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time)
-    update_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time, onupdate=time.time)
-    version: Mapped[int] = mapped_column(nullable=False, default=1)
-
-    __mapper_args__ = {
-        "version_id_col": version,
-    }
-
-    @classmethod
-    async def get_resources_by_id(
-        cls, session_factory: async_sessionmaker[AsyncSession], resources_id: str
-    ) -> Optional[ResourcesUpdate]:
-        async with session_factory() as session:
-            async with session.begin():
-                obj = await session.get(cls, resources_id)
-                if obj is None:
-                    return None
-                return obj.as_resources_update()
-
-    def as_resources_update(self) -> ResourcesUpdate:
-        return ResourcesUpdate(**self.model_dump())

agentlightning/store/database/orm/rollout.py

@@ -1,201 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import hashlib
-import logging
-import time
-import uuid
-from typing import Any, Dict, List, Optional, cast
-
-from sqlalchemy import JSON, Float, Integer, String, and_, select
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
-from sqlalchemy.orm import Mapped, mapped_column
-
-from agentlightning.types import AttemptedRollout, Rollout, RolloutConfig, RolloutStatus
-
-from ...base import is_finished, is_queuing
-from .attempt import AttemptInDB
-from .base import AttemptStatusUpdateMessage, PydanticInDB, SqlAlchemyBase
-
-logger = logging.getLogger(__name__)
-
-
-def _generate_rollout_id() -> str:
-    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
-    return "ro-" + short_id
-
-
-class RolloutConfigInDB(PydanticInDB):
-    """Custom SQLAlchemy type to store RolloutConfig as JSON in the database."""
-
-    target_type = RolloutConfig
-
-
-class RolloutInDB(SqlAlchemyBase):
-    __tablename__ = "rollouts"
-
-    input: Mapped[Any] = mapped_column(JSON, nullable=False)
-    rollout_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_rollout_id)
-    start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
-    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
-    mode: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
-    resources_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
-    status: Mapped[RolloutStatus] = mapped_column(String, default="queuing", nullable=False)
-    config: Mapped[RolloutConfig] = mapped_column(
-        RolloutConfigInDB, nullable=False, default_factory=RolloutConfig
-    )  # JSON serialized, convert to RolloutConfig when needed
-    rollout_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(
-        JSON, nullable=True, default=None
-    )  # JSON serialized, convert to Dict when needed
-
-    # Attempt-related helper methods can be added here if needed
-    num_attempts: Mapped[int] = mapped_column(
-        Integer, default=0, nullable=False
-    )  # number of attempts made for this rollout
-    enqueue_time: Mapped[Optional[float]] = mapped_column(
-        Float, nullable=True, default_factory=time.time
-    )  # time when the rollout was enqueued (for FIFO scheduling)
-    latest_attempt_id: Mapped[Optional[str]] = mapped_column(
-        String, nullable=True, default=None
-    )  # the attempt_id of the latest attempt
-
-    # use optimistic concurrency control
-    version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
-    __mapper_args__ = {
-        "version_id_col": version_id,
-    }
-
-    def __post_init__(self):
-        if self.status not in ["queuing", "running", "succeeded", "failed", "requeuing"]:
-            raise ValueError(f"Invalid rollout status: {self.status}")
-
-    def as_rollout(self) -> Rollout:
-        return Rollout(
-            **self.model_dump(
-                exclude={"rollout_metadata", "num_attempts", "enqueue_time", "latest_attempt_id", "version_id"},
-                mapper={
-                    "metadata": lambda obj: obj.rollout_metadata,  # type: ignore
-                    "config": lambda obj: obj.config if obj.config is not None else RolloutConfig(),  # type: ignore
-                },
-            )
-        )
-
-    def _validate_status_message(self, msg: Dict[str, str]) -> None:
-        """Validate the status update message.
-        Raises:
-            ValueError: If the message is invalid.
-        """
-        if "event" not in msg:
-            raise ValueError("Status update message must contain 'event' field.")
-        event = msg["event"]
-        if event not in [
-            "attempt_status_update",  # from attempt status update
-            "user_update",  # from user-initiated update
-        ]:
-            raise ValueError(f"Invalid event type in status update message: {event}")
-        if event == "user_update":
-            if "new_status" not in msg:
-                raise ValueError("Status update message for event 'user_update' must contain 'new_status' field.")
-        if event == "attempt_status_update":
-            # leverage AttemptStatusUpdateMessage for validation
-            pass
-
-    async def update_status(self, msg: Dict[str, Any] | AttemptStatusUpdateMessage) -> None:
-        """Update the rollout status based on the provided message.
-        Args:
-            msg (Dict[str, str]): The status update message. Refer to `_validate_status_message` for the expected format.
-            current_time (Optional[float]): The current time to set end_time or enqueue_time if needed.
-        """
-        if isinstance(msg, dict):
-            self._validate_status_message(msg)
-            event = msg["event"]
-            current_time = msg.get("timestamp", time.time())
-        else:
-            event = msg.event
-            current_time = msg.timestamp
-
-        old_status = self.status
-        new_status = self.status  # initialize new_status with old_status
-
-        # Step 1: Determine the new status based on the event
-        if event == "user_update":
-            assert isinstance(msg, dict)
-            new_status = msg["new_status"]
-        elif event == "attempt_status_update":
-            msg = AttemptStatusUpdateMessage(**msg) if isinstance(msg, dict) else msg
-            if msg.attempt_id == self.latest_attempt_id:
-                new_status = msg.new_status  # directly take the latest attempt status
-                if msg.is_succeeded:
-                    new_status = "succeeded"
-                elif msg.is_failed:
-                    # no other attempts running, decide whether to requeue or fail
-                    config = self.config
-                    if config.max_attempts > self.num_attempts and msg.new_status in config.retry_condition:
-                        new_status = "requeuing"
-                    else:
-                        new_status = "failed"
-                # elif msg.is_running and old_status in ["failed", "requeuing"]:
-                # new_status = "running"
-            else:
-                # ignore attempts from old attempts
-                new_status = old_status
-
-        # Step 2: Update the status if it has changed and handle follow-up actions
-        if new_status is None:
-            raise RuntimeError(
-                f"New status of `{old_status}` and `{self.latest_attempt_id}` could not be determined from the message {msg}."
-            )
-        if new_status == old_status:
-            return
-        self.status = cast(RolloutStatus, new_status)
-
-        if is_finished(self):  # type: ignore
-            self.end_time = current_time
-        if is_queuing(self):  # type: ignore
-            self.enqueue_time = current_time
-            # When requeuing, we do not reset latest_attempt_id or num_attempts,
-            # as they should persist across requeues.
-
-    @classmethod
-    async def get_rollout_by_id(
-        cls: type[RolloutInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
-    ) -> Optional[Rollout | AttemptedRollout]:
-        """Query a specific rollout from the database."""
-        async with session_factory() as session:
-            async with session.begin():
-                rollout_obj = await session.get(cls, rollout_id)
-                if rollout_obj is None:
-                    return None
-                if rollout_obj.latest_attempt_id is not None:
-                    attempt_obj = await session.get(AttemptInDB, rollout_obj.latest_attempt_id)
-                    if attempt_obj is not None:
-                        return AttemptedRollout(
-                            **rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt()
-                        )
-                return rollout_obj.as_rollout()
-
-    @classmethod
-    async def query_rollouts(
-        cls: type[RolloutInDB],
-        session_factory: async_sessionmaker[AsyncSession],
-        *,
-        statuses: Optional[List[str]] = None,
-        ids: Optional[List[str]] = None,
-    ) -> List[RolloutInDB]:
-        """
-        Query rollouts from the database with optional filters.
-        """
-        async with session_factory() as session:
-            async with session.begin():
-                conditions: list[Any] = []
-                if statuses is not None:
-                    conditions.append(cls.status.in_(statuses))
-                if ids is not None:
-                    conditions.append(cls.rollout_id.in_(ids))
-                query = select(cls)
-                if conditions:
-                    query = query.where(and_(*conditions))
-                result = await session.scalars(query)
-                rollout_objs = result.all()
-                return list(rollout_objs)

agentlightning/store/database/orm/span.py

@@ -1,101 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import logging
-from typing import Any, Dict, List, Optional
-
-from sqlalchemy import JSON, Float, Integer, String
-from sqlalchemy.orm import Mapped, mapped_column
-
-logger = logging.getLogger(__name__)
-
-from agentlightning.types.tracer import (
-    Attributes,
-    AttributeValue,
-    Event,
-    Link,
-    OtelResource,
-    Span,
-    SpanContext,
-    TraceStatus,
-)
-
-from .base import NamedDictBase, PydanticInDB, PydanticListInDB, SqlAlchemyBase
-
-
-class TraceStatusInDB(PydanticInDB):
-    target_type = TraceStatus
-
-
-class AttributesInDB(NamedDictBase):
-    target_alias = None  # type: ignore
-    value_type = AttributeValue
-
-
-class EventListInDB(PydanticListInDB):
-    value_type = Event
-
-
-class LinkListInDB(PydanticListInDB):
-    value_type = Link
-
-
-class SpanContextInDB(PydanticInDB):
-    target_type = SpanContext
-
-
-class OtelResourceInDB(PydanticInDB):
-    target_type = OtelResource
-
-
-class SpanInDB(SqlAlchemyBase):
-    __tablename__ = "spans"
-
-    rollout_id: Mapped[str] = mapped_column(String, nullable=False)  # The rollout which this span belongs to.
-    attempt_id: Mapped[str] = mapped_column(String, nullable=False)  # The attempt which this span belongs to.
-    sequence_id: Mapped[int] = mapped_column(
-        Integer, nullable=False
-    )  # The ID to make spans ordered within a single attempt.
-
-    # Current ID (in hex, formatted via trace_api.format_*)
-    trace_id: Mapped[str] = mapped_column(
-        String, nullable=False
-    )  # one rollout can have traces coming from multiple places
-
-    # FIXME: span_id may be not unique across different attempts/rollouts, use (rollout_id, attempt_id, sequence_id) as the primary key instead
-    span_id: Mapped[str] = mapped_column(
-        String, nullable=False
-    )  # The span ID of the span. This ID comes from the OpenTelemetry span ID generator.
-    parent_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)  # The parent span ID of the span.
-
-    # Core ReadableSpan fields
-    name: Mapped[str] = mapped_column(String, nullable=False)
-    status: Mapped[TraceStatus] = mapped_column(TraceStatusInDB, nullable=False)
-    attributes: Mapped[Attributes] = mapped_column(AttributesInDB, nullable=False)
-    events: Mapped[List[Event]] = mapped_column(EventListInDB, nullable=False)
-    links: Mapped[List[Link]] = mapped_column(LinkListInDB, nullable=False)
-
-    # Timestamps
-    start_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-    end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
-
-    # Other parsable fields
-    context: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
-    parent: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
-    resource: Mapped[OtelResource] = mapped_column(OtelResourceInDB, nullable=False)
-
-    # extra fields can be added here as needed
-    extra: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
-
-    __mapper_args__ = {
-        "primary_key": [rollout_id, attempt_id, sequence_id],
-    }
-
-    def as_span(self) -> Span:
-        return Span(
-            **self.model_dump(
-                exclude={"extra"},
-                mapper={"*": lambda obj: obj.extra or {}},  # type: ignore
-            )
-        )

agentlightning/store/database/retry_helper.py

@@ -1,316 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""This file contains a configurable async retry decorator based on exception type."""
-
-from __future__ import annotations
-
-import functools
-import importlib
-import logging
-import random
-from dataclasses import asdict, dataclass
-from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Optional, Type, TypeVar
-
-from tenacity import AsyncRetrying, RetryCallState, retry_if_exception
-
-# ----------------------------------------------------------------------
-# Logging setup
-# ----------------------------------------------------------------------
-logger = logging.getLogger("async_retry")
-logging.basicConfig(level=logging.INFO)
-
-# ----------------------------------------------------------------------
-# Type alias for async callable
-# ----------------------------------------------------------------------
-F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
-
-
-# ----------------------------------------------------------------------
-# Dataclass definition for retry configuration
-# ----------------------------------------------------------------------
-@dataclass
-class RetryStrategy:
-    """Configuration schema for retry behavior of a specific exception type.
-    The wait time before $n$-th retry is calculated as ($n$ starts from 1):
-        wait_time = wait_seconds * (backoff ** (n - 1)) * (1 + jitter * U(-1, 1))
-    where U(-1, 1) is a uniform random variable between -1 and 1.
-    Attributes:
-        max_attempts: Maximum number of attempts before giving up. Default is 1 (no retry). None means infinite retries.
-        max_retry_delay: Optional maximum delay between retries in seconds. Default is None (no limit).
-        wait_seconds: Base wait time in seconds before the first retry. Default is 0.0.
-        max_wait_seconds: Maximum wait time in seconds between retries. Default is None (no limit).
-        backoff: Exponential backoff multiplier. Default is 1.0 (no backoff).
-        jitter: Fractional (relative) jitter to apply to wait time. Default is 0.0 (no jitter).
-        log: Whether to log each retry attempt. Default is False.
-    """
-
-    max_attempts: Optional[int] = 1
-    max_retry_delay: Optional[float] = None
-    wait_seconds: float = 0.0
-    max_wait_seconds: Optional[float] = None
-    backoff: float = 1.0
-    jitter: float = 0.0
-    log: bool = False
-
-    def asdict(self) -> Dict[str, Any]:
-        return asdict(self)
-
-    def __post_init__(self):
-        if self.max_attempts is not None and self.max_attempts < 1:
-            raise ValueError("max_attempts must be at least 1 or None for infinite retries")
-        if self.wait_seconds < 0.0:
-            raise ValueError("wait_seconds must be non-negative")
-        if self.backoff < 1.0:
-            raise ValueError("backoff must be at least 1.0")
-        if not (0.0 <= self.jitter <= 1.0):
-            raise ValueError("jitter must be between 0.0 and 1.0")
-
-    def _get_wait_time(self, attempt_number: int) -> float:
-        """Calculate the wait time before the given attempt number."""
-        base_wait = self.wait_seconds * (self.backoff ** (attempt_number - 1))
-        if self.jitter > 0:
-            delta = base_wait * self.jitter
-            wait_time = random.uniform(base_wait - delta, base_wait + delta)
-        else:
-            wait_time = base_wait
-        wait_time = max(wait_time, 0.0)
-        if self.max_wait_seconds is not None:
-            wait_time = min(wait_time, self.max_wait_seconds)
-        return wait_time
-
-    def wait_func(self, retry_state: RetryCallState) -> float:
-        """Tenacity wait function based on the given strategy."""
-        return self._get_wait_time(retry_state.attempt_number)
-
-    def stop_func(self, retry_state: RetryCallState) -> bool:
-        """Tenacity stop function based on the given strategy."""
-        if self.max_attempts is not None:
-            if retry_state.attempt_number >= self.max_attempts:
-                return True
-        if self.max_retry_delay is not None:
-            time_since_start = retry_state.seconds_since_start
-            if time_since_start is None:
-                logger.warning("Cannot determine time since start for retry stop condition.")
-                return False
-            if time_since_start >= self.max_retry_delay:
-                return True
-        return False
-
-    async def before_sleep(self, retry_state: RetryCallState):
-        """Tenacity before_sleep callback to log retry attempts."""
-        if self.log:
-            exc = retry_state.outcome.exception() if retry_state.outcome else None
-            next_wait = self.wait_func(retry_state)
-            logger.warning(
-                f"[Retry] {exc.__class__.__name__}: attempt={retry_state.attempt_number}, "
-                f"next_wait={next_wait:.2f}s, message={exc}"
-            )
-
-
-# ----------------------------------------------------------------------
-# Exception Registry — shared, reusable, and extensible
-# ----------------------------------------------------------------------
-class ExceptionRegistry:
-    """
-    Global registry for mapping string keys to Exception classes.
-    Supports dynamic registration and fallback to importlib.
-    """
-
-    _registry: Dict[str, Type[BaseException]] = {}
-
-    @classmethod
-    def register(cls, name: str, exc_type: Type[BaseException] | None = None) -> None:
-        """Register an exception type under a given name."""
-        if name in cls._registry:
-            logger.warning(f"Overwriting existing exception registration for name '{name}'.")
-        if exc_type is None:
-            # Try to dynamically import the exception class
-            try:
-                module_name, class_name = name.rsplit(".", 1)
-                module = importlib.import_module(module_name)
-                exc_type = getattr(module, class_name)
-                if exc_type is None:
-                    raise TypeError(f"{name} is not an Exception type.")
-            except (ImportError, AttributeError, ValueError, TypeError) as e:
-                raise ValueError(f"Cannot resolve exception type for name '{name}': {e}")
-        cls._registry[name] = exc_type
-
-    @classmethod
-    def all_registered(cls) -> Dict[str, Type[BaseException]]:
-        """Return the current registry mapping."""
-        return dict(cls._registry)
-
-    @classmethod
-    def clear(cls):
-        """Clear all registered exception mappings."""
-        cls._registry.clear()
-
-
-# ----------------------------------------------------------------------
-# Async Retry Decorator
-# ----------------------------------------------------------------------
-class AsyncTypeBasedRetry:
-    """
-    A configurable async retry decorator based on exception type.
-
-    - Takes configuration as a Dict[str, RetryStrategy].
-    - Provides `from_json()` for quick loading.
-    - Uses a global ExceptionRegistry to resolve exception names.
-    """
-
-    def __init__(self, strategies: Dict[str, RetryStrategy], default_strategy: RetryStrategy | None = None):
-        self.exception_map = self._build_exception_map(strategies)
-        self.default_strategy = default_strategy or RetryStrategy()
-
-    # ------------------------------------------------------------------
-    # Build exception map
-    # ------------------------------------------------------------------
-    def _build_exception_map(self, strategies: Dict[str, RetryStrategy]) -> Dict[Type[BaseException], RetryStrategy]:
-        mapping: Dict[Type[BaseException], RetryStrategy] = {}
-        all_registered = ExceptionRegistry.all_registered()
-        for name, strat in strategies.items():
-            if name in all_registered:
-                exc_type = all_registered[name]
-            else:
-                raise ValueError(f"Exception type '{name}' is not registered in ExceptionRegistry.")
-            mapping[exc_type] = strat
-        return mapping
-
-    # ------------------------------------------------------------------
-    # Retry core logic
-    # ------------------------------------------------------------------
-    def get_exception(self, retry_state: RetryCallState) -> Optional[BaseException]:
-        """Get the exception from the given retry state, if any."""
-        return retry_state.outcome.exception() if retry_state.outcome else None
-
-    def get_strategy(self, retry_state: RetryCallState) -> Optional[RetryStrategy]:
-        """Get the RetryStrategy for the exception in the given retry state.
-        IF no matching exception type is found, return the default strategy.
-        IF no exception is found, return None.
-        """
-        exc = self.get_exception(retry_state)
-        if exc is None:
-            return None
-        for exc_type, strat in self.exception_map.items():
-            if isinstance(exc, exc_type):
-                return strat
-        return self.default_strategy
-
-    def should_retry(self, exc: BaseException) -> bool:
-        return any(isinstance(exc, t) for t in self.exception_map.keys())
-
-    def wait_func(self, retry_state: RetryCallState) -> float:
-        strat = self.get_strategy(retry_state)
-        if strat is None:
-            return 0.0
-        return strat.wait_func(retry_state)
-
-    def stop_func(self, retry_state: RetryCallState) -> bool:
-        strat = self.get_strategy(retry_state)
-        if strat is None:
-            return False
-        return strat.stop_func(retry_state)
-
-    async def before_sleep(self, retry_state: RetryCallState):
-        strat = self.get_strategy(retry_state)
-        if strat is None:
-            return
-        await strat.before_sleep(retry_state)
-
-    # ------------------------------------------------------------------
-    # Decorator entry point
-    # ------------------------------------------------------------------
-    def __call__(self, func: F) -> F:
-        @functools.wraps(func)
-        async def wrapper(*args, **kwargs):  # type: ignore
-            async for attempt in AsyncRetrying(
-                retry=retry_if_exception(lambda e: self.should_retry(e)),
-                wait=self.wait_func,
-                stop=self.stop_func,
-                before_sleep=self.before_sleep,
-                reraise=True,
-            ):
-                with attempt:
-                    return await func(*args, **kwargs)
-
-        return wrapper  # type: ignore
-
-
-# ----------------------------------------------------------------------
-# A configurable async retrier for any code block
-# ----------------------------------------------------------------------
-
-
-class AsyncRetryBlock:
-    """
-    Async retry helper for a single exception type and strategy.
-
-    Usage:
-        async with AsyncRetryBlock(strategy):
-            await some_async_function()
-    """
-
-    def __init__(self, strategy: RetryStrategy, **retry_kwargs):  # type: ignore
-        self.strategy = strategy
-        self._retryer = AsyncRetrying(
-            wait=self._wait_func,
-            stop=self._stop_func,
-            before_sleep=self._before_sleep,
-            **retry_kwargs,  # type: ignore
-        )
-
-    async def run(self, coro: Callable[..., Awaitable[Any]]) -> Any:
-        """Run the given coroutine with retries according to the strategy.
-        For example:
-            async def my_coro():
-                ...
-            retry_block = AsyncRetryBlock(strategy)
-            result = await retry_block.run(my_coro)
-        """
-        async for attempt in self._retryer:
-            with attempt:
-                return await coro()
-
-    # ------------------------------------------------------------------
-    # Core: async iterator interface
-    # ------------------------------------------------------------------
-    def __aiter__(self) -> AsyncIterator[Any]:
-        """Return an async iterator that yields retry attempts.
-        Usage:
-            async for attempt in retry_block:
-                with attempt:
-                    await some_async_function()
-        """
-        return self._retryer.__aiter__()
-
-    # ------------------------------------------------------------------
-    # Context manager entry
-    # ------------------------------------------------------------------
-    async def __aenter__(self):
-        self._aiter = self._retryer.__aiter__()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
-        # Consume the retry iterator
-        try:
-            # If exception occurred, let the retryer handle it
-            async for attempt in self._aiter:
-                with attempt:
-                    if exc_val:
-                        raise exc_val
-        except Exception:
-            # Allow exception to propagate if retries exhausted
-            pass
-        return False
-
-    # ------------------------------------------------------------------
-    # Strategy function
-    # ------------------------------------------------------------------
-    def _wait_func(self, retry_state: RetryCallState) -> float:
-        return self.strategy.wait_func(retry_state)
-
-    def _stop_func(self, retry_state: RetryCallState) -> bool:
-        return self.strategy.stop_func(retry_state)
-
-    async def _before_sleep(self, retry_state: RetryCallState):
-        await self.strategy.before_sleep(retry_state)

agentlightning/store/database/sqlite.py

@@ -1,685 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import asyncio
-import logging
-import os
-import time
-from datetime import datetime, timedelta
-from typing import Any, Dict, List, Literal, Optional, Sequence, Union
-
-from apscheduler.schedulers.background import BackgroundScheduler
-from apscheduler.triggers.interval import IntervalTrigger
-from opentelemetry.sdk.trace import ReadableSpan
-from pydantic import BaseModel
-from sqlalchemy import and_, or_, select
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
-from sqlalchemy.orm.exc import StaleDataError
-from tenacity import RetryError
-
-from agentlightning.types import (
-    Attempt,
-    AttemptedRollout,
-    AttemptStatus,
-    NamedResources,
-    ResourcesUpdate,
-    Rollout,
-    RolloutConfig,
-    RolloutStatus,
-    Span,
-    TaskInput,
-)
-
-from ..base import UNSET, LightningStore, Unset, is_finished
-from .orm import (
-    AttemptInDB,
-    ResourcesUpdateInDB,
-    RolloutInDB,
-    SpanInDB,
-    SpanSeqIdInDB,
-    SqlAlchemyBase,
-)
-from .retry_helper import AsyncRetryBlock, AsyncTypeBasedRetry, ExceptionRegistry, RetryStrategy
-
-logger = logging.getLogger(__name__)
-
-# TODO add periodic cleanup of old rollouts/attempts/spans
-
-ExceptionRegistry.register("sqlalchemy.orm.exc.StaleDataError")
-ExceptionRegistry.register("sqlalchemy.exc.OperationalError")
-
-db_retry = AsyncTypeBasedRetry(
-    {
-        "sqlalchemy.exc.OperationalError": RetryStrategy(
-            max_attempts=5, wait_seconds=1, backoff=1.5, jitter=0.3, log=True
-        ),
-        "sqlalchemy.orm.exc.StaleDataError": RetryStrategy(
-            max_attempts=100, wait_seconds=1e-3, backoff=1.0, jitter=0.1, log=True
-        ),
-    }
-)
-
-
-class _WaitForRolloutsCompleted(Exception):
-    """Internal exception to signal that not all rollouts have completed yet."""
-
-    pass
-
-
-class BackgroundTaskConfig(BaseModel):
-    name: str  # unique name for the task
-    method: str  # method name to call, currently only supports methods of SqlLightningStore
-    interval: Dict[Literal["seconds", "minutes", "hours"], float]  # interval for the task
-    is_async: bool = True  # whether the task method is async, default to True
-
-
-class SqlLightningStore(LightningStore):
-    """
-    A LightningStore implementation that uses a database backend to store and manage rollouts and attempts.
-    The database backend is expected to support asynchronous operations.
-    The store uses SQLAlchemy ORM models to interact with the database
-    Args:
-        database_url (string):
-            The database URL for connecting to the database.
-            If None, will read from the 'DATABASE_URL' environment variable.
-        retry_for_waiting (RetryStrategy):
-            Retry strategy for polling when waiting for rollouts to complete.
-            If None, a default strategy will be used.
-        wait_for_nonexistent_rollout (Bool):
-            If True, when waiting for rollouts, will wait for all specified rollouts to complete, including non-existing ones.
-            If False, will ignore non-existing rollouts as completed. (Default: False)
-        background_tasks_cfg (list[Dict[str, Any]]):
-            The configuration for in-process periodic tasks, following the definition of `BackgroundTaskConfig`.
-            IF not provided (None as default), the dbstore will incorporate a default set of periodic tasks as follows:
-            [
-                BackgroundTaskConfig(name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}),
-            ]
-            To disable all periodic tasks, provide an empty list `[]`.
-    Note:
-        Explicitly use async `start()` and `stop()` methods to manage the database connection lifecycle.
-    """
-
-    def __init__(
-        self,
-        database_url: Optional[str] = None,
-        *,
-        retry_for_waiting: Optional[dict[str, Any] | RetryStrategy] = None,
-        wait_for_nonexistent_rollout: bool = False,
-        background_tasks_cfg: list[Dict[str, Any]] | None = None,
-    ) -> None:
-        super().__init__()
-        if database_url is None:
-            database_url = os.getenv("DATABASE_URL", None)
-        if database_url is None:
-            raise ValueError(
-                "A database URL must be provided either via the 'database_url' parameter or the 'DATABASE_URL' environment variable."
-            )
-
-        self._engine = create_async_engine(database_url, echo=False)
-        self._async_session = async_sessionmaker(self._engine, expire_on_commit=False)
-
-        self._latest_resources_id = None
-
-        # special handling for retry strategy
-        retry_for_waiting = retry_for_waiting or RetryStrategy(
-            max_attempts=10,  # set a limit for retries if timeout is specified, otherwise will change to None later
-            max_retry_delay=None,  # set later
-            wait_seconds=10.0,  # poll every 10 seconds
-            max_wait_seconds=60.0,  # at most wait 60 seconds between retries
-            backoff=1.0,
-            jitter=0.0,
-            log=True,
-        )
-        self.retry_for_waiting = (
-            retry_for_waiting if isinstance(retry_for_waiting, RetryStrategy) else RetryStrategy(**retry_for_waiting)
-        )
-        self.wait_for_nonexistent_rollout = wait_for_nonexistent_rollout
-
-        # setup in-process periodic tasks
-        if background_tasks_cfg is None:
-            self.background_tasks_cfg = [
-                BackgroundTaskConfig(
-                    name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}
-                ),
-            ]
-        else:
-            self.background_tasks_cfg = [BackgroundTaskConfig(**cfg) for cfg in background_tasks_cfg]
-        self._background_scheduler = BackgroundScheduler()
-
-    async def start(self):
-        async with self._engine.begin() as conn:
-            await conn.run_sync(SqlAlchemyBase.metadata.create_all)
-        for task_cfg in self.background_tasks_cfg:
-            self.add_background_task(task_cfg, to_scheduler_only=True)
-        self._background_scheduler.start()  # type: ignore
-
-    async def stop(self):
-        await self._engine.dispose()
-        self._background_scheduler.shutdown()  # type: ignore
-
-    def add_background_task(
-        self, task_cfg: Dict[str, Any] | BackgroundTaskConfig, to_scheduler_only: bool = False
-    ) -> None:
-        """Add a new periodic background task to the scheduler.
-        Args:
-            task_cfg (Dict[str, Any] | BackgroundTaskConfig): The configuration for the background task.
-            to_scheduler_only (bool): If True, only add the task to the scheduler without updating the configuration list.
-        Raises:
-            ValueError: If the task method is not defined in SqlLightningStore.
-        """
-        config = task_cfg if isinstance(task_cfg, BackgroundTaskConfig) else BackgroundTaskConfig(**task_cfg)
-        if not to_scheduler_only:
-            # check existing tasks
-            for existing in self.background_tasks_cfg:
-                if existing.name == config.name:
-                    logger.warning(
-                        f"Background task {config.name} is already scheduled, will update its configuration."
-                    )
-            self.background_tasks_cfg.append(config)
-        delta_t = timedelta(**config.interval)
-        if not hasattr(self, config.method):
-            raise ValueError(f"Periodic task method {config.method} is not defined in SqlLightningStore.")
-        if config.is_async:
-            func = lambda: asyncio.run(getattr(self, config.method)())
-        else:
-            func = lambda: getattr(self, config.method)()
-
-        self._background_scheduler.add_job(  # type: ignore
-            func=func,
-            trigger=IntervalTrigger(**config.interval),  # type: ignore
-            name=f"SqlLightningStore.{config.name}",
-            replace_existing=True,
-            next_run_time=datetime.now() + delta_t,  # schedule the first run after the interval
-        )
-
-    # ------------------------------------------------------
-    # Public methods defined in LightningStore
-    # ------------------------------------------------------
-
-    @db_retry
-    async def start_rollout(
-        self,
-        input: TaskInput,
-        mode: Literal["train", "val", "test"] | None = None,
-        resources_id: str | None = None,
-        config: RolloutConfig | None = None,
-        metadata: Dict[str, Any] | None = None,
-    ) -> AttemptedRollout:
-        async with self._async_session() as session:
-            async with session.begin():
-                rollout_obj = RolloutInDB(
-                    input=input,
-                    mode=mode,
-                    resources_id=resources_id or self._latest_resources_id,
-                    status="queuing",
-                    config=config or RolloutConfig(),
-                    rollout_metadata=metadata,
-                )
-                session.add(rollout_obj)
-                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
-                await session.flush()  # ensure the object is written to the DB
-                return attempted_rollout
-
-    @db_retry
-    async def enqueue_rollout(
-        self,
-        input: TaskInput,
-        mode: Literal["train", "val", "test"] | None = None,
-        resources_id: str | None = None,
-        config: RolloutConfig | None = None,
-        metadata: Dict[str, Any] | None = None,
-    ) -> Rollout:
-        async with self._async_session() as session:
-            async with session.begin():
-                rollout_obj = RolloutInDB(
-                    input=input,
-                    mode=mode,
-                    resources_id=resources_id or self._latest_resources_id,
-                    status="queuing",
-                    config=config or RolloutConfig(),
-                    rollout_metadata=metadata,
-                )
-                session.add(rollout_obj)
-                await session.flush()  # ensure the object is written to the DB
-                return rollout_obj.as_rollout()
-
-    @db_retry
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        return await self._fifo_dequeue_rollout()
-
-    @db_retry
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
-        async with self._async_session() as session:
-            async with session.begin():
-                rollout_obj = await session.get(RolloutInDB, rollout_id)
-                if rollout_obj is None:
-                    raise ValueError(f"Rollout {rollout_id} not found")
-                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
-                await session.flush()  # ensure the object is written to the DB
-                return attempted_rollout
-
-    @db_retry
-    async def add_span(self, span: Span) -> Span:
-        seq_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, span.rollout_id, span.attempt_id)
-        return await self._add_span(span.model_dump(), seq_id=seq_id)
-
-    @db_retry
-    async def add_otel_span(
-        self,
-        rollout_id: str,
-        attempt_id: str,
-        readable_span: ReadableSpan,
-        sequence_id: int | None = None,
-    ) -> Span:
-        sequence_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id, sequence_id)
-        span = Span.from_opentelemetry(
-            src=readable_span,
-            rollout_id=rollout_id,
-            attempt_id=attempt_id,
-            sequence_id=sequence_id,
-        )
-        return await self._add_span(span.model_dump(), seq_id=sequence_id)
-
-    @db_retry
-    async def query_rollouts(
-        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Rollout]:
-        rollouts = await RolloutInDB.query_rollouts(self._async_session, statuses=status, ids=rollout_ids)  # type: ignore
-        attempt_ids = [r.latest_attempt_id for r in rollouts if r.latest_attempt_id is not None]
-        async with self._async_session() as session:
-            async with session.begin():
-                scalars = await session.scalars(select(AttemptInDB).where(AttemptInDB.attempt_id.in_(attempt_ids)))
-                attempts = scalars.all()
-                attempt_map = {a.attempt_id: a.as_attempt() for a in attempts}
-                return [
-                    (
-                        AttemptedRollout(**r.as_rollout().model_dump(), attempt=attempt_map[r.latest_attempt_id])
-                        if r.latest_attempt_id in attempt_map
-                        else r.as_rollout()
-                    )
-                    for r in rollouts
-                ]  # type: ignore
-
-    @db_retry
-    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
-        return await AttemptInDB.get_attempts_for_rollout(self._async_session, rollout_id)  # type: ignore
-
-    @db_retry
-    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
-        return await RolloutInDB.get_rollout_by_id(self._async_session, rollout_id)
-
-    @db_retry
-    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
-        return await AttemptInDB.get_latest_attempt_for_rollout(self._async_session, rollout_id)
-
-    @db_retry
-    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
-        return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, resources_id)
-
-    @db_retry
-    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
-        if self._latest_resources_id is None:
-            return None
-        return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, self._latest_resources_id)
-
-    @db_retry
-    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
-        return await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id)
-
-    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
-        # implementation the timeout via tenacity retry mechanism, by a `with` context
-        strategy = RetryStrategy(**self.retry_for_waiting.asdict())
-        if timeout is not None:
-            strategy.max_retry_delay = timeout
-            if strategy.max_attempts is not None:
-                strategy.wait_seconds = min(strategy.wait_seconds, timeout / (strategy.max_attempts + 1))
-        else:
-            strategy.max_attempts = None  # infinite retries
-
-        non_completed_ids, non_existing_ids = set(rollout_ids), set(rollout_ids)
-        completed_rollouts: Dict[str, Rollout] = {}
-        if len(non_completed_ids) < len(rollout_ids):
-            logger.warning("Duplicate rollout_ids found in wait_for_rollouts input. Duplicates will be ignored.")
-
-        try:
-            async for attempt in AsyncRetryBlock(
-                strategy,
-                reraise=True,
-            ):
-                with attempt:
-                    async with self._async_session() as session:
-                        async with session.begin():
-                            result = await session.scalars(
-                                select(RolloutInDB).where(RolloutInDB.rollout_id.in_(non_completed_ids))
-                            )
-                            rollouts = [r.as_rollout() for r in result.all()]
-                            for r in rollouts:
-                                if r.rollout_id in non_existing_ids:
-                                    non_existing_ids.discard(r.rollout_id)  # found existing rollout
-                                if is_finished(r):
-                                    completed_rollouts[r.rollout_id] = r
-                                    non_completed_ids.discard(r.rollout_id)
-                            # check termination conditions
-                            if self.wait_for_nonexistent_rollout:
-                                if len(non_completed_ids) == 0:
-                                    return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
-                                raise _WaitForRolloutsCompleted(
-                                    f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
-                                )
-                            else:
-                                if len(non_completed_ids) == len(non_existing_ids):
-                                    logger.warning(f"All remaining rollouts are non-existing: {non_existing_ids}.")
-                                    return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
-                                raise _WaitForRolloutsCompleted(
-                                    f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
-                                )
-
-        except (RetryError, _WaitForRolloutsCompleted):
-            return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
-        except Exception as e:
-            logger.error(f"Error while waiting for rollouts: {e}")
-            raise e
-
-        # Ensure a return value in case no rollouts are completed
-        return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
-
-    @db_retry
-    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
-        async with self._async_session() as session:
-            async with session.begin():
-                conditions: List[Any] = [SpanInDB.rollout_id == rollout_id]
-                if attempt_id is not None:
-                    if attempt_id == "latest":
-                        rollout_obj = await session.get(RolloutInDB, rollout_id)
-                        if rollout_obj is None:
-                            logger.warning(f"Rollout {rollout_id} does not exist. Cannot query latest attempt spans.")
-                            return []
-                        attempt_id = rollout_obj.latest_attempt_id
-                    conditions.append(SpanInDB.attempt_id == attempt_id)
-                query = select(SpanInDB).where(and_(*conditions)).order_by(SpanInDB.sequence_id.asc())
-                result = await session.scalars(query)
-                span_objs = result.all()
-                return [obj.as_span() for obj in span_objs]
-
-    @db_retry
-    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
-        async with self._async_session() as session:
-            async with session.begin():
-                current_time = time.time()
-                resource_obj = ResourcesUpdateInDB(
-                    resources=resources,
-                    create_time=current_time,
-                    update_time=current_time,
-                )
-                session.add(resource_obj)
-                await session.flush()  # ensure the object is written to the DB
-                self._latest_resources_id = resource_obj.resources_id
-                return resource_obj.as_resources_update()
-
-    @db_retry
-    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
-        async with self._async_session() as session:
-            async with session.begin():
-                obj = await session.get(ResourcesUpdateInDB, resources_id)
-                if obj is None:
-                    # raise ValueError(f"Failed to update resources {resources_id}. It may not exist.")
-                    # FIXME InMemoryLightningStore will create the resources if not exist, but the base method require to raise error
-                    # HACK here stick to the behavior of InMemoryLightningStore for compatibility
-                    current_time = time.time()
-                    obj = ResourcesUpdateInDB(
-                        resources_id=resources_id,
-                        resources=resources,
-                        create_time=current_time,
-                        update_time=current_time,
-                    )
-                    session.add(obj)
-                else:
-                    obj.resources = resources
-                await session.flush()
-                self._latest_resources_id = resources_id
-                return obj.as_resources_update()
-
-    @db_retry
-    async def query_resources(self) -> List[ResourcesUpdate]:
-        async with self._async_session() as session:
-            async with session.begin():
-                result = await session.scalars(
-                    select(ResourcesUpdateInDB).order_by(ResourcesUpdateInDB.create_time.asc())
-                )
-                resource_objs = result.all()
-                return [obj.as_resources_update() for obj in resource_objs]
-
-    @db_retry
-    async def update_rollout(
-        self,
-        rollout_id: str | None,
-        input: TaskInput | Unset = UNSET,
-        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
-        resources_id: Optional[str] | Unset = UNSET,
-        status: RolloutStatus | Unset = UNSET,
-        config: RolloutConfig | Unset = UNSET,
-        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> Rollout:
-        if rollout_id is None:
-            raise ValueError("rollout_id must be provided for updating a rollout.")
-
-        async with self._async_session() as session:
-            async with session.begin():
-                rollout_obj = await session.get(RolloutInDB, rollout_id)
-                if rollout_obj is None:
-                    raise ValueError(f"Rollout {rollout_id} not found")
-                # udpate fields
-                if not isinstance(input, Unset):
-                    rollout_obj.input = input
-                if not isinstance(mode, Unset):
-                    rollout_obj.mode = mode
-                if not isinstance(resources_id, Unset):
-                    rollout_obj.resources_id = resources_id
-                if not isinstance(status, Unset):
-                    await rollout_obj.update_status(dict(event="user_update", new_status=status))
-                if not isinstance(config, Unset):
-                    rollout_obj.config = config
-                if not isinstance(metadata, Unset):
-                    rollout_obj.rollout_metadata = metadata
-                await session.flush()  # ensure the object is written to the DB
-                return rollout_obj.as_rollout()
-
-    @db_retry
-    async def update_attempt(
-        self,
-        rollout_id: str,
-        attempt_id: str | Literal["latest"],
-        status: AttemptStatus | Unset = UNSET,
-        worker_id: str | Unset = UNSET,
-        last_heartbeat_time: float | Unset = UNSET,
-        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
-    ) -> Attempt:
-        async with self._async_session() as session:
-            async with session.begin():
-                rollout_obj = await session.get(RolloutInDB, rollout_id)
-                if rollout_obj is None:
-                    raise ValueError(f"Rollout {rollout_id} not found")
-                if attempt_id == "latest":
-                    if rollout_obj.latest_attempt_id is None:
-                        raise ValueError(f"Rollout {rollout_id} has no attempts. Cannot update latest attempt.")
-                    attempt_id = rollout_obj.latest_attempt_id
-                if attempt_id != rollout_obj.latest_attempt_id:
-                    logger.warning(
-                        f"Updating attempt {attempt_id} which is not the latest attempt for rollout {rollout_id}. Latest is {rollout_obj.latest_attempt_id}."
-                    )
-                attempt_obj = await session.get(AttemptInDB, attempt_id)
-                if attempt_obj is None:
-                    raise ValueError(f"No attempts found")
-                if attempt_obj.rollout_id != rollout_id:
-                    raise ValueError(f"Attempt {attempt_id} does not belong to rollout {rollout_id}.")
-                # update fields
-                if not isinstance(status, Unset):
-                    msg = attempt_obj.update_status(dict(event="user_update", new_status=status))
-                    if msg is not None:
-                        await rollout_obj.update_status(msg)
-                if not isinstance(worker_id, Unset):
-                    attempt_obj.worker_id = worker_id
-                if not isinstance(last_heartbeat_time, Unset):
-                    attempt_obj.last_heartbeat_time = last_heartbeat_time
-                if not isinstance(metadata, Unset):
-                    attempt_obj.attempt_metadata = metadata
-                await session.flush()  # ensure the object is written to the DB
-                return attempt_obj.as_attempt()
-
-    # ------------------------------------------------------
-    # periodic background tasks can be added here
-    # ------------------------------------------------------
-
-    async def check_attempt_timeout(self):
-        """Periodically check for attempts that have timed out and update their status accordingly."""
-        # use update with where condition to find and update timed-out attempts
-        current_time = time.time()
-
-        timed_out_results = await self._attempt_timeout_check(current_time)
-
-        # TODO run the tasks with a wrapper with asyncio semaphore to limit concurrency and handle exceptions
-        tasks = [self._process_timed_out_attempt(attempt, current_time) for attempt in timed_out_results]
-        await asyncio.gather(*tasks)
-
-    async def _process_timed_out_attempt(self, attempt_ref: AttemptInDB, current_time: float) -> None:
-        async with self._async_session() as session:
-            async with session.begin():
-                # Step 1: Update attempt status
-                attempt_obj = await session.get(
-                    AttemptInDB, attempt_ref.attempt_id
-                )  # refresh the object in the new session
-                if attempt_obj is None:
-                    raise ValueError(f"Attempt {attempt_ref.attempt_id} not found during timeout processing")
-                if attempt_obj.version_id != attempt_ref.version_id:
-                    # version mismatch, skip processing to avoid race conditions
-                    raise StaleDataError(f"Attempt {attempt_ref.attempt_id} version mismatch during timeout processing")
-                msg = {}
-                if attempt_obj.is_timed_out(current_time):
-                    msg = dict(event="overall_timeout", timestamp=current_time)
-                elif attempt_obj.is_unresponsive(current_time):
-                    msg = dict(event="single_step_timeout", timestamp=current_time)
-                else:
-                    raise ValueError(f"Attempt {attempt_ref.attempt_id} is not timed out during timeout processing")
-                msg2rollout = attempt_obj.update_status(msg)
-                if msg2rollout is None:
-                    return  # no further update needed
-
-                # Step 2: Update rollouts
-                rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
-                if rollout_obj is None:
-                    raise ValueError(f"Rollout {attempt_obj.rollout_id} not found during timeout processing")
-                await rollout_obj.update_status(msg2rollout)
-
-    # ------------------------------------------------------
-    # internal helper methods can be added here
-    # ------------------------------------------------------
-
-    async def _add_span(self, span: Dict[str, Any], seq_id: Optional[int] = None) -> Span:
-        """Add a new span to the database."""
-        if seq_id is not None:
-            span["sequence_id"] = seq_id
-        extra_dic: Dict[str, Any] = {}
-        for k in list(span.keys()):
-            if k not in SpanInDB.__table__.columns.keys():
-                extra_dic[k] = span.pop(k)
-        span["extra"] = extra_dic if extra_dic else None
-
-        async with self._async_session() as session:
-            async with session.begin():
-                # create SpanInDB object
-                span_obj = SpanInDB(**span)
-                session.add(span_obj)
-                # update attempt's last_heartbeat_time and status
-                attempt_obj = await session.get(AttemptInDB, span["attempt_id"])
-                if attempt_obj is None:
-                    raise ValueError(f"Attempt {span['attempt_id']} not found")
-                # ensure the attempt and rollout are in running status
-                msg = attempt_obj.update_status(dict(event="span_received"))
-                if msg is not None:
-                    rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
-                    if rollout_obj is None:
-                        raise ValueError(f"Rollout {attempt_obj.rollout_id} not found")
-                    await rollout_obj.update_status(msg)
-                await session.flush()  # ensure the object is written to the DB
-                return span_obj.as_span()
-
-    async def _fifo_dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        """Dequeue the next rollout in FIFO order (the one with the earliest enqueue_time).
-        Returns the RolloutInDB object if found, else None.
-        Note: This method does not update the status of the rollout. The caller should handle that.
-        """
-        async with self._async_session() as session:
-            async with session.begin():
-                # use the update...returning to atomically select the next rollout and claim it by updating its status to 'preparing'
-                result = await session.scalars(
-                    select(RolloutInDB)
-                    .where(RolloutInDB.status.in_(["queuing", "requeuing"]), RolloutInDB.enqueue_time.isnot(None))
-                    .order_by(RolloutInDB.enqueue_time.asc())
-                    .limit(1)
-                )
-                rollout_obj = result.one_or_none()
-                if rollout_obj is None:
-                    return None  # no rollout available
-                # update the status of the rollout to 'preparing' via Compare-and-Swap to avoid race
-                attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
-                await session.flush()  # ensure the object is written to the DB
-                return attempted_rollout
-
-    async def _start_attempt_for_rollout(self, session: AsyncSession, rollout_obj: RolloutInDB) -> AttemptedRollout:
-        """Create a new attempt for the given rollout and update the rollout's fields."""
-        # create a new attempt for this rollout
-        rollout_config = rollout_obj.config
-        attempt_obj = AttemptInDB(
-            rollout_id=rollout_obj.rollout_id,
-            sequence_id=rollout_obj.num_attempts + 1,
-            status="preparing",
-            max_duration=rollout_config.timeout_seconds,
-            max_heartbeat_interval=rollout_config.unresponsive_seconds,
-        )
-        session.add(attempt_obj)
-        # pre-update the rollout_obj fields for CAS
-        rollout_obj.status = attempt_obj.status  # type: ignore pre-update the status in the object for CAS
-        rollout_obj.enqueue_time = None  # pre-update the enqueue_time in the object for CAS
-        rollout_obj.num_attempts += 1  # pre-update the num_attempts in the object for CAS
-        rollout_obj.latest_attempt_id = attempt_obj.attempt_id  # pre-update the latest_attempt_id in the object for CAS
-
-        # create a sequence id tracker for each attempt
-        # FIXME currently InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
-        # create a sequence id tracker for this rollout, only if not exists
-        existing = await session.get(SpanSeqIdInDB, rollout_obj.rollout_id)
-        if existing is None:
-            seq_obj = SpanSeqIdInDB(
-                rollout_id=rollout_obj.rollout_id,
-                attempt_id=attempt_obj.attempt_id,
-            )
-            session.add(seq_obj)
-
-        return AttemptedRollout(**rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt())
-
-    async def _attempt_timeout_check(self, now: float) -> Sequence[AttemptInDB]:
-        """Scan the table for attempts that have timed out based on the given mode, and return them for further processing.
-        Returns:
-            list[AttemptInDB]:
-                A list of AttemptInDB objects that timed out.
-        """
-        async with self._async_session() as session:
-            async with session.begin():
-                scalars = await session.scalars(
-                    select(AttemptInDB).where(
-                        and_(
-                            AttemptInDB.status.in_(["preparing", "running"]),
-                            or_(
-                                and_(
-                                    AttemptInDB.max_duration.isnot(None),
-                                    (now - AttemptInDB.start_time) > AttemptInDB.max_duration,
-                                ),
-                                and_(
-                                    AttemptInDB.max_heartbeat_interval.isnot(None),
-                                    (now - AttemptInDB.last_heartbeat_time) > AttemptInDB.max_heartbeat_interval,
-                                ),
-                            ),
-                        )
-                    )
-                )
-                return scalars.all()

agentlightning/store/memory.py

@@ -46,7 +46,7 @@ from agentlightning.types import (
     TaskInput,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset, is_finished, is_queuing
+from .base import UNSET, LightningStore, Unset, is_finished, is_queuing
 from .utils import healthcheck, propagate_status
 
 T_callable = TypeVar("T_callable", bound=Callable[..., Any])
@@ -243,14 +243,6 @@ class InMemoryLightningStore(LightningStore):
         # Completion tracking for wait_for_rollouts (cross-loop safe)
         self._completion_events: Dict[str, threading.Event] = {}
 
-    def capabilities(self) -> LightningStoreCapabilities:
-        """Return the capabilities of the store."""
-        return LightningStoreCapabilities(
-            thread_safe=False,
-            async_safe=True,
-            zero_copy=False,
-        )
-
     @_healthcheck_wrapper
     async def start_rollout(
         self,
@@ -435,7 +427,7 @@ class InMemoryLightningStore(LightningStore):
     @_healthcheck_wrapper
     async def query_rollouts(
         self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
-    ) -> List[Union[Rollout, AttemptedRollout]]:
+    ) -> List[Rollout]:
         """Retrieves rollouts filtered by their status and rollout ids.
         If no status is provided, returns all rollouts.
 

agentlightning/store/sqlite.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# TODO: Implement this

agentlightning/store/threading.py

@@ -20,7 +20,7 @@ from agentlightning.types import (
     TaskInput,
 )
 
-from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
+from .base import UNSET, LightningStore, Unset
 
 
 class LightningStoreThreaded(LightningStore):
@@ -35,15 +35,6 @@ class LightningStoreThreaded(LightningStore):
         self.store = store
         self._lock = threading.Lock()
 
-    def capabilities(self) -> LightningStoreCapabilities:
-        """Return the capabilities of the store."""
-        capabilities = self.store.capabilities()
-        return {
-            **capabilities,
-            "async_safe": True,
-            "thread_safe": True,
-        }
-
     async def start_rollout(
         self,
         input: TaskInput,

agentlightning/types/core.py

@@ -117,7 +117,6 @@ AttemptStatus = Literal[
 ]
 """The status of an attempt."""
 
-
 RolloutMode = Literal["train", "val", "test"]
 """Possible rollout modes."""
 

agentlightning/utils/server_launcher.py

@@ -8,7 +8,6 @@ import logging
 import multiprocessing
 import queue
 import signal
-import socket
 import threading
 import time
 import traceback
@@ -25,21 +24,18 @@ from gunicorn.app.base import BaseApplication
 from gunicorn.arbiter import Arbiter
 from portpicker import pick_unused_port
 
-__all__ = ["PythonServerLauncher", "PythonServerLauncherArgs", "LaunchMode"]
+__all__ = ["PythonServerLauncher", "PythonServerLauncherArgs"]
 
 
 LaunchMode = Literal["asyncio", "thread", "mp"]
-"""The launch mode for the server."""
 
 
 @dataclass
 class PythonServerLauncherArgs:
     port: Optional[int] = None
     """The TCP port to listen on. If not provided, the server will use a random available port."""
-    host: Optional[str] = None
+    host: str = "127.0.0.1"
     """The hostname or IP address to bind the server to."""
-    access_host: Optional[str] = None
-    """The hostname or IP address to advertise to the client. If not provided, the server will use the default outbound IPv4 address for this machine."""
     launch_mode: LaunchMode = "asyncio"
     """The launch mode. `asyncio` is the default mode to runs the server in the current thread.
     `thread` runs the server in a separate thread. `mp` runs the server in a separate process."""
@@ -158,26 +154,19 @@ async def run_uvicorn_asyncio(
             # Normally, the program will not reach this point, as the server will throw the exception itself earlier.
             raise RuntimeError(f"Server did not start up within {timeout:.2f} seconds.") from server_start_exception
 
-        logger.info(f"Server started up in {time.time() - start_time:.2f} seconds.")
+        logger.debug(f"Server started up in {time.time() - start_time:.2f} seconds.")
 
         # Check for health endpoint status if provided
         if health_url is not None:
-            logger.info(f"Probing health endpoint {health_url}...")
             async with aiohttp.ClientSession() as session:
                 while time.time() < deadline:
-                    try:
+                    with suppress(Exception):
                         async with session.get(health_url) as resp:
                             if resp.status == 200:
-                                logger.info(
+                                logger.debug(
                                     f"Server is healthy at {health_url} in {time.time() - start_time:.2f} seconds."
                                 )
                                 return
-                            else:
-                                logger.debug(
-                                    f"Server is NOT healthy at {health_url} in {time.time() - start_time:.2f} seconds. Got status {resp.status}."
-                                )
-                    except Exception as e:
-                        logger.debug(f"Error probing health endpoint {health_url}: {str(e)}")
                     await asyncio.sleep(0.1)
 
             # If the server is not healthy, kill it if requested.
@@ -198,7 +187,7 @@ async def run_uvicorn_asyncio(
                 )
 
         else:
-            logger.info("Server does not provide a health check endpoint. Skipping health check.")
+            logger.debug("Server does not provide a health check endpoint. Skipping health check.")
 
     async def _serve_server() -> None:
         nonlocal server_start_exception
@@ -566,27 +555,6 @@ def run_gunicorn(
         watchdog_thread.join(timeout=5.0)
 
 
-def _get_default_ipv4_address() -> str:
-    """Determine the default outbound IPv4 address for this machine.
-
-    Implementation:
-        Opens a UDP socket and "connects" to a public address to force route
-        selection, then inspects the socket's local address. No packets are sent.
-
-    Returns:
-        str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
-    """
-    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    try:
-        # Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
-        s.connect(("8.8.8.8", 80))
-        return s.getsockname()[0]
-    except Exception:
-        return "127.0.0.1"
-    finally:
-        s.close()
-
-
 class PythonServerLauncher:
     """Unified launcher for FastAPI, using uvicorn or gunicorn per mode/worker count.
 
@@ -605,9 +573,7 @@ class PythonServerLauncher:
         self.app = app
         self.args = args
         self.serve_context = serve_context
-        self._host: Optional[str] = self.args.host
         self._port: Optional[int] = self.args.port
-        self._access_host: Optional[str] = self.args.access_host
 
         # uvicorn (in-proc asyncio)
         self._uvicorn_server: Optional[uvicorn.Server] = None
@@ -629,12 +595,14 @@ class PythonServerLauncher:
     @property
     def endpoint(self) -> str:
         """Return the externally advertised host:port pair regardless of accessibility."""
-        return f"http://{self._ensure_host()}:{self._ensure_port()}"
+        return f"http://{self.args.host}:{self._ensure_port()}"
 
     @property
-    def access_endpoint(self) -> str:
+    def access_url(self) -> str:
         """Return a loopback-friendly URL so health checks succeed even when binding to 0.0.0.0."""
-        return f"http://{self._ensure_access_host()}:{self._ensure_port()}"
+        # Probe host normalization for 0.0.0.0
+        host_for_probe = "127.0.0.1" if self.args.host in ("0.0.0.0", "::") else self.args.host
+        return f"http://{host_for_probe}:{self._ensure_port()}"
 
     @property
     def health_url(self) -> Optional[str]:
@@ -644,7 +612,7 @@ class PythonServerLauncher:
         path = self.args.healthcheck_url
         if not path.startswith("/"):
             path = "/" + path
-        return f"{self.access_endpoint}{path}"
+        return f"{self.access_url}{path}"
 
     async def start(self):
         """Starts the server according to launch_mode and n_workers."""
@@ -731,35 +699,16 @@ class PythonServerLauncher:
             return f"{module}:app"
         return "unknown:app"
 
-    def _ensure_host(self) -> str:
-        if self._host is None:
-            logger.warning("No host provided, using 0.0.0.0.")
-            self._host = "0.0.0.0"
-        return self._host
-
     def _ensure_port(self) -> int:
         if self._port is None:
             logger.warning("No port provided, using pick_unused_port to pick a random unused port.")
             self._port = pick_unused_port()
         return self._port
 
-    def _ensure_access_host(self) -> str:
-        if self.args.access_host is None:
-            if self._ensure_host() in ("0.0.0.0", "::"):
-                # Probe host normalization for 0.0.0.0
-                logger.warning("No access host provided, using default outbound IPv4 address for this machine.")
-                self._access_host = _get_default_ipv4_address()
-            else:
-                logger.warning("No access host provided, using the host provided.")
-                self._access_host = self._ensure_host()
-        else:
-            self._access_host = self.args.access_host
-        return self._access_host
-
     def _create_uvicorn_server(self) -> uvicorn.Server:
         config = uvicorn.Config(
             app=self.app,
-            host=self._ensure_host(),
+            host=self.args.host,
             port=self._ensure_port(),
             log_level=self.args.log_level,
             loop="asyncio",
@@ -833,9 +782,6 @@ class PythonServerLauncher:
         try:
             evt: ChildEvent = await asyncio.to_thread(self._thread_event_queue.get, True, timeout)
         except queue.Empty:
-            if not self._thread.is_alive():
-                logger.error("Threaded server failed to start and is not alive. No error event was received.")
-                return
             logger.error("Threaded server failed to start and sends no event. This should not happen.")
             await self._stop_uvicorn_thread()
             return
@@ -873,7 +819,6 @@ class PythonServerLauncher:
         if self.is_running():
             raise RuntimeError("Server process is already running. Stopping it first.")
 
-        host = self._ensure_host()
         port = self._ensure_port()
 
         try:
@@ -889,7 +834,7 @@ class PythonServerLauncher:
         if self.args.n_workers > 1:
             logger.info(f"Starting Gunicorn server...")
             options = {
-                "bind": f"{host}:{port}",
+                "bind": f"{self.args.host}:{port}",
                 "workers": int(self.args.n_workers),
                 "worker_class": "uvicorn_worker.UvicornWorker",
                 "loglevel": logging.getLevelName(self.args.log_level).lower(),
@@ -938,9 +883,6 @@ class PythonServerLauncher:
         try:
             evt: ChildEvent = await asyncio.to_thread(self._mp_event_queue.get, True, timeout)
         except queue.Empty:
-            if not self._proc.is_alive():
-                logger.error("Server process failed to start and is not alive. No error event was received.")
-                return
             logger.error("Server process failed to start and sends no event. This should not happen.")
             await self._stop_serving_process()
             return

agentlightning/verl/daemon.py

@@ -294,7 +294,7 @@ class AgentModeDaemon:
         self._proxy_thread.start()
         print(f"Proxy server running on port {self.proxy_port}")
 
-    async def _update_proxy_server_v1(self):
+    def _update_proxy_server_v1(self):
         model_name = self.train_information.get("model")
         if not model_name:
             raise ValueError("Model name is not set.")
@@ -313,7 +313,12 @@ class AgentModeDaemon:
             ],
         )
 
-        await self.llm_proxy.restart()
+        if self.llm_proxy.is_running():
+            # FIXME: Need to switch to a different port right now
+            # because the forked processes carried the old fd
+            self.llm_proxy.restart(_port=_find_available_port())
+        else:
+            self.llm_proxy.start()
 
     def start(self):
         """Starts the main AgentLightningServer and the proxy server."""
@@ -347,7 +352,7 @@ class AgentModeDaemon:
         if server_addresses != self.backend_llm_server_addresses:
             self.backend_llm_server_addresses = server_addresses
             if self.mode == "v1" and not self.llm_proxy.is_running():
-                await self._update_proxy_server_v1()
+                self._update_proxy_server_v1()
         self.is_train = is_train
 
         # 1. Update resources on the server for clients to use

dashboard/.storybook/main.ts

@@ -10,7 +10,7 @@ const config: StorybookConfig = {
   },
   stories: ['../src/**/*.mdx', '../src/**/*.story.@(js|jsx|ts|tsx)'],
   staticDirs: ['../static'],
-  addons: ['@storybook/addon-themes', '@storybook/addon-vitest'],
+  addons: ['@storybook/addon-themes'],
   framework: {
     name: '@storybook/react-vite',
     options: {},

dashboard/.storybook/vitest.setup.ts

@@ -1,8 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-import { setProjectAnnotations } from '@storybook/react-vite';
-import * as projectAnnotations from './preview';
-
-// This is an important step to apply the right configuration when testing your stories.
-// More info at: https://storybook.js.org/docs/api/portable-stories/portable-stories-vitest#setprojectannotations
-setProjectAnnotations([projectAnnotations]);

dashboard/package-lock.json

@@ -25,7 +25,6 @@
         "@eslint/js": "^9.37.0",
         "@ianvs/prettier-plugin-sort-imports": "^4.7.0",
         "@storybook/addon-themes": "^9.1.10",
-        "@storybook/addon-vitest": "^9.1.16",
         "@storybook/react": "^9.1.10",
         "@storybook/react-vite": "^9.1.10",
         "@stylistic/eslint-plugin": "^5.5.0",
@@ -37,8 +36,6 @@
         "@types/react": "^19.2.2",
         "@types/react-dom": "^19.2.1",
         "@vitejs/plugin-react": "^5.0.4",
-        "@vitest/browser-playwright": "4.0.4",
-        "@vitest/coverage-v8": "4.0.4",
         "chromatic": "^13.3.3",
         "eslint": "^9.37.0",
         "eslint-config-mantine": "^4.0.3",
@@ -48,7 +45,6 @@
         "jsdom": "^27.0.0",
         "msw": "^2.11.6",
         "msw-storybook-addon": "^2.0.6",
-        "playwright": "^1.56.1",
         "postcss": "^8.5.6",
         "postcss-preset-mantine": "1.18.0",
         "postcss-simple-vars": "^7.0.1",
@@ -437,16 +433,6 @@
         "node": ">=6.9.0"
       }
     },
-    "node_modules/@bcoe/v8-coverage": {
-      "version": "1.0.2",
-      "resolved": "https://registry.npmjs.org/@bcoe/v8-coverage/-/v8-coverage-1.0.2.tgz",
-      "integrity": "sha512-6zABk/ECA/QYSCQ1NGiVwwbQerUCZ+TQbp64Q3AgmfNvurHH0j8TtXa1qbShXA6qqkpAj4V5W8pP6mLe1mcMqA==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=18"
-      }
-    },
     "node_modules/@cacheable/memoize": {
       "version": "2.0.3",
       "resolved": "https://registry.npmjs.org/@cacheable/memoize/-/memoize-2.0.3.tgz",
@@ -1816,13 +1802,6 @@
         "node": ">=14"
       }
     },
-    "node_modules/@polka/url": {
-      "version": "1.0.0-next.29",
-      "resolved": "https://registry.npmjs.org/@polka/url/-/url-1.0.0-next.29.tgz",
-      "integrity": "sha512-wwQAWhWSuHaag8c4q/KN/vCoeOJYshAIvMQwD4GpSb3OiZklFfvAgmj0VCBBImRpuF/aFgIRzllXlVX93Jevww==",
-      "dev": true,
-      "license": "MIT"
-    },
     "node_modules/@reduxjs/toolkit": {
       "version": "2.9.2",
       "resolved": "https://registry.npmjs.org/@reduxjs/toolkit/-/toolkit-2.9.2.tgz",
@@ -2216,44 +2195,6 @@
         "storybook": "^9.1.16"
       }
     },
-    "node_modules/@storybook/addon-vitest": {
-      "version": "9.1.16",
-      "resolved": "https://registry.npmjs.org/@storybook/addon-vitest/-/addon-vitest-9.1.16.tgz",
-      "integrity": "sha512-X0rOOUMb5UHbfekcjnTeiDTarZdsg5irXXPxxL//8QQCFyCLF6Bdm1YNlCdF560PtwaaQPXzlxByD0FfGbtdWA==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@storybook/global": "^5.0.0",
-        "@storybook/icons": "^1.4.0",
-        "prompts": "^2.4.0",
-        "ts-dedent": "^2.2.0"
-      },
-      "funding": {
-        "type": "opencollective",
-        "url": "https://opencollective.com/storybook"
-      },
-      "peerDependencies": {
-        "@vitest/browser": "^3.0.0 || ^4.0.0",
-        "@vitest/browser-playwright": "^4.0.0",
-        "@vitest/runner": "^3.0.0 || ^4.0.0",
-        "storybook": "^9.1.16",
-        "vitest": "^3.0.0 || ^4.0.0"
-      },
-      "peerDependenciesMeta": {
-        "@vitest/browser": {
-          "optional": true
-        },
-        "@vitest/browser-playwright": {
-          "optional": true
-        },
-        "@vitest/runner": {
-          "optional": true
-        },
-        "vitest": {
-          "optional": true
-        }
-      }
-    },
     "node_modules/@storybook/builder-vite": {
       "version": "9.1.16",
       "resolved": "https://registry.npmjs.org/@storybook/builder-vite/-/builder-vite-9.1.16.tgz",
@@ -2297,20 +2238,6 @@
       "dev": true,
       "license": "MIT"
     },
-    "node_modules/@storybook/icons": {
-      "version": "1.6.0",
-      "resolved": "https://registry.npmjs.org/@storybook/icons/-/icons-1.6.0.tgz",
-      "integrity": "sha512-hcFZIjW8yQz8O8//2WTIXylm5Xsgc+lW9ISLgUk1xGmptIJQRdlhVIXCpSyLrQaaRiyhQRaVg7l3BD9S216BHw==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=14.0.0"
-      },
-      "peerDependencies": {
-        "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0-beta",
-        "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0-beta"
-      }
-    },
     "node_modules/@storybook/react": {
       "version": "9.1.16",
       "resolved": "https://registry.npmjs.org/@storybook/react/-/react-9.1.16.tgz",
@@ -2943,263 +2870,6 @@
         "vite": "^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0"
       }
     },
-    "node_modules/@vitest/browser": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/browser/-/browser-4.0.4.tgz",
-      "integrity": "sha512-1ZXztcBtRd3maKliHzWbQohsyRjam0ws6OPRWNWfGxFUOHTlNBtDnJAm8z1x7IzVkZ6JcOAumHJAbxNJh4tkDw==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/mocker": "4.0.4",
-        "@vitest/utils": "4.0.4",
-        "magic-string": "^0.30.19",
-        "pixelmatch": "7.1.0",
-        "pngjs": "^7.0.0",
-        "sirv": "^3.0.2",
-        "tinyrainbow": "^3.0.3",
-        "ws": "^8.18.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      },
-      "peerDependencies": {
-        "vitest": "4.0.4"
-      }
-    },
-    "node_modules/@vitest/browser-playwright": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/browser-playwright/-/browser-playwright-4.0.4.tgz",
-      "integrity": "sha512-jGKnGZ5ZKXuwQ1Ldwll/rZxk3webz4gz3kvoTYX2NH2ASPiwFGck8D09Sf2wVjCuDqebPXXd69zUIt1o4yQ5tA==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/browser": "4.0.4",
-        "@vitest/mocker": "4.0.4",
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      },
-      "peerDependencies": {
-        "playwright": "*",
-        "vitest": "4.0.4"
-      },
-      "peerDependenciesMeta": {
-        "playwright": {
-          "optional": false
-        }
-      }
-    },
-    "node_modules/@vitest/browser-playwright/node_modules/@vitest/mocker": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.0.4.tgz",
-      "integrity": "sha512-UTtKgpjWj+pvn3lUM55nSg34098obGhSHH+KlJcXesky8b5wCUgg7s60epxrS6yAG8slZ9W8T9jGWg4PisMf5Q==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/spy": "4.0.4",
-        "estree-walker": "^3.0.3",
-        "magic-string": "^0.30.19"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      },
-      "peerDependencies": {
-        "msw": "^2.4.9",
-        "vite": "^6.0.0 || ^7.0.0-0"
-      },
-      "peerDependenciesMeta": {
-        "msw": {
-          "optional": true
-        },
-        "vite": {
-          "optional": true
-        }
-      }
-    },
-    "node_modules/@vitest/browser-playwright/node_modules/@vitest/spy": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.0.4.tgz",
-      "integrity": "sha512-G9L13AFyYECo40QG7E07EdYnZZYCKMTSp83p9W8Vwed0IyCG1GnpDLxObkx8uOGPXfDpdeVf24P1Yka8/q1s9g==",
-      "dev": true,
-      "license": "MIT",
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/browser-playwright/node_modules/estree-walker": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/estree-walker/-/estree-walker-3.0.3.tgz",
-      "integrity": "sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@types/estree": "^1.0.0"
-      }
-    },
-    "node_modules/@vitest/browser-playwright/node_modules/tinyrainbow": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.0.3.tgz",
-      "integrity": "sha512-PSkbLUoxOFRzJYjjxHJt9xro7D+iilgMX/C9lawzVuYiIdcihh9DXmVibBe8lmcFrRi/VzlPjBxbN7rH24q8/Q==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=14.0.0"
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/@vitest/mocker": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.0.4.tgz",
-      "integrity": "sha512-UTtKgpjWj+pvn3lUM55nSg34098obGhSHH+KlJcXesky8b5wCUgg7s60epxrS6yAG8slZ9W8T9jGWg4PisMf5Q==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/spy": "4.0.4",
-        "estree-walker": "^3.0.3",
-        "magic-string": "^0.30.19"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      },
-      "peerDependencies": {
-        "msw": "^2.4.9",
-        "vite": "^6.0.0 || ^7.0.0-0"
-      },
-      "peerDependenciesMeta": {
-        "msw": {
-          "optional": true
-        },
-        "vite": {
-          "optional": true
-        }
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/@vitest/pretty-format": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.0.4.tgz",
-      "integrity": "sha512-lHI2rbyrLVSd1TiHGJYyEtbOBo2SDndIsN3qY4o4xe2pBxoJLD6IICghNCvD7P+BFin6jeyHXiUICXqgl6vEaQ==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/@vitest/spy": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.0.4.tgz",
-      "integrity": "sha512-G9L13AFyYECo40QG7E07EdYnZZYCKMTSp83p9W8Vwed0IyCG1GnpDLxObkx8uOGPXfDpdeVf24P1Yka8/q1s9g==",
-      "dev": true,
-      "license": "MIT",
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/@vitest/utils": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.0.4.tgz",
-      "integrity": "sha512-4bJLmSvZLyVbNsYFRpPYdJViG9jZyRvMZ35IF4ymXbRZoS+ycYghmwTGiscTXduUg2lgKK7POWIyXJNute1hjw==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/pretty-format": "4.0.4",
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/estree-walker": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/estree-walker/-/estree-walker-3.0.3.tgz",
-      "integrity": "sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@types/estree": "^1.0.0"
-      }
-    },
-    "node_modules/@vitest/browser/node_modules/tinyrainbow": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.0.3.tgz",
-      "integrity": "sha512-PSkbLUoxOFRzJYjjxHJt9xro7D+iilgMX/C9lawzVuYiIdcihh9DXmVibBe8lmcFrRi/VzlPjBxbN7rH24q8/Q==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=14.0.0"
-      }
-    },
-    "node_modules/@vitest/coverage-v8": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-4.0.4.tgz",
-      "integrity": "sha512-YM7gDj2TX2AXyGLz0p/B7hvTsTfaQc+kSV/LU0nEnKlep/ZfbdCDppPND4YQiQC43OXyrhkG3y8ZSTqYb2CKqQ==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@bcoe/v8-coverage": "^1.0.2",
-        "@vitest/utils": "4.0.4",
-        "ast-v8-to-istanbul": "^0.3.5",
-        "debug": "^4.4.3",
-        "istanbul-lib-coverage": "^3.2.2",
-        "istanbul-lib-report": "^3.0.1",
-        "istanbul-lib-source-maps": "^5.0.6",
-        "istanbul-reports": "^3.2.0",
-        "magicast": "^0.3.5",
-        "std-env": "^3.9.0",
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      },
-      "peerDependencies": {
-        "@vitest/browser": "4.0.4",
-        "vitest": "4.0.4"
-      },
-      "peerDependenciesMeta": {
-        "@vitest/browser": {
-          "optional": true
-        }
-      }
-    },
-    "node_modules/@vitest/coverage-v8/node_modules/@vitest/pretty-format": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.0.4.tgz",
-      "integrity": "sha512-lHI2rbyrLVSd1TiHGJYyEtbOBo2SDndIsN3qY4o4xe2pBxoJLD6IICghNCvD7P+BFin6jeyHXiUICXqgl6vEaQ==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/coverage-v8/node_modules/@vitest/utils": {
-      "version": "4.0.4",
-      "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.0.4.tgz",
-      "integrity": "sha512-4bJLmSvZLyVbNsYFRpPYdJViG9jZyRvMZ35IF4ymXbRZoS+ycYghmwTGiscTXduUg2lgKK7POWIyXJNute1hjw==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@vitest/pretty-format": "4.0.4",
-        "tinyrainbow": "^3.0.3"
-      },
-      "funding": {
-        "url": "https://opencollective.com/vitest"
-      }
-    },
-    "node_modules/@vitest/coverage-v8/node_modules/tinyrainbow": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.0.3.tgz",
-      "integrity": "sha512-PSkbLUoxOFRzJYjjxHJt9xro7D+iilgMX/C9lawzVuYiIdcihh9DXmVibBe8lmcFrRi/VzlPjBxbN7rH24q8/Q==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=14.0.0"
-      }
-    },
     "node_modules/@vitest/expect": {
       "version": "3.2.4",
       "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.4.tgz",
@@ -3655,35 +3325,6 @@
       "dev": true,
       "license": "MIT"
     },
-    "node_modules/ast-v8-to-istanbul": {
-      "version": "0.3.8",
-      "resolved": "https://registry.npmjs.org/ast-v8-to-istanbul/-/ast-v8-to-istanbul-0.3.8.tgz",
-      "integrity": "sha512-szgSZqUxI5T8mLKvS7WTjF9is+MVbOeLADU73IseOcrqhxr/VAvy6wfoVE39KnKzA7JRhjF5eUagNlHwvZPlKQ==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@jridgewell/trace-mapping": "^0.3.31",
-        "estree-walker": "^3.0.3",
-        "js-tokens": "^9.0.1"
-      }
-    },
-    "node_modules/ast-v8-to-istanbul/node_modules/estree-walker": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/estree-walker/-/estree-walker-3.0.3.tgz",
-      "integrity": "sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@types/estree": "^1.0.0"
-      }
-    },
-    "node_modules/ast-v8-to-istanbul/node_modules/js-tokens": {
-      "version": "9.0.1",
-      "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-9.0.1.tgz",
-      "integrity": "sha512-mxa9E9ITFOt0ban3j6L5MpjwegGz6lBQmM1IJkWeBZGcMxto50+eWdjC/52xDbS2vy0k7vIMK0Fe2wfL9OQSpQ==",
-      "dev": true,
-      "license": "MIT"
-    },
     "node_modules/astral-regex": {
       "version": "2.0.0",
       "resolved": "https://registry.npmjs.org/astral-regex/-/astral-regex-2.0.0.tgz",
@@ -5891,13 +5532,6 @@
         "node": ">=18"
       }
     },
-    "node_modules/html-escaper": {
-      "version": "2.0.2",
-      "resolved": "https://registry.npmjs.org/html-escaper/-/html-escaper-2.0.2.tgz",
-      "integrity": "sha512-H2iMtd0I4Mt5eYiapRdIDjp+XzelXQ0tFE4JS7YFwFevXXMmOp9myNrUvCg0D6ws8iqkRPBfKHgbwig1SmlLfg==",
-      "dev": true,
-      "license": "MIT"
-    },
     "node_modules/html-tags": {
       "version": "3.3.1",
       "resolved": "https://registry.npmjs.org/html-tags/-/html-tags-3.3.1.tgz",
@@ -6520,60 +6154,6 @@
       "dev": true,
       "license": "ISC"
     },
-    "node_modules/istanbul-lib-coverage": {
-      "version": "3.2.2",
-      "resolved": "https://registry.npmjs.org/istanbul-lib-coverage/-/istanbul-lib-coverage-3.2.2.tgz",
-      "integrity": "sha512-O8dpsF+r0WV/8MNRKfnmrtCWhuKjxrq2w+jpzBL5UZKTi2LeVWnWOmWRxFlesJONmc+wLAGvKQZEOanko0LFTg==",
-      "dev": true,
-      "license": "BSD-3-Clause",
-      "engines": {
-        "node": ">=8"
-      }
-    },
-    "node_modules/istanbul-lib-report": {
-      "version": "3.0.1",
-      "resolved": "https://registry.npmjs.org/istanbul-lib-report/-/istanbul-lib-report-3.0.1.tgz",
-      "integrity": "sha512-GCfE1mtsHGOELCU8e/Z7YWzpmybrx/+dSTfLrvY8qRmaY6zXTKWn6WQIjaAFw069icm6GVMNkgu0NzI4iPZUNw==",
-      "dev": true,
-      "license": "BSD-3-Clause",
-      "dependencies": {
-        "istanbul-lib-coverage": "^3.0.0",
-        "make-dir": "^4.0.0",
-        "supports-color": "^7.1.0"
-      },
-      "engines": {
-        "node": ">=10"
-      }
-    },
-    "node_modules/istanbul-lib-source-maps": {
-      "version": "5.0.6",
-      "resolved": "https://registry.npmjs.org/istanbul-lib-source-maps/-/istanbul-lib-source-maps-5.0.6.tgz",
-      "integrity": "sha512-yg2d+Em4KizZC5niWhQaIomgf5WlL4vOOjZ5xGCmF8SnPE/mDWWXgvRExdcpCgh9lLRRa1/fSYp2ymmbJ1pI+A==",
-      "dev": true,
-      "license": "BSD-3-Clause",
-      "dependencies": {
-        "@jridgewell/trace-mapping": "^0.3.23",
-        "debug": "^4.1.1",
-        "istanbul-lib-coverage": "^3.0.0"
-      },
-      "engines": {
-        "node": ">=10"
-      }
-    },
-    "node_modules/istanbul-reports": {
-      "version": "3.2.0",
-      "resolved": "https://registry.npmjs.org/istanbul-reports/-/istanbul-reports-3.2.0.tgz",
-      "integrity": "sha512-HGYWWS/ehqTV3xN10i23tkPkpH46MLCIMFNCaaKNavAXTF1RkqxawEPtnjnGZ6XKSInBKkiOA5BKS+aZiY3AvA==",
-      "dev": true,
-      "license": "BSD-3-Clause",
-      "dependencies": {
-        "html-escaper": "^2.0.0",
-        "istanbul-lib-report": "^3.0.0"
-      },
-      "engines": {
-        "node": ">=8"
-      }
-    },
     "node_modules/iterator.prototype": {
       "version": "1.1.5",
       "resolved": "https://registry.npmjs.org/iterator.prototype/-/iterator.prototype-1.1.5.tgz",
@@ -6758,16 +6338,6 @@
         "node": ">=0.10.0"
       }
     },
-    "node_modules/kleur": {
-      "version": "3.0.3",
-      "resolved": "https://registry.npmjs.org/kleur/-/kleur-3.0.3.tgz",
-      "integrity": "sha512-eTIzlVOSUR+JxdDFepEYcBMtZ9Qqdef+rnzWdRZuMbOywu5tO2w2N7rqjoANZ5k9vywhL6Br1VRjUIgTQx4E8w==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=6"
-      }
-    },
     "node_modules/known-css-properties": {
       "version": "0.37.0",
       "resolved": "https://registry.npmjs.org/known-css-properties/-/known-css-properties-0.37.0.tgz",
@@ -6896,34 +6466,6 @@
         "@jridgewell/sourcemap-codec": "^1.5.5"
       }
     },
-    "node_modules/magicast": {
-      "version": "0.3.5",
-      "resolved": "https://registry.npmjs.org/magicast/-/magicast-0.3.5.tgz",
-      "integrity": "sha512-L0WhttDl+2BOsybvEOLK7fW3UA0OQ0IQ2d6Zl2x/a6vVRs3bAY0ECOSHHeL5jD+SbOpOCUEi0y1DgHEn9Qn1AQ==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@babel/parser": "^7.25.4",
-        "@babel/types": "^7.25.4",
-        "source-map-js": "^1.2.0"
-      }
-    },
-    "node_modules/make-dir": {
-      "version": "4.0.0",
-      "resolved": "https://registry.npmjs.org/make-dir/-/make-dir-4.0.0.tgz",
-      "integrity": "sha512-hXdUTZYIVOt1Ex//jAQi+wTZZpUpwBj/0QsOzqegb3rGMMeJiSEu5xLHnYfBrRV4RH2+OCSOO95Is/7x1WJ4bw==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "semver": "^7.5.3"
-      },
-      "engines": {
-        "node": ">=10"
-      },
-      "funding": {
-        "url": "https://github.com/sponsors/sindresorhus"
-      }
-    },
     "node_modules/mantine-datatable": {
       "version": "8.2.0",
       "resolved": "https://registry.npmjs.org/mantine-datatable/-/mantine-datatable-8.2.0.tgz",
@@ -7086,16 +6628,6 @@
         "marked": "14.0.0"
       }
     },
-    "node_modules/mrmime": {
-      "version": "2.0.1",
-      "resolved": "https://registry.npmjs.org/mrmime/-/mrmime-2.0.1.tgz",
-      "integrity": "sha512-Y3wQdFg2Va6etvQ5I82yUhGdsKrcYox6p7FfL1LbK2J4V01F9TGlepTIhnK24t7koZibmg82KGglhA1XK5IsLQ==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=10"
-      }
-    },
     "node_modules/ms": {
       "version": "2.1.3",
       "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz",
@@ -7572,76 +7104,6 @@
         "url": "https://github.com/sponsors/jonschlinkert"
       }
     },
-    "node_modules/pixelmatch": {
-      "version": "7.1.0",
-      "resolved": "https://registry.npmjs.org/pixelmatch/-/pixelmatch-7.1.0.tgz",
-      "integrity": "sha512-1wrVzJ2STrpmONHKBy228LM1b84msXDUoAzVEl0R8Mz4Ce6EPr+IVtxm8+yvrqLYMHswREkjYFaMxnyGnaY3Ng==",
-      "dev": true,
-      "license": "ISC",
-      "dependencies": {
-        "pngjs": "^7.0.0"
-      },
-      "bin": {
-        "pixelmatch": "bin/pixelmatch"
-      }
-    },
-    "node_modules/playwright": {
-      "version": "1.56.1",
-      "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.56.1.tgz",
-      "integrity": "sha512-aFi5B0WovBHTEvpM3DzXTUaeN6eN0qWnTkKx4NQaH4Wvcmc153PdaY2UBdSYKaGYw+UyWXSVyxDUg5DoPEttjw==",
-      "dev": true,
-      "license": "Apache-2.0",
-      "dependencies": {
-        "playwright-core": "1.56.1"
-      },
-      "bin": {
-        "playwright": "cli.js"
-      },
-      "engines": {
-        "node": ">=18"
-      },
-      "optionalDependencies": {
-        "fsevents": "2.3.2"
-      }
-    },
-    "node_modules/playwright-core": {
-      "version": "1.56.1",
-      "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.56.1.tgz",
-      "integrity": "sha512-hutraynyn31F+Bifme+Ps9Vq59hKuUCz7H1kDOcBs+2oGguKkWTU50bBWrtz34OUWmIwpBTWDxaRPXrIXkgvmQ==",
-      "dev": true,
-      "license": "Apache-2.0",
-      "bin": {
-        "playwright-core": "cli.js"
-      },
-      "engines": {
-        "node": ">=18"
-      }
-    },
-    "node_modules/playwright/node_modules/fsevents": {
-      "version": "2.3.2",
-      "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz",
-      "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==",
-      "dev": true,
-      "hasInstallScript": true,
-      "license": "MIT",
-      "optional": true,
-      "os": [
-        "darwin"
-      ],
-      "engines": {
-        "node": "^8.16.0 || ^10.6.0 || >=11.0.0"
-      }
-    },
-    "node_modules/pngjs": {
-      "version": "7.0.0",
-      "resolved": "https://registry.npmjs.org/pngjs/-/pngjs-7.0.0.tgz",
-      "integrity": "sha512-LKWqWJRhstyYo9pGvgor/ivk2w94eSjE3RGVuzLGlr3NmD8bf7RcYGze1mNdEHRP6TRP6rMuDHk5t44hnTRyow==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=14.19.0"
-      }
-    },
     "node_modules/possible-typed-array-names": {
       "version": "1.1.0",
       "resolved": "https://registry.npmjs.org/possible-typed-array-names/-/possible-typed-array-names-1.1.0.tgz",
@@ -7936,20 +7398,6 @@
         "url": "https://github.com/chalk/ansi-styles?sponsor=1"
       }
     },
-    "node_modules/prompts": {
-      "version": "2.4.2",
-      "resolved": "https://registry.npmjs.org/prompts/-/prompts-2.4.2.tgz",
-      "integrity": "sha512-NxNv/kLguCA7p3jE8oL2aEBsrJWgAakBpgmgK6lpPWV+WuOmY6r2/zbAVnP+T8bQlA0nzHXSJSJW0Hq7ylaD2Q==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "kleur": "^3.0.3",
-        "sisteransi": "^1.0.5"
-      },
-      "engines": {
-        "node": ">= 6"
-      }
-    },
     "node_modules/prop-types": {
       "version": "15.8.1",
       "resolved": "https://registry.npmjs.org/prop-types/-/prop-types-15.8.1.tgz",
@@ -8772,28 +8220,6 @@
         "url": "https://github.com/sponsors/isaacs"
       }
     },
-    "node_modules/sirv": {
-      "version": "3.0.2",
-      "resolved": "https://registry.npmjs.org/sirv/-/sirv-3.0.2.tgz",
-      "integrity": "sha512-2wcC/oGxHis/BoHkkPwldgiPSYcpZK3JU28WoMVv55yHJgcZ8rlXvuG9iZggz+sU1d4bRgIGASwyWqjxu3FM0g==",
-      "dev": true,
-      "license": "MIT",
-      "dependencies": {
-        "@polka/url": "^1.0.0-next.24",
-        "mrmime": "^2.0.0",
-        "totalist": "^3.0.0"
-      },
-      "engines": {
-        "node": ">=18"
-      }
-    },
-    "node_modules/sisteransi": {
-      "version": "1.0.5",
-      "resolved": "https://registry.npmjs.org/sisteransi/-/sisteransi-1.0.5.tgz",
-      "integrity": "sha512-bLGGlR1QxBcynn2d5YmDX4MGjlZvy2MRBDRNHLJ8VI6l6+9FUiyTFNJ0IveOSP0bcXgVDPRcfGqA0pjaqUpfVg==",
-      "dev": true,
-      "license": "MIT"
-    },
     "node_modules/slash": {
       "version": "3.0.0",
       "resolved": "https://registry.npmjs.org/slash/-/slash-3.0.0.tgz",
@@ -9707,16 +9133,6 @@
         "node": ">=8.0"
       }
     },
-    "node_modules/totalist": {
-      "version": "3.0.1",
-      "resolved": "https://registry.npmjs.org/totalist/-/totalist-3.0.1.tgz",
-      "integrity": "sha512-sf4i37nQ2LBx4m3wB74y+ubopq6W/dIzXg0FDGjsYnZHVa1Da8FH853wlL2gtUhg+xJXjfk3kUZS3BRoQeoQBQ==",
-      "dev": true,
-      "license": "MIT",
-      "engines": {
-        "node": ">=6"
-      }
-    },
     "node_modules/tough-cookie": {
       "version": "6.0.0",
       "resolved": "https://registry.npmjs.org/tough-cookie/-/tough-cookie-6.0.0.tgz",

dashboard/package.json

@@ -10,8 +10,7 @@
     "eslint": "eslint .",
     "stylelint": "stylelint '**/*.css'",
     "prettier": "prettier --check \"**/*.{ts,tsx,mjs,cjs}\"",
-    "vitest": "vitest run --project unit",
-    "vitest-storybook": "vitest run --project storybook",
+    "vitest": "vitest run",
     "storybook": "storybook dev -p 6006",
     "build-storybook": "storybook build",
     "chromatic": "chromatic"
@@ -34,7 +33,6 @@
     "@eslint/js": "^9.37.0",
     "@ianvs/prettier-plugin-sort-imports": "^4.7.0",
     "@storybook/addon-themes": "^9.1.10",
-    "@storybook/addon-vitest": "^9.1.16",
     "@storybook/react": "^9.1.10",
     "@storybook/react-vite": "^9.1.10",
     "@stylistic/eslint-plugin": "^5.5.0",
@@ -67,9 +65,6 @@
     "typescript-eslint": "^8.46.0",
     "vite": "^7.1.9",
     "vite-tsconfig-paths": "^5.1.4",
-    "vitest": "^4.0.0",
-    "playwright": "^1.56.1",
-    "@vitest/browser-playwright": "4.0.4",
-    "@vitest/coverage-v8": "4.0.4"
+    "vitest": "^4.0.0"
   }
 }

dashboard/public/index.html

@@ -2,7 +2,7 @@
 <html lang="en">
   <head>
     <meta charset="UTF-8" />
-    <link rel="icon" type="image/svg+xml" href="../src/favicon.svg" />
+    <link rel="icon" type="image/svg+xml" href="/src/favicon.svg" />
     <meta name="viewport" content="minimum-scale=1, initial-scale=1, width=device-width, user-scalable=no" />
     <title>Agent-lightning Dashboard</title>
   </head>

dashboard/src/components/AppDrawer.component.tsx

@@ -1,13 +1,11 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from 'react';
+import { useCallback, useEffect, useMemo, useState, type ReactNode } from 'react';
 import { Editor } from '@monaco-editor/react';
 import { IconCheck, IconCopy } from '@tabler/icons-react';
 import type { DataTableSortStatus } from 'mantine-datatable';
-import { createSearchParams, Link, useInRouterContext, useLocation } from 'react-router-dom';
 import {
   ActionIcon,
-  Anchor,
   Badge,
   Box,
   CopyButton,
@@ -271,7 +269,7 @@ export function JsonEditor({ value }: JsonEditorProps) {
   const editorTheme = colorScheme === 'dark' ? 'vs-dark' : 'vs-light';
 
   return (
-    <Box data-testid='json-editor-container' style={{ flex: 1, minHeight: 0 }}>
+    <Box style={{ flex: 1, minHeight: 0 }}>
       <Editor
         height='100%'
         language='json'
@@ -324,15 +322,6 @@ function RolloutTracesDrawerBody({ rollout, attempt, onShowRollout, onShowSpanDe
   const { data, isFetching, isError, error, refetch } = useGetSpansQuery(queryArgs);
   const spans = data?.items ?? [];
   const totalRecords = data?.total ?? 0;
-  const tracesLinkSearch = useMemo(() => {
-    const params = createSearchParams({
-      rolloutId: rollout.rolloutId,
-      ...(attempt?.attemptId ? { attemptId: attempt.attemptId } : {}),
-    });
-    return params.toString();
-  }, [attempt?.attemptId, rollout.rolloutId]);
-  const tracesLinkHref = tracesLinkSearch ? `/traces?${tracesLinkSearch}` : '/traces';
-  const isWithinRouter = useInRouterContext();
 
   const handleSortStatusChange = useCallback((status: DataTableSortStatus<TracesTableRecord>) => {
     setSort({
@@ -352,38 +341,11 @@ function RolloutTracesDrawerBody({ rollout, attempt, onShowRollout, onShowSpanDe
 
   return (
     <Stack gap='md' style={{ flex: 1, minHeight: 0 }}>
-      <Group justify='space-between' align='center' gap='sm' wrap='nowrap'>
-        <Text size='sm' style={{ flex: 1, minWidth: 0 }}>
-          Showing spans for{' '}
-          <Text component='span' fw={600}>
-            {rollout.rolloutId}
-            {attempt ? ` · Attempt ${attempt.sequenceId} (${attempt.attemptId})` : ' · Latest attempt'}
-          </Text>
-        </Text>
-        {isWithinRouter ? (
-          <Anchor
-            component={Link}
-            to={tracesLinkHref}
-            size='sm'
-            aria-label={`Open traces page for rollout ${rollout.rolloutId}${
-              attempt ? ` attempt ${attempt.sequenceId}` : ''
-            }`}
-          >
-            View full traces
-          </Anchor>
-        ) : (
-          <Anchor
-            href={tracesLinkHref}
-            size='sm'
-            aria-label={`Open traces page for rollout ${rollout.rolloutId}${
-              attempt ? ` attempt ${attempt.sequenceId}` : ''
-            }`}
-          >
-            View full traces
-          </Anchor>
-        )}
-      </Group>
-      <Box data-testid='traces-drawer-table-container' style={{ flex: 1, minHeight: 0, overflow: 'auto' }}>
+      <Text size='sm' c='dimmed'>
+        Showing spans for rollout {rollout.rolloutId}
+        {attempt ? ` · Attempt ${attempt.sequenceId} (${attempt.attemptId})` : ' · Latest attempt'}
+      </Text>
+      <Box style={{ flex: 1, minHeight: 0 }}>
         <TracesTable
           spans={spans}
           totalRecords={totalRecords}
@@ -412,16 +374,10 @@ export function AppDrawerContainer() {
   const dispatch = useAppDispatch();
   const isOpen = useAppSelector(selectDrawerIsOpen);
   const content = useAppSelector(selectDrawerContent);
-  const isRouterAvailable = useInRouterContext();
 
   const handleClose = useCallback(() => {
     dispatch(closeDrawer());
   }, [dispatch]);
-  const handleNavigation = useCallback(() => {
-    if (isOpen) {
-      dispatch(closeDrawer());
-    }
-  }, [dispatch, isOpen]);
 
   const derivedContent = useMemo(() => {
     if (!content) {
@@ -488,29 +444,5 @@ export function AppDrawerContainer() {
 
   const { title, body } = derivedContent;
 
-  return (
-    <>
-      {isRouterAvailable ? <DrawerLocationWatcher onNavigation={handleNavigation} /> : null}
-      <AppDrawer opened={isOpen} onClose={handleClose} title={title} body={body} />
-    </>
-  );
-}
-
-type DrawerLocationWatcherProps = {
-  onNavigation: () => void;
-};
-
-function DrawerLocationWatcher({ onNavigation }: DrawerLocationWatcherProps) {
-  const location = useLocation();
-  const lastLocationKeyRef = useRef(location.key);
-
-  useEffect(() => {
-    if (lastLocationKeyRef.current === location.key) {
-      return;
-    }
-    lastLocationKeyRef.current = location.key;
-    onNavigation();
-  }, [location.key, onNavigation]);
-
-  return null;
+  return <AppDrawer opened={isOpen} onClose={handleClose} title={title} body={body} />;
 }

dashboard/src/components/ResourcesTable.component.tsx

@@ -4,8 +4,7 @@ import { useCallback, useEffect, useMemo, useState, type ReactNode, type SetStat
 import { IconCheck, IconCopy, IconRefresh } from '@tabler/icons-react';
 import { DataTable, type DataTableColumn, type DataTableSortStatus } from 'mantine-datatable';
 import { ActionIcon, Box, Button, CopyButton, Group, Stack, Text, Tooltip } from '@mantine/core';
-import { useElementSize, useViewportSize } from '@mantine/hooks';
-import { getLayoutAwareWidth } from '@/layouts/helper';
+import { useElementSize } from '@mantine/hooks';
 import type { Resources } from '@/types';
 import { getErrorDescriptor } from '@/utils/error';
 import { formatDateTime, safeStringify } from '@/utils/format';
@@ -163,12 +162,6 @@ export function ResourcesTable({
 }: ResourcesTableProps) {
   const [expandedRecordIds, setExpandedRecordIds] = useState<string[]>([]);
   const { ref: tableContainerRef, width: containerWidth } = useElementSize();
-  const { width: viewportWidth } = useViewportSize();
-
-  const layoutAwareContainerWidth = useMemo(
-    () => getLayoutAwareWidth(containerWidth, viewportWidth),
-    [containerWidth, viewportWidth],
-  );
 
   const resourcesRecords = useMemo<ResourcesTableRecord[]>(() => {
     if (!resourcesList) {
@@ -180,8 +173,8 @@ export function ResourcesTable({
   const columns = useMemo(() => createResourcesColumns({}), []);
 
   const responsiveColumns = useMemo(
-    () => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
-    [columns, layoutAwareContainerWidth],
+    () => createResponsiveColumns(columns, containerWidth, COLUMN_VISIBILITY),
+    [columns, containerWidth],
   );
 
   const totalPages = useMemo(

dashboard/src/components/RolloutTable.component.tsx

@@ -24,7 +24,7 @@ import {
   Text,
   Tooltip,
 } from '@mantine/core';
-import { useElementSize, useViewportSize } from '@mantine/hooks';
+import { useElementSize } from '@mantine/hooks';
 import {
   type Attempt,
   type AttemptStatus,
@@ -33,7 +33,6 @@ import {
   type RolloutsSortState,
   type RolloutStatus,
 } from '@/features/rollouts';
-import { getLayoutAwareWidth } from '@/layouts/helper';
 import {
   clampToNow,
   formatDateTime,
@@ -79,14 +78,14 @@ const ROLLOUT_MODE_OPTIONS: RolloutMode[] = ['train', 'val', 'test'];
 const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
 
 const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
-  rolloutId: { fixedWidth: 12.5, priority: 0 },
+  rolloutId: { fixedWidth: 10, priority: 0 },
   actionsPlaceholder: { fixedWidth: 6.5, priority: 0 },
   inputText: { minWidth: 14, priority: 1 },
   statusValue: { fixedWidth: 10, priority: 1 },
   startTimestamp: { fixedWidth: 12, priority: 2 },
   durationSeconds: { fixedWidth: 10, priority: 2 },
   attemptId: { fixedWidth: 12, priority: 3 },
-  resourcesId: { fixedWidth: 10, priority: 3 },
+  resourcesId: { fixedWidth: 8, priority: 3 },
   mode: { fixedWidth: 8, priority: 3 },
   lastHeartbeatTimestamp: { fixedWidth: 10, priority: 3 },
   workerId: { fixedWidth: 10, priority: 3 },
@@ -294,14 +293,7 @@ function createRolloutColumns({
       accessor: 'inputText',
       title: 'Input',
       render: ({ inputText }) => (
-        <Text
-          size='sm'
-          ff='monospace'
-          c='dimmed'
-          lineClamp={1}
-          title={inputText}
-          style={{ width: '100%', wordBreak: 'break-all', overflow: 'hidden' }}
-        >
+        <Text size='sm' ff='monospace' c='dimmed' lineClamp={1} style={{ width: '100%' }}>
           {inputText}
         </Text>
       ),
@@ -541,11 +533,6 @@ export function RolloutTable({
 }: RolloutTableProps) {
   const [expandedRecordIds, setExpandedRecordIds] = useState<string[]>([]);
   const { ref: tableContainerRef, width: containerWidth } = useElementSize();
-  const { width: viewportWidth } = useViewportSize();
-
-  const layoutAwareContainerWidth = useMemo(() => {
-    return getLayoutAwareWidth(containerWidth, viewportWidth);
-  }, [containerWidth, viewportWidth]);
 
   const rolloutRecords = useMemo<RolloutTableRecord[]>(() => {
     if (!rollouts) {
@@ -579,8 +566,8 @@ export function RolloutTable({
   );
 
   const responsiveColumns = useMemo(
-    () => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
-    [columns, layoutAwareContainerWidth],
+    () => createResponsiveColumns(columns, containerWidth, COLUMN_VISIBILITY),
+    [columns, containerWidth],
   );
 
   const totalPages = useMemo(
@@ -666,7 +653,7 @@ export function RolloutTable({
   );
 
   return (
-    <Box ref={tableContainerRef} data-testid='rollouts-table-container'>
+    <Box ref={tableContainerRef}>
       <DataTable<RolloutTableRecord>
         classNames={{ root: 'rollouts-table' }}
         withTableBorder

dashboard/src/components/TracesTable.component.tsx

@@ -11,25 +11,23 @@ import {
 } from '@tabler/icons-react';
 import { DataTable, type DataTableColumn, type DataTableSortStatus } from 'mantine-datatable';
 import { ActionIcon, Badge, Box, Button, CopyButton, Group, Stack, Text, Tooltip } from '@mantine/core';
-import { useElementSize, useViewportSize } from '@mantine/hooks';
-import { getLayoutAwareWidth } from '@/layouts/helper';
+import { useElementSize } from '@mantine/hooks';
 import type { Span } from '@/types';
 import { getErrorDescriptor } from '@/utils/error';
-import { formatDateTimeWithMilliseconds, formatDuration, toTimestamp } from '@/utils/format';
+import { formatDateTime, formatDuration, toTimestamp } from '@/utils/format';
 import { createResponsiveColumns, type ColumnVisibilityConfig } from '@/utils/table';
 
 const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
 
 const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
   name: { minWidth: 12.5, priority: 0 },
-  spanId: { fixedWidth: 14, priority: 1 },
-  traceId: { fixedWidth: 24, priority: 3 },
+  spanId: { fixedWidth: 12, priority: 1 },
+  traceId: { fixedWidth: 12, priority: 2 },
   parentId: { fixedWidth: 12, priority: 2 },
   statusCode: { fixedWidth: 8, priority: 2 },
   attributeKeys: { minWidth: 12.5, priority: 2 },
-  startTime: { fixedWidth: 15, priority: 1 },
-  endTime: { fixedWidth: 15, priority: 1 },
-  duration: { fixedWidth: 10, priority: 3 },
+  startTime: { fixedWidth: 12, priority: 1 },
+  duration: { fixedWidth: 10, priority: 2 },
   actionsPlaceholder: { fixedWidth: 6, priority: 0 },
 };
 
@@ -213,14 +211,7 @@ function createTracesColumns({
       title: 'Start Time',
       sortable: true,
       textAlign: 'left',
-      render: ({ startTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(startTime))}</Text>,
-    },
-    {
-      accessor: 'endTime',
-      title: 'End Time',
-      sortable: true,
-      textAlign: 'left',
-      render: ({ endTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(endTime))}</Text>,
+      render: ({ startTime }) => <Text size='sm'>{formatDateTime(toTimestamp(startTime))}</Text>,
     },
     {
       accessor: 'duration',
@@ -310,7 +301,6 @@ export function TracesTable({
   recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
 }: TracesTableProps) {
   const { ref: tableContainerRef, width: containerWidth } = useElementSize();
-  const { width: viewportWidth } = useViewportSize();
 
   const traceRecords = useMemo<TracesTableRecord[]>(() => {
     if (!spans) {
@@ -334,14 +324,9 @@ export function TracesTable({
     [onShowRollout, onShowSpanDetail, onParentIdClick, spanIds],
   );
 
-  const layoutAwareContainerWidth = useMemo(
-    () => getLayoutAwareWidth(containerWidth, viewportWidth),
-    [containerWidth, viewportWidth],
-  );
-
   const responsiveColumns = useMemo(
-    () => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
-    [columns, layoutAwareContainerWidth],
+    () => createResponsiveColumns(columns, containerWidth, COLUMN_VISIBILITY),
+    [columns, containerWidth],
   );
 
   const totalPages = useMemo(

dashboard/src/features/rollouts/api.ts

@@ -94,21 +94,15 @@ const normalizeSpan = (value: unknown): Span => {
     };
   };
   const rawStatus = camelized.status ?? { status_code: 'UNSET', description: null };
-  const result = {
+  return {
     ...camelized,
     parentId: camelized.parentId ?? null,
-    // The following fields does not need to be normalized to camel case
-    // For example, gen_ai.xxx should not become genAi.xxx
-    attributes: (value as any).attributes ?? {},
-    context: (value as any).context ?? {},
-    parent: (value as any).parent ?? null,
-    resource: (value as any).resource ?? {},
+    attributes: camelized.attributes ?? {},
     status: {
       status_code: rawStatus.status_code ?? rawStatus.statusCode ?? 'UNSET',
       description: rawStatus.description ?? null,
     },
   };
-  return result;
 };
 
 const normalizeResources = (value: unknown): Resources => {
@@ -127,20 +121,20 @@ const normalizePaginatedResponse = <T>(value: unknown, normalizer: (item: unknow
     throw new Error('Expected paginated response payload');
   }
 
-  const converted = value as {
+  const camelized = camelCaseKeys(value) as {
     items?: unknown;
     limit?: number;
     offset?: number;
     total?: number;
   };
 
-  const itemsSource = Array.isArray(converted.items) ? converted.items : [];
+  const itemsSource = Array.isArray(camelized.items) ? camelized.items : [];
 
   return {
     items: itemsSource.map((item) => normalizer(item)),
-    limit: typeof converted.limit === 'number' ? converted.limit : itemsSource.length,
-    offset: typeof converted.offset === 'number' ? converted.offset : 0,
-    total: typeof converted.total === 'number' ? converted.total : itemsSource.length,
+    limit: typeof camelized.limit === 'number' ? camelized.limit : itemsSource.length,
+    offset: typeof camelized.offset === 'number' ? camelized.offset : 0,
+    total: typeof camelized.total === 'number' ? camelized.total : itemsSource.length,
   };
 };
 

dashboard/src/features/rollouts/rollouts.test.ts

@@ -53,24 +53,6 @@ describe('rollouts feature integration', () => {
     expect(data.items[0].status).toBeDefined();
   });
 
-  it('includes attempts directly on rollout payloads when they exist', async () => {
-    const store = createServerBackedStore();
-    const queryArgs = selectRolloutsQueryArgs(store.getState());
-
-    const subscription = store.dispatch(rolloutsApi.endpoints.getRollouts.initiate(queryArgs));
-    const data = await subscription.unwrap();
-    subscription.unsubscribe();
-
-    const rolloutWithAttempt = data.items.find((rollout) => rollout.rolloutId === 'ro-story-002');
-    expect(rolloutWithAttempt).toBeDefined();
-    expect(rolloutWithAttempt?.attempt).not.toBeNull();
-    expect(rolloutWithAttempt?.attempt?.attemptId).toBe('at-story-022');
-
-    const rolloutWithoutAttempt = data.items.find((rollout) => rollout.rolloutId === 'ro-story-004');
-    expect(rolloutWithoutAttempt).toBeDefined();
-    expect(rolloutWithoutAttempt?.attempt).toBeNull();
-  });
-
   it('retrieves attempts for a rollout from the Python server', async () => {
     const store = createServerBackedStore();
     const subscription = store.dispatch(

dashboard/src/features/traces/slice.ts

@@ -69,34 +69,6 @@ const tracesSlice = createSlice({
       state.page = initialTracesUiState.page;
       state.sort = initialTracesUiState.sort;
     },
-    hydrateTracesStateFromQuery(
-      state,
-      action: PayloadAction<{ rolloutId?: string | null; attemptId?: string | null }>,
-    ) {
-      const payload = action.payload;
-      if (Object.hasOwn(payload, 'rolloutId')) {
-        const nextRolloutId = payload.rolloutId ?? null;
-        if (state.rolloutId !== nextRolloutId) {
-          state.rolloutId = nextRolloutId;
-          state.page = 1;
-          state.attemptId = null;
-        } else if (nextRolloutId === null) {
-          state.attemptId = null;
-        }
-      }
-
-      if (Object.hasOwn(payload, 'attemptId')) {
-        if (state.rolloutId === null) {
-          state.attemptId = null;
-          return;
-        }
-        const nextAttemptId = payload.attemptId ?? null;
-        if (state.attemptId !== nextAttemptId) {
-          state.attemptId = nextAttemptId;
-          state.page = 1;
-        }
-      }
-    },
   },
 });
 
@@ -109,7 +81,6 @@ export const {
   setTracesSort,
   setTracesViewMode,
   resetTracesFilters,
-  hydrateTracesStateFromQuery,
 } = tracesSlice.actions;
 
 export const tracesReducer = tracesSlice.reducer;

dashboard/src/layouts/AppLayout.tsx

@@ -99,7 +99,6 @@ function useServerConnection({ baseUrl, autoRefreshMs }: ConnectionOptions) {
 
     check();
 
-    // FIXME: autorefresh only refresh server status, not the data
     if (autoRefreshMs && autoRefreshMs > 0) {
       intervalId = window.setInterval(check, autoRefreshMs);
     }

dashboard/src/layouts/helper.ts

@@ -1,89 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-function parseCssNumber(value: string | null | undefined): number {
-  if (!value) {
-    return 0;
-  }
-
-  const parsed = Number.parseFloat(value);
-  return Number.isFinite(parsed) ? parsed : 0;
-}
-
-function getElementWidth(selectors: string | string[]): number {
-  if (typeof window === 'undefined') {
-    return 0;
-  }
-
-  const selectorList = Array.isArray(selectors) ? selectors : [selectors];
-  for (const selector of selectorList) {
-    const element = window.document.querySelector<HTMLElement>(selector);
-    if (element) {
-      return element.getBoundingClientRect().width || 0;
-    }
-  }
-
-  return 0;
-}
-
-function getAppShellContentWidth(): number {
-  if (typeof window === 'undefined') {
-    return 0;
-  }
-
-  const selectors = ['.mantine-AppShell-main', '[data-mantine-component="AppShellMain"]'];
-  return getElementWidth(selectors);
-}
-
-export function getAppShellOffsets(): number {
-  if (typeof window === 'undefined' || !window.document?.documentElement) {
-    return 0;
-  }
-
-  const rootStyles = window.getComputedStyle(window.document.documentElement);
-
-  const navbarOffset =
-    parseCssNumber(rootStyles.getPropertyValue('--app-shell-navbar-offset')) ||
-    getElementWidth(['.mantine-AppShell-navbar', '[data-mantine-component="AppShellNavbar"]']);
-
-  const asideOffset =
-    parseCssNumber(rootStyles.getPropertyValue('--app-shell-aside-offset')) ||
-    getElementWidth(['.mantine-AppShell-aside', '[data-mantine-component="AppShellAside"]']);
-
-  const paddingVar = parseCssNumber(rootStyles.getPropertyValue('--app-shell-padding'));
-  let paddingTotal = paddingVar ? paddingVar * 2 : 0;
-
-  if (paddingTotal === 0) {
-    const selectors = ['.mantine-AppShell-main', '[data-mantine-component="AppShellMain"]'];
-    for (const selector of selectors) {
-      const main = window.document.querySelector<HTMLElement>(selector);
-      if (!main) {
-        continue;
-      }
-      const mainStyles = window.getComputedStyle(main);
-      const computedPadding = parseCssNumber(mainStyles.paddingLeft) + parseCssNumber(mainStyles.paddingRight);
-      if (computedPadding > 0) {
-        paddingTotal = computedPadding;
-        break;
-      }
-    }
-  }
-
-  return navbarOffset + asideOffset + paddingTotal;
-}
-
-export function getLayoutAwareWidth(containerWidth: number, viewportWidth: number): number {
-  const appShellContentWidth = getAppShellContentWidth();
-  const layoutOffsets = getAppShellOffsets();
-  const viewportAvailable = viewportWidth && viewportWidth > 0 ? Math.max(viewportWidth - layoutOffsets, 0) : undefined;
-  const effectiveAvailable = appShellContentWidth > 0 ? appShellContentWidth : viewportAvailable;
-
-  if (!containerWidth && effectiveAvailable) {
-    return effectiveAvailable;
-  }
-
-  if (!effectiveAvailable) {
-    return containerWidth;
-  }
-
-  return Math.min(containerWidth, effectiveAvailable);
-}

dashboard/src/pages/Resources.page.story.tsx

@@ -1,17 +1,12 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 import type { Meta, StoryObj } from '@storybook/react';
-import { waitFor, within } from '@testing-library/dom';
-import userEvent from '@testing-library/user-event';
 import { delay, http, HttpResponse } from 'msw';
 import { Provider } from 'react-redux';
-import { createMemoryRouter, RouterProvider } from 'react-router-dom';
 import { AppAlertBanner } from '@/components/AppAlertBanner';
-import { AppDrawerContainer } from '@/components/AppDrawer.component';
 import { initialConfigState } from '@/features/config/slice';
 import { initialResourcesUiState } from '@/features/resources/slice';
 import { initialRolloutsUiState } from '@/features/rollouts/slice';
-import { AppLayout } from '@/layouts/AppLayout';
 import { createAppStore } from '@/store';
 import type { Resources } from '@/types';
 import { createResourcesHandlers } from '@/utils/mock';
@@ -153,8 +148,8 @@ const sampleResources: Resources[] = [
 
 const defaultHandlers = createResourcesHandlers(sampleResources);
 
-function createStoryStore(configOverrides?: Partial<typeof initialConfigState>) {
-  return createAppStore({
+function renderWithStore(configOverrides?: Partial<typeof initialConfigState>) {
+  const store = createAppStore({
     config: {
       ...initialConfigState,
       baseUrl: STORY_BASE_URL,
@@ -164,53 +159,11 @@ function createStoryStore(configOverrides?: Partial<typeof initialConfigState>)
     rollouts: initialRolloutsUiState,
     resources: initialResourcesUiState,
   });
-}
-
-function renderWithStore(configOverrides?: Partial<typeof initialConfigState>) {
-  const store = createStoryStore(configOverrides);
 
   return (
     <Provider store={store}>
-      <>
-        <ResourcesPage />
-        <AppAlertBanner />
-        <AppDrawerContainer />
-      </>
-    </Provider>
-  );
-}
-
-function renderWithAppLayout(configOverrides?: Partial<typeof initialConfigState>) {
-  const store = createStoryStore(configOverrides);
-  const router = createMemoryRouter(
-    [
-      {
-        path: '/',
-        element: (
-          <AppLayout
-            config={{
-              baseUrl: store.getState().config.baseUrl,
-              autoRefreshMs: store.getState().config.autoRefreshMs,
-            }}
-          />
-        ),
-        children: [
-          {
-            path: '/resources',
-            element: <ResourcesPage />,
-          },
-        ],
-      },
-    ],
-    { initialEntries: ['/resources'] },
-  );
-
-  return (
-    <Provider store={store}>
-      <>
-        <RouterProvider router={router} />
-        <AppDrawerContainer />
-      </>
+      <ResourcesPage />
+      <AppAlertBanner />
     </Provider>
   );
 }
@@ -224,41 +177,6 @@ export const Default: Story = {
   },
 };
 
-export const WithSidebarLayout: Story = {
-  name: 'Within AppLayout',
-  render: () => renderWithAppLayout(),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-};
-
-export const Search: Story = {
-  render: () => renderWithStore(),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    await canvas.findByText('rs-a1b2c3d4e5f6');
-
-    const searchInput = canvas.getByPlaceholderText('Search by Resources ID');
-    await userEvent.type(searchInput, 'rs-abcdef123456');
-
-    await waitFor(() => {
-      if (canvas.queryByText('rs-a1b2c3d4e5f6')) {
-        throw new Error('Expected search to filter out non-matching resources');
-      }
-      if (!canvas.queryByText('rs-abcdef123456')) {
-        throw new Error('Expected matching resource to remain visible');
-      }
-    });
-  },
-};
-
 export const EmptyState: Story = {
   render: () => renderWithStore(),
   parameters: {

dashboard/src/pages/Rollouts.page.story.tsx

@@ -5,10 +5,8 @@ import { waitFor, within } from '@testing-library/dom';
 import userEvent from '@testing-library/user-event';
 import { delay, http, HttpResponse } from 'msw';
 import { Provider } from 'react-redux';
-import { createMemoryRouter, RouterProvider } from 'react-router-dom';
 import { AppAlertBanner } from '@/components/AppAlertBanner';
 import { AppDrawerContainer } from '@/components/AppDrawer.component';
-import { AppLayout } from '@/layouts/AppLayout';
 import { createMockHandlers } from '@/utils/mock';
 import { STORY_BASE_URL, STORY_DATE_NOW_SECONDS } from '../../.storybook/constants';
 import { allModes } from '../../.storybook/modes';
@@ -297,105 +295,6 @@ const sampleSpansByAttempt: Record<string, Span[]> = {
   ],
 };
 
-const overflowDrawerSpans: Span[] = Array.from({ length: 160 }, (_, index) => ({
-  rolloutId: 'ro-7fa3b6e2',
-  attemptId: 'at-9001',
-  sequenceId: index + 1,
-  traceId: `tr-overflow-${Math.floor(index / 5)}`,
-  spanId: `sp-overflow-${index + 1}`,
-  parentId: index === 0 ? null : `sp-overflow-${index}`,
-  name: `Overflow span ${index + 1}`,
-  status: { status_code: 'OK', description: null },
-  attributes: { step: `overflow-${index + 1}`, duration_ms: 20 + (index % 5) },
-  startTime: now - 1_200 - index * 20,
-  endTime: now - 1_180 - index * 20,
-  events: [],
-  links: [],
-  context: {},
-  parent: null,
-  resource: {},
-}));
-
-const overflowSpansByAttempt: Record<string, Span[]> = {
-  ...sampleSpansByAttempt,
-  'ro-7fa3b6e2:at-9001': overflowDrawerSpans,
-};
-
-const longJsonLogs = Array.from({ length: 200 }, (_, index) => ({
-  id: index + 1,
-  detail: `Log entry ${index + 1} ${'x'.repeat(32)}`,
-  timestamp: now - index * 2,
-}));
-
-const jsonOverflowAttempt: Attempt = {
-  rolloutId: 'ro-json-overflow',
-  attemptId: 'at-json-overflow',
-  sequenceId: 1,
-  status: 'succeeded',
-  startTime: now - 600,
-  endTime: now - 300,
-  workerId: 'worker-scroll',
-  lastHeartbeatTime: now - 300,
-  metadata: { notes: 'Completed with a very large JSON payload' },
-};
-
-const jsonOverflowAttemptEndTime = jsonOverflowAttempt.endTime ?? jsonOverflowAttempt.startTime + 1;
-
-const jsonOverflowRollout: Rollout = {
-  rolloutId: 'ro-json-overflow',
-  input: {
-    task: 'Render large JSON',
-    payload: longJsonLogs,
-    summary: 'This rollout includes many log lines to test scroll behavior.',
-  },
-  status: 'succeeded',
-  mode: 'train',
-  resourcesId: 'rs-json-overflow',
-  startTime: jsonOverflowAttempt.startTime,
-  endTime: jsonOverflowAttemptEndTime,
-  attempt: jsonOverflowAttempt,
-  config: {
-    retries: 0,
-    parameters: { max_steps: 200, batch: 5 },
-  },
-  metadata: {
-    owner: 'scroll-tester',
-    description: 'Synthetic rollout with oversized JSON payload for storybook validation.',
-    tags: Array.from({ length: 40 }, (_, index) => `tag-${index + 1}`),
-  },
-};
-
-const jsonOverflowSpans: Span[] = [
-  {
-    rolloutId: jsonOverflowRollout.rolloutId,
-    attemptId: jsonOverflowAttempt.attemptId,
-    sequenceId: 1,
-    traceId: 'tr-json-overflow',
-    spanId: 'sp-json-root',
-    parentId: null,
-    name: 'json-overflow-root',
-    status: { status_code: 'OK', description: null },
-    attributes: { detail: 'root span' },
-    startTime: jsonOverflowAttempt.startTime,
-    endTime: jsonOverflowAttemptEndTime,
-    events: [],
-    links: [],
-    context: {},
-    parent: null,
-    resource: {},
-  },
-];
-
-const jsonOverflowRollouts = [jsonOverflowRollout, ...sampleRollouts];
-const jsonOverflowAttemptsByRollout: Record<string, Attempt[]> = {
-  ...attemptsByRollout,
-  [jsonOverflowRollout.rolloutId]: [jsonOverflowAttempt],
-};
-const jsonOverflowSpansByAttempt: Record<string, Span[]> = {
-  ...sampleSpansByAttempt,
-  [`${jsonOverflowRollout.rolloutId}:${jsonOverflowAttempt.attemptId}`]: jsonOverflowSpans,
-};
-
 const longDurationRollouts: Rollout[] = [
   {
     rolloutId: 'ro-long-duration',
@@ -656,12 +555,8 @@ const autoExpandAttempts: Record<string, Attempt[]> = {
     },
   ],
 };
-
-function createStoryStore(
-  uiOverrides?: Partial<RolloutsUiState>,
-  configOverrides?: Partial<typeof initialConfigState>,
-) {
-  return createAppStore({
+function renderWithStore(uiOverrides?: Partial<RolloutsUiState>, configOverrides?: Partial<typeof initialConfigState>) {
+  const store = createAppStore({
     config: {
       ...initialConfigState,
       baseUrl: STORY_BASE_URL,
@@ -674,10 +569,6 @@ function createStoryStore(
     },
     resources: initialResourcesUiState,
   });
-}
-
-function renderWithStore(uiOverrides?: Partial<RolloutsUiState>, configOverrides?: Partial<typeof initialConfigState>) {
-  const store = createStoryStore(uiOverrides, configOverrides);
 
   return (
     <Provider store={store}>
@@ -690,51 +581,7 @@ function renderWithStore(uiOverrides?: Partial<RolloutsUiState>, configOverrides
   );
 }
 
-function renderWithAppLayout(
-  uiOverrides?: Partial<RolloutsUiState>,
-  configOverrides?: Partial<typeof initialConfigState>,
-) {
-  const store = createStoryStore(uiOverrides, configOverrides);
-  const router = createMemoryRouter(
-    [
-      {
-        path: '/',
-        element: (
-          <AppLayout
-            config={{
-              baseUrl: store.getState().config.baseUrl,
-              autoRefreshMs: store.getState().config.autoRefreshMs,
-            }}
-          />
-        ),
-        children: [
-          {
-            path: '/rollouts',
-            element: <RolloutsPage />,
-          },
-        ],
-      },
-    ],
-    { initialEntries: ['/rollouts'] },
-  );
-
-  return (
-    <Provider store={store}>
-      <>
-        <RouterProvider router={router} />
-        <AppDrawerContainer />
-      </>
-    </Provider>
-  );
-}
-
 const defaultHandlers = createMockHandlers(sampleRollouts, attemptsByRollout, sampleSpansByAttempt);
-const overflowHandlers = createMockHandlers(sampleRollouts, attemptsByRollout, overflowSpansByAttempt);
-const jsonOverflowHandlers = createMockHandlers(
-  jsonOverflowRollouts,
-  jsonOverflowAttemptsByRollout,
-  jsonOverflowSpansByAttempt,
-);
 
 export const Default: Story = {
   render: () => renderWithStore(),
@@ -745,40 +592,6 @@ export const Default: Story = {
   },
 };
 
-export const WithSidebarLayout: Story = {
-  name: 'Within AppLayout',
-  render: () => renderWithAppLayout(),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-};
-
-export const WithSidebarStatusFilter: Story = {
-  name: 'Within AppLayout (Status Filter)',
-  render: () => renderWithAppLayout({ statusFilters: ['running'] }),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-  play: async () => {
-    await waitFor(() => {
-      const container = document.querySelector<HTMLElement>('[data-testid="rollouts-table-container"]');
-      const main = document.querySelector<HTMLElement>('.mantine-AppShell-main');
-      if (!container || !main) {
-        throw new Error('Unable to locate rollout table container or AppShell main region');
-      }
-      const containerRect = container.getBoundingClientRect();
-      const mainRect = main.getBoundingClientRect();
-      if (containerRect.right > mainRect.right + 1) {
-        throw new Error('Rollouts table extends beyond the AppShell content area');
-      }
-    });
-  },
-};
-
 export const DarkTheme: Story = {
   render: () => renderWithStore(undefined, { theme: 'dark' }),
   parameters: {
@@ -908,7 +721,7 @@ export const AutoExpandedAttempt: Story = {
   },
 };
 
-export const Search: Story = {
+export const RawJsonDrawer: Story = {
   render: () => renderWithStore(),
   parameters: {
     msw: {
@@ -918,99 +731,21 @@ export const Search: Story = {
   play: async ({ canvasElement }) => {
     const canvas = within(canvasElement);
     await canvas.findByText('ro-7fa3b6e2');
+    const rolloutCell = canvas.getByText('ro-7fa3b6e2');
+    const rolloutRow = rolloutCell.closest('tr');
 
-    const searchInput = canvas.getByPlaceholderText('Search by Rollout ID');
-    await userEvent.type(searchInput, 'ro-116eab45');
-
-    await waitFor(() => {
-      if (canvas.queryByText('ro-7fa3b6e2')) {
-        throw new Error('Expected search to filter out non-matching rollouts');
-      }
-      if (!canvas.queryByText('ro-116eab45')) {
-        throw new Error('Expected search to keep the matching rollout visible');
-      }
-    });
-  },
-};
-
-async function openSampleTracesDrawer(canvasElement: HTMLElement) {
-  const canvas = within(canvasElement);
-  await canvas.findByText('ro-7fa3b6e2');
-  const rolloutCell = canvas.getByText('ro-7fa3b6e2');
-  const rolloutRow = rolloutCell.closest('tr');
-
-  if (!rolloutRow) {
-    throw new Error('Unable to locate rollout row for traces drawer');
-  }
-
-  const rowScope = within(rolloutRow);
-  const traceButtons = rowScope.getAllByRole('button', { name: 'View traces' });
-  const tracesButton = traceButtons[0];
-  await userEvent.click(tracesButton);
-
-  return within(document.body).findByRole('dialog');
-}
-
-async function openRawJsonDrawer(canvasElement: HTMLElement, rolloutId = 'ro-7fa3b6e2') {
-  const canvas = within(canvasElement);
-  await canvas.findByText(rolloutId);
-  const rolloutCell = canvas.getByText(rolloutId);
-  const rolloutRow = rolloutCell.closest('tr');
-
-  if (!rolloutRow) {
-    throw new Error(`Unable to locate rollout row for ${rolloutId}`);
-  }
-
-  const rowScope = within(rolloutRow);
-  const rawButtons = rowScope.getAllByRole('button', { name: 'View raw JSON' });
-  const rawButton = rawButtons[0];
-  await userEvent.click(rawButton);
-
-  return within(document.body).findByRole('dialog');
-}
-
-export const RawJsonDrawer: Story = {
-  render: () => renderWithStore(),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const drawer = await openRawJsonDrawer(canvasElement);
-    await waitFor(
-      async () => {
-        await within(drawer).findByText('Attempt');
-        await within(drawer).findByText(/worker-alpha/);
-      },
-      { timeout: 3_000 },
-    );
-  },
-};
-
-export const RawJsonDrawerScrollable: Story = {
-  name: 'Raw JSON Drawer Scrollable',
-  render: () => renderWithStore(),
-  parameters: {
-    msw: {
-      handlers: jsonOverflowHandlers,
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const drawer = await openRawJsonDrawer(canvasElement, 'ro-json-overflow');
-    const editorContainer = drawer.querySelector('[data-testid="json-editor-container"]') as HTMLElement | null;
-    if (!editorContainer) {
-      throw new Error('Unable to locate JSON editor container');
+    if (!rolloutRow) {
+      throw new Error('Unable to locate rollout row for raw JSON drawer');
     }
-    await waitFor(() => {
-      const scrollable = editorContainer.querySelector('.monaco-scrollable-element') as HTMLElement | null;
-      if (!scrollable) {
-        throw new Error('Monaco editor not ready yet');
-      }
-      if (scrollable.scrollHeight <= scrollable.clientHeight) {
-        throw new Error('Expected JSON content to overflow and allow scrolling');
-      }
-    });
+
+    const rowScope = within(rolloutRow);
+    const rawButtons = rowScope.getAllByRole('button', { name: 'View raw JSON' });
+    const rawButton = rawButtons[0];
+    await userEvent.click(rawButton);
+
+    const drawer = await within(document.body).findByRole('dialog');
+    await within(drawer).findByText('Attempt');
+    await within(drawer).findByText(/worker-alpha/);
   },
 };
 
@@ -1022,56 +757,20 @@ export const TracesDrawer: Story = {
     },
   },
   play: async ({ canvasElement }) => {
-    await openSampleTracesDrawer(canvasElement);
-  },
-};
+    const canvas = within(canvasElement);
+    await canvas.findByText('ro-7fa3b6e2');
+    const rolloutCell = canvas.getByText('ro-7fa3b6e2');
+    const rolloutRow = rolloutCell.closest('tr');
 
-export const TracesDrawerLink: Story = {
-  name: 'Traces Drawer Link',
-  render: () => renderWithStore(),
-  parameters: {
-    msw: {
-      handlers: defaultHandlers,
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const drawer = await openSampleTracesDrawer(canvasElement);
-    const link = await within(drawer).findByText('View full traces');
-    const href = link.getAttribute('href');
-    if (!href) {
-      throw new Error('Expected traces drawer to render a link to the traces page');
+    if (!rolloutRow) {
+      throw new Error('Unable to locate rollout row for traces drawer');
     }
-    if (!href.includes('rolloutId=ro-7fa3b6e2')) {
-      throw new Error(`Link href ${href} is missing rolloutId query parameter`);
-    }
-    if (!href.includes('attemptId=at-9001')) {
-      throw new Error(`Link href ${href} is missing attemptId query parameter`);
-    }
-  },
-};
 
-export const TracesDrawerScrollableTable: Story = {
-  name: 'Traces Drawer Scrollable Table',
-  render: () => renderWithStore(),
-  parameters: {
-    msw: {
-      handlers: overflowHandlers,
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const drawer = await openSampleTracesDrawer(canvasElement);
-    const container = drawer.querySelector('[data-testid="traces-drawer-table-container"]') as HTMLElement | null;
-    if (!container) {
-      throw new Error('Unable to locate traces table container inside drawer');
-    }
-    const overflowStyle = window.getComputedStyle(container).overflowY;
-    if (overflowStyle !== 'auto' && overflowStyle !== 'scroll') {
-      throw new Error('Expected traces table container to allow vertical scrolling');
-    }
-    await waitFor(() => {
-      if (container.scrollHeight <= container.clientHeight) {
-        throw new Error('Expected traces table content to overflow and enable scrolling');
-      }
-    });
+    const rowScope = within(rolloutRow);
+    const traceButtons = rowScope.getAllByRole('button', { name: 'View traces' });
+    const tracesButton = traceButtons[0];
+    await userEvent.click(tracesButton);
+
+    await within(document.body).findByRole('dialog');
   },
 };

dashboard/src/pages/Settings.page.tsx

@@ -8,11 +8,10 @@ import { useAppDispatch, useAppSelector } from '../store/hooks';
 
 const AUTO_REFRESH_OPTIONS = [
   { label: 'Off', value: '0' },
-  // TODO: Support real auto-refresh
-  { label: 'Every 5 seconds', value: '5000', disabled: true },
-  { label: 'Every 15 seconds', value: '15000', disabled: true },
-  { label: 'Every 60 seconds', value: '60000', disabled: true },
-  { label: 'Every 5 minutes', value: '300000', disabled: true },
+  { label: 'Every 5 seconds', value: '5000' },
+  { label: 'Every 15 seconds', value: '15000' },
+  { label: 'Every 60 seconds', value: '60000' },
+  { label: 'Every 5 minutes', value: '300000' },
 ];
 
 const THEME_OPTIONS = [

dashboard/src/pages/Traces.page.story.tsx

@@ -1,14 +1,10 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 import type { Meta, StoryObj } from '@storybook/react';
-import { waitFor, within } from '@testing-library/dom';
-import userEvent from '@testing-library/user-event';
 import { delay, http, HttpResponse } from 'msw';
 import { Provider } from 'react-redux';
-import { createMemoryRouter, MemoryRouter, RouterProvider } from 'react-router-dom';
 import { AppAlertBanner } from '@/components/AppAlertBanner';
 import { AppDrawerContainer } from '@/components/AppDrawer.component';
-import { AppLayout } from '@/layouts/AppLayout';
 import { createMockHandlers } from '@/utils/mock';
 import { STORY_BASE_URL, STORY_DATE_NOW_SECONDS } from '../../.storybook/constants';
 import { allModes } from '../../.storybook/modes';
@@ -235,105 +231,73 @@ const singleSpansByAttempt = Object.fromEntries(
   Object.entries(spansByAttempt).filter(([key]) => key.startsWith(`${singleRollout.rolloutId}:`)),
 ) as Record<string, Span[]>;
 
-const rolloutWithoutAttempt: Rollout = {
-  rolloutId: 'ro-traces-no-attempt',
-  input: { task: 'Legacy rollout without attempts' },
-  status: 'failed',
-  mode: 'train',
-  resourcesId: 'rs-traces-no-attempt',
-  startTime: now - 7200,
-  endTime: now - 7000,
-  attempt: null,
-  config: { retries: 0 },
-  metadata: { owner: 'casey' },
-};
-const noAttemptRollouts: Rollout[] = [rolloutWithoutAttempt];
-const noAttemptAttemptsByRollout: Record<string, Attempt[]> = {
-  [rolloutWithoutAttempt.rolloutId]: [],
-};
-const noAttemptSpansByAttempt: Record<string, Span[]> = {};
-
 const owners = ['ava', 'ben', 'carla', 'diego'] as const;
 
-function createSyntheticRollouts(prefix: string, count: number) {
-  const attemptsByRollout: Record<string, Attempt[]> = {};
-  const spansByAttemptByRollout: Record<string, Span[]> = {};
-  const rollouts: Rollout[] = Array.from({ length: count }, (_, index) => {
-    const rolloutId = `ro-${prefix}-${String(index + 1).padStart(3, '0')}`;
-    const statusOptions = ['running', 'succeeded', 'failed'] as const;
-    const modeOptions = ['train', 'val', 'test'] as const;
-    const status = statusOptions[index % statusOptions.length];
-    const mode = modeOptions[index % modeOptions.length];
-    const startTime = now - (index + 1) * 420;
-    const endTime = status === 'running' ? null : startTime + 240;
-    const attemptId = `${rolloutId}-attempt`;
-    const attemptStatus: Attempt['status'] =
-      status === 'failed' ? 'failed' : status === 'succeeded' ? 'succeeded' : 'running';
-    const attempt: Attempt = {
+const manyAttemptsByRollout: Record<string, Attempt[]> = {};
+const manySpansByAttempt: Record<string, Span[]> = {};
+
+const manyRollouts: Rollout[] = Array.from({ length: 24 }, (_, index) => {
+  const rolloutId = `ro-many-${String(index + 1).padStart(3, '0')}`;
+  const statusOptions = ['running', 'succeeded', 'failed'] as const;
+  const modeOptions = ['train', 'val', 'test'] as const;
+  const status = statusOptions[index % statusOptions.length];
+  const mode = modeOptions[index % modeOptions.length];
+  const startTime = now - (index + 1) * 420;
+  const endTime = status === 'running' ? null : startTime + 240;
+  const attemptId = `${rolloutId}-attempt`;
+  const attemptStatus: Attempt['status'] =
+    status === 'failed' ? 'failed' : status === 'succeeded' ? 'succeeded' : 'running';
+  const attempt: Attempt = {
+    rolloutId,
+    attemptId,
+    sequenceId: 1,
+    status: attemptStatus,
+    startTime,
+    endTime,
+    workerId: `worker-${String.fromCharCode(97 + (index % 26))}`,
+    lastHeartbeatTime: endTime ?? startTime + 180,
+    metadata: { region: index % 2 === 0 ? 'us-east-1' : 'eu-west-1' },
+  };
+  manyAttemptsByRollout[rolloutId] = [attempt];
+  manySpansByAttempt[`${rolloutId}:${attemptId}`] = [
+    {
       rolloutId,
       attemptId,
       sequenceId: 1,
-      status: attemptStatus,
-      startTime,
-      endTime,
-      workerId: `worker-${String.fromCharCode(97 + (index % 26))}`,
-      lastHeartbeatTime: endTime ?? startTime + 180,
-      metadata: { region: index % 2 === 0 ? 'us-east-1' : 'eu-west-1' },
-    };
-    attemptsByRollout[rolloutId] = [attempt];
-    spansByAttemptByRollout[`${rolloutId}:${attemptId}`] = [
-      {
-        rolloutId,
-        attemptId,
-        sequenceId: 1,
-        traceId: `tr-${prefix}-${index + 1}`,
-        spanId: `sp-${prefix}-${index + 1}-root`,
-        parentId: null,
-        name: `Synthetic root span ${prefix} ${index + 1}`,
-        status: {
-          status_code: status === 'failed' ? 'ERROR' : 'OK',
-          description: status === 'failed' ? 'Synthetic failure' : null,
-        },
-        attributes: {
-          'trace.sample': index + 1,
-          'duration_ms': 240,
-        },
-        startTime,
-        endTime: endTime ?? startTime + 240,
-        events: [],
-        links: [],
-        context: {},
-        parent: null,
-        resource: {},
+      traceId: `tr-many-${index + 1}`,
+      spanId: `sp-many-${index + 1}-root`,
+      parentId: null,
+      name: 'Synthetic root span',
+      status: {
+        status_code: status === 'failed' ? 'ERROR' : 'OK',
+        description: status === 'failed' ? 'Synthetic failure' : null,
+      },
+      attributes: {
+        'trace.sample': index + 1,
+        'duration_ms': 240,
       },
-    ];
-    return {
-      rolloutId,
-      input: { task: `Synthetic trace ${index + 1}` },
-      status,
-      mode,
-      resourcesId: `rs-${prefix}-${(index % 7) + 1}`,
       startTime,
-      endTime,
-      attempt,
-      config: { retries: index % 3 },
-      metadata: { owner: owners[index % owners.length] },
-    };
-  });
-  return { rollouts, attemptsByRollout, spansByAttempt: spansByAttemptByRollout };
-}
-
-const {
-  rollouts: manyRollouts,
-  attemptsByRollout: manyAttemptsByRollout,
-  spansByAttempt: manySpansByAttempt,
-} = createSyntheticRollouts('many', 24);
-
-const {
-  rollouts: vastRollouts,
-  attemptsByRollout: vastAttemptsByRollout,
-  spansByAttempt: vastSpansByAttempt,
-} = createSyntheticRollouts('vast', 160);
+      endTime: endTime ?? startTime + 240,
+      events: [],
+      links: [],
+      context: {},
+      parent: null,
+      resource: {},
+    },
+  ];
+  return {
+    rolloutId,
+    input: { task: `Synthetic trace ${index + 1}` },
+    status,
+    mode,
+    resourcesId: `rs-many-${(index % 7) + 1}`,
+    startTime,
+    endTime,
+    attempt,
+    config: { retries: index % 3 },
+    metadata: { owner: owners[index % owners.length] },
+  };
+});
 
 function createHandlers(delayMs?: number) {
   return createMockHandlers(sampleRollouts, attemptsByRollout, spansByAttempt, delayMs);
@@ -361,11 +325,11 @@ function createRequestTimeoutHandlers() {
 
 const rolloutsAndAttemptsHandlers = createMockHandlers(sampleRollouts, attemptsByRollout);
 
-function createStoryStore(
+function renderTracesPage(
   preloadedTracesState?: Partial<TracesUiState>,
   configOverrides?: Partial<typeof initialConfigState>,
 ) {
-  return createAppStore({
+  const store = createAppStore({
     config: {
       ...initialConfigState,
       baseUrl: STORY_BASE_URL,
@@ -375,60 +339,12 @@ function createStoryStore(
     resources: initialResourcesUiState,
     traces: { ...initialTracesUiState, ...preloadedTracesState },
   });
-}
-
-function renderTracesPage(
-  preloadedTracesState?: Partial<TracesUiState>,
-  configOverrides?: Partial<typeof initialConfigState>,
-) {
-  const store = createStoryStore(preloadedTracesState, configOverrides);
 
   return (
     <Provider store={store}>
-      <MemoryRouter initialEntries={['/traces']}>
-        <TracesPage />
-        <AppAlertBanner />
-        <AppDrawerContainer />
-      </MemoryRouter>
-    </Provider>
-  );
-}
-
-function renderTracesPageWithAppLayout(
-  preloadedTracesState?: Partial<TracesUiState>,
-  configOverrides?: Partial<typeof initialConfigState>,
-  initialEntry: string = '/traces',
-) {
-  const store = createStoryStore(preloadedTracesState, configOverrides);
-  const router = createMemoryRouter(
-    [
-      {
-        path: '/',
-        element: (
-          <AppLayout
-            config={{
-              baseUrl: store.getState().config.baseUrl,
-              autoRefreshMs: store.getState().config.autoRefreshMs,
-            }}
-          />
-        ),
-        children: [
-          {
-            path: '/traces',
-            element: <TracesPage />,
-          },
-        ],
-      },
-    ],
-    { initialEntries: [initialEntry] },
-  );
-
-  return (
-    <Provider store={store}>
-      <>
-        <RouterProvider router={router} />
-        <AppDrawerContainer />
-      </>
+      <TracesPage />
+      <AppAlertBanner />
+      <AppDrawerContainer />
     </Provider>
   );
 }
@@ -442,66 +358,6 @@ export const DefaultView: Story = {
   },
 };
 
-export const WithSidebarLayout: Story = {
-  name: 'Within AppLayout',
-  render: () => renderTracesPageWithAppLayout(),
-  parameters: {
-    msw: {
-      handlers: createHandlers(),
-    },
-  },
-};
-
-export const QueryParams: Story = {
-  name: 'Loads From Query Params',
-  render: () =>
-    renderTracesPageWithAppLayout(undefined, undefined, '/traces?rolloutId=ro-traces-002&attemptId=at-traces-004'),
-  parameters: {
-    msw: {
-      handlers: createHandlers(),
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    const rolloutInput = (await canvas.findByLabelText('Select rollout')) as HTMLInputElement;
-    await waitFor(() => {
-      if (rolloutInput.value !== 'ro-traces-002') {
-        throw new Error('Expected rollout select to use value from query string');
-      }
-    });
-    const attemptInput = (await canvas.findByLabelText('Select attempt')) as HTMLInputElement;
-    await waitFor(() => {
-      if (attemptInput.value.indexOf('at-traces-004') === -1) {
-        throw new Error('Expected attempt select to use value from query string');
-      }
-    });
-  },
-};
-
-export const MissingRolloutQuery: Story = {
-  name: 'Missing Rollout From Query Params',
-  render: () => renderTracesPageWithAppLayout(undefined, undefined, '/traces?rolloutId=ro-missing-999'),
-  parameters: {
-    msw: {
-      handlers: createMockHandlers(manyRollouts, manyAttemptsByRollout, manySpansByAttempt),
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    const rolloutInput = (await canvas.findByLabelText('Select rollout')) as HTMLInputElement;
-    await waitFor(() => {
-      if (rolloutInput.value !== '') {
-        throw new Error('Expected rollout select to remain empty when the query rollout does not exist');
-      }
-    });
-    await waitFor(() => {
-      if (!canvas.getByText('Select a rollout and attempt to view traces.')) {
-        throw new Error('Expected empty selection message when rollout query param is invalid');
-      }
-    });
-  },
-};
-
 export const DarkTheme: Story = {
   render: () => renderTracesPage(undefined, { theme: 'dark' }),
   parameters: {
@@ -530,24 +386,6 @@ export const SingleResult: Story = {
   },
 };
 
-export const NoAttemptPlaceholder: Story = {
-  render: () => renderTracesPage(),
-  parameters: {
-    msw: {
-      handlers: createMockHandlers(noAttemptRollouts, noAttemptAttemptsByRollout, noAttemptSpansByAttempt),
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    const attemptInput = (await canvas.findByLabelText('Select attempt')) as HTMLInputElement;
-    await waitFor(() => {
-      if (attemptInput.placeholder !== 'No Attempt') {
-        throw new Error('Expected attempt select placeholder to read "No Attempt" when no attempts are available');
-      }
-    });
-  },
-};
-
 export const ManyResults: Story = {
   render: () => renderTracesPage(),
   parameters: {
@@ -557,67 +395,6 @@ export const ManyResults: Story = {
   },
 };
 
-export const LargeDatasetSearch: Story = {
-  render: () => renderTracesPage(),
-  parameters: {
-    msw: {
-      handlers: createMockHandlers(vastRollouts, vastAttemptsByRollout, vastSpansByAttempt),
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    const rolloutTrigger = (await canvas.findByLabelText('Select rollout')) as HTMLInputElement;
-    await userEvent.click(rolloutTrigger);
-    for (let i = 0; i < 'ro-vast-001'.length + 1; i++) {
-      await userEvent.type(rolloutTrigger, '{backspace}');
-    }
-    await userEvent.type(rolloutTrigger, 'ro-vast-150');
-
-    await waitFor(() => {
-      const option = within(document.body).queryByText('ro-vast-150');
-      if (!option) {
-        throw new Error('Expected remote rollout search to return IDs beyond the initial list');
-      }
-    });
-
-    const option = within(document.body).getByText('ro-vast-150');
-    await userEvent.click(option);
-
-    await waitFor(() => {
-      if (rolloutTrigger.value !== 'ro-vast-150') {
-        throw new Error('Expected rollout select to use the searched rollout ID');
-      }
-    });
-
-    await canvas.findByText('Synthetic root span vast 150');
-  },
-};
-
-export const Search: Story = {
-  render: () => renderTracesPage(),
-  parameters: {
-    msw: {
-      handlers: createHandlers(),
-    },
-  },
-  play: async ({ canvasElement }) => {
-    const canvas = within(canvasElement);
-    await canvas.findByLabelText('Search spans');
-
-    const searchInput = canvas.getByLabelText('Search spans');
-    await userEvent.type(searchInput, 'Fetch');
-
-    await waitFor(() => {
-      if (!canvas.queryByText('Fetch resources')) {
-        throw new Error('Expected matching span to be displayed after searching');
-      }
-      if (canvas.queryByText('Initialize rollout')) {
-        throw new Error('Expected non-matching spans to be filtered out');
-      }
-    });
-  },
-};
-
 export const LoadingState: Story = {
   render: () => renderTracesPage(),
   parameters: {

dashboard/src/pages/Traces.page.tsx

@@ -1,12 +1,10 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-import { useCallback, useEffect, useMemo, useState } from 'react';
+import { useCallback, useEffect, useMemo } from 'react';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { IconCheck, IconChevronDown, IconSearch } from '@tabler/icons-react';
 import type { DataTableSortStatus } from 'mantine-datatable';
-import { useSearchParams } from 'react-router-dom';
 import { Button, Group, Menu, Select, Skeleton, Stack, TextInput, Title } from '@mantine/core';
-import { useDebouncedValue } from '@mantine/hooks';
 import { TracesTable, type TracesTableRecord } from '@/components/TracesTable.component';
 import { selectAutoRefreshMs } from '@/features/config';
 import {
@@ -16,7 +14,6 @@ import {
   type GetRolloutsQueryArgs,
 } from '@/features/rollouts';
 import {
-  hydrateTracesStateFromQuery,
   resetTracesFilters,
   selectTracesAttemptId,
   selectTracesPage,
@@ -56,38 +53,16 @@ function getLatestAttempt(attempts: Attempt[]): Attempt | null {
   return [...attempts].sort((a, b) => a.sequenceId - b.sequenceId).at(-1) ?? null;
 }
 
-function mergeRolloutCache(cache: Record<string, Rollout>, items: Rollout[]): Record<string, Rollout> {
-  if (!items.length) {
-    return cache;
+function findRollout(rollouts: Rollout[] | undefined, rolloutId: string | null): Rollout | null {
+  if (!rollouts || !rolloutId) {
+    return null;
   }
-  let changed = false;
-  const next = { ...cache };
-  for (const item of items) {
-    const existing = next[item.rolloutId];
-    if (!existing || existing !== item) {
-      next[item.rolloutId] = item;
-      changed = true;
-    }
-  }
-  return changed ? next : cache;
+  return rollouts.find((rollout) => rollout.rolloutId === rolloutId) ?? null;
 }
 
 export function TracesPage() {
   const dispatch = useAppDispatch();
   const autoRefreshMs = useAppSelector(selectAutoRefreshMs);
-  const [searchParams, setSearchParams] = useSearchParams();
-  const searchParamsKey = searchParams.toString();
-  const [hydratedSearchParamsKey, setHydratedSearchParamsKey] = useState<string | null>(null);
-  const [rolloutSearchValue, setRolloutSearchValue] = useState('');
-  const [debouncedRolloutSearchValue] = useDebouncedValue(rolloutSearchValue, 300);
-  const normalizedRolloutSearchValue = debouncedRolloutSearchValue.trim();
-  const rolloutSearchActive = normalizedRolloutSearchValue.length > 0;
-  const [rolloutLookup, setRolloutLookup] = useState<Record<string, Rollout>>({});
-  const hasRolloutQueryParam = searchParams.has('rolloutId');
-  const hasAttemptQueryParam = searchParams.has('attemptId');
-  const [initialHasRolloutQueryParam] = useState(hasRolloutQueryParam);
-  const rolloutIdFromQuery = hasRolloutQueryParam ? searchParams.get('rolloutId') || null : undefined;
-  const attemptIdFromQuery = hasAttemptQueryParam ? searchParams.get('attemptId') || null : undefined;
   const rolloutId = useAppSelector(selectTracesRolloutId);
   const attemptId = useAppSelector(selectTracesAttemptId);
   const searchTerm = useAppSelector(selectTracesSearchTerm);
@@ -97,7 +72,7 @@ export function TracesPage() {
   const viewMode = useAppSelector(selectTracesViewMode);
   const spansQueryArgs = useAppSelector(selectTracesQueryArgs);
 
-  const baseRolloutsQueryArgs = useMemo<GetRolloutsQueryArgs>(
+  const rolloutsQueryArgs = useMemo<GetRolloutsQueryArgs>(
     () => ({
       limit: 100,
       offset: 0,
@@ -113,68 +88,13 @@ export function TracesPage() {
     isFetching: rolloutsFetching,
     isError: rolloutsIsError,
     error: rolloutsError,
-  } = useGetRolloutsQuery(baseRolloutsQueryArgs, {
+  } = useGetRolloutsQuery(rolloutsQueryArgs, {
     pollingInterval: autoRefreshMs > 0 ? autoRefreshMs : undefined,
   });
 
-  const baseRolloutItems = rolloutsData?.items ?? [];
+  const rolloutItems = rolloutsData?.items ?? [];
 
-  const rolloutSearchQueryArgs = useMemo<GetRolloutsQueryArgs | null>(() => {
-    if (!rolloutSearchActive) {
-      return null;
-    }
-    return {
-      limit: 20,
-      offset: 0,
-      sortBy: 'start_time',
-      sortOrder: 'desc',
-      rolloutIdContains: normalizedRolloutSearchValue,
-    };
-  }, [normalizedRolloutSearchValue, rolloutSearchActive]);
-
-  const { data: rolloutSearchData, isFetching: rolloutSearchFetching } = useGetRolloutsQuery(
-    rolloutSearchQueryArgs ?? skipToken,
-  );
-
-  const rolloutByIdQueryArgs = useMemo<GetRolloutsQueryArgs | null>(() => {
-    if (!rolloutId || rolloutLookup[rolloutId]) {
-      return null;
-    }
-    return {
-      limit: 20,
-      offset: 0,
-      sortBy: 'start_time',
-      sortOrder: 'desc',
-      rolloutIdContains: rolloutId,
-    };
-  }, [rolloutId, rolloutLookup]);
-
-  const { data: rolloutByIdData, isFetching: rolloutByIdFetching } = useGetRolloutsQuery(
-    rolloutByIdQueryArgs ?? skipToken,
-  );
-
-  const searchRolloutItems = rolloutSearchData?.items ?? [];
-  const rolloutByIdItems = rolloutByIdData?.items ?? [];
-
-  useEffect(() => {
-    if (baseRolloutItems.length > 0) {
-      setRolloutLookup((prev) => mergeRolloutCache(prev, baseRolloutItems));
-    }
-  }, [baseRolloutItems]);
-
-  useEffect(() => {
-    if (searchRolloutItems.length > 0) {
-      setRolloutLookup((prev) => mergeRolloutCache(prev, searchRolloutItems));
-    }
-  }, [searchRolloutItems]);
-
-  useEffect(() => {
-    if (rolloutByIdItems.length > 0) {
-      setRolloutLookup((prev) => mergeRolloutCache(prev, rolloutByIdItems));
-    }
-  }, [rolloutByIdItems]);
-
-  const selectedRollout = rolloutId ? (rolloutLookup[rolloutId] ?? null) : null;
+  const selectedRollout = useMemo(() => findRollout(rolloutItems, rolloutId), [rolloutItems, rolloutId]);
 
   const attemptsQueryArgs =
     rolloutId !== null
@@ -205,61 +125,21 @@ export function TracesPage() {
     pollingInterval: autoRefreshMs > 0 ? autoRefreshMs : undefined,
   });
 
-  const shouldResolveRollout = rolloutId !== null && !rolloutLookup[rolloutId];
-
   useEffect(() => {
     if (!rolloutsData) {
       return;
     }
-
-    if (rolloutsData.total === 0) {
+    if (rolloutItems.length === 0) {
       if (rolloutId !== null) {
         dispatch(setTracesRolloutId(null));
       }
       return;
     }
-
-    if (rolloutId === null) {
-      if (!initialHasRolloutQueryParam && baseRolloutItems[0]) {
-        dispatch(setTracesRolloutId(baseRolloutItems[0].rolloutId));
-      }
-      return;
+    const rolloutExists = rolloutId ? rolloutItems.some((rollout) => rollout.rolloutId === rolloutId) : false;
+    if (!rolloutExists) {
+      dispatch(setTracesRolloutId(rolloutItems[0].rolloutId));
     }
-
-    if (!shouldResolveRollout) {
-      return;
-    }
-
-    if (rolloutSearchActive && normalizedRolloutSearchValue === rolloutId) {
-      return;
-    }
-
-    if (rolloutByIdQueryArgs) {
-      if (rolloutByIdFetching || !rolloutByIdData || rolloutsLoading) {
-        return;
-      }
-      if (rolloutByIdData.items.length === 0) {
-        dispatch(setTracesRolloutId(null));
-      }
-      return;
-    }
-
-    dispatch(setTracesRolloutId(null));
-  }, [
-    baseRolloutItems,
-    dispatch,
-    initialHasRolloutQueryParam,
-    normalizedRolloutSearchValue,
-    rolloutByIdData,
-    rolloutByIdFetching,
-    rolloutByIdQueryArgs,
-    rolloutId,
-    rolloutLookup,
-    rolloutSearchActive,
-    rolloutsData,
-    rolloutsLoading,
-    shouldResolveRollout,
-  ]);
+  }, [dispatch, rolloutsData, rolloutId, rolloutItems]);
 
   useEffect(() => {
     if (!rolloutId) {
@@ -280,29 +160,20 @@ export function TracesPage() {
       return;
     }
 
-    if (attemptId === null) {
-      const fallbackAttemptId = selectedRollout?.attempt?.attemptId ?? null;
-      if (fallbackAttemptId !== attemptId) {
-        dispatch(setTracesAttemptId(fallbackAttemptId));
-      }
+    const fallbackAttemptId = selectedRollout?.attempt?.attemptId ?? null;
+    if (fallbackAttemptId !== attemptId) {
+      dispatch(setTracesAttemptId(fallbackAttemptId));
     }
   }, [attemptsData, attemptId, dispatch, rolloutId, selectedRollout]);
 
-  const visibleRolloutItems = rolloutSearchActive ? searchRolloutItems : baseRolloutItems;
-  const rolloutSelectIsFetching =
-    (rolloutSearchActive ? rolloutSearchFetching : rolloutsFetching) ||
-    Boolean(rolloutByIdQueryArgs && rolloutByIdFetching);
-
-  const rolloutOptions = useMemo(() => {
-    const options = visibleRolloutItems.map((rollout) => ({
-      value: rollout.rolloutId,
-      label: rollout.rolloutId,
-    }));
-    if (rolloutId && !visibleRolloutItems.some((rollout) => rollout.rolloutId === rolloutId)) {
-      options.push({ value: rolloutId, label: rolloutId });
-    }
-    return options;
-  }, [rolloutId, visibleRolloutItems]);
+  const rolloutOptions = useMemo(
+    () =>
+      rolloutItems.map((rollout) => ({
+        value: rollout.rolloutId,
+        label: rollout.rolloutId,
+      })),
+    [rolloutItems],
+  );
 
   const attemptOptions = useMemo(() => {
     if (attemptsData && attemptsData.items.length > 0) {
@@ -325,21 +196,11 @@ export function TracesPage() {
     return [];
   }, [attemptsData, selectedRollout]);
 
-  const attemptPlaceholder = useMemo(() => {
-    if (!rolloutId) {
-      return 'Select Attempt';
-    }
-    if (attemptOptions.length === 0) {
-      return 'No Attempt';
-    }
-    return 'Latest Attempt';
-  }, [attemptOptions.length, rolloutId]);
-
   const rawSpansData = spansData as any as { items?: Span[]; total?: number } | undefined;
   const spans = rawSpansData?.items ?? [];
   const spansTotal = rawSpansData?.total ?? 0;
   const recordsPerPageOptions = [50, 100, 200, 500];
-  const isInitialLoading = rolloutsLoading && baseRolloutItems.length === 0;
+  const isInitialLoading = rolloutsLoading && rolloutItems.length === 0;
   const isFetching = spansFetching || rolloutsFetching || attemptsFetching;
 
   const selectionMessage = useMemo<string | undefined>(() => {
@@ -403,52 +264,6 @@ export function TracesPage() {
     [dispatch],
   );
 
-  useEffect(() => {
-    const payload: { rolloutId?: string | null; attemptId?: string | null } = {};
-    if (hasRolloutQueryParam) {
-      payload.rolloutId = rolloutIdFromQuery;
-    }
-    if (hasAttemptQueryParam) {
-      payload.attemptId = attemptIdFromQuery;
-    }
-    if (Object.keys(payload).length > 0) {
-      dispatch(hydrateTracesStateFromQuery(payload));
-    }
-    setHydratedSearchParamsKey((prev) => (prev === searchParamsKey ? prev : searchParamsKey));
-  }, [attemptIdFromQuery, dispatch, hasAttemptQueryParam, hasRolloutQueryParam, rolloutIdFromQuery, searchParamsKey]);
-
-  useEffect(() => {
-    if (hydratedSearchParamsKey !== searchParamsKey) {
-      return;
-    }
-    const next = new URLSearchParams(searchParams);
-    let changed = false;
-
-    if (rolloutId) {
-      if (next.get('rolloutId') !== rolloutId) {
-        next.set('rolloutId', rolloutId);
-        changed = true;
-      }
-    } else if (next.has('rolloutId')) {
-      next.delete('rolloutId');
-      changed = true;
-    }
-
-    if (rolloutId && attemptId) {
-      if (next.get('attemptId') !== attemptId) {
-        next.set('attemptId', attemptId);
-        changed = true;
-      }
-    } else if (next.has('attemptId')) {
-      next.delete('attemptId');
-      changed = true;
-    }
-
-    if (changed) {
-      setSearchParams(next, { replace: true });
-    }
-  }, [attemptId, hydratedSearchParamsKey, rolloutId, searchParams, searchParamsKey, setSearchParams]);
-
   const handleSearchTermChange = useCallback(
     (value: string) => {
       dispatch(setTracesSearchTerm(value));
@@ -494,7 +309,10 @@ export function TracesPage() {
 
   const handleShowRollout = useCallback(
     (record: TracesTableRecord) => {
-      const rollout = rolloutLookup[record.rolloutId];
+      if (rolloutItems.length === 0) {
+        return;
+      }
+      const rollout = rolloutItems.find((item) => item.rolloutId === record.rolloutId);
       if (!rollout) {
         return;
       }
@@ -512,12 +330,13 @@ export function TracesPage() {
         }),
       );
     },
-    [attemptsData, dispatch, rolloutLookup],
+    [attemptsData, dispatch, rolloutItems],
   );
 
   const handleShowSpanDetail = useCallback(
     (record: TracesTableRecord) => {
-      const rolloutForSpan = rolloutLookup[record.rolloutId] ?? null;
+      const rolloutForSpan =
+        rolloutItems.length > 0 ? (rolloutItems.find((item) => item.rolloutId === record.rolloutId) ?? null) : null;
       const attempts = attemptsData?.items ?? [];
       const attemptForSpan =
         attempts.find((attempt) => attempt.attemptId === record.attemptId) ?? rolloutForSpan?.attempt ?? null;
@@ -531,7 +350,7 @@ export function TracesPage() {
         }),
       );
     },
-    [attemptsData, dispatch, rolloutLookup],
+    [attemptsData, dispatch, rolloutItems],
   );
 
   const handleParentIdClick = useCallback(
@@ -563,27 +382,14 @@ export function TracesPage() {
                 if (value !== rolloutId) {
                   dispatch(setTracesRolloutId(value));
                 }
-                setRolloutSearchValue('');
               }}
               searchable
-              searchValue={rolloutSearchValue}
-              onSearchChange={(value) => {
-                setRolloutSearchValue(value ?? '');
-              }}
               placeholder='Select rollout'
               aria-label='Select rollout'
-              nothingFoundMessage={
-                rolloutSelectIsFetching ? 'Loading...' : rolloutSearchActive ? 'No matching rollouts' : 'No rollouts'
-              }
+              nothingFoundMessage={rolloutsFetching ? 'Loading...' : 'No rollouts'}
               comboboxProps={{ withinPortal: true }}
-              onDropdownOpen={() => {
-                setRolloutSearchValue('');
-              }}
-              onDropdownClose={() => {
-                setRolloutSearchValue('');
-              }}
               w={260}
-              disabled={rolloutOptions.length === 0 && !rolloutSelectIsFetching}
+              disabled={rolloutOptions.length === 0}
             />
             <Select
               data={attemptOptions}
@@ -594,7 +400,7 @@ export function TracesPage() {
                 }
               }}
               searchable
-              placeholder={attemptPlaceholder}
+              placeholder='Latest attempt'
               aria-label='Select attempt'
               nothingFoundMessage={attemptsFetching ? 'Loading...' : 'No attempts'}
               comboboxProps={{ withinPortal: true }}

dashboard/src/utils/format.ts

@@ -27,14 +27,6 @@ export function formatDateTime(timestamp: number | null): string {
   return dayjs(timestamp * 1000).format('YYYY-MM-DD HH:mm:ss');
 }
 
-export function formatDateTimeWithMilliseconds(timestamp: number | null): string {
-  if (timestamp == null) {
-    return '—';
-  }
-
-  return dayjs(timestamp * 1000).format('YYYY-MM-DD HH:mm:ss.SSS');
-}
-
 export function formatDuration(seconds: number | null): string {
   if (seconds == null) {
     return '—';

dashboard/src/utils/table.ts

@@ -48,21 +48,7 @@ export function compareRecords<T, K extends keyof T>(a: T, b: T, key: K): number
  * Create responsive columns based on container width
  * Columns with priority 0 are always shown, others are shown based on available space
  */
-const FALLBACK_EM_IN_PIXELS = 16;
-
-function getEmInPixels(): number {
-  if (typeof window === 'undefined' || !window.document?.documentElement) {
-    return FALLBACK_EM_IN_PIXELS;
-  }
-
-  const rootFontSize = window.getComputedStyle(window.document.documentElement).fontSize;
-  const parsed = Number.parseFloat(rootFontSize);
-  if (Number.isFinite(parsed) && parsed > 0) {
-    return parsed;
-  }
-
-  return FALLBACK_EM_IN_PIXELS;
-}
+const EM_IN_PIXELS = 16;
 
 function resolveWidth(config: ColumnVisibilityConfig): { widthEm: number; fixed: boolean } {
   if ('fixedWidth' in config && typeof config.fixedWidth === 'number') {
@@ -76,8 +62,7 @@ export function createResponsiveColumns<T>(
   containerWidth: number,
   columnVisibilityConfig: Record<string, ColumnVisibilityConfig>,
 ): DataTableColumn<T>[] {
-  const measuredWidth = containerWidth ? Math.max(containerWidth, 0) : Number.POSITIVE_INFINITY;
-  const emInPixels = getEmInPixels();
+  const measuredWidth = containerWidth ? Math.max(containerWidth - 48, 0) : Number.POSITIVE_INFINITY;
 
   const columnEntries = columns.map((column, index) => {
     const accessorKey = String(column.accessor);
@@ -94,7 +79,7 @@ export function createResponsiveColumns<T>(
       accessorKey,
       ...config,
       widthEm,
-      widthPx: widthEm * emInPixels,
+      widthPx: widthEm * EM_IN_PIXELS,
       fixed,
     };
   });

dashboard/tsconfig.json

@@ -21,5 +21,5 @@
       "@test-utils": ["./test-utils"]
     }
   },
-  "include": ["src", "public", "test-utils", ".storybook/main.ts", ".storybook/preview.tsx", ".storybook/modes.ts", ".storybook/constants.ts", ".storybook/vitest.setup.ts"]
+  "include": ["src", "public", "test-utils", ".storybook/main.ts", ".storybook/preview.tsx", ".storybook/modes.ts", ".storybook/constants.ts"]
 }

dashboard/vite.config.mjs

@@ -1,22 +1,16 @@
 // Copyright (c) Microsoft. All rights reserved.
 
-/// <reference types="vitest/config" />
-
 import path from 'node:path';
 import { fileURLToPath } from 'node:url';
-import { storybookTest } from '@storybook/addon-vitest/vitest-plugin';
 import react from '@vitejs/plugin-react';
-import { playwright } from '@vitest/browser-playwright';
 import { defineConfig } from 'vite';
 import tsconfigPaths from 'vite-tsconfig-paths';
 
-const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
-
-// More info at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
 const projectRoot = __dirname;
 const appRoot = path.resolve(projectRoot, 'public');
+
 export default defineConfig({
   root: appRoot,
   plugins: [react(), tsconfigPaths()],
@@ -27,47 +21,14 @@ export default defineConfig({
     },
   },
   build: {
-    outDir: path.resolve(projectRoot, '../agentlightning/dashboard'),
+    outDir: path.resolve(projectRoot, 'dist'),
     emptyOutDir: true,
   },
   test: {
     globals: true,
     environment: 'jsdom',
     setupFiles: './vitest.setup.mjs',
+    globalSetup: './vitest.global-setup.mjs',
     root: projectRoot,
-    projects: [
-      {
-        extends: true,
-        test: {
-          name: 'unit',
-          globalSetup: './vitest.global-setup.mjs',
-        },
-      },
-      {
-        // TODO: vitest for storybook has been setup but it's not working yet.
-        extends: true,
-        plugins: [
-          // The plugin will run tests for the stories defined in your Storybook config
-          // See options at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon#storybooktest
-          storybookTest({
-            configDir: path.join(dirname, '.storybook'),
-          }),
-        ],
-        test: {
-          name: 'storybook',
-          browser: {
-            enabled: true,
-            headless: true,
-            provider: playwright({}),
-            instances: [
-              {
-                browser: 'chromium',
-              },
-            ],
-          },
-          setupFiles: ['.storybook/vitest.setup.ts'],
-        },
-      },
-    ],
   },
 });

dashboard/vitest.shims.d.ts

@@ -1,3 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-/// <reference types="@vitest/browser-playwright" />

docs/how-to/unsloth-sft.md

@@ -81,7 +81,7 @@ model_list = [
 ]
 llm_proxy.update_model_list(model_list)
 # If the proxy is not running, it will start automatically.
-await llm_proxy.restart()
+llm_proxy.restart()
 # Add the proxy as a resource to the store so that the runners can access it via URL.
 resource_update = await store.add_resources({"main_llm": llm_proxy.as_resource()})
 ```

docs/reference/internal.md

@@ -34,12 +34,6 @@
 
 ::: agentlightning.tracer.agentops.LightningSpanProcessor
 
-::: agentlightning.utils.server_launcher.PythonServerLauncher
-
-::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
-
-::: agentlightning.utils.server_launcher.LaunchMode
-
 ## Deprecated APIs
 
 ::: agentlightning.server.AgentLightningServer

docs/tutorials/debug.md

@@ -79,43 +79,24 @@ vllm serve Qwen/Qwen2.5-0.5B-Instruct --port 8080
 Then start the LLM proxy via the following script:
 
 ```python
-import asyncio
-import aiohttp
 import agentlightning as agl
 
-async def serve_llm_proxy():
-    store = agl.InMemoryLightningStore()
-    store_server = agl.LightningStoreServer(store, "127.0.0.1", 8081)
-    await store_server.start()
+llm_proxy = agl.LLMProxy(
+    port=8081,
+    model_list=[
+        {
+            "model_name": "Qwen/Qwen2.5-0.5B-Instruct",
+            "litellm_params": {
+                "model": "hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",
+                "api_base": "http://localhost:8080/v1",
+            },
+        }
+    ],
+    store=agl.InMemoryLightningStore(),
+)
 
-    llm_proxy = agl.LLMProxy(
-        port=8082,
-        model_list=[
-            {
-                "model_name": "Qwen/Qwen2.5-0.5B-Instruct",
-                "litellm_params": {
-                    "model": "hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",
-                    "api_base": "http://localhost:8080/v1",
-                },
-            }
-        ],
-        store=store_server,
-    )
-
-    await llm_proxy.start()
-    await asyncio.sleep(1000000)
-```
-
-Test the served LLM proxy with a client like:
-
-```python
-async def test_llm_proxy():
-    async with aiohttp.ClientSession() as session:
-        async with session.post("http://localhost:8082/v1/chat/completions", json={
-            "model": "Qwen/Qwen2.5-0.5B-Instruct",
-            "messages": [{"role": "user", "content": "Hello, world!"}],
-        }) as response:
-            print(await response.json())
+llm_proxy.start()
+time.sleep(1000000)
 ```
 
 You can now use the LLM proxy by specifying environment variables:

docs/tutorials/installation.md

@@ -152,18 +152,6 @@ uv sync --frozen \
 
 Read more about Agent-lightning managed dependency groups [here]({{ src("pyproject.toml") }}).
 
-### Building the Dashboard
-
-The Agent-Lightning dashboard is built using [Vite](https://vite.dev/). To build the dashboard, run the following command:
-
-```bash
-cd dashboard
-npm ci
-npm run build
-```
-
-Some HTML and JavaScript assets will be generated in the `agentlightning/dashboard` directory.
-
 ### Activating Your Environment
 
 After syncing dependencies, `uv` automatically creates a virtual environment inside the `.venv/` directory.

examples/tinker/agl_tinker/algo.py

@@ -81,7 +81,6 @@ class Tinker(Algorithm):
                 port=config.llm_proxy_port,
                 model_list=[],
                 store=store,
-                launch_mode="thread",
             )
 
         await main_training_loop(config, store, adapter, llm_proxy)  # type: ignore

examples/tinker/agl_tinker/llm.py

@@ -303,7 +303,5 @@ def create_llm_proxy(
         store=store,
         model_list=tinker_llm.as_model_list(),
         num_retries=2,
-        # Must use thread mode here because otherwise the Tinker sampling client will hang.
-        launch_mode="thread",
         _add_return_token_ids=_add_return_token_ids,
     )

examples/tinker/agl_tinker/train.py

@@ -168,7 +168,7 @@ async def do_sync_training(
         t_start = time.time()
 
         llm_proxy.update_model_list(tinker_llm.as_model_list())
-        await llm_proxy.restart()
+        llm_proxy.restart()
 
         logger.info(f"[Batch {i_batch}] LiteLLM model list: {llm_proxy.model_list}")
         llm_resource = llm_proxy.as_resource()
@@ -221,7 +221,7 @@ async def do_sync_training(
         ml_logger.log_metrics(metrics, step=i_batch)
         logger.info(f"[Batch {i_batch}] Sampling and training completed")
 
-    await llm_proxy.stop()
+    llm_proxy.stop()
 
 
 @scope
@@ -343,8 +343,6 @@ async def main(config: Config) -> None:
         model_list=[],
         store=store,
         num_retries=config.llm_proxy_retry_attempts,
-        # Must use thread mode here because otherwise the Tinker sampling client will hang.
-        launch_mode="thread",
     )
 
     await main_training_loop(config, store, adapter, llm_proxy)

examples/tinker/hello.py

@@ -173,12 +173,7 @@ def oneclick():
     )
     trainer = agl.Trainer(
         algorithm=Tinker(config),
-        llm_proxy=agl.LLMProxy(
-            port=12306,
-            num_retries=3,
-            # Must use thread mode here because otherwise the Tinker sampling client will hang.
-            launch_mode="thread",
-        ),
+        llm_proxy=agl.LLMProxy(port=12306, num_retries=3),
         n_runners=8,
         port=_find_available_port(),
     )

examples/tinker/q20_evaluate.py

@@ -33,7 +33,6 @@ game statistics after the run.
 from __future__ import annotations
 
 import argparse
-import asyncio
 import json
 import traceback
 from pathlib import Path
@@ -45,12 +44,12 @@ from crewai import LLM as CrewLLM
 from q20_agent import AnswererResponse, SearchTool, TwentyQuestionsFlow
 from rich.console import Console
 
-from agentlightning import InMemoryLightningStore, LightningStoreThreaded, LLMProxy
+from agentlightning import InMemoryLightningStore, LLMProxy
 
 console = Console()
 
 
-async def evaluate_q20(
+def evaluate_q20(
     model_name: str,
     search: bool,
     port: int,
@@ -69,7 +68,7 @@ async def evaluate_q20(
         seed: Optional random seed for shuffling the dataset; ``None`` disables deterministic shuffling.
     """
 
-    store = LightningStoreThreaded(InMemoryLightningStore())
+    store = InMemoryLightningStore()
     df = pd.read_csv(dataset_path)  # type: ignore
     if df.empty:
         console.print(f"[bold yellow]Dataset '{dataset_path}' is empty. Nothing to evaluate.[/bold yellow]")
@@ -89,7 +88,6 @@ async def evaluate_q20(
                 {"model_name": model_name, "litellm_params": {"model": "openai/" + model_name}},
             ],
             num_retries=2,
-            launch_mode="thread",
             _add_return_token_ids=False,
         )
 
@@ -110,7 +108,7 @@ async def evaluate_q20(
     console.print("Model list:", llm_proxy.model_list)
 
     try:
-        await llm_proxy.start()
+        llm_proxy.start()
         player_llm = CrewLLM(
             model="openai/" + model_name, base_url=f"http://localhost:{port}/v1", api_key="dummy", timeout=60.0
         )
@@ -145,7 +143,7 @@ async def evaluate_q20(
 
             flow = TwentyQuestionsFlow(player_llm=player_llm, answer_llm=answer_llm, search_tool=search_tool)
             try:
-                await flow.kickoff_async(
+                flow.kickoff(
                     {
                         "answer": row["answer"],
                         "category": row["category"],
@@ -163,7 +161,7 @@ async def evaluate_q20(
             with output_path.open("a") as f:
                 f.write(json.dumps(result_json) + "\n")
     finally:
-        await llm_proxy.stop()
+        llm_proxy.stop()
 
 
 def main(argv: Optional[List[str]] = None) -> None:
@@ -207,15 +205,13 @@ def main(argv: Optional[List[str]] = None) -> None:
     )
 
     args = parser.parse_args(argv)
-    asyncio.run(
-        evaluate_q20(
-            model_name=args.model,
-            search=args.search,
-            port=args.port,
-            output_file=args.output_file,
-            dataset_path=args.dataset,
-            seed=None if args.seed == -1 else args.seed,
-        )
+    evaluate_q20(
+        model_name=args.model,
+        search=args.search,
+        port=args.port,
+        output_file=args.output_file,
+        dataset_path=args.dataset,
+        seed=None if args.seed == -1 else args.seed,
     )
 
 

examples/tinker/q20_train.py

@@ -133,7 +133,7 @@ def dry_run():
 
     Uses in-memory store and processes 4 sample tasks to verify the setup works.
     """
-    store = agl.LightningStoreThreaded(agl.InMemoryLightningStore())
+    store = agl.InMemoryLightningStore()
     llm_proxy = create_llm_proxy("Qwen/Qwen3-30B-A3B-Instruct-2507", "qwen3_instruct", store=store)
     trainer = agl.Trainer(
         n_runners=2,
@@ -141,13 +141,13 @@ def dry_run():
         store=store,
     )
     try:
-        asyncio.run(llm_proxy.start())
+        llm_proxy.start()
         sampled_csv = pd.read_csv("q20_nouns.csv").sample(n=4, random_state=42)  # type: ignore
         sampled_csv["search_enabled"] = False
         dataset = sampled_csv.to_dict(orient="records")  # type: ignore
         trainer.dev(q20_agent, cast(agl.Dataset[Q20Task], dataset))
     finally:
-        asyncio.run(llm_proxy.stop())
+        llm_proxy.stop()
 
 
 async def algo(search: bool, model: Literal["qwen4b", "qwen30b"], port: int):

examples/tinker/tests/test_tinker_llm.py

@@ -27,7 +27,6 @@ from agentlightning import (
     configure_logger,
     emit_reward,
 )
-from agentlightning.store import LightningStoreThreaded
 
 configure_logger(name="agentlightning")
 configure_logger(name="agl_tinker", level=logging.INFO)
@@ -62,7 +61,7 @@ async def test_tracer():
 
         # init tracer before llm_proxy to avoid tracer provider being not active.
         console.print("Starting LLM proxy...")
-        await llm_proxy.start()
+        llm_proxy.start()
         console.print("LLM proxy started")
 
         # client = openai.OpenAI(
@@ -100,7 +99,7 @@ async def test_tracer():
         print(trajectory)
     finally:
         console.print("Stopping LLM proxy...")
-        await llm_proxy.stop()
+        llm_proxy.stop()
         console.print("LLM proxy stopped")
 
 
@@ -118,20 +117,19 @@ async def test_llm_proxy():
     )
     tinker_llm.rewrite_litellm_custom_providers()
 
-    store = LightningStoreThreaded(InMemoryLightningStore())
+    store = InMemoryLightningStore()
     rollout = await store.start_rollout("dummy", "train")
     llm_proxy = LLMProxy(
         port=4000,
         store=store,
         model_list=tinker_llm.as_model_list(),
         num_retries=0,
-        launch_mode="thread",
     )
 
     try:
         # init tracer before llm_proxy to avoid tracer provider being not active.
         console.print("Starting LLM proxy...")
-        await llm_proxy.start()
+        llm_proxy.start()
         console.print("LLM proxy started")
 
         client = openai.OpenAI(
@@ -159,7 +157,7 @@ async def test_llm_proxy():
         print(trajectory)
     finally:
         console.print("Stopping LLM proxy...")
-        await llm_proxy.stop()
+        llm_proxy.stop()
         console.print("LLM proxy stopped")
 
 

examples/unsloth/sft_algorithm.py

@@ -185,7 +185,7 @@ async def sft_one_iter(
         llm_proxy.update_model_list(model_list)
         # Restart the LLM proxy after backend model list update
         # If LLM proxy has never been started, it will be started
-        await llm_proxy.restart()
+        llm_proxy.restart()
 
         # Put the LLM proxy address into the store as an address
         resources_update = await store.add_resources(

pyproject.toml

@@ -21,9 +21,6 @@ dependencies = [
   "pydantic>=2.11",
   "openai",
   "rich",
-  "sqlalchemy[asyncio]",
-  "aiosqlite",
-  "tenacity",
   "portpicker",
   "gunicorn",
   "uvicorn_worker",
@@ -263,15 +260,7 @@ build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
 packages = ["agentlightning"]
-include = [
-  "agentlightning/**/*.yaml",
-  "agentlightning/**/*.yml",
-  "agentlightning/**/*.poml",
-  "agentlightning/**/*.html",
-  "agentlightning/**/*.js",
-  "agentlightning/**/*.css",
-  "agentlightning/**/*.svg",
-]
+include = ["agentlightning/**/*.yaml", "agentlightning/**/*.yml", "agentlightning/**/*.poml"]
 
 [tool.hatch.build.targets.sdist]
 exclude = [
@@ -279,7 +268,6 @@ exclude = [
   "tests/**",
   "docs/**",
   "scripts/**",
-  "dashboard/**",
 ]
 
 [tool.pytest.ini_options]

tests/llm_proxy/test_llm_proxy_cpu.py

@@ -15,9 +15,7 @@ from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
 from agentlightning.store.memory import InMemoryLightningStore
-from agentlightning.store.threading import LightningStoreThreaded
 from agentlightning.types import Span
-from agentlightning.utils.server_launcher import PythonServerLauncherArgs
 
 from ..common.network import get_free_port
 from ..common.tracer import clear_tracer_provider
@@ -123,10 +121,10 @@ def test_exporter_helpers():
     # TODO: add more complex tests for the exporter helper
 
 
-@pytest.mark.asyncio
-async def test_update_model_list():
+def test_update_model_list():
     store = InMemoryLightningStore()
     proxy = LLMProxy(
+        port=get_free_port(),
         model_list=[
             {
                 "model_name": "gpt-4o-arbitrary",
@@ -135,11 +133,9 @@ async def test_update_model_list():
                 },
             }
         ],
-        launch_mode="asyncio",
-        port=get_free_port(),
         store=store,
     )
-    await proxy.start()
+    proxy.start()
     assert proxy.is_running()
     assert proxy.model_list == [
         {
@@ -168,11 +164,10 @@ async def test_update_model_list():
         }
     ]
     assert proxy.is_running()
-    await proxy.stop()
+    proxy.stop()
 
 
-@pytest.mark.asyncio
-async def test_restart_resets_litellm_logging_worker() -> None:
+def test_restart_resets_litellm_logging_worker() -> None:
     """LLMProxy.start() should recreate LiteLLM's logging worker on each run."""
 
     try:
@@ -182,6 +177,7 @@ async def test_restart_resets_litellm_logging_worker() -> None:
 
     store = InMemoryLightningStore()
     proxy = LLMProxy(
+        port=get_free_port(),
         model_list=[
             {
                 "model_name": "dummy-model",
@@ -190,24 +186,17 @@ async def test_restart_resets_litellm_logging_worker() -> None:
             }
         ],
         store=store,
-        launcher_args=PythonServerLauncherArgs(
-            port=get_free_port(),
-            launch_mode="asyncio",
-            healthcheck_url="/health",
-            startup_timeout=10.0,
-            process_join_timeout=10.0,
-        ),
     )
 
     try:
-        await proxy.start()
+        proxy.start()
         first_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
-        await proxy.stop()
+        proxy.stop()
 
-        await proxy.start()
+        proxy.start()
         second_worker = litellm_logging_worker.GLOBAL_LOGGING_WORKER
     finally:
-        await proxy.stop()
+        proxy.stop()
 
     assert first_worker is not second_worker, "LiteLLM logging worker should be refreshed after restart"
 
@@ -232,18 +221,18 @@ class TestLLM(CustomLLM):
         )
 
 
-@pytest.mark.asyncio
-async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixture) -> None:
+def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixture) -> None:
     clear_tracer_provider()
 
     restart_times: int = 30
 
-    store = LightningStoreThreaded(InMemoryLightningStore())
+    store = InMemoryLightningStore()
     caplog.set_level(logging.WARNING)
 
     port = get_free_port()
     try:
         llm_proxy = LLMProxy(
+            port=port,
             model_list=[
                 {
                     "model_name": "gpt-4o-arbitrary",
@@ -254,27 +243,22 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
                     },
                 }
             ],
-            launcher_args=PythonServerLauncherArgs(
-                launch_mode="thread",
-                healthcheck_url="/health",
-                port=port,
-            ),
             store=store,
         )
         for restart_idx in range(restart_times):
             llm_instance = TestLLM(f"Hi! {restart_idx}")
             litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
             custom_llm_setup()
-            await llm_proxy.restart()
+            llm_proxy.restart()
             assert llm_proxy.is_running()
 
-            openai_client = openai.AsyncOpenAI(
-                base_url=llm_proxy.server_launcher.access_endpoint,
+            openai_client = openai.OpenAI(
+                base_url=f"http://localhost:{port}",
                 api_key="token-abc123",
                 timeout=5,
                 max_retries=0,
             )
-            response = await openai_client.chat.completions.create(
+            response = openai_client.chat.completions.create(
                 model="gpt-4o-arbitrary",
                 messages=[{"role": "user", "content": "Hello world"}],
                 stream=False,
@@ -282,11 +266,10 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
             assert response.choices[0].message.content == f"Hi! {restart_idx}"
 
             error_logs = [record.message for record in caplog.records if record.levelno >= logging.ERROR]
-            error_logs = [message for message in error_logs if "Task was destroyed but it is pending!" not in message]
             assert not error_logs, f"Found error logs: {error_logs}"
             assert not any("Cannot add callback" in record.message for record in caplog.records)
 
-        await llm_proxy.stop()
+        llm_proxy.stop()
     finally:
         litellm.custom_provider_map = []
         custom_llm_setup()

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -22,7 +22,6 @@ import openai
 import pytest
 
 from agentlightning.llm_proxy import LLMProxy, _reset_litellm_logging_worker  # pyright: ignore[reportPrivateUsage]
-from agentlightning.store import LightningStore, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.types import LLM, Span
 
@@ -69,9 +68,7 @@ def test_qwen25_model_sanity(qwen25_model: RemoteOpenAIServer):
 @pytest.mark.asyncio
 async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
     clear_tracer_provider()
-    inmemory_store = InMemoryLightningStore()
-    store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=get_free_port())
-    await store.start()
+    store = InMemoryLightningStore()
     proxy = LLMProxy(
         port=get_free_port(),
         model_list=[
@@ -88,10 +85,12 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
 
     rollout = await store.start_rollout(None)
 
-    await proxy.start()
+    proxy.start()
 
     resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
 
+    import openai
+
     client = openai.OpenAI(base_url=resource.endpoint, api_key="token-abc123")
     response = client.chat.completions.create(
         model="gpt-4o-arbitrary",
@@ -101,12 +100,10 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
     assert response.choices[0].message.content is not None
     assert "hello, world" in response.choices[0].message.content.lower()
 
-    await proxy.stop()
+    proxy.stop()
 
     spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
 
-    await store.stop()
-
     # Verify all spans have correct rollout_id, attempt_id, and sequence_id
     assert len(spans) > 0, "Should have captured spans"
     for span in spans:
@@ -171,14 +168,12 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
     assert "gen_ai.completion.0.finish_reason" in litellm_span.attributes, "gen_ai.completion.0.finish_reason not found"
 
 
-async def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: int = 0, gunicorn: bool = False):
+def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: int = 0):
     clear_tracer_provider()
     _reset_litellm_logging_worker()  # type: ignore
     store = InMemoryLightningStore()
-    store_server = LightningStoreServer(store=store, host="127.0.0.1", port=get_free_port())
-    # When the server is forked into subprocess, it automatically becomes a client of the store
-    await store_server.start()
     proxy = LLMProxy(
+        port=get_free_port(),
         model_list=[
             {
                 "model_name": "gpt-4o-arbitrary",
@@ -188,16 +183,14 @@ async def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: in
                 },
             }
         ],
-        port=get_free_port(),
-        num_workers=4 if gunicorn else 1,
-        store=store_server,
+        store=store,
         num_retries=retries,
     )
-    await proxy.start()
-    return proxy, store_server
+    proxy.start()
+    return proxy, store
 
 
-async def _new_resource(proxy: LLMProxy, store: LightningStore):
+async def _new_resource(proxy: LLMProxy, store: InMemoryLightningStore):
     rollout = await store.start_rollout(None)
     return proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id), rollout
 
@@ -220,7 +213,7 @@ def _attr(s: Span, key: str, default: Any = None):  # type: ignore
 
 @pytest.mark.asyncio
 async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+    proxy, store = _make_proxy_and_store(qwen25_model)
     try:
         resource, rollout = await _new_resource(proxy, store)
         client = _get_client_for_resource(resource)
@@ -241,14 +234,12 @@ async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
         assert len(_find_span(spans, "raw_gen_ai_request")) == 3
         # TODO: Check response contents and token ids for the 3 requests respectively
     finally:
-        await proxy.stop()
-        await store.stop()
+        proxy.stop()
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("gunicorn", [False, True])
-async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicorn: bool):
-    proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=gunicorn)
+async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer):
+    proxy, store = _make_proxy_and_store(qwen25_model)
     try:
         resource, rollout = await _new_resource(proxy, store)
         aclient = _get_async_client_for_resource(resource)
@@ -263,21 +254,19 @@ async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicor
 
         outs = await asyncio.gather(*[_one(i) for i in range(10)])
         assert len([o for o in outs if o]) == 10
-        await asyncio.sleep(1.0)  # Allow some extra time for the spans to be recorded
 
         spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
         assert len(_find_span(spans, "raw_gen_ai_request")) == 10
         assert {s.sequence_id for s in spans} == {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
         # TODO: Check whether the sequence ids get mixed up or not
     finally:
-        await proxy.stop()
-        await store.stop()
+        proxy.stop()
 
 
 @pytest.mark.asyncio
 async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
     # litellm proxy accepts Anthropic schema and forwards to OpenAI backend
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+    proxy, store = _make_proxy_and_store(qwen25_model)
     try:
         resource, rollout = await _new_resource(proxy, store)
 
@@ -294,13 +283,12 @@ async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
         spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
         assert len(spans) > 0
     finally:
-        await proxy.stop()
-        await store.stop()
+        proxy.stop()
 
 
 @pytest.mark.asyncio
 async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+    proxy, store = _make_proxy_and_store(qwen25_model)
     try:
         resource, rollout = await _new_resource(proxy, store)
         client = _get_client_for_resource(resource)
@@ -363,14 +351,13 @@ async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
 
         # TODO: Check response contents and token ids for the 2 requests respectively
     finally:
-        await proxy.stop()
-        await store.stop()
+        proxy.stop()
 
 
 @pytest.mark.skip(reason="Streaming is not supported yet")
 @pytest.mark.asyncio
 async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
-    proxy, store = await _make_proxy_and_store(qwen25_model)
+    proxy, store = _make_proxy_and_store(qwen25_model)
     try:
         resource, rollout = await _new_resource(proxy, store)
         client = _get_client_for_resource(resource)
@@ -392,5 +379,4 @@ async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
         assert len(spans) > 0
         # TODO: didn't test the token ids in streaming chunks here
     finally:
-        await proxy.stop()
-        await store.stop()
+        proxy.stop()

tests/runner/test_agent_integration.py

@@ -1,8 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import asyncio
+import multiprocessing
 import os
-from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Dict, Optional, cast
+import time
+from multiprocessing.synchronize import Event as MpEvent
+from typing import Any, Dict, Optional, cast
 
 import openai
 import pytest
@@ -154,17 +157,6 @@ def server():
         yield server
 
 
-class LLMProxyWithClearedTracerProvider(LLMProxy):
-    """LLMProxy that clears the tracer provider before serving."""
-
-    @asynccontextmanager
-    async def _serve_context(self) -> AsyncGenerator[None, None]:
-        # This will be run inside the LLM proxy's own process
-        clear_tracer_provider()
-        async with super()._serve_context():
-            yield
-
-
 @pytest.mark.asyncio
 async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenAIServer) -> None:
     torch = pytest.importorskip("torch")
@@ -193,7 +185,7 @@ async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenA
     await server_store.start()
     client_store = LightningStoreClient(server_store.endpoint)
 
-    proxy = LLMProxyWithClearedTracerProvider(
+    proxy = LLMProxy(
         port=get_free_port(),
         model_list=[
             {
@@ -207,7 +199,16 @@ async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenA
         store=client_store,
     )
 
-    await proxy.start()
+    def run_proxy_server(proxy: LLMProxy, event: MpEvent):
+        clear_tracer_provider()  # clear once more before the proxy starts
+        proxy.start()
+        event.set()
+        time.sleep(3600)  # Keep the server running
+
+    event = multiprocessing.Event()
+    process = multiprocessing.Process(target=run_proxy_server, args=(proxy, event))
+    process.start()
+    event.wait(timeout=30)
 
     try:
         await runner.step("Say hello to Agent Lightning", resources={"llm": proxy.as_resource()})
@@ -233,6 +234,9 @@ async def test_runner_integration_with_spawned_litellm_proxy(server: RemoteOpenA
         assert last_spans[0].attributes.get("reward") == 0.5
     finally:
         teardown_runner(runner)
-        await proxy.stop()
+        process.terminate()
         await client_store.close()
         await server_store.stop()
+        await asyncio.to_thread(process.join, timeout=1)
+        if process.is_alive():
+            process.kill()

tests/store/conftest.py

@@ -1,18 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import os
 import time
-import typing
-import uuid
 from unittest.mock import Mock
 
 import pytest
-import pytest_asyncio
 from opentelemetry.sdk.trace import ReadableSpan
 from pytest import FixtureRequest
 
-from agentlightning.store import InMemoryLightningStore, SqlLightningStore
 from agentlightning.store.base import LightningStore
+from agentlightning.store.memory import InMemoryLightningStore
 
 __all__ = [
     "inmemory_store",
@@ -26,35 +22,15 @@ def inmemory_store() -> InMemoryLightningStore:
     return InMemoryLightningStore()
 
 
-@pytest_asyncio.fixture
-async def sql_store() -> typing.AsyncGenerator[SqlLightningStore, None]:
+@pytest.fixture
+def sql_store():
     """Placeholder fixture for SQL store implementation. Returns None until SQL store is ready."""
-    """Helper generator to create a SqlLightningStore using a SQLite file for testing."""
-    tmp_path = ".pytest_cache"
-    # Ensure the directory exists and create a random file in it
-    os.makedirs(tmp_path, exist_ok=True)
-    db_path = os.path.join(tmp_path, f"test_db_{uuid.uuid4().hex}.sqlite3")
-    database_url = f"sqlite+aiosqlite:///{db_path}"
-    store = SqlLightningStore(database_url=database_url)
-    store.retry_for_waiting.wait_seconds = 0.2  # Set polling interval to 0.2s for test
-
-    # Config db_store with a short time interval for healthcheck
-    store.add_background_task(
-        {"name": "test_healthcheck", "method": "check_attempt_timeout", "interval": {"seconds": 0.1}}
-    )
-
-    await store.start()
-    try:
-        yield store
-    finally:
-        await store.stop()
-        if os.path.exists(db_path):
-            os.remove(db_path)
+    return None
 
 
 # Uncomment this when sql store is ready
-@pytest.fixture(params=["inmemory_store", "sql_store"])
-# @pytest.fixture(params=["inmemory_store"])
+# @pytest.fixture(params=["inmemory_store", "sql_store"])
+@pytest.fixture(params=["inmemory_store"])
 def store_fixture(request: FixtureRequest) -> LightningStore:
     """Parameterized fixture that provides different store implementations for testing.
     Currently supports InMemoryLightningStore, with SQL store support planned.

tests/store/test_client_server.py

@@ -292,16 +292,12 @@ async def test_client_server_end_to_end(
     all_rollouts = await client.query_rollouts()
     assert any(r.rollout_id == enqueued.rollout_id for r in all_rollouts)
     assert await client.query_rollouts(rollout_ids=[enqueued.rollout_id])
-    # Test that attempt is present in the rollout
-    assert any(hasattr(r, "attempt") and r.attempt is not None for r in all_rollouts)  # type: ignore
     attempts = await client.query_attempts(dequeued_client.rollout_id)
     assert attempts
     assert await client.get_latest_attempt(dequeued_client.rollout_id) is not None
     stored_client_rollout = await client.get_rollout_by_id(dequeued_client.rollout_id)
     assert stored_client_rollout is not None
     assert stored_client_rollout.config.unresponsive_seconds == 6.0
-    # Test that attempt is present in the rollout
-    assert hasattr(stored_client_rollout, "attempt") and stored_client_rollout.attempt is not None  # type: ignore
 
     client_span = _make_span(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id, 101, "client-span")
     stored_span = await client.add_span(client_span)

tests/store/test_implementation.py

@@ -897,17 +897,10 @@ async def test_span_triggers_status_transition(store_fixture: LightningStore, mo
     # Get the attempt
     attempts = await store_fixture.query_attempts(rollout.rollout_id)
     attempt_id = attempts[0].attempt_id
-    assert attempts[0].status == "preparing"
 
     # Add first span
     await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span)
 
-    # Attempt status should be changed
-    attempt_v2 = await store_fixture.get_latest_attempt(rollout.rollout_id)
-    assert attempt_v2 is not None
-    assert attempt_v2.attempt_id == attempt_id
-    assert attempt_v2.status == "running"
-
     # Status should transition to running
     rollouts = await store_fixture.query_rollouts(status=["running"])
     assert len(rollouts) == 1
@@ -1845,7 +1838,7 @@ async def test_healthcheck_timeout_behavior(store_fixture: LightningStore, mock_
     assert len(running_rollouts) == 1
 
     # Wait for timeout to occur
-    await asyncio.sleep(0.3)  # Wait longer than timeout_seconds
+    await asyncio.sleep(0.15)  # Wait longer than timeout_seconds
 
     # Trigger healthcheck by calling any decorated method
     # Verify the attempt was marked as timeout and rollout was requeued
@@ -1883,7 +1876,7 @@ async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore,
     assert running_attempts[0].last_heartbeat_time is not None
 
     # Wait for unresponsive timeout
-    await asyncio.sleep(0.3)  # Wait longer than unresponsive_seconds
+    await asyncio.sleep(0.15)  # Wait longer than unresponsive_seconds
 
     # Verify attempt was marked as unresponsive
     attempts_after = await store_fixture.query_attempts(rollout.rollout_id)

tests/utils/test_server_launcher.py

@@ -1308,17 +1308,17 @@ async def test_launcher_assigns_random_port_when_none(launch_mode: LaunchMode):
 
     await launcher.start()
     assert launcher.is_running()
-    # Port is chosen and embedded in endpoint/access_endpoint
+    # Port is chosen and embedded in endpoint/access_url
     assert launcher.endpoint.startswith(f"http://{host}:")
-    assert launcher.access_endpoint.startswith(f"http://{host}:")
-    await _probe_json(f"{launcher.access_endpoint}/", {"hello": "world"})
+    assert launcher.access_url.startswith(f"http://{host}:")
+    await _probe_json(f"{launcher.access_url}/", {"hello": "world"})
 
     await launcher.stop()
 
 
 @pytest.mark.asyncio
-async def test_launcher_endpoint_access_endpoint_health_url_normalization():
-    host = "0.0.0.0"  # should flip to 127.0.0.1 for access_endpoint
+async def test_launcher_endpoint_access_url_health_url_normalization():
+    host = "0.0.0.0"  # should flip to 127.0.0.1 for access_url
     port = portpicker.pick_unused_port()
     app = _make_app_health()
 
@@ -1335,14 +1335,10 @@ async def test_launcher_endpoint_access_endpoint_health_url_normalization():
     await launcher.start()
     try:
         assert launcher.endpoint == f"http://{host}:{port}"
-        # access_endpoint should map 0.0.0.0 -> 127.0.0.1
-        assert launcher.access_endpoint.startswith("http://") and launcher.access_endpoint.endswith(f":{port}")
-        assert (
-            launcher.health_url
-            and launcher.health_url.startswith("http://")
-            and launcher.health_url.endswith(f"/health")
-        )
-        await _probe_json(f"{launcher.access_endpoint}/", {"hello": "world"})
+        # access_url should map 0.0.0.0 -> 127.0.0.1
+        assert launcher.access_url == f"http://127.0.0.1:{port}"
+        assert launcher.health_url == f"http://127.0.0.1:{port}/health"
+        await _probe_json(f"{launcher.access_url}/", {"hello": "world"})
     finally:
         await launcher.stop()
         _free_port("127.0.0.1", port)

uv.lock

@@ -128,7 +128,6 @@ dependencies = [
     { name = "agentops", version = "0.4.18", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'group-14-agentlightning-core-legacy') or (sys_platform == 'linux' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (sys_platform == 'linux' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable')" },
     { name = "agentops", version = "0.4.21", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'group-14-agentlightning-core-stable') or (sys_platform == 'linux' and extra == 'group-14-agentlightning-tinker') or (sys_platform == 'linux' and extra == 'group-14-agentlightning-torch-gpu-stable') or (sys_platform == 'linux' and extra != 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-cpu' and extra != 'group-14-agentlightning-torch-legacy') or (sys_platform == 'linux' and extra != 'group-14-agentlightning-core-legacy' and extra != 'group-14-agentlightning-torch-gpu-legacy' and extra != 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "aiohttp", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
-    { name = "aiosqlite", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "fastapi", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "flask", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "graphviz", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
@@ -147,8 +146,6 @@ dependencies = [
     { name = "pydantic", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "rich", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "setproctitle", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
-    { name = "sqlalchemy", extra = ["asyncio"], marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
-    { name = "tenacity", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "uvicorn", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
     { name = "uvicorn-worker", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
 ]
@@ -349,7 +346,6 @@ trl = [
 requires-dist = [
     { name = "agentops", specifier = ">=0.4.13" },
     { name = "aiohttp" },
-    { name = "aiosqlite" },
     { name = "fastapi" },
     { name = "flask" },
     { name = "graphviz" },
@@ -366,8 +362,6 @@ requires-dist = [
     { name = "pydantic", specifier = ">=2.11" },
     { name = "rich" },
     { name = "setproctitle" },
-    { name = "sqlalchemy", extras = ["asyncio"] },
-    { name = "tenacity" },
     { name = "uvicorn" },
     { name = "uvicorn-worker" },
     { name = "verl", marker = "extra == 'verl'", specifier = ">=0.5.0" },
@@ -824,18 +818,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" },
 ]
 
-[[package]]
-name = "aiosqlite"
-version = "0.21.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/13/7d/8bca2bf9a247c2c5dfeec1d7a5f40db6518f88d314b8bca9da29670d2671/aiosqlite-0.21.0.tar.gz", hash = "sha256:131bb8056daa3bc875608c631c678cda73922a2d4ba8aec373b19f18c17e7aa3", size = 13454, upload-time = "2025-02-03T07:30:16.235Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f5/10/6c25ed6de94c49f88a91fa5018cb4c0f3625f31d5be9f771ebe5cc7cd506/aiosqlite-0.21.0-py3-none-any.whl", hash = "sha256:2549cf4057f95f53dcba16f2b64e8e2791d7e1adedb13197dd8ed77bb226d7d0", size = 15792, upload-time = "2025-02-03T07:30:13.6Z" },
-]
-
 [[package]]
 name = "airportsdata"
 version = "20250909"
@@ -10317,11 +10299,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/9c/5e/6a29fa884d9fb7ddadf6b69490a9d45fded3b38541713010dad16b77d015/sqlalchemy-2.0.44-py3-none-any.whl", hash = "sha256:19de7ca1246fbef9f9d1bff8f1ab25641569df226364a0e40457dc5457c54b05", size = 1928718, upload-time = "2025-10-10T15:29:45.32Z" },
 ]
 
-[package.optional-dependencies]
-asyncio = [
-    { name = "greenlet", marker = "sys_platform == 'linux' or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-core-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-tinker') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-core-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-core-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-tinker' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-cu128') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-legacy') or (extra == 'group-14-agentlightning-torch-cpu' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-gpu-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-gpu-legacy' and extra == 'group-14-agentlightning-trl') or (extra == 'group-14-agentlightning-torch-gpu-stable' and extra == 'group-14-agentlightning-torch-legacy') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-torch-stable') or (extra == 'group-14-agentlightning-torch-legacy' and extra == 'group-14-agentlightning-trl')" },
-]
-
 [[package]]
 name = "sqlparse"
 version = "0.5.3"