.

Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>

.github/workflows/badge-compat.yml

@@ -0,0 +1,29 @@
+name: Badge - Compatibility
+
+on:
+  workflow_run:
+    workflows:
+      - Examples - Backward Compatibility
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'examples-compat.yml', label: 'examples-compat', variants: ['legacy', 'stable'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/badge-unit.yml

@@ -0,0 +1,31 @@
+name: Badge - Unit Test
+
+on:
+  workflow_run:
+    workflows:
+      - CPU Test
+      - GPU Test
+    types: [completed]
+
+  workflow_dispatch:
+
+permissions:
+  actions: read
+  contents: read
+
+jobs:
+  badge:
+    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const badgeAggregation = require('./scripts/badge_aggregation.js');
+            const dependencies = [
+              { workflow: 'tests-full.yml', label: 'tests-full', variants: ['legacy', 'stable'] },
+              { workflow: 'tests.yml', label: 'tests', variants: ['legacy', 'stable', 'Lint', 'documentation', 'JavaScript'] },
+            ];
+            await badgeAggregation({ github, context, core, dependencies });

.github/workflows/examples-calc-x.yml

@@ -22,12 +22,12 @@ run-name: >-
       || format('Calc-X - {0}', github.event_name) }}
 
 jobs:
-  calc-x:
+  calc-x-perf:
     if: >
       github.event_name != 'repository_dispatch' ||
       github.event.action == 'ci-calc-x' ||
       github.event.action == 'ci-all'
-    name: Calc-X (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    name: Calc-X Performance (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
     runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
     timeout-minutes: 90
     strategy:
@@ -74,7 +74,7 @@ jobs:
       - name: Upload dependencies artifact
         uses: actions/upload-artifact@v4
         with:
-          name: dependencies-calc-x-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          name: dependencies-calc-x-performance-${{ matrix.python-version }}-${{ matrix.setup-script }}
           path: requirements-freeze.txt
           compression-level: 0
 
@@ -116,13 +116,11 @@ jobs:
       # Don't ask why. Don't touch this.
       - name: Calc-X training
         run: |
-          set -ex
           source .venv/bin/activate
           cd examples/calc_x
           ../../scripts/restart_ray.sh
           sleep 5
-          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci
-          sleep 10
+          python train_calc_agent.py --val-file data/test_mini.parquet --ci
         shell: bash
         env:
           WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
@@ -137,14 +135,118 @@ jobs:
           WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
           WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
 
-      - name: Calc-X training LLM Proxy
+  calc-x-variants:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-calc-x' ||
+      github.event.action == 'ci-all'
+    name: Calc-X Variants (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 90
+    strategy:
+      matrix:
+        include:
+        - python-version: '3.10'
+          setup-script: 'legacy'
+        - python-version: '3.12'
+          setup-script: 'stable'
+        - python-version: '3.13'
+          setup-script: 'latest'
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - name: Check disk space
+        run: df -h
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade dependencies (latest)
+        run: uv lock --upgrade
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (latest)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group torch-gpu-stable
+        if: matrix.setup-script == 'latest'
+      - name: Sync dependencies (stable & legacy)
+        run: |
+          uv sync --frozen --no-default-groups --extra verl \
+            --group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
+        if: matrix.setup-script != 'latest'
+      - name: Freeze dependencies
+        run: |
+          set -ex
+          uv pip freeze | tee requirements-freeze.txt
+          echo "UV_LOCKED=1" >> $GITHUB_ENV
+          echo "UV_NO_SYNC=1" >> $GITHUB_ENV
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-calc-x-variants-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          ./scripts/litellm_run.sh
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
+
+      - name: Prepare Calc-X dataset
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
+          unzip calc-x-data.zip -d data
+          rm calc-x-data.zip
+
+      - name: Calc-X MCP sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run tests/test_mcp_calculator.py
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: Calc-X sanity check
+        run: |
+          set -ex
+          cd examples/calc_x
+          uv run legacy_calc_agent_debug.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Training with local model
         run: |
           set -ex
           source .venv/bin/activate
           cd examples/calc_x
           ../../scripts/restart_ray.sh
           sleep 5
-          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci --llm-proxy
+          hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir data/qwen_model
+          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --model $(realpath data/qwen_model)
+          sleep 10
+        shell: bash
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        id: calc_x_train_local_model
+
+      - name: Training with LLM Proxy
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --llm-proxy
           sleep 10
         shell: bash
         env:
@@ -152,7 +254,7 @@ jobs:
           WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
         id: calc_x_train_llm_proxy
 
-      - name: Calc-X training with external store
+      - name: Training with external store
         run: |
           set -euo pipefail
           source .venv/bin/activate
@@ -182,7 +284,7 @@ jobs:
           WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
         id: calc_x_train_external_store
 
-      - name: Calc-X training with role-based environment variables
+      - name: Training with role-based environment variables
         run: |
           set -euo pipefail
           source .venv/bin/activate

README.md

@@ -4,7 +4,7 @@
 
 # Agent Lightning⚡
 
-[![Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml)
+[![Unit Tests](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
 [![Documentation](https://img.shields.io/badge/GitHub%20Pages-Documentation-blue)](https://microsoft.github.io/agent-lightning/)
 [![PyPI version](https://badge.fury.io/py/agentlightning.svg)](https://badge.fury.io/py/agentlightning)
 [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
@@ -34,6 +34,12 @@ Read more on our [documentation website](https://microsoft.github.io/agent-light
 pip install agentlightning
 ```
 
+For the latest nightly build (cutting-edge features), you can install from Test PyPI:
+
+```bash
+pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
+```
+
 Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
 
 To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).
@@ -69,10 +75,11 @@ No rewrites, no lock-in, just a clear path from first rollout to steady improvem
 | Workflow | Status |
 |----------|--------|
 | CPU Tests | [![tests workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
-| GPU Tests | [![tests-full workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml) |
+| Full Tests | [![tests summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
+| UI Tests | [![UI Tests](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml) |
 | Examples Integration | [![examples summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml) |
 | Latest Dependency Compatibility | [![latest summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml) |
-| Legacy Examples Compatibility | [![examples compatibility workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml) |
+| Legacy Examples Compatibility | [![compat summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml) |
 
 ## ⚡ Citation
 

agentlightning/__init__.py

@@ -10,7 +10,9 @@ from .emitter import *
 from .execution import *
 from .litagent import *
 from .llm_proxy import *
-from .logging import *
+from .logging import configure_logger  # deprecated  # type: ignore
+from .logging import setup as setup_logging  # type: ignore
+from .logging import setup_module as setup_module_logging  # type: ignore
 from .runner import *
 from .server import AgentLightningServer  # deprecated  # type: ignore
 from .store import *

agentlightning/cli/store.py

@@ -9,7 +9,7 @@ import asyncio
 import logging
 from typing import Iterable
 
-from agentlightning.logging import configure_logger
+from agentlightning import setup_logging
 from agentlightning.store.client_server import LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 
@@ -27,7 +27,7 @@ def main(argv: Iterable[str] | None = None) -> int:
     )
     args = parser.parse_args(list(argv) if argv is not None else None)
 
-    configure_logger()
+    setup_logging()
 
     store = InMemoryLightningStore()
     server = LightningStoreServer(
@@ -35,6 +35,7 @@ def main(argv: Iterable[str] | None = None) -> int:
         host="0.0.0.0",
         port=args.port,
         cors_allow_origins=args.cors_origins,
+        launch_mode="asyncio",
     )
     try:
         asyncio.run(server.run_forever())

agentlightning/execution/client_server.py

@@ -67,8 +67,8 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
         server_host: str | None = None,
         server_port: int | None = None,
         n_runners: int = 1,
-        graceful_timeout: float = 5.0,
-        terminate_timeout: float = 5.0,
+        graceful_timeout: float = 10.0,
+        terminate_timeout: float = 10.0,
         main_process: Literal["algorithm", "runner"] = "algorithm",
         managed_store: bool | None = None,
     ) -> None:

agentlightning/llm_proxy.py

@@ -4,12 +4,15 @@ from __future__ import annotations
 
 import ast
 import asyncio
+import json
 import logging
 import os
 import re
 import tempfile
 import threading
+import time
 from contextlib import asynccontextmanager
+from datetime import datetime
 from typing import (
     Any,
     AsyncGenerator,
@@ -18,8 +21,11 @@ from typing import (
     Dict,
     Iterable,
     List,
+    Literal,
     Optional,
     Sequence,
+    Tuple,
+    Type,
     TypedDict,
     Union,
     cast,
@@ -29,12 +35,15 @@ import litellm
 import opentelemetry.trace as trace_api
 import yaml
 from fastapi import Request, Response
+from fastapi.responses import StreamingResponse
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
 from litellm.proxy.proxy_server import app, save_worker_config  # pyright: ignore[reportUnknownVariableType]
+from litellm.types.utils import CallTypes
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import Scope
 
 from agentlightning.types import LLM, ProxyLLM
 from agentlightning.utils.server_launcher import (
@@ -352,7 +361,9 @@ class LightningSpanExporter(SpanExporter):
                 headers_merged.update(cast(Dict[str, Any], headers))
 
             if not headers_merged:
-                logger.warning(f"No headers found in {len(subtree_spans)} subtree spans. Cannot log to store.")
+                logger.warning(
+                    f"No headers found in {len(subtree_spans)} subtree spans of root {root_span_id}. Cannot log to store."
+                )
                 continue
 
             # Validate and normalize required header fields.
@@ -454,6 +465,23 @@ class LightningOpenTelemetry(OpenTelemetry):
 
         super().__init__(config=config)  # pyright: ignore[reportUnknownMemberType]
 
+    async def async_pre_call_deployment_hook(
+        self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] = None
+    ) -> Optional[Dict[str, Any]]:
+        """The root span is sometimes missing (e.g., when Anthropic endpoint is used).
+        It is created in an auth module in LiteLLM. If it's missing, we create it here.
+        """
+        if "metadata" not in kwargs or "litellm_parent_otel_span" not in kwargs["metadata"]:
+            parent_otel_span = self.create_litellm_proxy_request_started_span(  # type: ignore
+                start_time=datetime.now(),
+                headers=kwargs.get("headers", {}),
+            )
+            updated_metadata = {**kwargs.get("metadata", {}), "litellm_parent_otel_span": parent_otel_span}
+
+            return {**kwargs, "metadata": updated_metadata}
+        else:
+            return kwargs
+
 
 class RolloutAttemptMiddleware(BaseHTTPMiddleware):
     """
@@ -499,6 +527,433 @@ class RolloutAttemptMiddleware(BaseHTTPMiddleware):
         return response
 
 
+class MessageInspectionMiddleware(BaseHTTPMiddleware):
+    """Middleware to inspect the request and response bodies.
+
+    It's for debugging purposes. Add it via "message_inspection" middleware alias.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        ti = time.time()
+        logger.info(f"Received request with scope: {request.scope}")
+        logger.info(f"Received request with body: {await request.body()}")
+        response = await call_next(request)
+        elapsed = time.time() - ti
+        logger.info(f"Response to request took {elapsed} seconds")
+        logger.info(f"Received response with status code: {response.status_code}")
+        logger.info(f"Received response with body: {response.body}")
+        return response
+
+
+class StreamConversionMiddleware(BaseHTTPMiddleware):
+    """Middleware to convert streaming responses to non-streaming responses.
+
+    Useful for backend that only supports non-streaming responses.
+
+    LiteLLM's OpenTelemetry is also buggy with streaming responses.
+    The conversion will hopefully bypass the bug.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        # Only process POST requests to completion endpoints
+        if request.method != "POST":
+            return await call_next(request)
+
+        # Check if it's a chat completions or messages endpoint
+        endpoint_format: Literal["openai", "anthropic", "unknown"] = "unknown"
+        if request.url.path.endswith("/chat/completions") or "/chat/completions?" in request.url.path:
+            endpoint_format = "openai"
+        elif request.url.path.endswith("/messages") or "/messages?" in request.url.path:
+            endpoint_format = "anthropic"
+        else:
+            endpoint_format = "unknown"
+
+        if endpoint_format == "unknown":
+            # Directly bypass the middleware
+            return await call_next(request)
+
+        # Read the request body
+        try:
+            json_body = await request.json()
+        except json.JSONDecodeError:
+            logger.warning(f"Request body is not valid JSON: {request.body}")
+            return await call_next(request)
+
+        # Check if streaming is requested
+        is_streaming = json_body.get("stream", False)
+
+        # Simple case: no streaming requested, just return the response
+        if not is_streaming:
+            return await call_next(request)
+
+        # Now the stream case
+        return await self._handle_stream_case(request, json_body, endpoint_format, call_next)
+
+    async def _handle_stream_case(
+        self,
+        request: Request,
+        json_body: Dict[str, Any],
+        endpoint_format: Literal["openai", "anthropic"],
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        # 1) Modify the request body to force stream=False
+        modified_json = dict(json_body)
+        modified_json["stream"] = False
+        modified_body = json.dumps(modified_json).encode("utf-8")
+
+        # 2) Build a new scope + receive that yields our modified body
+        scope: Scope = dict(request.scope)
+        # rewrite headers for accept/content-length
+        new_headers: List[Tuple[bytes, bytes]] = []
+        saw_accept = False
+        for k, v in scope["headers"]:
+            kl = k.lower()
+            if kl == b"accept":
+                saw_accept = True
+                new_headers.append((k, b"application/json"))
+            elif kl == b"content-length":
+                # replace with new length
+                continue
+            else:
+                new_headers.append((k, v))
+        if not saw_accept:
+            new_headers.append((b"accept", b"application/json"))
+        new_headers.append((b"content-length", str(len(modified_body)).encode("ascii")))
+        scope["headers"] = new_headers
+
+        # Directly modify the request body
+        # Creating a new request won't work because request is cached in the base class
+        request._body = modified_body  # type: ignore
+
+        response = await call_next(request)
+
+        buffered: Optional[bytes] = None
+        # 4) If OK, buffer the response body (it should be JSON because we forced stream=False)
+        if 200 <= response.status_code < 300:
+            try:
+                if hasattr(response, "body_iterator"):
+                    # Buffer body safely
+                    body_chunks: List[bytes] = []
+                    async for chunk in response.body_iterator:  # type: ignore
+                        body_chunks.append(chunk)  # type: ignore
+                    buffered = b"".join(body_chunks)
+                else:
+                    buffered = response.body  # type: ignore
+
+                data = json.loads(buffered or b"{}")
+
+                if endpoint_format == "anthropic":
+                    return StreamingResponse(
+                        self.anthropic_stream_generator(data),
+                        media_type="text/event-stream",
+                        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+                    )
+                else:
+                    # openai format
+                    return StreamingResponse(
+                        self.openai_stream_generator(data),
+                        media_type="text/event-stream",
+                        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+                    )
+            except Exception as e:
+                # If anything goes wrong, fall back to non-streaming JSON
+                logger.exception(f"Error converting to stream; returning non-stream response: {e}")
+                # Rebuild the consumed response
+                return Response(
+                    content=buffered if buffered is not None else b"",
+                    status_code=response.status_code,
+                    headers=dict(response.headers),
+                    media_type=response.media_type,
+                    background=response.background,
+                )
+        else:
+            return response
+
+    async def anthropic_stream_generator(self, original_response: Dict[str, Any]):
+        """Generate Anthropic SSE-formatted chunks from complete content blocks
+
+        This is a dirty hack for Anthropic-style streaming from non-streaming response.
+        The sse format is subject to change based on Anthropic's implementation.
+        If so, try to use `MessageInspectionMiddleware` to inspect the update and fix accordingly.
+        """
+        # Anthropic format - handle multiple content blocks (text + tool_use)
+        content_blocks: List[Dict[str, Any]] = original_response.get("content", [])
+        message_id = original_response.get("id", f"msg_{int(time.time() * 1000)}")
+        model = original_response.get("model", "claude")
+
+        # Send message_start event
+        message_start: Dict[str, Any] = {
+            "type": "message_start",
+            "message": {
+                "id": message_id,
+                "type": "message",
+                "role": "assistant",
+                "content": [],
+                "model": model,
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": original_response.get("usage", {"input_tokens": 0, "output_tokens": 0}),
+            },
+        }
+        yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n"
+
+        # Send ping to keep connection alive
+        ping = {"type": "ping"}
+        yield f"event: ping\ndata: {json.dumps(ping)}\n\n"
+
+        # Process each content block
+        for block_index, block in enumerate(content_blocks):
+            block_type = block.get("type", "text")
+
+            if block_type == "text":
+                # Handle text block
+                content = block.get("text", "")
+
+                # Send content_block_start event
+                content_block_start = {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "text", "text": ""},
+                }
+                yield f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n"
+
+                # Stream text content in chunks
+                if content:
+                    words = content.split()
+                    chunk_size = 5
+
+                    for i in range(0, len(words), chunk_size):
+                        chunk_words = words[i : i + chunk_size]
+                        text_chunk = " ".join(chunk_words)
+
+                        # Add space after chunk unless it's the last one
+                        if i + chunk_size < len(words):
+                            text_chunk += " "
+
+                        content_block_delta = {
+                            "type": "content_block_delta",
+                            "index": block_index,
+                            "delta": {"type": "text_delta", "text": text_chunk},
+                        }
+                        yield f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n"
+                        await asyncio.sleep(0.02)
+
+                # Send content_block_stop event
+                content_block_stop = {"type": "content_block_stop", "index": block_index}
+                yield f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n"
+
+            elif block_type == "tool_use":
+                # Handle tool_use block
+                tool_name = block.get("name", "")
+                tool_input = block.get("input", {})
+                tool_id = block.get("id", f"toolu_{int(time.time() * 1000)}")
+
+                # Send content_block_start event for tool use
+                content_block_start: Dict[str, Any] = {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "tool_use", "id": tool_id, "name": tool_name, "input": {}},
+                }
+                yield f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n"
+
+                # Stream tool input as JSON string chunks
+                input_json = json.dumps(tool_input)
+                chunk_size = 20  # characters per chunk for JSON
+
+                for i in range(0, len(input_json), chunk_size):
+                    json_chunk = input_json[i : i + chunk_size]
+
+                    content_block_delta = {
+                        "type": "content_block_delta",
+                        "index": block_index,
+                        "delta": {"type": "input_json_delta", "partial_json": json_chunk},
+                    }
+                    yield f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n"
+                    await asyncio.sleep(0.01)
+
+                # Send content_block_stop event
+                content_block_stop = {"type": "content_block_stop", "index": block_index}
+                yield f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n"
+
+        # Send message_delta event with stop reason
+        message_delta = {
+            "type": "message_delta",
+            "delta": {"stop_reason": original_response.get("stop_reason", "end_turn"), "stop_sequence": None},
+            "usage": {"output_tokens": original_response.get("usage", {}).get("output_tokens", 0)},
+        }
+        yield f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n"
+
+        # Send message_stop event
+        message_stop = {"type": "message_stop"}
+        yield f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n"
+
+    async def openai_stream_generator(self, response_json: Dict[str, Any]) -> AsyncGenerator[str, Any]:
+        """
+        Convert a *complete* OpenAI chat.completions choice into a stream of
+        OpenAI-compatible SSE chunks.
+
+        This emits:
+
+          - an initial delta with the role ("assistant"),
+          - a sequence of deltas for message.content (split into small chunks),
+          - deltas for any tool_calls (including id/name and chunked arguments),
+          - a terminal chunk with finish_reason,
+          - and finally the literal '[DONE]'.
+
+        Notes:
+
+        - We only handle a *single* choice (index 0 typically).
+        - We purposefully don't attempt to stream logprobs.
+        - Chunking strategy is simple and conservative to avoid splitting
+          multi-byte characters: we slice on spaces where possible, then fall
+          back to fixed-size substrings.
+        """
+        choice = cast(Dict[str, Any], (response_json.get("choices") or [{}])[0])
+        model = response_json.get("model", "unknown")
+        created: int = int(time.time())
+        index: int = choice.get("index", 0)
+
+        message: Dict[str, Any] = choice.get("message", {}) or {}
+        role: str = message.get("role", "assistant")
+        content: str = message.get("content") or ""
+        tool_calls: List[Any] = message.get("tool_calls") or []
+        finish_reason: Optional[str] = choice.get(
+            "finish_reason"
+        )  # e.g., "stop", "length", "tool_calls", "content_filter"
+
+        def sse_chunk(obj: Dict[str, Any]) -> str:
+            print("sse_chunk: ", obj)
+            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"
+
+        # 1) initial chunk with the role
+        yield sse_chunk(
+            {
+                "id": f"chatcmpl-{created}",
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model,
+                "choices": [{"index": index, "delta": {"role": role}, "finish_reason": None}],
+            }
+        )
+
+        # 2) stream textual content as small deltas
+        async def stream_content(text: str):
+            if not text:
+                return
+            # prefer splitting on spaces in ~20–40 char pieces
+            approx = 28
+            start = 0
+            n = len(text)
+            while start < n:
+                end = min(start + approx, n)
+                if end < n:
+                    # try to break on a space going forward
+                    space = text.rfind(" ", start, end)
+                    if space > start:
+                        end = space + 1
+                delta_text = text[start:end]
+                start = end
+                if not delta_text:
+                    break
+                yield sse_chunk(
+                    {
+                        "id": f"chatcmpl-{created}",
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": model,
+                        "choices": [{"index": index, "delta": {"content": delta_text}, "finish_reason": None}],
+                    }
+                )
+                # tiny pause helps some UIs animate smoothly; keep very small
+                await asyncio.sleep(0.0)
+
+        async for piece in stream_content(content):  # type: ignore[misc]
+            yield piece  # pass through the produced chunks
+
+        # 3) stream tool_calls if present (id/name first, then arguments piecemeal)
+        for tc_index, tc in enumerate(tool_calls):
+            tc_type = tc.get("type", "function")
+            tc_id = tc.get("id") or f"call_{created}_{tc_index}"
+            fn: Dict[str, Any] = (tc.get("function") or {}) if tc_type == "function" else {}
+            fn_name: str = fn.get("name", "")
+            fn_args: str = fn.get("arguments", "") or ""
+
+            # (a) delta that announces the tool call id/type/name
+            yield sse_chunk(
+                {
+                    "id": f"chatcmpl-{created}",
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model,
+                    "choices": [
+                        {
+                            "index": index,
+                            "delta": {
+                                "tool_calls": [
+                                    {"index": tc_index, "id": tc_id, "type": tc_type, "function": {"name": fn_name}}
+                                ]
+                            },
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+
+            # (b) stream arguments in small substrings
+            arg_chunk_size = 40
+            for pos in range(0, len(fn_args), arg_chunk_size):
+                partial = fn_args[pos : pos + arg_chunk_size]
+                yield sse_chunk(
+                    {
+                        "id": f"chatcmpl-{created}",
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": model,
+                        "choices": [
+                            {
+                                "index": index,
+                                "delta": {"tool_calls": [{"index": tc_index, "function": {"arguments": partial}}]},
+                                "finish_reason": None,
+                            }
+                        ],
+                    }
+                )
+                await asyncio.sleep(0.0)
+
+        # 4) terminal chunk with finish_reason (default to "stop" if missing)
+        yield sse_chunk(
+            {
+                "id": f"chatcmpl-{created}",
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model,
+                "choices": [
+                    {
+                        "index": index,
+                        "delta": {},
+                        "finish_reason": finish_reason or ("tool_calls" if tool_calls else "stop"),
+                    }
+                ],
+            }
+        )
+
+        # 5) literal DONE sentinel
+        yield "data: [DONE]\n\n"
+
+
+_MIDDLEWARE_REGISTRY: Dict[str, Type[BaseHTTPMiddleware]] = {
+    "rollout_attempt": RolloutAttemptMiddleware,
+    "stream_conversion": StreamConversionMiddleware,
+    "message_inspection": MessageInspectionMiddleware,
+}
+
+
+_CALLBACK_REGISTRY = {
+    "return_token_ids": AddReturnTokenIds,
+    "opentelemetry": LightningOpenTelemetry,
+}
+
+
 class LLMProxy:
     """Host a LiteLLM OpenAI-compatible proxy bound to a LightningStore.
 
@@ -522,7 +977,10 @@ class LLMProxy:
 
     !!! warning
 
-        The LLM Proxy does support streaming, but the tracing is still problematic when streaming is enabled.
+        By default (or when "stream_conversion" middleware is enabled), the LLM Proxy will convert OpenAI and Anthropic requests with `stream=True`
+        to a non-streaming request before going through the LiteLLM proxy. This is because the OpenTelemetry tracer provided by
+        LiteLLM is buggy with streaming responses. You can disable this by removing the "stream_conversion" middleware.
+        In that case, you might lose some tracing information like token IDs.
 
     !!! danger
 
@@ -544,6 +1002,13 @@ class LLMProxy:
             `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.
+        middlewares: List of FastAPI middleware classes or strings to register. You can specify the class aliases or classes that have been imported.
+            If not provided, the default middlewares (RolloutAttemptMiddleware and StreamConversionMiddleware) will be used.
+            Available middleware aliases are: "rollout_attempt", "stream_conversion", "message_inspection".
+            Middlewares are the **first layer** of request processing. They are applied to all requests before the LiteLLM proxy.
+        callbacks: List of LiteLLM callback classes or strings to register. You can specify the class aliases or classes that have been imported.
+            If not provided, the default callbacks (AddReturnTokenIds and LightningOpenTelemetry) will be used.
+            Available callback aliases are: "return_token_ids", "opentelemetry".
     """
 
     def __init__(
@@ -557,7 +1022,8 @@ class LLMProxy:
         num_workers: int = 1,
         launch_mode: LaunchMode = "mp",
         launcher_args: PythonServerLauncherArgs | None = None,
-        _add_return_token_ids: bool = True,
+        middlewares: List[Union[Type[BaseHTTPMiddleware], str]] | None = None,
+        callbacks: List[Union[Type[CustomLogger], str]] | None = None,
     ):
         self.store = store
 
@@ -589,7 +1055,33 @@ class LLMProxy:
 
         self._config_file = None
 
-        self._add_return_token_ids = _add_return_token_ids
+        self.middlewares: List[Type[BaseHTTPMiddleware]] = []
+        if middlewares is None:
+            middlewares = ["rollout_attempt", "stream_conversion"]
+        for middleware in middlewares:
+            if isinstance(middleware, str):
+                if middleware not in _MIDDLEWARE_REGISTRY:
+                    raise ValueError(
+                        f"Invalid middleware alias: {middleware}. Available aliases are: {list(_MIDDLEWARE_REGISTRY.keys())}"
+                    )
+                middleware = _MIDDLEWARE_REGISTRY[middleware]
+                self.middlewares.append(middleware)
+            else:
+                self.middlewares.append(middleware)
+
+        self.callbacks: List[Type[CustomLogger]] = []
+        if callbacks is None:
+            callbacks = ["return_token_ids", "opentelemetry"]
+        for callback in callbacks:
+            if isinstance(callback, str):
+                if callback not in _CALLBACK_REGISTRY:
+                    raise ValueError(
+                        f"Invalid callback alias: {callback}. Available aliases are: {list(_CALLBACK_REGISTRY.keys())}"
+                    )
+                callback = _CALLBACK_REGISTRY[callback]
+                self.callbacks.append(callback)
+            else:
+                self.callbacks.append(callback)
 
     def get_store(self) -> Optional[LightningStore]:
         """Get the store used by the proxy.
@@ -641,21 +1133,18 @@ class LLMProxy:
         set_active_llm_proxy(self)
 
         # Install middleware if it's not already installed.
-        installed: bool = False
+        installation_status: Dict[Any, bool] = {}
         for mw in app.user_middleware:
-            if mw.cls is RolloutAttemptMiddleware:
-                # Check whether the middleware is installed.
-                # It could be installed by other LLM Proxy instances, but it doesn't matter.
-                logger.info("Found existing RolloutAttemptMiddleware installed. Will not install a new one.")
-                installed = True
-                break
+            installation_status[mw.cls] = True
 
-        if not installed:
-            # Fallback to adding a new middleware
-            logger.info("Adding a new middleware to the FastAPI app.")
-            app.add_middleware(RolloutAttemptMiddleware)
+        for mw in self.middlewares:
+            if mw not in installation_status:
+                logger.info(f"Adding middleware {mw} to the FastAPI app.")
+                app.add_middleware(mw)
+            else:
+                logger.info(f"Middleware {mw} is already installed. Will not install a new one.")
 
-        if not initialize_llm_callbacks(self._add_return_token_ids):
+        if not initialize_llm_callbacks(self.callbacks):
             # If it's not the first time to initialize the callbacks, also
             # reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
             _reset_litellm_logging_worker()
@@ -720,7 +1209,7 @@ class LLMProxy:
         if self.store is None:
             raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
 
-        store_capabilities = self.store.capabilities()
+        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."
@@ -849,7 +1338,7 @@ def set_active_llm_proxy(proxy: LLMProxy) -> None:
     _global_llm_proxy = proxy
 
 
-def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
+def initialize_llm_callbacks(callback_classes: List[Type[CustomLogger]]) -> bool:
     """Restore `litellm.callbacks` to a state that is just initialized by agent-lightning.
 
     When litellm is restarted multiple times in the same process, more and more callbacks
@@ -857,8 +1346,7 @@ def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
     This function remembers the initial state of `litellm.callbacks` and always restore to that state.
 
     Args:
-        _add_return_token_ids: Whether to add the return token ids callback. Internal use only.
-        Ideally the callback should automatically be enabled when the backend supports it.
+        callback_classes: List of callback classes to register.
 
     Returns:
         Whether the callbacks are initialized for the first time.
@@ -866,31 +1354,31 @@ def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
     global _callbacks_before_litellm_start
 
     if _callbacks_before_litellm_start is None:
-        litellm.callbacks.extend(  # type: ignore
-            [
-                AddReturnTokenIds(),
-                LightningOpenTelemetry(),
-            ]
-            if _add_return_token_ids
-            else [
-                LightningOpenTelemetry(),
-            ]
-        )
+        litellm.callbacks.extend([cls() for cls in callback_classes])  # type: ignore
         _callbacks_before_litellm_start = [*litellm.callbacks]  # type: ignore
         return True
 
+    else:
+        # Put whatever is missing in the new callback classes to the existing callbacks.
+        for cls in callback_classes:
+            if not any(isinstance(cb, cls) for cb in _callbacks_before_litellm_start):
+                logger.info(f"Adding missing callback {cls} to the existing callbacks.")
+                _callbacks_before_litellm_start.append(cls())
+
     _reset_litellm_logging_callback_manager()
 
-    # Check if tracer provider is malformed due to global tracer clear in tests.
-    if not _check_tracer_provider():
-        logger.warning(
-            "Global tracer provider might have been cleared outside. Re-initializing OpenTelemetry callback."
-        )
-        _callbacks_before_litellm_start = [
-            cb for cb in _callbacks_before_litellm_start if not isinstance(cb, LightningOpenTelemetry)
-        ] + [LightningOpenTelemetry()]
-    else:
-        logger.debug("Global tracer provider is valid. Reusing existing OpenTelemetry callback.")
+    if LightningOpenTelemetry in callback_classes:
+        # Check if tracer provider is malformed due to global tracer clear in tests.
+        if not _check_tracer_provider():
+            logger.warning(
+                "Global tracer provider might have been cleared outside. Re-initializing OpenTelemetry callback."
+            )
+            _callbacks_before_litellm_start = [
+                cb for cb in _callbacks_before_litellm_start if not isinstance(cb, LightningOpenTelemetry)
+            ] + [LightningOpenTelemetry()]
+        else:
+            logger.debug("Global tracer provider is valid. Reusing existing OpenTelemetry callback.")
+    # Otherwise, we just skip the check for opentelemetry and use the existing callback.
 
     litellm.callbacks.clear()  # type: ignore
     litellm.callbacks.extend(_callbacks_before_litellm_start)  # type: ignore

agentlightning/logging.py

@@ -1,10 +1,18 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import logging
 import os
 import platform
+import sys
+import warnings
+from logging.config import dictConfig
+from typing import Any, Dict, Optional
 
-__all__ = ["configure_logger"]
+from rich.console import Console
+
+__all__ = ["setup", "configure_logger", "setup_module"]
 
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
@@ -15,6 +23,10 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
     not propagate to the root logger, preventing duplicate log emission when
     applications compose multiple logging configurations.
 
+    !!! danger
+
+        This function is deprecated in favor of [`setup_logging`][agentlightning.setup_logging].
+
     Args:
         level: Logging level applied both to the logger and the installed
             handler. Defaults to `logging.INFO`.
@@ -32,23 +44,327 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
         logger.info("agent-lightning is ready!")
         ```
     """
+    warnings.warn("This function is deprecated in favor of `setup_logging`.", DeprecationWarning, stacklevel=2)
 
+    return setup_module(level=level, name=name, console=True, color=True, propagate=False)
+
+
+DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s"
+DATE_FORMAT = "%H:%M:%S"
+
+
+def _to_level_value(lvl: int | str) -> int:
+    if isinstance(lvl, int):
+        return lvl
+    val = getattr(logging, str(lvl).upper(), None)
+    if val is None:
+        raise ValueError(f"Invalid log level: {lvl}")
+    return val
+
+
+def _ensure_file_handler(
+    logger: logging.Logger,
+    filename: str,
+    *,
+    level: int,
+    formatter: Optional[logging.Formatter],
+) -> None:
+    """Attach a FileHandler to `logger` for `filename` if it doesn't already exist."""
+    abspath = os.path.abspath(filename)
+
+    # Avoid duplicates
+    for h in logger.handlers:
+        if isinstance(h, logging.FileHandler) and getattr(h, "baseFilename", None) == abspath:
+            return
+
+    # Ensure directory exists
+    dirname = os.path.dirname(abspath)
+    if dirname:
+        os.makedirs(dirname, exist_ok=True)
+
+    fh = logging.FileHandler(abspath, encoding="utf-8")
+    fh.setLevel(level)
+    if formatter is not None:
+        fh.setFormatter(formatter)
+    else:
+        fh.setFormatter(logging.Formatter(DEFAULT_FORMAT, DATE_FORMAT))
+
+    logger.addHandler(fh)
+
+
+def setup(
+    level: int | str = "INFO",
+    *,
+    console: bool = True,
+    color: bool | Dict[str, Any] = True,
+    propagate: bool = False,
+    disable_existing_loggers: bool = False,
+    capture_warnings: bool = False,
+    submodule_levels: Optional[dict[str, int | str]] = None,
+    extra_handlers: Optional[list[logging.Handler]] = None,
+    formatter: Optional[logging.Formatter] = None,
+    apply_to: Optional[list[str]] = None,
+    files: Optional[str | dict[str, str]] = None,
+) -> None:
+    """Configures logging for the `agentlightning` logger hierarchy.
+
+    This function provides a one-stop setup utility for configuring the
+    `agentlightning` root logger and optionally its submodules or external
+    loggers. It supports console logging, colored rich output, per-submodule
+    log levels, and optional handler/formatter injection.
+
+    The setup is intentionally isolated: it does not modify the global root
+    logger or loggers belonging to other libraries unless explicitly directed
+    via `apply_to`.
+
+    Args:
+        level:
+            Logging level for the base `agentlightning` logger. Accepts either
+            an integer (e.g., `logging.DEBUG`) or a string level name
+            (e.g., `"INFO"`). Defaults to `"INFO"`.
+        console:
+            Whether to attach a console handler to the logger. Defaults to
+            `True`.
+        color:
+            Enables rich-formatted output using `RichHandler` when `True`
+            or a configuration dict. If `False`, a plain text formatter is
+            used instead. Defaults to `True`.
+        propagate:
+            Whether `agentlightning` logs should propagate to ancestor
+            loggers. Defaults to `False`.
+        disable_existing_loggers:
+            Passed to `logging.config.dictConfig`. If `True`, disables all
+            existing configured loggers before applying this configuration.
+            Defaults to `False`.
+        capture_warnings:
+            If `True`, redirects Python `warnings` emitted via the `warnings`
+            module into the logging system. Defaults to `False`.
+        submodule_levels:
+            Mapping of submodule logger names to logging levels. If a specified
+            submodule level is more verbose than the base level, a warning is emitted.
+        extra_handlers:
+            A list of user-provided handlers to attach to the `agentlightning` logger.
+            Handlers are added idempotently; duplicates are not reattached.
+        formatter:
+            A formatter to apply to any handler under `agentlightning` that does not
+            already have one assigned. Useful for customizing output without overwriting
+            formatters on custom handlers.
+        apply_to:
+            A list of additional logger names to configure identically to
+            `agentlightning` base logger. Their handlers are replaced with copies of the base
+            handlers, and propagation is disabled to avoid duplicate log emission.
+        files:
+            If a string, attach a FileHandler to the base `agentlightning` logger.
+            If a dict, for each `(logger_name, filename)` pair, attach a FileHandler
+            directly to that logger.
+            Each file handler should use the logger's effective level at creation.
+
+    Notes:
+        * On Windows, this function forces UTF-8 mode in the console to prevent
+          issues with rich output or special characters.
+        * Submodule loggers can generate records below the handler's emission
+          threshold. Whether such records appear depends on both the logger's
+          level and the handler's level.
+        * `apply_to` loggers inherit the same handlers but do not propagate
+          upward, yielding isolated, consistent behavior.
+
+    Examples:
+        Basic setup:
+
+        >>> setup()
+
+        Enabling debug mode with no color:
+
+        >>> setup(level="DEBUG", color=False)
+
+        Overriding specific submodule levels:
+
+        >>> setup(submodule_levels={"agentlightning.io": "DEBUG"})
+
+        Attaching an additional file handler:
+
+        >>> fh = logging.FileHandler("app.log")
+        >>> setup(extra_handlers=[fh])
+    """
     # Ensure UTF-8 encoding on Windows consoles
-    # Note: This change does not fully represent support for execution under the windown system.
+    # Note: This change does not fully represent support for execution under the windows system.
     # It only fixes console printing issues caused by special characters.
     # TODO: More comprehensive Windows support may be needed in the future.
     if platform.system() == "Windows":
         os.environ["PYTHONUTF8"] = "1"
 
-    logger = logging.getLogger(name)
-    logger.handlers.clear()  # clear existing handlers
+    base_logger = setup_module(
+        level,
+        name="agentlightning",
+        console=console,
+        color=color,
+        propagate=propagate,
+        disable_existing_loggers=disable_existing_loggers,
+    )
 
-    # log to stdout
-    handler = logging.StreamHandler()
-    handler.setLevel(level)
-    formatter = logging.Formatter("%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s)   %(message)s")
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-    logger.setLevel(level)
-    logger.propagate = False  # prevent double logging
-    return logger
+    base_level_value = base_logger.level
+
+    # Apply user-provided formatter (only to handlers without one,
+    # so we don't clobber custom extra_handlers)
+    if formatter is not None:
+        for h in base_logger.handlers:
+            if h.formatter is None:
+                h.setFormatter(formatter)
+
+    # Attach user-provided handler(s) if any, idempotently
+    if extra_handlers:
+        for h in extra_handlers:
+            if h not in base_logger.handlers:
+                base_logger.addHandler(h)
+
+    # Per-submodule levels
+    if submodule_levels:
+        for name, lvl in submodule_levels.items():
+            sub_level = _to_level_value(lvl)
+
+            # Emit a warning if submodule level is lower (more verbose) than the global/base level
+            if sub_level < base_level_value:
+                base_logger.warning(
+                    "Submodule logger '%s' level %s (%s) is more verbose than base "
+                    "logger level %s (%s). Records below the base level may still be "
+                    "filtered out by handlers depending on their own levels.",
+                    name,
+                    lvl,
+                    sub_level,
+                    logging.getLevelName(base_level_value),
+                    base_level_value,
+                )
+
+            # The logger will *create* records down to the logger's level, but a handler
+            # with a higher level will still drop anything below its own threshold.
+            # Effective emission is gated by both: record.level >= logger.level AND handler.level.
+            logging.getLogger(name).setLevel(lvl)
+
+    # Attach file handlers if requested
+    if files is not None:
+        if isinstance(files, str):
+            # Single file for the entire `agentlightning` hierarchy.
+            _ensure_file_handler(
+                logger=base_logger,
+                filename=files,
+                level=base_level_value,
+                formatter=formatter,
+            )
+        else:
+            # Per-logger files
+            for logger_name, filename in files.items():
+                lg = logging.getLogger(logger_name)
+                # Use the logger's *effective* level at creation time
+                effective_level = lg.getEffectiveLevel()
+                _ensure_file_handler(
+                    logger=lg,
+                    filename=filename,
+                    level=effective_level,
+                    formatter=formatter,
+                )
+
+    # Optionally apply the same handler setup to other loggers outside this module
+    if apply_to:
+        for name in apply_to:
+            lg = logging.getLogger(name)
+            # This removes any existing handlers so we don't duplicate output
+            # and ensures these loggers share exactly the same handlers as base_logger.
+            lg.handlers.clear()
+            for h in base_logger.handlers:
+                lg.addHandler(h)
+            lg.setLevel(base_logger.level)
+            # We've attached handlers directly to these loggers; if propagate
+            # stayed True, records would bubble up to ancestor loggers and could be
+            # emitted twice (here and on the parent/root). Setting False isolates them.
+            lg.propagate = False
+
+    # Optionally capture warnings
+    if capture_warnings:
+        logging.captureWarnings(True)
+
+
+def setup_module(
+    level: int | str = "INFO",
+    *,
+    name: str = "agentlightning",
+    console: bool = True,
+    color: bool | Dict[str, Any] = True,
+    propagate: bool = False,
+    disable_existing_loggers: bool = False,
+) -> logging.Logger:
+    """Initializes and returns the base logger for `agentlightning`.
+
+    This function constructs and applies a `dictConfig` configuration for the
+    logger hierarchy rooted at `name`. It supports either rich console
+    formatting (via `RichHandler`) or plain text formatting, based on the
+    `color` argument.
+
+    Unlike [`setup_logging`][agentlightning.setup_logging], this function configures only a single logger namespace
+    and does not attach extra handlers or submodule levels. It is primarily used
+    internally by [`setup_logging`][agentlightning.setup_logging] but is also suitable for direct integration in
+    custom logging workflows.
+    """
+    root_cfg: Dict[str, Any] = {
+        "version": 1,
+        "disable_existing_loggers": disable_existing_loggers,
+        "loggers": {
+            name: {
+                "handlers": [],
+                "level": level,
+                "propagate": propagate,
+            }
+        },
+        "handlers": {},
+        "formatters": {},
+    }
+
+    # Choose formatter / handler definition
+    if color is not False and console:
+        # Console must be true to display colored outputs
+        if isinstance(color, dict):
+            rich_handler_config = color
+        else:
+            rich_handler_config: Dict[str, Any] = {
+                "rich_tracebacks": False,
+                "markup": False,
+                "show_time": True,
+                "show_path": True,
+            }
+
+            if not _has_width():
+                # e.g., in a CI environment.
+                rich_handler_config["console"] = Console(width=200)
+
+        root_cfg["handlers"]["console"] = {
+            "class": "rich.logging.RichHandler",
+            "level": level,
+            **rich_handler_config,
+        }
+        # RichHandler manages its own style; keep formatter None
+    else:
+        fmt_name = "plain"
+        root_cfg["formatters"][fmt_name] = {
+            "format": DEFAULT_FORMAT,
+            "datefmt": DATE_FORMAT,
+        }
+
+        if console:
+            root_cfg["handlers"]["console"] = {
+                "class": "logging.StreamHandler",
+                "level": level,
+                "formatter": fmt_name,
+            }
+
+    # Attach selected handlers to agentlightning
+    handler_names = list(root_cfg["handlers"].keys())
+    root_cfg["loggers"][name]["handlers"] = handler_names
+
+    # Apply dictConfig (this resets the logger handlers)
+    dictConfig(root_cfg)
+
+    return logging.getLogger(name)
+
+
+def _has_width() -> bool:
+    """Automatically determine whether the terminal has a width."""
+    return sys.stdout.isatty()

agentlightning/runner/agent.py

@@ -11,8 +11,21 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import threading
 import time
-from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
+from contextlib import suppress
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypeVar,
+    cast,
+)
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -30,6 +43,7 @@ from agentlightning.types import (
     RolloutRawResult,
     Span,
 )
+from agentlightning.utils.system_snapshot import system_snapshot
 
 if TYPE_CHECKING:
     from agentlightning.execution.events import ExecutionEvent
@@ -52,7 +66,14 @@ class LitAgentRunner(Runner[T_task]):
         worker_id: Identifier for the active worker process, if any.
     """
 
-    def __init__(self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
+    def __init__(
+        self,
+        tracer: Tracer,
+        max_rollouts: Optional[int] = None,
+        poll_interval: float = 5.0,
+        heartbeat_interval: float = 10.0,
+        heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
+    ) -> None:
         """Initialize the agent runner.
 
         Args:
@@ -60,11 +81,16 @@ class LitAgentRunner(Runner[T_task]):
             max_rollouts: Optional cap on iterations processed by
                 [`iter`][agentlightning.LitAgentRunner.iter].
             poll_interval: Seconds to wait between store polls when no work is available.
+            heartbeat_interval: Seconds to wait between sending heartbeats to the store.
+            heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
+                "asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
         """
         super().__init__()
         self._tracer = tracer
         self._max_rollouts = max_rollouts
         self._poll_interval = poll_interval
+        self._heartbeat_interval = heartbeat_interval
+        self._heartbeat_launch_mode = heartbeat_launch_mode
 
         # Set later
         self._agent: Optional[LitAgent[T_task]] = None
@@ -304,6 +330,67 @@ class LitAgentRunner(Runner[T_task]):
 
         return trace_spans
 
+    async def _emit_heartbeat(self, store: LightningStore) -> None:
+        """Send a heartbeat tick to the store."""
+        worker_id = self.get_worker_id()
+
+        try:
+            await store.update_worker(worker_id, system_snapshot())
+        except asyncio.CancelledError:
+            # bypass the exception
+            raise
+        except Exception:
+            logger.exception("%s Unable to update worker heartbeat.", self._log_prefix())
+
+    def _start_heartbeat_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
+        """Start a background heartbeat loop and return an async stopper."""
+
+        if self._heartbeat_interval <= 0:
+            return None
+
+        if self.worker_id is None:
+            logger.warning("%s Cannot start heartbeat loop without worker_id.", self._log_prefix())
+            return None
+
+        if self._heartbeat_launch_mode == "asyncio":
+            stop_event = asyncio.Event()
+
+            async def heartbeat_loop() -> None:
+                while not stop_event.is_set():
+                    await self._emit_heartbeat(store)
+                    with suppress(asyncio.TimeoutError):
+                        await asyncio.wait_for(stop_event.wait(), timeout=self._heartbeat_interval)
+
+            task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
+
+            async def stop() -> None:
+                stop_event.set()
+                with suppress(asyncio.CancelledError):
+                    await task
+
+            return stop
+
+        if self._heartbeat_launch_mode == "thread":
+            stop_evt = threading.Event()
+
+            def thread_worker() -> None:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                while not stop_evt.is_set():
+                    loop.run_until_complete(self._emit_heartbeat(store))
+                    stop_evt.wait(self._heartbeat_interval)
+
+            thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
+            thread.start()
+
+            async def stop() -> None:
+                stop_evt.set()
+                await asyncio.to_thread(thread.join)
+
+            return stop
+
+        raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
+
     async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
         """Sleep until the next poll interval, with optional event-based interruption.
 
@@ -450,39 +537,49 @@ class LitAgentRunner(Runner[T_task]):
         logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
         store = self.get_store()
 
-        while not (event is not None and event.is_set()) and (
-            self._max_rollouts is None or num_tasks_processed < self._max_rollouts
-        ):
-            # Retrieve the next rollout
-            next_rollout: Optional[Rollout] = None
-            while not (event is not None and event.is_set()):
-                logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
-                next_rollout = await store.dequeue_rollout()
+        stop_heartbeat = self._start_heartbeat_loop(store)
+
+        try:
+            while not (event is not None and event.is_set()) and (
+                self._max_rollouts is None or num_tasks_processed < self._max_rollouts
+            ):
+                # Retrieve the next rollout
+                next_rollout: Optional[Rollout] = None
+                while not (event is not None and event.is_set()):
+                    logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
+                    next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
+                    if next_rollout is None:
+                        logger.debug(
+                            f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."
+                        )
+                        await self._sleep_until_next_poll(event)
+                    else:
+                        break
+
                 if next_rollout is None:
-                    logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
-                    await self._sleep_until_next_poll(event)
-                else:
-                    break
+                    return
 
-            if next_rollout is None:
-                return
+                try:
+                    # Claim the rollout but updating the current worker id
+                    await store.update_attempt(
+                        next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
+                    )
+                except Exception:
+                    # This exception could happen if the rollout is dequeued and the other end died for some reason
+                    logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
+                    continue
 
-            try:
-                # Claim the rollout but updating the current worker id
-                await store.update_attempt(
-                    next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
-                )
-            except Exception:
-                # This exception could happen if the rollout is dequeued and the other end died for some reason
-                logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
-                continue
+                # Execute the step
+                await self._step_impl(next_rollout)
 
-            # Execute the step
-            await self._step_impl(next_rollout)
-
-            num_tasks_processed += 1
-            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
-                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}")
+                num_tasks_processed += 1
+                if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                    logger.info(
+                        f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}"
+                    )
+        finally:
+            if stop_heartbeat is not None:
+                await stop_heartbeat()
 
         logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
 
@@ -526,6 +623,12 @@ class LitAgentRunner(Runner[T_task]):
             resources_id = None
 
         attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        # Register the attempt as running by the current worker
+        await self.get_store().update_attempt(
+            attempted_rollout.rollout_id,
+            attempted_rollout.attempt.attempt_id,
+            worker_id=self.get_worker_id(),
+        )
         rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
 
         completed_rollout = await store.get_rollout_by_id(rollout_id)

agentlightning/store/base.py

@@ -17,6 +17,7 @@ from agentlightning.types import (
     RolloutStatus,
     Span,
     TaskInput,
+    Worker,
 )
 
 
@@ -85,6 +86,7 @@ class LightningStore:
     Unless stated otherwise, missing identifiers should result in a `ValueError`.
     """
 
+    @property
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
@@ -167,7 +169,7 @@ class LightningStore:
         """
         raise NotImplementedError()
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """Claim the oldest queued rollout and transition it to `preparing`.
 
         This function do not block.
@@ -180,6 +182,8 @@ class LightningStore:
           the number of attempts already registered for the rollout plus one.
         * Return an [`AttemptedRollout`][agentlightning.AttemptedRollout] snapshot so the
           runner knows both rollout metadata and the attempt identifier.
+        * Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
+          (e.g., `last_dequeue_time`) when `worker_id` is provided.
 
         Returns:
             The next attempt to execute, or `None` when no eligible rollouts are queued.
@@ -527,6 +531,12 @@ class LightningStore:
         Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],
         parameters also default to the sentinel [`UNSET`][agentlightning.store.base.UNSET].
 
+        If `worker_id` is present, the worker status will be updated following the rules:
+
+        1. If attempt status is "succeeded" or "failed", the corresponding worker status will be set to "idle".
+        2. If attempt status is "unresponsive" or "timeout", the corresponding worker status will be set to "unknown".
+        3. Otherwise, the worker status will be set to "busy".
+
         Args:
             rollout_id: Identifier of the rollout whose attempt will be updated.
             attempt_id: Attempt identifier or `"latest"` as a convenience.
@@ -543,3 +553,46 @@ class LightningStore:
             ValueError: Implementations must raise when the rollout or attempt is unknown.
         """
         raise NotImplementedError()
+
+    async def query_workers(
+        self,
+    ) -> List[Worker]:
+        """Query all workers in the system.
+
+        Returns:
+            A list of all workers.
+        """
+        raise NotImplementedError()
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        """Retrieve a single worker by identifier.
+
+        Args:
+            worker_id: Identifier of the worker.
+
+        Returns:
+            The worker record if it exists, otherwise `None`.
+
+        Raises:
+            NotImplementedError: Subclasses must implement lookup semantics.
+        """
+        raise NotImplementedError()
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        """Record a heartbeat for `worker_id` and refresh telemetry.
+
+        Implementations must treat this API as heartbeat-only: it should snapshot
+        the latest stats when provided, stamp `last_heartbeat_time` with the
+        current wall clock, and rely on other store mutations (`dequeue_rollout`,
+        `update_attempt`, etc.) to drive the worker's busy/idle status,
+        assignment, and activity timestamps.
+
+        Args:
+            worker_id: Identifier of the worker to update.
+            heartbeat_stats: Replacement worker heartbeat statistics (non-null when provided).
+        """
+        raise NotImplementedError()

agentlightning/store/client_server.py

@@ -8,12 +8,10 @@ import os
 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
 
 import aiohttp
-import uvicorn
 from fastapi import APIRouter, Body, Depends, FastAPI, HTTPException
 from fastapi import Query as FastAPIQuery
 from fastapi import Request, Response
@@ -34,11 +32,15 @@ from agentlightning.types import (
     RolloutStatus,
     Span,
     TaskInput,
+    Worker,
+    WorkerStatus,
 )
+from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncher, PythonServerLauncherArgs
 
 from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
 
-logger = logging.getLogger(__name__)
+server_logger = logging.getLogger("agentlightning.store.server")
+client_logger = logging.getLogger("agentlightning.store.client")
 
 API_V1_PREFIX = "/v1"
 API_AGL_PREFIX = "/agl"
@@ -62,6 +64,10 @@ class RolloutRequest(BaseModel):
     metadata: Optional[Dict[str, Any]] = None
 
 
+class DequeueRolloutRequest(BaseModel):
+    worker_id: Optional[str] = None
+
+
 class QueryRolloutsRequest(BaseModel):
     status_in: Optional[List[RolloutStatus]] = Field(FastAPIQuery(default=None))
     rollout_id_in: Optional[List[str]] = Field(FastAPIQuery(default=None))
@@ -106,6 +112,10 @@ class UpdateAttemptRequest(BaseModel):
     metadata: Optional[Dict[str, Any]] = None
 
 
+class UpdateWorkerRequest(BaseModel):
+    heartbeat_stats: Optional[Dict[str, Any]] = None
+
+
 class QueryAttemptsRequest(BaseModel):
     # Pagination
     limit: int = -1
@@ -148,6 +158,19 @@ class QuerySpansRequest(BaseModel):
     sort_order: Literal["asc", "desc"] = "asc"
 
 
+class QueryWorkersRequest(BaseModel):
+    status_in: Optional[List[WorkerStatus]] = Field(FastAPIQuery(default=None))
+    worker_id_contains: Optional[str] = None
+    # Pagination
+    limit: int = -1
+    offset: int = 0
+    # Sorting
+    sort_by: Optional[str] = None
+    sort_order: Literal["asc", "desc"] = "asc"
+    # Filtering logic
+    filter_logic: Literal["and", "or"] = "and"
+
+
 def _apply_filters_sort_paginate(
     items: List[T],
     filters: Dict[str, Any],
@@ -257,36 +280,70 @@ class LightningStoreServer(LightningStore):
 
     `agl store` is a convenient CLI to start a store server.
 
+    When the server is executed in a subprocess, the store will discover itself having a different PID
+    and automatically delegate to an HTTP client instead of using the local store.
+    This ensures one single copy of the store will be shared across all processes.
+
     Args:
         store: The underlying store to delegate operations to.
         host: The hostname or IP address to bind the server to.
         port: The TCP port to listen on.
         cors_allow_origins: A list of CORS origins to allow. Use '*' to allow all origins.
+        launch_mode: The launch mode to use for the server. Defaults to "thread",
+            which runs the server in a separate thread.
+        launcher_args: The arguments to use for the server launcher.
+            It's not allowed to set `host`, `port`, `launch_mode` together with `launcher_args`.
     """
 
     def __init__(
         self,
         store: LightningStore,
-        host: str,
-        port: int,
+        host: str | None = None,
+        port: int | None = None,
         cors_allow_origins: Sequence[str] | str | None = None,
+        launch_mode: LaunchMode = "thread",
+        launcher_args: PythonServerLauncherArgs | None = None,
     ):
         super().__init__()
         self.store = store
-        self._lock = threading.Lock()
-        self.host = host
-        self.port = port
-        self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
+
+        if launcher_args is not None:
+            if host is not None or port is not None or launch_mode != "thread":
+                raise ValueError("host, port, and launch_mode cannot be set when launcher_args is provided.")
+            self.launcher_args = launcher_args
+        else:
+            if port is None:
+                server_logger.warning("No port provided, using default port 4747.")
+                port = 4747
+            self.launcher_args = PythonServerLauncherArgs(
+                host=host,
+                port=port,
+                launch_mode=launch_mode,
+                healthcheck_url=API_V1_AGL_PREFIX + "/health",
+            )
+
+        store_capabilities = self.store.capabilities
+        if not store_capabilities["async_safe"]:
+            raise ValueError("The store is not async-safe. Please use another store for the server.")
+        if self.launcher_args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
+            raise ValueError(
+                "The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
+            )
+        if self.launcher_args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
+            server_logger.warning(
+                "The store is not thread-safe. Please be careful when using the store server and the underlying store in different threads."
+            )
+
         self.app: FastAPI | None = FastAPI(title="LightningStore Server")
+        self.server_launcher = PythonServerLauncher(
+            app=self.app,
+            args=self.launcher_args,
+        )
+
+        self._lock: threading.Lock = threading.Lock()
+        self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
         self._apply_cors()
         self._setup_routes()
-        self._uvicorn_config: uvicorn.Config | None = uvicorn.Config(
-            self.app, host="0.0.0.0", port=self.port, log_level="error"
-        )
-        self._uvicorn_server: uvicorn.Server | None = uvicorn.Server(self._uvicorn_config)
-
-        self._serving_thread: Optional[threading.Thread] = None
-        self._server_start_exception: Optional[BaseException] = None
 
         # Process-awareness:
         # LightningStoreServer holds a plain Python object (self.store) in one process
@@ -298,30 +355,30 @@ class LightningStoreServer(LightningStore):
         self._owner_pid = os.getpid()
         self._client: Optional[LightningStoreClient] = None
 
+    @property
     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
+        return LightningStoreCapabilities(
+            async_safe=True,
+            thread_safe=True,
+            zero_copy=True,
+        )
 
     def __getstate__(self):
         """
         Control pickling to prevent server state from being sent to subprocesses.
 
         When LightningStoreServer is pickled (e.g., passed to a subprocess), we only
-        serialize the underlying store and connection details. The FastAPI app and
-        uvicorn server are excluded as they should not be transferred between processes.
+        serialize the underlying store and connection details. The client instance
+        and process-awareness state are excluded as they should not be transferred between processes.
 
         The subprocess should create its own server instance if needed.
         """
+        # server-launcher is needed for the host/port address are propagated to the subprocess
         return {
-            "store": self.store,
-            "host": self.host,
-            "port": self.port,
+            "launcher_args": self.launcher_args,
+            "server_launcher": self.server_launcher,
             "_owner_pid": self._owner_pid,
-            "_cors_allow_origins": self._cors_allow_origins,
         }
 
     def __setstate__(self, state: Dict[str, Any]):
@@ -331,10 +388,13 @@ class LightningStoreServer(LightningStore):
         Note: This creates a new server instance without FastAPI/uvicorn initialized.
         Call __init__() pattern or create a new LightningStoreServer if you need
         a fully functional server in the subprocess.
+        The unpickled server will also have no app and store attributes,
+        this is to make sure there is only one copy of the server in the whole system.
         """
-        self.store = state["store"]
-        self.host = state["host"]
-        self.port = state["port"]
+        self.app = None
+        self.store = None
+        self.launcher_args = state["launcher_args"]
+        self.server_launcher = state["server_launcher"]
         self._owner_pid = state["_owner_pid"]
         self._cors_allow_origins = state.get("_cors_allow_origins")
         self._client = None
@@ -380,151 +440,38 @@ class LightningStoreServer(LightningStore):
 
     @property
     def endpoint(self) -> str:
-        return f"http://{self.host}:{self.port}"
+        """Endpoint is the address that the client will use to connect to the server."""
+        return self.server_launcher.access_endpoint
 
     async def start(self):
         """Starts the FastAPI server in the background.
 
         You need to call this method in the same process as the server was created in.
         """
-        assert self._uvicorn_server is not None
-        logger.info(f"Starting server at {self.endpoint}")
+        server_logger.info(
+            f"Serving the lightning store at {self.server_launcher.endpoint}, accessible at {self.server_launcher.access_endpoint}"
+        )
 
-        uvicorn_server = self._uvicorn_server
-        self._server_start_exception = None
-
-        def run_server_forever():
-            try:
-                asyncio.run(uvicorn_server.serve())
-            except (SystemExit, Exception) as exc:
-                logger.debug("LightningStore server thread exiting due to %s", exc, exc_info=exc)
-                self._server_start_exception = exc
-
-        serving_thread = threading.Thread(target=run_server_forever, daemon=True)
-        self._serving_thread = serving_thread
-        serving_thread.start()
-
-        # Wait for uvicorn to report that it has started before pinging /health.
-        start_deadline = time.time() + 10
-        while time.time() < start_deadline:
-            if uvicorn_server.started:
-                break
-            if self._server_start_exception is not None or not serving_thread.is_alive():
-                self._handle_failed_start()
-                raise RuntimeError(self._format_start_failure_reason())
-            await asyncio.sleep(0.05)
-        else:
-            self._handle_failed_start()
-            raise RuntimeError("Server failed to start within the 10 seconds.")
-
-        # Wait for /health to be available once uvicorn reports started.
-        if not await self._server_health_check():
-            self._handle_failed_start()
-            raise RuntimeError("Server failed to start within the 10 seconds.")
-
-        # If startup failed (e.g. port already in use), uvicorn never flips `started`
-        # and the worker thread stops immediately. Guard against latching on to a
-        # different process that happened to satisfy the health check.
-        if not uvicorn_server.started or not serving_thread.is_alive() or self._server_start_exception is not None:
-            self._handle_failed_start()
-            failure_reason = self._format_start_failure_reason()
-            raise RuntimeError(failure_reason)
-
-    async def _server_health_check(self) -> bool:
-        """Checks if the server is healthy."""
-        current_time = time.time()
-        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:
-                        if response.status == 200:
-                            return True
-            await asyncio.sleep(0.1)
-        return False
-
-    def _handle_failed_start(self) -> None:
-        """Clean up thread state when startup fails."""
-        if self._uvicorn_server is not None:
-            self._uvicorn_server.should_exit = True
-        if self._serving_thread is not None:
-            # Thread already exited in most failure scenarios; join defensively.
-            self._serving_thread.join(timeout=0.1)
-            self._serving_thread = None
-
-    def _format_start_failure_reason(self) -> str:
-        base_message = f"LightningStore server failed to start on {self.endpoint}."
-        if isinstance(self._server_start_exception, SystemExit):
-            return f"{base_message} Another process may already be using this port."
-        if isinstance(self._server_start_exception, OSError):
-            return f"{base_message} {self._server_start_exception.strerror}."
-        if self._server_start_exception is not None:
-            return f"{base_message} Reason: {self._server_start_exception}."
-        return f"{base_message} Another process may already be using this port."
+        start_time = time.time()
+        await self.server_launcher.start()
+        end_time = time.time()
+        server_logger.info(f"Lightning store server started in {end_time - start_time:.2f} seconds")
 
     async def run_forever(self):
-        """Runs the FastAPI server indefinitely.
-
-        You need to call this method in the same process as the server was created in.
-        """
-        assert self._uvicorn_server is not None
-        uvicorn_server = self._uvicorn_server
-
-        async def _wait_till_healthy():
-            health = await self._server_health_check()
-            if not health:
-                raise RuntimeError("Server did not become healthy within the 10 seconds.")
-            logger.info("Store server is online at %s", self.endpoint)
-
-        async def _serve_capture():
-            try:
-                await uvicorn_server.serve()
-            except KeyboardInterrupt:
-                raise
-            except (SystemExit, Exception) as exc:
-                logger.debug("LightningStore server serve() raised %s", exc, exc_info=exc)
-                self._server_start_exception = exc
-                raise RuntimeError("LightningStore server failed to serve") from exc
-
-        # We run _wait_till_healthy and self._uvicorn_server.serve in parallel
-        # until one of them raises an exception.
-        try:
-            await asyncio.gather(_wait_till_healthy(), _serve_capture())
-        except BaseException as exc:
-            if isinstance(exc, KeyboardInterrupt):
-                raise
-            startup_failed = not uvicorn_server.started or isinstance(
-                self._server_start_exception, (SystemExit, OSError)
-            )
-            if startup_failed:
-                self._handle_failed_start()
-                raise RuntimeError(self._format_start_failure_reason())
-            raise
+        """Runs the FastAPI server indefinitely."""
+        server_logger.info(
+            f"Running the lightning store server at {self.server_launcher.endpoint}, accessible at {self.server_launcher.access_endpoint}"
+        )
+        await self.server_launcher.run_forever()
 
     async def stop(self):
         """Gracefully stops the running FastAPI server.
 
         You need to call this method in the same process as the server was created in.
         """
-        assert self._uvicorn_server is not None
-        if self._uvicorn_server.started:
-            logger.info("Stopping server...")
-            self._uvicorn_server.should_exit = True
-            if self._serving_thread is not None:
-                self._serving_thread.join(timeout=10)
-            self._serving_thread = None
-            logger.info("Server stopped.")
-
-    def _backend(self) -> LightningStore:
-        """Returns the object to delegate to in *this* process.
-
-        - In the owner process: delegate to the in-process store.
-        - In a different process: delegate to a HTTP client talking to the server.
-        """
-        if os.getpid() == self._owner_pid:
-            return self.store
-        if self._client is None:
-            self._client = LightningStoreClient(self.endpoint)
-        return self._client
+        server_logger.info("Stopping the lightning store server...")
+        await self.server_launcher.stop()
+        server_logger.info("Lightning store server stopped.")
 
     def _setup_routes(self):
         """Set up FastAPI routes for all store operations."""
@@ -549,7 +496,7 @@ class LightningStoreServer(LightningStore):
             except Exception as exc:
                 # decide whether to convert this into your 400 JSONResponse
                 if request.url.path.startswith(API_V1_AGL_PREFIX):
-                    logger.exception("Unhandled application error", exc_info=exc)
+                    server_logger.exception("Unhandled application error", exc_info=exc)
                     payload = {
                         "detail": "Internal server error",
                         "error_type": type(exc).__name__,
@@ -576,7 +523,7 @@ class LightningStoreServer(LightningStore):
                 client_address = "unknown"
             else:
                 client_address = f"{client.host}:{client.port}"
-            logger.info(
+            server_logger.debug(
                 f"{client_address} - "
                 f'"{request.method} {request.url.path} HTTP/{request.scope["http_version"]}" '
                 f"{response.status_code} in {duration:.2f} ms"
@@ -600,8 +547,11 @@ class LightningStoreServer(LightningStore):
             )
 
         @api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
-        async def dequeue_rollout():  # pyright: ignore[reportUnusedFunction]
-            return await self.dequeue_rollout()
+        async def dequeue_rollout(  # pyright: ignore[reportUnusedFunction]
+            request: DequeueRolloutRequest | None = Body(None),
+        ):
+            worker_id = request.worker_id if request else None
+            return await self.dequeue_rollout(worker_id=worker_id)
 
         @api.post(API_AGL_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
         async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
@@ -641,9 +591,11 @@ class LightningStoreServer(LightningStore):
         async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
             return await self.get_rollout_by_id(rollout_id)
 
-        def _get_mandatory_field_or_unset(request: BaseModel, field: str) -> Any:
+        def _get_mandatory_field_or_unset(request: BaseModel | None, field: str) -> Any:
             # If some fields are mandatory by the underlying store, but optional in the FastAPI,
             # we make sure it's set to non-null value or UNSET via this function.
+            if request is None:
+                return UNSET
             if field in request.model_fields_set:
                 value = getattr(request, field)
                 if value is None:
@@ -683,6 +635,39 @@ class LightningStoreServer(LightningStore):
                 metadata=_get_mandatory_field_or_unset(request, "metadata"),
             )
 
+        @api.get(API_AGL_PREFIX + "/workers", response_model=PaginatedResponse[Worker])
+        async def query_workers(params: QueryWorkersRequest = Depends()):  # pyright: ignore[reportUnusedFunction]
+            all_workers = await self.query_workers()
+
+            filters: Dict[str, Any] = {}
+            if params.status_in:
+                filters["status_in"] = params.status_in
+            if params.worker_id_contains is not None:
+                filters["worker_id_contains"] = params.worker_id_contains
+
+            return _apply_filters_sort_paginate(
+                all_workers,
+                filters,
+                params.filter_logic,
+                params.sort_by,
+                params.sort_order,
+                params.limit,
+                params.offset,
+            )
+
+        @api.get(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Optional[Worker])
+        async def get_worker(worker_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.get_worker_by_id(worker_id)
+
+        @api.post(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Worker)
+        async def update_worker(  # pyright: ignore[reportUnusedFunction]
+            worker_id: str, request: UpdateWorkerRequest | None = Body(None)
+        ):
+            return await self.update_worker(
+                worker_id=worker_id,
+                heartbeat_stats=_get_mandatory_field_or_unset(request, "heartbeat_stats"),
+            )
+
         @api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", response_model=PaginatedResponse[Attempt])
         async def query_attempts(  # pyright: ignore[reportUnusedFunction]
             rollout_id: str, params: QueryAttemptsRequest = Depends()
@@ -815,19 +800,19 @@ class LightningStoreServer(LightningStore):
 
         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)
+            server_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(
+            server_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)
+            server_logger.error("Dashboard index file not found at %s. Please build the dashboard first.", index_file)
             return
 
         # Mount the static files in dashboard/assets
@@ -844,20 +829,25 @@ class LightningStoreServer(LightningStore):
             # Let the frontend router handle it
             return FileResponse(index_file)
 
-        logger.info("Agent-lightning dashboard will be available at %s", self.endpoint)
+        server_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()
-        method = getattr(backend, method_name)
-        if backend is self.store:
+        """First decide what store to delegate to in *this* process, and then call the method on it.
+
+        - In the owner process: delegate to the in-process store.
+        - In a different process: delegate to a HTTP client talking to the server.
+        """
+        if os.getpid() == self._owner_pid:
             if method_name == "wait_for_rollouts":
                 # wait_for_rollouts can block for a long time; avoid holding the lock
                 # so other requests can make progress while we wait.
-                return await method(*args, **kwargs)
+                return await getattr(self.store, method_name)(*args, **kwargs)
             with self._lock:
-                return await method(*args, **kwargs)
-        return await method(*args, **kwargs)
+                return await getattr(self.store, method_name)(*args, **kwargs)
+        if self._client is None:
+            self._client = LightningStoreClient(self.endpoint)
+        return await getattr(self._client, method_name)(*args, **kwargs)
 
     async def start_rollout(
         self,
@@ -893,8 +883,8 @@ class LightningStoreServer(LightningStore):
             metadata,
         )
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        return await self._call_store_method("dequeue_rollout")
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
+        return await self._call_store_method("dequeue_rollout", worker_id)
 
     async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         return await self._call_store_method("start_attempt", rollout_id)
@@ -999,6 +989,23 @@ class LightningStoreServer(LightningStore):
             metadata,
         )
 
+    async def query_workers(self) -> List[Worker]:
+        return await self._call_store_method("query_workers")
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        return await self._call_store_method("get_worker_by_id", worker_id)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        return await self._call_store_method(
+            "update_worker",
+            worker_id,
+            heartbeat_stats,
+        )
+
 
 class LightningStoreClient(LightningStore):
     """HTTP client that talks to a remote LightningStoreServer.
@@ -1008,8 +1015,12 @@ class LightningStoreClient(LightningStore):
         retry_delays:
             Backoff schedule (seconds) used when the initial request fails for a
             non-application reason. Each entry is a retry attempt.
+            Setting to an empty sequence to disable retries.
         health_retry_delays:
             Delays between /health probes while waiting for the server to come back.
+            Setting to an empty sequence to disable health checks.
+        request_timeout: Timeout (seconds) for each request.
+        connection_timeout: Timeout (seconds) for establishing connection.
     """
 
     def __init__(
@@ -1018,19 +1029,26 @@ class LightningStoreClient(LightningStore):
         *,
         retry_delays: Sequence[float] = (1.0, 2.0, 5.0),
         health_retry_delays: Sequence[float] = (0.1, 0.2, 0.5),
+        request_timeout: float = 30.0,
+        connection_timeout: float = 5.0,
     ):
         self.server_address = server_address.rstrip("/") + API_V1_AGL_PREFIX
         self._sessions: Dict[int, aiohttp.ClientSession] = {}  # id(loop) -> ClientSession
-        self._lock = threading.RLock()
+        self._lock = threading.Lock()
 
         # retry config
         self._retry_delays = tuple(float(d) for d in retry_delays)
         self._health_retry_delays = tuple(float(d) for d in health_retry_delays)
 
+        # Timeouts
+        self._request_timeout = request_timeout
+        self._connection_timeout = connection_timeout
+
         # Store whether the dequeue was successful in history
         self._dequeue_was_successful: bool = False
         self._dequeue_first_unsuccessful: bool = True
 
+    @property
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
@@ -1049,6 +1067,8 @@ class LightningStoreClient(LightningStore):
             "server_address": self.server_address,
             "_retry_delays": self._retry_delays,
             "_health_retry_delays": self._health_retry_delays,
+            "_request_timeout": self._request_timeout,
+            "_connection_timeout": self._connection_timeout,
         }
 
     def __setstate__(self, state: Dict[str, Any]):
@@ -1059,9 +1079,11 @@ class LightningStoreClient(LightningStore):
         """
         self.server_address = state["server_address"]
         self._sessions = {}
-        self._lock = threading.RLock()
+        self._lock = threading.Lock()
         self._retry_delays = state["_retry_delays"]
         self._health_retry_delays = state["_health_retry_delays"]
+        self._request_timeout = state["_request_timeout"]
+        self._connection_timeout = state["_connection_timeout"]
         self._dequeue_was_successful = False
         self._dequeue_first_unsuccessful = True
 
@@ -1087,7 +1109,12 @@ class LightningStoreClient(LightningStore):
         with self._lock:
             sess = self._sessions.get(key)
             if sess is None or sess.closed:
-                timeout = aiohttp.ClientTimeout(total=30.0, connect=5.0, sock_connect=5.0, sock_read=30.0)
+                timeout = aiohttp.ClientTimeout(
+                    total=self._request_timeout,
+                    connect=self._connection_timeout,
+                    sock_connect=self._connection_timeout,
+                    sock_read=self._request_timeout,
+                )
                 sess = aiohttp.ClientSession(timeout=timeout)
                 self._sessions[key] = sess
         return sess
@@ -1097,20 +1124,24 @@ class LightningStoreClient(LightningStore):
         Probe the server's /health until it responds 200 or retries are exhausted.
         Returns True if healthy, False otherwise.
         """
-        logger.info(f"Waiting for server to be healthy at {self.server_address}/health")
+        if not self._health_retry_delays:
+            client_logger.info("No health retry delays configured; skipping health checks.")
+            return True
+
+        client_logger.info(f"Waiting for server to be healthy at {self.server_address}/health")
         for delay in [*self._health_retry_delays, 0.0]:
             try:
                 async with session.get(f"{self.server_address}/health") as r:
                     if r.status == 200:
-                        logger.info(f"Server is healthy at {self.server_address}/health")
+                        client_logger.info(f"Server is healthy at {self.server_address}/health")
                         return True
             except Exception:
                 # swallow and retry
                 if delay > 0.0:
-                    logger.warning(f"Server is not healthy yet. Retrying in {delay} seconds.")
+                    client_logger.warning(f"Server is not healthy yet. Retrying in {delay} seconds.")
             if delay > 0.0:
                 await asyncio.sleep(delay)
-        logger.error(
+        client_logger.error(
             f"Server is not healthy at {self.server_address}/health after {len(self._health_retry_delays)} retry attempts"
         )
         return False
@@ -1143,7 +1174,7 @@ class LightningStoreClient(LightningStore):
 
         for delay in attempts:
             if delay:
-                logger.info(f"Waiting {delay} seconds before retrying {method}: {path}")
+                client_logger.info(f"Waiting {delay} seconds before retrying {method}: {path}")
                 await asyncio.sleep(delay)
             try:
                 http_call = getattr(session, method)
@@ -1153,12 +1184,12 @@ class LightningStoreClient(LightningStore):
             except aiohttp.ClientResponseError as cre:
                 # Respect app-level 4xx as final
                 # 4xx => application issue; do not retry (except 408 which is transient)
-                logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
+                client_logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
                 if 400 <= cre.status < 500 and cre.status != 408:
                     raise
                 # 5xx and others will be retried below if they raise
                 last_exc = cre
-                logger.info(f"5xx and other status codes will be retried. Retrying the request {method}: {path}")
+                client_logger.info(f"5xx and other status codes will be retried. Retrying the request {method}: {path}")
                 # before next retry, ensure server is healthy
                 if not await self._wait_until_healthy(session):
                     break  # server is not healthy, do not retry
@@ -1169,9 +1200,9 @@ class LightningStoreClient(LightningStore):
                 asyncio.TimeoutError,
             ) as net_exc:
                 # Network/session issue: probe health before retrying
-                logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
+                client_logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
                 last_exc = net_exc
-                logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
+                client_logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
                 if not await self._wait_until_healthy(session):
                     break  # server is not healthy, do not retry
 
@@ -1242,7 +1273,7 @@ class LightningStoreClient(LightningStore):
         )
         return Rollout.model_validate(data)
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """
         Dequeue a rollout from the server queue.
 
@@ -1255,8 +1286,11 @@ class LightningStoreClient(LightningStore):
         """
         session = await self._get_session()
         url = f"{self.server_address}/queues/rollouts/dequeue"
+        request_kwargs: Dict[str, Any] = {}
+        if worker_id is not None:
+            request_kwargs["json"] = {"worker_id": worker_id}
         try:
-            async with session.post(url) as resp:
+            async with session.post(url, **request_kwargs) as resp:
                 resp.raise_for_status()
                 data = await resp.json()
                 self._dequeue_was_successful = True
@@ -1264,9 +1298,9 @@ class LightningStoreClient(LightningStore):
         except Exception as e:
             if self._dequeue_was_successful:
                 if self._dequeue_first_unsuccessful:
-                    logger.warning(f"dequeue_rollout failed with exception: {e}")
+                    client_logger.warning(f"dequeue_rollout failed with exception: {e}")
                     self._dequeue_first_unsuccessful = False
-            logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
+            client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
             # Else ignore the exception because the server is not ready yet
             return None
 
@@ -1320,7 +1354,9 @@ class LightningStoreClient(LightningStore):
             data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts/latest")
             return Attempt.model_validate(data) if data else None
         except Exception as e:
-            logger.error(f"get_latest_attempt failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
+            client_logger.error(
+                f"get_latest_attempt failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True
+            )
             return None
 
     async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
@@ -1344,7 +1380,9 @@ class LightningStoreClient(LightningStore):
             else:
                 return Rollout.model_validate(data)
         except Exception as e:
-            logger.error(f"get_rollout_by_id failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
+            client_logger.error(
+                f"get_rollout_by_id failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True
+            )
             return None
 
     async def query_resources(self) -> List[ResourcesUpdate]:
@@ -1385,7 +1423,7 @@ class LightningStoreClient(LightningStore):
             data = await self._request_json("get", f"/resources/{resources_id}")
             return ResourcesUpdate.model_validate(data) if data else None
         except Exception as e:
-            logger.error(
+            client_logger.error(
                 f"get_resources_by_id failed after all retries for resources_id={resources_id}: {e}", exc_info=True
             )
             return None
@@ -1405,7 +1443,7 @@ class LightningStoreClient(LightningStore):
             data = await self._request_json("get", "/resources/latest")
             return ResourcesUpdate.model_validate(data) if data else None
         except Exception as e:
-            logger.error(f"get_latest_resources failed after all retries: {e}", exc_info=True)
+            client_logger.error(f"get_latest_resources failed after all retries: {e}", exc_info=True)
             return None
 
     async def add_span(self, span: Span) -> Span:
@@ -1525,3 +1563,27 @@ class LightningStoreClient(LightningStore):
             json=payload,
         )
         return Attempt.model_validate(data)
+
+    async def query_workers(self) -> List[Worker]:
+        data = await self._request_json("get", "/workers")
+        items = data.get("items", [])
+        return [Worker.model_validate(item) for item in items]
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        data = await self._request_json("get", f"/workers/{worker_id}")
+        if data is None:
+            return None
+        return Worker.model_validate(data)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        payload: Dict[str, Any] = {}
+        if not isinstance(heartbeat_stats, Unset):
+            payload["heartbeat_stats"] = heartbeat_stats
+        json_payload = payload if payload else None
+
+        data = await self._request_json("post", f"/workers/{worker_id}", json=json_payload)
+        return Worker.model_validate(data)

agentlightning/store/memory.py

@@ -44,6 +44,7 @@ from agentlightning.types import (
     RolloutStatus,
     Span,
     TaskInput,
+    Worker,
 )
 
 from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset, is_finished, is_queuing
@@ -242,7 +243,50 @@ class InMemoryLightningStore(LightningStore):
 
         # Completion tracking for wait_for_rollouts (cross-loop safe)
         self._completion_events: Dict[str, threading.Event] = {}
+        # Worker tracking
+        self._workers: Dict[str, Worker] = {}
 
+        # Running rollouts cache, including preparing and running rollouts
+        self._running_rollout_ids: Set[str] = set()
+
+    def _get_or_create_worker(self, worker_id: str) -> Worker:
+        worker = self._workers.get(worker_id)
+        if worker is None:
+            worker = Worker(worker_id=worker_id)
+            self._workers[worker_id] = worker
+        return worker
+
+    def _sync_worker_with_attempt(self, attempt: Attempt) -> None:
+        worker_id = attempt.worker_id
+        if not worker_id:
+            return
+
+        worker = self._get_or_create_worker(worker_id)
+        now = time.time()
+
+        if attempt.status in ("succeeded", "failed"):
+            if worker.status != "idle":
+                worker.last_idle_time = now
+            worker.status = "idle"
+            worker.current_rollout_id = None
+            worker.current_attempt_id = None
+        elif attempt.status in ("timeout", "unresponsive"):
+            if worker.status != "unknown":
+                worker.last_idle_time = now
+            worker.status = "unknown"
+            worker.current_rollout_id = None
+            worker.current_attempt_id = None
+        else:
+            transitioned = worker.status != "busy" or worker.current_attempt_id != attempt.attempt_id
+            if transitioned:
+                worker.last_busy_time = now
+            worker.status = "busy"
+            worker.current_rollout_id = attempt.rollout_id
+            worker.current_attempt_id = attempt.attempt_id
+
+        Worker.model_validate(worker.model_dump())
+
+    @property
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
         return LightningStoreCapabilities(
@@ -281,6 +325,7 @@ class InMemoryLightningStore(LightningStore):
                 config=rollout_config,
                 metadata=rollout_metadata,
             )
+            self._running_rollout_ids.add(rollout.rollout_id)
 
             # Create the initial attempt
             attempt_id = _generate_attempt_id()
@@ -338,7 +383,7 @@ class InMemoryLightningStore(LightningStore):
             return rollout
 
     @_healthcheck_wrapper
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         """Retrieves the next task from the queue without blocking.
         Returns `None` if the queue is empty.
 
@@ -347,6 +392,11 @@ class InMemoryLightningStore(LightningStore):
         See [`LightningStore.dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] for semantics.
         """
         async with self._lock:
+            if worker_id is not None:
+                worker = self._get_or_create_worker(worker_id)
+                worker.last_dequeue_time = time.time()
+                worker.status = "idle"
+
             # Keep looking until we find a rollout that's still in queuing status
             # or the queue is empty
             while self._task_queue:
@@ -357,6 +407,7 @@ class InMemoryLightningStore(LightningStore):
                 if is_queuing(rollout):
                     # Update status to preparing
                     rollout.status = "preparing"
+                    self._running_rollout_ids.add(rollout.rollout_id)
 
                     # Create a new attempt (could be first attempt or retry)
                     attempt_id = _generate_attempt_id()
@@ -655,6 +706,7 @@ class InMemoryLightningStore(LightningStore):
         if current_attempt == latest_attempt:
             if rollout.status == "preparing":
                 rollout.status = "running"
+                self._running_rollout_ids.add(rollout.rollout_id)
             elif rollout.status in ["queuing", "requeuing"]:
                 try:
                     self._task_queue.remove(rollout)
@@ -663,6 +715,7 @@ class InMemoryLightningStore(LightningStore):
                         f"Trying to remove rollout {rollout.rollout_id} from the queue but it's not in the queue."
                     )
                 rollout.status = "running"
+                self._running_rollout_ids.add(rollout.rollout_id)
 
         return span
 
@@ -908,6 +961,12 @@ class InMemoryLightningStore(LightningStore):
         elif is_queuing(rollout) and rollout not in self._task_queue:
             self._task_queue.append(rollout)
 
+        # Updating running rollouts cache
+        if rollout.status in ["preparing", "running"]:
+            self._running_rollout_ids.add(rollout.rollout_id)
+        else:
+            self._running_rollout_ids.discard(rollout.rollout_id)
+
         # If the rollout is no longer in a queueing state, remove it from the queue.
         if not isinstance(status, Unset) and not is_queuing(rollout) and rollout in self._task_queue:
             try:
@@ -951,19 +1010,26 @@ class InMemoryLightningStore(LightningStore):
             if not attempt:
                 raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
 
+        worker_sync_required = False
+
         # Update fields if they are not UNSET
+        if not isinstance(worker_id, Unset):
+            attempt.worker_id = worker_id
+            worker_sync_required = worker_sync_required or bool(worker_id)
         if not isinstance(status, Unset):
             attempt.status = status
             # Also update end_time if the status indicates completion
             if status in ["failed", "succeeded"]:
                 attempt.end_time = time.time()
-        if not isinstance(worker_id, Unset):
-            attempt.worker_id = worker_id
+            worker_sync_required = worker_sync_required or bool(attempt.worker_id)
         if not isinstance(last_heartbeat_time, Unset):
             attempt.last_heartbeat_time = last_heartbeat_time
         if not isinstance(metadata, Unset):
             attempt.metadata = metadata
 
+        if worker_sync_required and attempt.worker_id:
+            self._sync_worker_with_attempt(attempt)
+
         # Re-validate the attempt to ensure legality
         Attempt.model_validate(attempt.model_dump())
 
@@ -981,12 +1047,40 @@ class InMemoryLightningStore(LightningStore):
 
         return attempt
 
+    @_healthcheck_wrapper
+    async def query_workers(self) -> List[Worker]:
+        """Return the current snapshot of all workers."""
+        async with self._lock:
+            return list(self._workers.values())
+
+    @_healthcheck_wrapper
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        async with self._lock:
+            return self._workers.get(worker_id)
+
+    @_healthcheck_wrapper
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        """Create or update a worker entry."""
+        async with self._lock:
+            worker = self._get_or_create_worker(worker_id)
+            if not isinstance(heartbeat_stats, Unset):
+                worker.heartbeat_stats = dict(heartbeat_stats)
+            worker.last_heartbeat_time = time.time()
+
+            Worker.model_validate(worker.model_dump())
+            return worker
+
     async def _healthcheck(self) -> None:
         """Perform healthcheck against all running rollouts in the store."""
         async with self._lock:
             running_rollouts: List[AttemptedRollout] = []
-            for rollout in self._rollouts.values():
-                if rollout.status in ["preparing", "running"]:
+            for rollout_id in self._running_rollout_ids:
+                rollout = self._rollouts.get(rollout_id)
+                if rollout is not None and rollout.status in ["preparing", "running"]:
                     all_attempts = self._attempts.get(rollout.rollout_id, [])
                     if not all_attempts:
                         # The rollout is running but has no attempts, this should not happen

agentlightning/store/threading.py

@@ -18,6 +18,7 @@ from agentlightning.types import (
     RolloutStatus,
     Span,
     TaskInput,
+    Worker,
 )
 
 from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
@@ -35,9 +36,10 @@ class LightningStoreThreaded(LightningStore):
         self.store = store
         self._lock = threading.Lock()
 
+    @property
     def capabilities(self) -> LightningStoreCapabilities:
         """Return the capabilities of the store."""
-        capabilities = self.store.capabilities()
+        capabilities = self.store.capabilities
         return {
             **capabilities,
             "async_safe": True,
@@ -66,9 +68,9 @@ class LightningStoreThreaded(LightningStore):
         with self._lock:
             return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
         with self._lock:
-            return await self.store.dequeue_rollout()
+            return await self.store.dequeue_rollout(worker_id=worker_id)
 
     async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
         with self._lock:
@@ -180,3 +182,22 @@ class LightningStoreThreaded(LightningStore):
                 last_heartbeat_time=last_heartbeat_time,
                 metadata=metadata,
             )
+
+    async def query_workers(self) -> List[Worker]:
+        with self._lock:
+            return await self.store.query_workers()
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        with self._lock:
+            return await self.store.get_worker_by_id(worker_id)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        with self._lock:
+            return await self.store.update_worker(
+                worker_id=worker_id,
+                heartbeat_stats=heartbeat_stats,
+            )

agentlightning/types/core.py

@@ -49,6 +49,8 @@ __all__ = [
     "Attempt",
     "AttemptedRollout",
     "Hook",
+    "Worker",
+    "WorkerStatus",
 ]
 
 T_co = TypeVar("T_co", covariant=True)
@@ -200,6 +202,32 @@ class AttemptedRollout(Rollout):
         return self
 
 
+WorkerStatus = Literal["idle", "busy", "unknown"]
+
+
+class Worker(BaseModel):
+    """Worker information. This is actually the same as Runner info."""
+
+    worker_id: str
+    """The ID of the worker."""
+    status: WorkerStatus = "unknown"
+    """The status of the worker."""
+    heartbeat_stats: Optional[Dict[str, Any]] = None
+    """Statistics about the worker's heartbeat."""
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported the stats."""
+    last_dequeue_time: Optional[float] = None
+    """The last time when the worker has tried to dequeue a rollout."""
+    last_busy_time: Optional[float] = None
+    """The last time when the worker has started an attempt and became busy."""
+    last_idle_time: Optional[float] = None
+    """The last time when the worker has triggered the end of an attempt and became idle."""
+    current_rollout_id: Optional[str] = None
+    """The ID of the current rollout that the worker is processing."""
+    current_attempt_id: Optional[str] = None
+    """The ID of the current attempt that the worker is processing."""
+
+
 TaskInput = Any
 """Task input type. Accepts arbitrary payloads."""
 

agentlightning/utils/server_launcher.py

@@ -53,6 +53,8 @@ class PythonServerLauncherArgs:
     """
     log_level: int = logging.INFO
     """The log level to use."""
+    access_log: bool = False
+    """Whether to turn on access logs."""
     startup_timeout: float = 60.0
     """The timeout to wait for the server to start up."""
     kill_unhealthy_server: bool = True
@@ -156,7 +158,9 @@ async def run_uvicorn_asyncio(
 
         if not uvicorn_server.started:
             # 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
+            raise RuntimeError(
+                f"Server did not start up within {time.time() - start_time:.2f} seconds."
+            ) from server_start_exception
 
         logger.info(f"Server started up in {time.time() - start_time:.2f} seconds.")
 
@@ -608,6 +612,13 @@ class PythonServerLauncher:
         self._host: Optional[str] = self.args.host
         self._port: Optional[int] = self.args.port
         self._access_host: Optional[str] = self.args.access_host
+        self.initialize()
+
+    def initialize(self):
+        # ensure the host/port/access_host are set
+        self._ensure_host()
+        self._ensure_port()
+        self._ensure_access_host()
 
         # uvicorn (in-proc asyncio)
         self._uvicorn_server: Optional[uvicorn.Server] = None
@@ -626,6 +637,26 @@ class PythonServerLauncher:
         # is_running flag
         self._is_running: bool = False
 
+    def __getstate__(self):
+        """Control pickling to prevent server state from being sent to subprocesses."""
+        return {
+            "app": self.app,
+            "args": self.args,
+            "serve_context": self.serve_context,
+            "_host": self._host,
+            "_port": self._port,
+            "_access_host": self._access_host,
+        }
+
+    def __setstate__(self, state: Dict[str, Any]):
+        self.app = state["app"]
+        self.args = state["args"]
+        self.serve_context = state["serve_context"]
+        self._host = state["_host"]
+        self._port = state["_port"]
+        self._access_host = state["_access_host"]
+        self.initialize()
+
     @property
     def endpoint(self) -> str:
         """Return the externally advertised host:port pair regardless of accessibility."""
@@ -744,17 +775,18 @@ class PythonServerLauncher:
         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()
+        if self._access_host is None:
+            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:
-                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
+                self._access_host = self.args.access_host
+        return self._access_host  # type: ignore
 
     def _create_uvicorn_server(self) -> uvicorn.Server:
         config = uvicorn.Config(
@@ -762,6 +794,7 @@ class PythonServerLauncher:
             host=self._ensure_host(),
             port=self._ensure_port(),
             log_level=self.args.log_level,
+            access_log=self.args.access_log,
             loop="asyncio",
         )
         return uvicorn.Server(config)
@@ -834,17 +867,19 @@ class PythonServerLauncher:
             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.")
+                raise RuntimeError("Threaded server failed to start and is not alive. No error event was received.")
+            logger.error(
+                "Threaded server failed to start and sends no event. This should not happen. Shutting down server."
+            )
             await self._stop_uvicorn_thread()
-            return
+            raise RuntimeError("Threaded server failed to start and sends no event. This should not happen.")
 
         if evt.kind == "error":
             logger.error("Threaded server failed to start (%s): %s\n%s", evt.exc_type, evt.message, evt.traceback)
             await asyncio.to_thread(self._thread.join, self.args.thread_join_timeout)
             if self._thread.is_alive():
-                raise RuntimeError(evt.message or "Threaded server failed to start and refused to shut down.")
+                logger.error("Threaded server failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
         else:
             logger.info("Threaded server started successfully.")
             self._is_running = True
@@ -893,7 +928,7 @@ class PythonServerLauncher:
                 "workers": int(self.args.n_workers),
                 "worker_class": "uvicorn_worker.UvicornWorker",
                 "loglevel": logging.getLevelName(self.args.log_level).lower(),
-                "accesslog": None,
+                "accesslog": "-" if self.args.access_log else None,
                 "errorlog": "-",
                 "preload_app": True,
                 "graceful_timeout": int(
@@ -939,11 +974,12 @@ class PythonServerLauncher:
             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.")
+                raise RuntimeError("Server process failed to start and is not alive. No error event was received.")
+            logger.error(
+                "Server process failed to start and sends no event. This should not happen. Shutting down server."
+            )
             await self._stop_serving_process()
-            return
+            raise RuntimeError("Server process failed to start and sends no event. This should not happen.")
 
         if evt.kind == "error":
             logger.error(
@@ -955,7 +991,8 @@ class PythonServerLauncher:
             )
             await asyncio.to_thread(self._proc.join, self.args.process_join_timeout)
             if self._proc.is_alive():
-                raise RuntimeError(evt.message or "Server process failed to start and refused to shut down.")
+                logger.error("Server process failed to start and refused to shut down.")
+            raise RuntimeError(evt.message)
         else:
             logger.info("Subprocess server started successfully.")
             self._is_running = True

agentlightning/utils/system_snapshot.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import platform
+import socket
+from contextlib import suppress
+from datetime import datetime
+from typing import Any, Dict, List, cast
+
+import psutil
+from gpustat import GPUStat, GPUStatCollection
+
+
+def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
+    # CPU
+    cpu = {
+        "cpu_name": platform.processor(),
+        "cpu_cores": psutil.cpu_count(logical=False),
+        "cpu_threads": psutil.cpu_count(logical=True),
+        "cpu_usage_pct": psutil.cpu_percent(0.05),
+    }
+
+    # Memory
+    vm = psutil.virtual_memory()
+    mem = {
+        "mem_used_gb": round(vm.used / (2**30), 2),
+        "mem_total_gb": round(vm.total / (2**30), 2),
+        "mem_pct": vm.percent,
+    }
+
+    # Disk
+    du = psutil.disk_usage("/")
+    disk = {
+        "disk_used_gb": round(du.used / (2**30), 2),
+        "disk_total_gb": round(du.total / (2**30), 2),
+        "disk_pct": du.percent,
+    }
+
+    # GPU
+    gpus: List[Dict[str, Any]] = []
+    with suppress(Exception):
+        for g in GPUStatCollection.new_query().gpus:  # type: ignore
+            g = cast(GPUStat, g)
+            gpus.append(
+                {
+                    "gpu": g.name,  # type: ignore
+                    "util_pct": g.utilization,
+                    "mem_used_mb": g.memory_used,
+                    "mem_total_mb": g.memory_total,
+                    "temp_c": g.temperature,
+                }
+            )
+
+    # Network
+    net = psutil.net_io_counters()
+    netinfo = {
+        "bytes_sent_mb": round(net.bytes_sent / (2**20), 2),
+        "bytes_recv_mb": round(net.bytes_recv / (2**20), 2),
+    }
+
+    # OS / meta
+    return {
+        "timestamp": datetime.now().isoformat(timespec="seconds"),
+        "host": socket.gethostname(),
+        "os": platform.platform(),
+        **cpu,
+        **mem,
+        **disk,
+        **netinfo,
+        **({"gpus": gpus} if include_gpu else {}),
+    }

agentlightning/verl/daemon.py

@@ -18,13 +18,13 @@ from flask import Flask, Response, abort, request
 from tensordict import TensorDict
 from verl import DataProto
 
-from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, configure_logger
+from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, setup_logging
 from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store.base import LightningStore
 from agentlightning.types import Rollout, RolloutConfig, Task
 
-configure_logger()
+setup_logging()
 
 __all__ = [
     "AgentModeDaemon",

agentlightning/verl/trainer.py

@@ -12,6 +12,7 @@ from typing import Dict, Tuple
 
 import numpy as np
 import torch
+import verl
 from codetiming import Timer
 from omegaconf import OmegaConf
 from tqdm import tqdm
@@ -403,14 +404,20 @@ class AgentLightningTrainer(RayPPOTrainer):
         assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
         if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
             raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
+        verl_version = verl.__version__
+        if verl_version == "0.5.0":
+            # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
+            # However, it is possible that verl updates the naming and causes incompatibility.
+            # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
+            model = "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:])
+        else:
+            # For other versions (e.g., 0.6.0), we use the full path to the model.
+            model = self.config.actor_rollout_ref.model.path
         self.agent_mode_daemon = AgentModeDaemon(
             self.config.agentlightning.port,
             self.config.actor_rollout_ref.rollout.n,
             train_information={
-                # Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
-                # However, it is possible that verl updates the naming and causes incompatibility.
-                # Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
-                "model": "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:]),
+                "model": model,
                 "temperature": self.config.actor_rollout_ref.rollout.temperature,
             },
             tokenizer=self.tokenizer,

dashboard/src/Router.tsx

@@ -6,6 +6,7 @@ import { ResourcesPage } from './pages/Resources.page';
 import { RolloutsPage } from './pages/Rollouts.page';
 import { SettingsPage } from './pages/Settings.page';
 import { TracesPage } from './pages/Traces.page';
+import { WorkersPage } from './pages/Workers.page';
 
 const router = createBrowserRouter([
   {
@@ -28,6 +29,10 @@ const router = createBrowserRouter([
         path: 'traces',
         element: <TracesPage />,
       },
+      {
+        path: 'runners',
+        element: <WorkersPage />,
+      },
       {
         path: 'settings',
         element: <SettingsPage />,

dashboard/src/components/AppDrawer.component.tsx

@@ -21,7 +21,7 @@ import {
 import { useGetSpansQuery } from '@/features/rollouts';
 import { closeDrawer, openDrawer, selectDrawerContent, selectDrawerIsOpen } from '@/features/ui/drawer';
 import { useAppDispatch, useAppSelector } from '@/store/hooks';
-import type { Attempt, AttemptStatus, Rollout, RolloutStatus, Span } from '@/types';
+import type { Attempt, AttemptStatus, Rollout, RolloutStatus, Span, Worker } from '@/types';
 import { formatStatusLabel } from '@/utils/format';
 import { TracesTable, type TracesTableRecord } from './TracesTable.component';
 
@@ -50,6 +50,12 @@ const SPAN_STATUS_COLORS: Record<Span['status']['status_code'], string> = {
   ERROR: 'red',
 };
 
+const WORKER_STATUS_COLORS: Record<Worker['status'], string> = {
+  busy: 'orange',
+  idle: 'teal',
+  unknown: 'gray',
+};
+
 const TRACES_SORT_FIELD_MAP: Record<string, string> = {
   name: 'name',
   traceId: 'trace_id',
@@ -408,6 +414,60 @@ function RolloutTracesDrawerBody({ rollout, attempt, onShowRollout, onShowSpanDe
   );
 }
 
+type WorkerDrawerTitleProps = {
+  worker: Worker;
+};
+
+function WorkerDrawerTitle({ worker }: WorkerDrawerTitleProps) {
+  const badgeColor = WORKER_STATUS_COLORS[worker.status] ?? 'gray';
+  return (
+    <Stack gap={3}>
+      <Group gap={6} align='center'>
+        <Text fw={600}>{worker.workerId}</Text>
+        <CopyButton value={worker.workerId}>
+          {({ copied, copy }) => (
+            <Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
+              <ActionIcon
+                aria-label={`Copy worker ID ${worker.workerId}`}
+                variant='subtle'
+                color={copied ? 'teal' : 'gray'}
+                size='sm'
+                onClick={(event) => {
+                  event.stopPropagation();
+                  copy();
+                }}
+              >
+                {copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
+              </ActionIcon>
+            </Tooltip>
+          )}
+        </CopyButton>
+        <Badge size='sm' variant='light' color={badgeColor}>
+          {formatStatusLabel(worker.status)}
+        </Badge>
+      </Group>
+      <Group gap='xl'>
+        <Group gap={4}>
+          <Text size='sm' c='dimmed' fw={500}>
+            Rollout
+          </Text>
+          <Text size='sm' c='dimmed'>
+            {worker.currentRolloutId ?? '—'}
+          </Text>
+        </Group>
+        <Group gap={4}>
+          <Text size='sm' c='dimmed' fw={500}>
+            Attempt
+          </Text>
+          <Text size='sm' c='dimmed'>
+            {worker.currentAttemptId ?? '—'}
+          </Text>
+        </Group>
+      </Group>
+    </Stack>
+  );
+}
+
 export function AppDrawerContainer() {
   const dispatch = useAppDispatch();
   const isOpen = useAppSelector(selectDrawerIsOpen);
@@ -428,6 +488,13 @@ export function AppDrawerContainer() {
       return null;
     }
 
+    if (content.type === 'worker-detail') {
+      const { worker } = content;
+      const title = <WorkerDrawerTitle worker={worker} />;
+      const body = <JsonEditor value={worker} />;
+      return { title, body };
+    }
+
     if (content.type === 'trace-detail') {
       const { span } = content;
       const title = <TraceDrawerTitle span={span} />;

dashboard/src/components/TracesTable.component.tsx

@@ -22,6 +22,7 @@ const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
 
 const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
   name: { minWidth: 12.5, priority: 0 },
+  sequenceId: { fixedWidth: 6, priority: 1 },
   spanId: { fixedWidth: 14, priority: 1 },
   traceId: { fixedWidth: 24, priority: 3 },
   parentId: { fixedWidth: 12, priority: 2 },
@@ -86,6 +87,12 @@ function createTracesColumns({
         </Text>
       ),
     },
+    {
+      accessor: 'sequenceId',
+      title: 'Seq.',
+      sortable: true,
+      render: ({ sequenceId }) => <Text size='sm'>{sequenceId}</Text>,
+    },
     {
       accessor: 'traceId',
       title: 'Trace ID',

dashboard/src/components/WorkersTable.component.tsx

@@ -0,0 +1,362 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { useCallback, useEffect, useMemo } from 'react';
+import { IconCheck, IconCopy, IconInfoCircle, IconRefresh } 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 type { Worker } from '@/types';
+import { getErrorDescriptor } from '@/utils/error';
+import { formatDateTime, formatRelativeTime, formatStatusLabel } 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> = {
+  workerId: { fixedWidth: 12, priority: 0 },
+  status: { fixedWidth: 6, priority: 1 },
+  currentRolloutId: { fixedWidth: 14, priority: 3 },
+  currentAttemptId: { fixedWidth: 14, priority: 3 },
+  lastHeartbeatTime: { fixedWidth: 10, priority: 2 },
+  lastBusyTime: { fixedWidth: 10, priority: 3 },
+  lastIdleTime: { fixedWidth: 10, priority: 3 },
+  lastDequeueTime: { fixedWidth: 10, priority: 1 },
+  actions: { fixedWidth: 5, priority: 0 },
+};
+
+export type WorkersTableRecord = Worker & {
+  timestamps: Record<
+    'lastHeartbeatTime' | 'lastBusyTime' | 'lastIdleTime' | 'lastDequeueTime',
+    { absolute: string; relative: string }
+  >;
+};
+
+const buildTimestampMeta = (value: Worker['lastHeartbeatTime']) => ({
+  absolute: formatDateTime(value),
+  relative: formatRelativeTime(value),
+});
+
+function buildWorkerRecord(worker: Worker): WorkersTableRecord {
+  return {
+    ...worker,
+    timestamps: {
+      lastHeartbeatTime: buildTimestampMeta(worker.lastHeartbeatTime),
+      lastBusyTime: buildTimestampMeta(worker.lastBusyTime),
+      lastIdleTime: buildTimestampMeta(worker.lastIdleTime),
+      lastDequeueTime: buildTimestampMeta(worker.lastDequeueTime),
+    },
+  };
+}
+
+type WorkersColumnsOptions = {
+  onShowDetails: (worker: Worker) => void;
+};
+
+const STATUS_COLORS: Record<Worker['status'], string> = {
+  busy: 'orange',
+  idle: 'teal',
+  unknown: 'gray',
+};
+
+function createWorkersColumns({ onShowDetails }: WorkersColumnsOptions): DataTableColumn<WorkersTableRecord>[] {
+  return [
+    {
+      accessor: 'workerId',
+      title: 'Runner ID',
+      sortable: true,
+      render: ({ workerId }) => (
+        <Group gap={2} wrap='nowrap'>
+          <Text fw={500} size='sm'>
+            {workerId}
+          </Text>
+          <CopyButton value={workerId}>
+            {({ copied, copy }) => (
+              <Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
+                <ActionIcon
+                  aria-label={`Copy worker ID ${workerId}`}
+                  variant='subtle'
+                  color={copied ? 'teal' : 'gray'}
+                  size='sm'
+                  onClick={(event) => {
+                    event.stopPropagation();
+                    copy();
+                  }}
+                >
+                  {copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
+                </ActionIcon>
+              </Tooltip>
+            )}
+          </CopyButton>
+        </Group>
+      ),
+    },
+    {
+      accessor: 'status',
+      title: 'Status',
+      sortable: true,
+      render: ({ status }) => {
+        const color = STATUS_COLORS[status] ?? 'gray';
+        return (
+          <Badge size='sm' variant='light' color={color} radius='sm'>
+            {formatStatusLabel(status)}
+          </Badge>
+        );
+      },
+    },
+    {
+      accessor: 'currentRolloutId',
+      title: 'Current Rollout',
+      sortable: true,
+      render: ({ currentRolloutId }) => <Text size='sm'>{currentRolloutId ?? '—'}</Text>,
+    },
+    {
+      accessor: 'currentAttemptId',
+      title: 'Current Attempt',
+      sortable: true,
+      render: ({ currentAttemptId }) => <Text size='sm'>{currentAttemptId ?? '—'}</Text>,
+    },
+    {
+      accessor: 'lastHeartbeatTime',
+      title: 'Heartbeat',
+      sortable: true,
+      render: ({ timestamps }) => (
+        <Stack gap={0} justify='center'>
+          <Text size='sm'>{timestamps.lastHeartbeatTime.relative}</Text>
+          {timestamps.lastHeartbeatTime.absolute !== '—' && (
+            <Text size='xs' c='dimmed'>
+              {timestamps.lastHeartbeatTime.absolute}
+            </Text>
+          )}
+        </Stack>
+      ),
+    },
+    {
+      accessor: 'lastBusyTime',
+      title: 'Last Busy',
+      sortable: true,
+      render: ({ timestamps }) => (
+        <Stack gap={0} justify='center'>
+          <Text size='sm'>{timestamps.lastBusyTime.relative}</Text>
+          {timestamps.lastBusyTime.absolute !== '—' && (
+            <Text size='xs' c='dimmed'>
+              {timestamps.lastBusyTime.absolute}
+            </Text>
+          )}
+        </Stack>
+      ),
+    },
+    {
+      accessor: 'lastIdleTime',
+      title: 'Last Idle',
+      sortable: true,
+      render: ({ timestamps }) => (
+        <Stack gap={0} justify='center'>
+          <Text size='sm'>{timestamps.lastIdleTime.relative}</Text>
+          {timestamps.lastIdleTime.absolute !== '—' && (
+            <Text size='xs' c='dimmed'>
+              {timestamps.lastIdleTime.absolute}
+            </Text>
+          )}
+        </Stack>
+      ),
+    },
+    {
+      accessor: 'lastDequeueTime',
+      title: 'Last Dequeue',
+      sortable: true,
+      render: ({ timestamps }) => (
+        <Stack gap={0} justify='center'>
+          <Text size='sm'>{timestamps.lastDequeueTime.relative}</Text>
+          {timestamps.lastDequeueTime.absolute !== '—' && (
+            <Text size='xs' c='dimmed'>
+              {timestamps.lastDequeueTime.absolute}
+            </Text>
+          )}
+        </Stack>
+      ),
+    },
+    {
+      accessor: 'actions',
+      title: 'Actions',
+      textAlign: 'left',
+      render: (record) => (
+        <Tooltip label='Show runner detail' withArrow disabled={!onShowDetails}>
+          <ActionIcon
+            aria-label='Show runner detail'
+            variant='subtle'
+            color='gray'
+            onClick={(event) => {
+              event.stopPropagation();
+              onShowDetails(record);
+            }}
+          >
+            <IconInfoCircle size={16} />
+          </ActionIcon>
+        </Tooltip>
+      ),
+    },
+  ];
+}
+
+export type WorkersTableProps = {
+  workers: Worker[] | undefined;
+  totalRecords: number;
+  isFetching: boolean;
+  isError: boolean;
+  error: unknown;
+  searchTerm: string;
+  sort: { column: string; direction: 'asc' | 'desc' };
+  page: number;
+  recordsPerPage: number;
+  onSortStatusChange: (status: DataTableSortStatus<WorkersTableRecord>) => void;
+  onPageChange: (page: number) => void;
+  onRecordsPerPageChange: (value: number) => void;
+  onResetFilters: () => void;
+  onRefetch: () => void;
+  onShowDetails: (worker: Worker) => void;
+  recordsPerPageOptions?: number[];
+};
+
+export function WorkersTable({
+  workers,
+  totalRecords,
+  isFetching,
+  isError,
+  error,
+  searchTerm,
+  sort,
+  page,
+  recordsPerPage,
+  onSortStatusChange,
+  onPageChange,
+  onRecordsPerPageChange,
+  onResetFilters,
+  onRefetch,
+  onShowDetails,
+  recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
+}: WorkersTableProps) {
+  const { ref: tableContainerRef, width: containerWidth } = useElementSize();
+  const { width: viewportWidth } = useViewportSize();
+
+  const layoutAwareContainerWidth = useMemo(
+    () => getLayoutAwareWidth(containerWidth, viewportWidth),
+    [containerWidth, viewportWidth],
+  );
+
+  const workerRecords = useMemo<WorkersTableRecord[]>(() => {
+    if (!workers) {
+      return [];
+    }
+    return workers.map((worker) => buildWorkerRecord(worker));
+  }, [workers]);
+
+  const columns = useMemo(() => createWorkersColumns({ onShowDetails }), [onShowDetails]);
+  const responsiveColumns = useMemo(
+    () => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
+    [columns, layoutAwareContainerWidth],
+  );
+
+  const totalPages = useMemo(
+    () => Math.max(1, Math.ceil(Math.max(0, totalRecords) / Math.max(1, recordsPerPage))),
+    [recordsPerPage, totalRecords],
+  );
+
+  useEffect(() => {
+    if (page > totalPages) {
+      onPageChange(totalPages);
+    }
+  }, [onPageChange, page, totalPages]);
+
+  const hasActiveFilters = searchTerm.trim().length > 0;
+
+  const sortStatus: DataTableSortStatus<WorkersTableRecord> = {
+    columnAccessor: sort.column,
+    direction: sort.direction,
+  };
+
+  const handleSortStatusChange = useCallback(
+    (status: DataTableSortStatus<WorkersTableRecord>) => {
+      onSortStatusChange(status);
+    },
+    [onSortStatusChange],
+  );
+
+  const errorDescriptor = isError ? getErrorDescriptor(error) : null;
+  const errorMessage = isError
+    ? `Workers are temporarily unavailable${errorDescriptor ? ` (${errorDescriptor})` : ''}.`
+    : 'Workers are temporarily unavailable.';
+
+  const emptyState = (
+    <Stack gap='sm' align='center' py='lg'>
+      {isError ? (
+        <>
+          <Text fw={600} size='sm'>
+            {errorMessage}
+          </Text>
+          <Text size='sm' c='dimmed' ta='center'>
+            Use the retry button to try again, or adjust the search to broaden the results.
+          </Text>
+          <Group gap='xs'>
+            <Button size='xs' variant='light' color='gray' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
+              Retry
+            </Button>
+            {hasActiveFilters ? (
+              <Button size='xs' variant='subtle' onClick={onResetFilters}>
+                Clear filters
+              </Button>
+            ) : null}
+          </Group>
+        </>
+      ) : (
+        <>
+          <Text fw={600} size='sm'>
+            No workers found
+          </Text>
+          <Text size='sm' c='dimmed' ta='center'>
+            {hasActiveFilters
+              ? 'Try adjusting the search to see more results.'
+              : 'Try refreshing to fetch the latest worker status.'}
+          </Text>
+          <Group gap='xs'>
+            <Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
+              Refresh
+            </Button>
+            {hasActiveFilters ? (
+              <Button size='xs' variant='subtle' onClick={onResetFilters}>
+                Clear filters
+              </Button>
+            ) : null}
+          </Group>
+        </>
+      )}
+    </Stack>
+  );
+
+  return (
+    <Box ref={tableContainerRef}>
+      <DataTable<WorkersTableRecord>
+        classNames={{ root: 'workers-table' }}
+        withTableBorder
+        withColumnBorders
+        highlightOnHover
+        verticalAlign='center'
+        minHeight={workerRecords.length === 0 ? 400 : undefined}
+        idAccessor='workerId'
+        records={workerRecords}
+        columns={responsiveColumns}
+        totalRecords={totalRecords}
+        recordsPerPage={recordsPerPage}
+        page={page}
+        onPageChange={onPageChange}
+        onRecordsPerPageChange={onRecordsPerPageChange}
+        recordsPerPageOptions={recordsPerPageOptions}
+        sortStatus={sortStatus}
+        onSortStatusChange={handleSortStatusChange}
+        fetching={isFetching}
+        loaderSize='sm'
+        emptyState={workerRecords.length === 0 ? emptyState : undefined}
+      />
+    </Box>
+  );
+}

dashboard/src/components/WorkersTable.story.tsx

@@ -0,0 +1,154 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { useMemo, useState } from 'react';
+import type { Meta, StoryObj } from '@storybook/react';
+import { IconSearch } from '@tabler/icons-react';
+import { Box, Stack, TextInput, Title } from '@mantine/core';
+import type { Worker } from '@/types';
+import { WorkersTable } from './WorkersTable.component';
+
+const meta: Meta<typeof WorkersTable> = {
+  title: 'Components/WorkersTable',
+  component: WorkersTable,
+  parameters: {
+    layout: 'fullscreen',
+  },
+};
+
+export default meta;
+
+type Story = StoryObj<typeof WorkersTable>;
+
+const now = Math.floor(Date.now() / 1000);
+
+const sampleWorkers: Worker[] = [
+  {
+    workerId: 'worker-east',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 2, gpuUtilization: 0.82 },
+    lastHeartbeatTime: now - 20,
+    lastDequeueTime: now - 60,
+    lastBusyTime: now - 120,
+    lastIdleTime: now - 600,
+    currentRolloutId: 'ro-story-001',
+    currentAttemptId: 'at-story-010',
+  },
+  {
+    workerId: 'worker-west',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 1 },
+    lastHeartbeatTime: now - 45,
+    lastDequeueTime: now - 300,
+    lastBusyTime: now - 200,
+    lastIdleTime: now - 4800,
+    currentRolloutId: 'ro-story-003',
+    currentAttemptId: 'at-story-033',
+  },
+  {
+    workerId: 'worker-north',
+    status: 'idle',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 90,
+    lastDequeueTime: now - 3600,
+    lastBusyTime: now - 5400,
+    lastIdleTime: now - 5400,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-south',
+    status: 'idle',
+    heartbeatStats: null,
+    lastHeartbeatTime: now - 900,
+    lastDequeueTime: now - 7200,
+    lastBusyTime: now - 8600,
+    lastIdleTime: now - 8600,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-standby',
+    status: 'unknown',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 15,
+    lastDequeueTime: now - 4000,
+    lastBusyTime: null,
+    lastIdleTime: null,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+];
+
+type WorkersTableStoryWrapperProps = {
+  maxWidth: number;
+  initialSort?: { column: string; direction: 'asc' | 'desc' };
+};
+
+function WorkersTableStoryWrapper({ maxWidth, initialSort }: WorkersTableStoryWrapperProps) {
+  const [searchTerm, setSearchTerm] = useState('');
+  const [page, setPage] = useState(1);
+  const [recordsPerPage, setRecordsPerPage] = useState(5);
+  const [sort, setSort] = useState<{ column: string; direction: 'asc' | 'desc' }>(
+    () => initialSort ?? { column: 'lastHeartbeatTime', direction: 'desc' },
+  );
+
+  const filteredWorkers = useMemo(() => {
+    const normalized = searchTerm.trim().toLowerCase();
+    if (!normalized) {
+      return sampleWorkers;
+    }
+    return sampleWorkers.filter((worker) => worker.workerId.toLowerCase().includes(normalized));
+  }, [searchTerm]);
+
+  return (
+    <Stack gap='md' p='lg'>
+      <Title order={2}>Workers ({maxWidth}px max width)</Title>
+      <TextInput
+        placeholder='Search'
+        leftSection={<IconSearch size={16} />}
+        value={searchTerm}
+        onChange={(event) => setSearchTerm(event.currentTarget.value)}
+        w='100%'
+        style={{ maxWidth: 360 }}
+      />
+      <Box style={{ maxWidth }}>
+        <WorkersTable
+          workers={filteredWorkers}
+          totalRecords={filteredWorkers.length}
+          isFetching={false}
+          isError={false}
+          error={null}
+          searchTerm={searchTerm}
+          sort={sort}
+          page={page}
+          recordsPerPage={recordsPerPage}
+          onSortStatusChange={(status) => {
+            return setSort({ column: status.columnAccessor as string, direction: status.direction });
+          }}
+          onPageChange={setPage}
+          onRecordsPerPageChange={setRecordsPerPage}
+          onResetFilters={() => {
+            setSearchTerm('');
+            setPage(1);
+          }}
+          onRefetch={() => {}}
+          onShowDetails={() => {}}
+        />
+      </Box>
+    </Stack>
+  );
+}
+
+export const Wide: Story = {
+  render: () => <WorkersTableStoryWrapper maxWidth={1600} />,
+};
+
+export const Narrow: Story = {
+  render: () => <WorkersTableStoryWrapper maxWidth={780} />,
+};
+
+export const SortedByCurrentRollout: Story = {
+  render: () => (
+    <WorkersTableStoryWrapper maxWidth={1200} initialSort={{ column: 'currentRolloutId', direction: 'asc' }} />
+  ),
+};

dashboard/src/features/rollouts/api.ts

@@ -13,6 +13,8 @@ import type {
   RolloutStatus,
   Span,
   Timestamp,
+  Worker,
+  WorkerStatus,
 } from '../../types';
 
 const rawBaseQuery = fetchBaseQuery({
@@ -122,6 +124,21 @@ const normalizeResources = (value: unknown): Resources => {
   };
 };
 
+const normalizeWorker = (value: unknown): Worker => {
+  const camelized = camelCaseKeys(value) as Worker;
+  return {
+    workerId: camelized.workerId,
+    status: camelized.status,
+    heartbeatStats: camelized.heartbeatStats ?? null,
+    lastHeartbeatTime: camelized.lastHeartbeatTime ?? null,
+    lastDequeueTime: camelized.lastDequeueTime ?? null,
+    lastBusyTime: camelized.lastBusyTime ?? null,
+    lastIdleTime: camelized.lastIdleTime ?? null,
+    currentRolloutId: camelized.currentRolloutId ?? null,
+    currentAttemptId: camelized.currentAttemptId ?? null,
+  };
+};
+
 const normalizePaginatedResponse = <T>(value: unknown, normalizer: (item: unknown) => T): PaginatedResponse<T> => {
   if (!value || typeof value !== 'object') {
     throw new Error('Expected paginated response payload');
@@ -183,6 +200,15 @@ export type GetResourcesQueryArgs = {
   resourcesIdContains?: string | null;
 };
 
+export type GetWorkersQueryArgs = {
+  limit: number;
+  offset: number;
+  sortBy?: string | null;
+  sortOrder?: 'asc' | 'desc';
+  workerIdContains?: string | null;
+  statusIn?: WorkerStatus[];
+};
+
 export type GetRolloutAttemptsQueryArgs = {
   rolloutId: string;
   limit?: number;
@@ -208,7 +234,7 @@ export type GetSpansQueryArgs = {
 export const rolloutsApi = createApi({
   reducerPath: 'rolloutsApi',
   baseQuery: dynamicBaseQuery,
-  tagTypes: ['Rollout', 'Span', 'Resources'],
+  tagTypes: ['Rollout', 'Span', 'Resources', 'Worker'],
   endpoints: (builder) => ({
     getResources: builder.query<PaginatedResponse<Resources>, GetResourcesQueryArgs>({
       query: ({ limit, offset, sortBy, sortOrder, resourcesIdContains }) => {
@@ -238,6 +264,37 @@ export const rolloutsApi = createApi({
             ]
           : [{ type: 'Resources' as const, id: 'LIST' }],
     }),
+    getWorkers: builder.query<PaginatedResponse<Worker>, GetWorkersQueryArgs>({
+      query: ({ limit, offset, sortBy, sortOrder, workerIdContains, statusIn }) => {
+        const searchParams = new URLSearchParams();
+        searchParams.set('limit', String(typeof limit === 'number' ? limit : -1));
+        searchParams.set('offset', String(typeof offset === 'number' ? offset : 0));
+        if (sortBy) {
+          searchParams.set('sort_by', sortBy);
+        }
+        if (sortOrder) {
+          searchParams.set('sort_order', sortOrder);
+        }
+        if (workerIdContains && workerIdContains.trim().length > 0) {
+          searchParams.set('worker_id_contains', workerIdContains.trim());
+        }
+        if (statusIn && statusIn.length > 0) {
+          statusIn.forEach((status) => searchParams.append('status_in', status));
+        }
+
+        const queryString = searchParams.toString();
+        const url = queryString.length > 0 ? `v1/agl/workers?${queryString}` : 'v1/agl/workers';
+        return { url, method: 'GET' };
+      },
+      transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeWorker),
+      providesTags: (result) =>
+        result
+          ? [
+              { type: 'Worker' as const, id: 'LIST' },
+              ...result.items.map((worker) => ({ type: 'Worker' as const, id: worker.workerId })),
+            ]
+          : [{ type: 'Worker' as const, id: 'LIST' }],
+    }),
     getRollouts: builder.query<PaginatedResponse<Rollout>, GetRolloutsQueryArgs>({
       query: ({ limit, offset, sortBy, sortOrder, statusIn, rolloutIdContains, modeIn }) => {
         const searchParams = new URLSearchParams();
@@ -343,4 +400,10 @@ export const rolloutsApi = createApi({
   }),
 });
 
-export const { useGetResourcesQuery, useGetRolloutsQuery, useGetRolloutAttemptsQuery, useGetSpansQuery } = rolloutsApi;
+export const {
+  useGetResourcesQuery,
+  useGetWorkersQuery,
+  useGetRolloutsQuery,
+  useGetRolloutAttemptsQuery,
+  useGetSpansQuery,
+} = rolloutsApi;

dashboard/src/features/ui/drawer/slice.ts

@@ -1,9 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
-import type { Attempt, Rollout, Span } from '@/types';
+import type { Attempt, Rollout, Span, Worker } from '@/types';
 
-export type DrawerType = 'rollout-json' | 'rollout-traces' | 'trace-detail';
+export type DrawerType = 'rollout-json' | 'rollout-traces' | 'trace-detail' | 'worker-detail';
 
 export type DrawerContent =
   | {
@@ -17,6 +17,10 @@ export type DrawerContent =
       span: Span;
       rollout: Rollout | null;
       attempt: Attempt | null;
+    }
+  | {
+      type: 'worker-detail';
+      worker: Worker;
     };
 
 export type DrawerState = {

dashboard/src/features/workers/index.ts

@@ -0,0 +1,5 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+export * from './slice';
+export * from './selectors';
+export { useGetWorkersQuery } from '../rollouts';

dashboard/src/features/workers/selectors.ts

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { createSelector } from '@reduxjs/toolkit';
+import type { GetWorkersQueryArgs } from '@/features/rollouts';
+import type { RootState } from '@/store';
+import type { WorkersSortState } from './slice';
+
+const WORKERS_SORT_FIELD_MAP: Record<string, string> = {
+  workerId: 'worker_id',
+  status: 'status',
+  currentRolloutId: 'current_rollout_id',
+  currentAttemptId: 'current_attempt_id',
+  lastHeartbeatTime: 'last_heartbeat_time',
+  lastDequeueTime: 'last_dequeue_time',
+  lastBusyTime: 'last_busy_time',
+  lastIdleTime: 'last_idle_time',
+};
+
+const resolveWorkersSortField = (sort: WorkersSortState): string =>
+  WORKERS_SORT_FIELD_MAP[sort.column] ?? 'last_heartbeat_time';
+
+export const selectWorkersUiState = (state: RootState) => state.workers;
+
+export const selectWorkersSearchTerm = (state: RootState) => selectWorkersUiState(state).searchTerm;
+export const selectWorkersPage = (state: RootState) => selectWorkersUiState(state).page;
+export const selectWorkersRecordsPerPage = (state: RootState) => selectWorkersUiState(state).recordsPerPage;
+export const selectWorkersSort = (state: RootState) => selectWorkersUiState(state).sort;
+
+export const selectWorkersQueryArgs = createSelector(
+  [selectWorkersSearchTerm, selectWorkersPage, selectWorkersRecordsPerPage, selectWorkersSort],
+  (searchTerm, page, recordsPerPage, sort): GetWorkersQueryArgs => {
+    const normalizedSearch = searchTerm.trim();
+    const limit = Math.max(1, recordsPerPage);
+    const offset = Math.max(0, (page - 1) * limit);
+    const sortBy = resolveWorkersSortField(sort);
+
+    return {
+      limit,
+      offset,
+      sortBy,
+      sortOrder: sort.direction,
+      workerIdContains: normalizedSearch.length > 0 ? normalizedSearch : undefined,
+    };
+  },
+);

dashboard/src/features/workers/slice.ts

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
+
+export type SortDirection = 'asc' | 'desc';
+
+export type WorkersSortState = {
+  column: string;
+  direction: SortDirection;
+};
+
+export type WorkersUiState = {
+  searchTerm: string;
+  page: number;
+  recordsPerPage: number;
+  sort: WorkersSortState;
+};
+
+export const initialWorkersUiState: WorkersUiState = {
+  searchTerm: '',
+  page: 1,
+  recordsPerPage: 50,
+  sort: {
+    column: 'lastHeartbeatTime',
+    direction: 'desc',
+  },
+};
+
+const workersSlice = createSlice({
+  name: 'workers',
+  initialState: initialWorkersUiState,
+  reducers: {
+    setWorkersSearchTerm(state, action: PayloadAction<string>) {
+      state.searchTerm = action.payload;
+      state.page = 1;
+    },
+    setWorkersPage(state, action: PayloadAction<number>) {
+      state.page = action.payload;
+    },
+    setWorkersRecordsPerPage(state, action: PayloadAction<number>) {
+      state.recordsPerPage = action.payload;
+      state.page = 1;
+    },
+    setWorkersSort(state, action: PayloadAction<WorkersSortState>) {
+      state.sort = action.payload;
+    },
+    resetWorkersFilters(state) {
+      state.searchTerm = initialWorkersUiState.searchTerm;
+      state.page = initialWorkersUiState.page;
+      state.recordsPerPage = initialWorkersUiState.recordsPerPage;
+      state.sort = initialWorkersUiState.sort;
+    },
+  },
+});
+
+export const { setWorkersSearchTerm, setWorkersPage, setWorkersRecordsPerPage, setWorkersSort, resetWorkersFilters } =
+  workersSlice.actions;
+
+export const workersReducer = workersSlice.reducer;

dashboard/src/features/workers/workers.test.ts

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { createServerBackedStore } from '@test-utils';
+import { describe, expect, it } from 'vitest';
+import { rolloutsApi } from '@/features/rollouts';
+import type { Worker } from '@/types';
+import { selectWorkersQueryArgs } from './selectors';
+import {
+  resetWorkersFilters,
+  setWorkersPage,
+  setWorkersRecordsPerPage,
+  setWorkersSearchTerm,
+  setWorkersSort,
+} from './slice';
+
+const extractWorkerIds = (workers: Worker[]): string[] => workers.map((worker) => worker.workerId);
+
+describe('workers feature integration', () => {
+  it('builds default query arguments from the UI state', () => {
+    const store = createServerBackedStore();
+    const queryArgs = selectWorkersQueryArgs(store.getState());
+
+    expect(queryArgs).toMatchObject({
+      limit: 50,
+      offset: 0,
+      sortBy: 'last_heartbeat_time',
+      sortOrder: 'desc',
+      workerIdContains: undefined,
+    });
+  });
+
+  it('fetches workers from the Python LightningStore server', async () => {
+    const store = createServerBackedStore();
+    const queryArgs = selectWorkersQueryArgs(store.getState());
+
+    const subscription = store.dispatch(rolloutsApi.endpoints.getWorkers.initiate(queryArgs));
+    const data = await subscription.unwrap();
+    subscription.unsubscribe();
+
+    expect(data.total).toBeGreaterThanOrEqual(4);
+    expect(data.items).toHaveLength(Math.min(queryArgs.limit, data.total));
+
+    const workerIds = extractWorkerIds(data.items);
+    expect(workerIds).toEqual(expect.arrayContaining(['worker-east', 'worker-west']));
+
+    const heartbeatTimes = data.items.map((worker) => worker.lastHeartbeatTime ?? 0);
+    const sortedHeartbeatTimes = [...heartbeatTimes].sort((a, b) => b - a);
+    expect(heartbeatTimes).toEqual(sortedHeartbeatTimes);
+  });
+
+  it('paginates worker results based on UI state', async () => {
+    const store = createServerBackedStore();
+    store.dispatch(setWorkersRecordsPerPage(2));
+    store.dispatch(setWorkersPage(2));
+
+    const queryArgs = selectWorkersQueryArgs(store.getState());
+    expect(queryArgs).toMatchObject({ limit: 2, offset: 2 });
+
+    const subscription = store.dispatch(rolloutsApi.endpoints.getWorkers.initiate(queryArgs));
+    const data = await subscription.unwrap();
+    subscription.unsubscribe();
+
+    expect(data.items).toHaveLength(2);
+    expect(data.total).toBeGreaterThanOrEqual(4);
+  });
+
+  it('applies search and sorting preferences', async () => {
+    const store = createServerBackedStore();
+    store.dispatch(resetWorkersFilters());
+    store.dispatch(setWorkersSearchTerm('worker-west'));
+    store.dispatch(setWorkersSort({ column: 'workerId', direction: 'asc' }));
+
+    const queryArgs = selectWorkersQueryArgs(store.getState());
+    expect(queryArgs).toMatchObject({
+      limit: 50,
+      offset: 0,
+      sortBy: 'worker_id',
+      sortOrder: 'asc',
+      workerIdContains: 'worker-west',
+    });
+
+    const subscription = store.dispatch(rolloutsApi.endpoints.getWorkers.initiate(queryArgs));
+    const data = await subscription.unwrap();
+    subscription.unsubscribe();
+
+    expect(data.items).toHaveLength(1);
+    expect(data.items[0].workerId).toBe('worker-west');
+  });
+
+  it('maps current rollout/attempt sorting to backend fields', () => {
+    const store = createServerBackedStore();
+    store.dispatch(setWorkersSort({ column: 'currentRolloutId', direction: 'desc' }));
+    let queryArgs = selectWorkersQueryArgs(store.getState());
+    expect(queryArgs.sortBy).toBe('current_rollout_id');
+    expect(queryArgs.sortOrder).toBe('desc');
+
+    store.dispatch(setWorkersSort({ column: 'currentAttemptId', direction: 'asc' }));
+    queryArgs = selectWorkersQueryArgs(store.getState());
+    expect(queryArgs.sortBy).toBe('current_attempt_id');
+    expect(queryArgs.sortOrder).toBe('asc');
+  });
+});

dashboard/src/layouts/AppLayout.story.tsx

@@ -39,6 +39,10 @@ const ROUTES = [
         path: 'traces',
         element: <Placeholder title='Traces' description='Browse telemetry spans across attempts.' />,
       },
+      {
+        path: 'runners',
+        element: <Placeholder title='Runners' description='Monitor runner activity and status.' />,
+      },
       {
         path: 'settings',
         element: (

dashboard/src/layouts/AppLayout.tsx

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 import { useEffect, useMemo, useState, type ReactNode } from 'react';
-import { IconCpu, IconRouteSquare, IconSettings, IconTimeline } from '@tabler/icons-react';
+import { IconCpu, IconRouteSquare, IconRun, IconSettings, IconTimeline } from '@tabler/icons-react';
 import { Outlet, NavLink as RouterNavLink, useLocation, useNavigate } from 'react-router-dom';
 import { AppShell, Badge, Group, Image, NavLink as MantineNavLink, Stack, Text, UnstyledButton } from '@mantine/core';
 import { AppAlertBanner } from '@/components/AppAlertBanner';
@@ -22,6 +22,7 @@ const NAV_ITEMS: NavItem[] = [
   { label: 'Rollouts', to: '/rollouts', icon: <IconRouteSquare size={16} /> },
   { label: 'Resources', to: '/resources', icon: <IconCpu size={16} /> },
   { label: 'Traces', to: '/traces', icon: <IconTimeline size={16} /> },
+  { label: 'Runners', to: '/runners', icon: <IconRun size={16} /> },
   { label: 'Settings', to: '/settings', icon: <IconSettings size={16} /> },
 ];
 

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

@@ -0,0 +1,270 @@
+// 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 { 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 { initialWorkersUiState } from '@/features/workers/slice';
+import { AppLayout } from '@/layouts/AppLayout';
+import { createAppStore } from '@/store';
+import type { Worker } from '@/types';
+import { createWorkersHandlers } from '@/utils/mock';
+import { STORY_BASE_URL, STORY_DATE_NOW_SECONDS } from '../../.storybook/constants';
+import { allModes } from '../../.storybook/modes';
+import { WorkersPage } from './Workers.page';
+
+const meta: Meta<typeof WorkersPage> = {
+  title: 'Pages/WorkersPage',
+  component: WorkersPage,
+  parameters: {
+    layout: 'fullscreen',
+    chromatic: {
+      modes: allModes,
+    },
+  },
+};
+
+export default meta;
+
+type Story = StoryObj<typeof WorkersPage>;
+
+const now = STORY_DATE_NOW_SECONDS;
+
+const sampleWorkers: Worker[] = [
+  {
+    workerId: 'worker-east',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 2, gpuUtilization: 0.82 },
+    lastHeartbeatTime: now - 20,
+    lastDequeueTime: now - 120,
+    lastBusyTime: now - 60,
+    lastIdleTime: now - 600,
+    currentRolloutId: 'ro-story-001',
+    currentAttemptId: 'at-story-010',
+  },
+  {
+    workerId: 'worker-west',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 1 },
+    lastHeartbeatTime: now - 45,
+    lastDequeueTime: now - 300,
+    lastBusyTime: now - 120,
+    lastIdleTime: now - 4800,
+    currentRolloutId: 'ro-story-003',
+    currentAttemptId: 'at-story-033',
+  },
+  {
+    workerId: 'worker-north',
+    status: 'idle',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 120,
+    lastDequeueTime: now - 3600,
+    lastBusyTime: now - 5400,
+    lastIdleTime: now - 180,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-south',
+    status: 'idle',
+    heartbeatStats: null,
+    lastHeartbeatTime: now - 900,
+    lastDequeueTime: now - 7200,
+    lastBusyTime: now - 8600,
+    lastIdleTime: now - 8600,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-central',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 3, cpuUtilization: 0.55 },
+    lastHeartbeatTime: now - 8,
+    lastDequeueTime: now - 45,
+    lastBusyTime: now - 10,
+    lastIdleTime: now - 900,
+    currentRolloutId: 'ro-story-005',
+    currentAttemptId: 'at-story-013',
+  },
+  {
+    workerId: 'worker-standby',
+    status: 'idle',
+    heartbeatStats: { queueDepth: 0, threads: 32 },
+    lastHeartbeatTime: now - 300,
+    lastDequeueTime: now - 10800,
+    lastBusyTime: now - 14400,
+    lastIdleTime: now - 200,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-observer',
+    status: 'unknown',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 30,
+    lastDequeueTime: now - 6400,
+    lastBusyTime: null,
+    lastIdleTime: null,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+];
+
+const defaultHandlers = createWorkersHandlers(sampleWorkers);
+
+function createStoryStore(configOverrides?: Partial<typeof initialConfigState>) {
+  return createAppStore({
+    config: {
+      ...initialConfigState,
+      baseUrl: STORY_BASE_URL,
+      autoRefreshMs: 0,
+      ...configOverrides,
+    },
+    workers: initialWorkersUiState,
+  });
+}
+
+function renderWithStore(configOverrides?: Partial<typeof initialConfigState>) {
+  const store = createStoryStore(configOverrides);
+
+  return (
+    <Provider store={store}>
+      <>
+        <WorkersPage />
+        <AppAlertBanner />
+        <AppDrawerContainer />
+      </>
+    </Provider>
+  );
+}
+
+function renderWithinAppLayout(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: '/runners',
+            element: <WorkersPage />,
+          },
+        ],
+      },
+    ],
+    { initialEntries: ['/runners'] },
+  );
+
+  return (
+    <Provider store={store}>
+      <>
+        <RouterProvider router={router} />
+        <AppDrawerContainer />
+      </>
+    </Provider>
+  );
+}
+
+const manyWorkers = Array.from({ length: 80 }, (_, index) => {
+  const suffix = (index + 1).toString().padStart(3, '0');
+  const busy = index % 2 === 0;
+  return {
+    workerId: `worker-batch-${suffix}`,
+    status: busy ? 'busy' : 'idle',
+    heartbeatStats: busy ? { queueDepth: (index % 5) + 1 } : { queueDepth: 0 },
+    lastHeartbeatTime: now - (index * 5 + 15),
+    lastDequeueTime: now - (index * 20 + 60),
+    lastBusyTime: busy ? now - (index * 10 + 30) : null,
+    lastIdleTime: busy ? null : now - (index * 10 + 45),
+    currentRolloutId: busy ? `ro-many-${suffix}` : null,
+    currentAttemptId: busy ? `at-many-${suffix}` : null,
+  } satisfies Worker;
+});
+
+export const Default: Story = {
+  render: () => renderWithinAppLayout(),
+  parameters: {
+    msw: {
+      handlers: defaultHandlers,
+    },
+  },
+};
+
+export const Search: Story = {
+  render: () => renderWithStore(),
+  parameters: {
+    msw: {
+      handlers: defaultHandlers,
+    },
+  },
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    await canvas.findByText('worker-east');
+
+    const searchInput = canvas.getByPlaceholderText('Search by Runner ID');
+    await userEvent.type(searchInput, 'worker-west');
+
+    await waitFor(() => {
+      if (canvas.queryByText('worker-east')) {
+        throw new Error('Expected filtered table to hide worker-east');
+      }
+      if (!canvas.queryByText('worker-west')) {
+        throw new Error('Expected worker-west to remain visible');
+      }
+    });
+  },
+};
+
+export const DrawerOpen: Story = {
+  render: () => renderWithStore(),
+  parameters: {
+    msw: {
+      handlers: defaultHandlers,
+    },
+  },
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    await canvas.findByText('worker-east');
+
+    const detailsButtons = await canvas.findAllByRole('button', { name: /detail/i });
+    await userEvent.click(detailsButtons[0]);
+
+    const body = within(document.body);
+    await waitFor(() => {
+      if (!body.queryByTestId('json-editor-container')) {
+        throw new Error('Expected worker detail drawer with JSON view');
+      }
+    });
+  },
+};
+
+export const ManyWorkers: Story = {
+  render: () => renderWithStore(),
+  parameters: {
+    msw: {
+      handlers: createWorkersHandlers(manyWorkers),
+    },
+  },
+};
+
+export const DarkTheme: Story = {
+  render: () => renderWithStore({ theme: 'dark' }),
+  parameters: {
+    theme: 'dark',
+    msw: {
+      handlers: defaultHandlers,
+    },
+  },
+};

dashboard/src/pages/Workers.page.tsx

@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+import { useCallback, useEffect } from 'react';
+import { IconSearch } from '@tabler/icons-react';
+import type { DataTableSortStatus } from 'mantine-datatable';
+import { Skeleton, Stack, TextInput, Title } from '@mantine/core';
+import { WorkersTable, type WorkersTableRecord } from '@/components/WorkersTable.component';
+import { selectAutoRefreshMs } from '@/features/config';
+import { hideAlert, showAlert } from '@/features/ui/alert';
+import { openDrawer } from '@/features/ui/drawer';
+import {
+  resetWorkersFilters,
+  selectWorkersPage,
+  selectWorkersQueryArgs,
+  selectWorkersRecordsPerPage,
+  selectWorkersSearchTerm,
+  selectWorkersSort,
+  setWorkersPage,
+  setWorkersRecordsPerPage,
+  setWorkersSearchTerm,
+  setWorkersSort,
+  useGetWorkersQuery,
+} from '@/features/workers';
+import { useAppDispatch, useAppSelector } from '@/store/hooks';
+import type { PaginatedResponse, Worker } from '@/types';
+import { getErrorDescriptor } from '@/utils/error';
+
+export function WorkersPage() {
+  const dispatch = useAppDispatch();
+  const autoRefreshMs = useAppSelector(selectAutoRefreshMs);
+  const searchTerm = useAppSelector(selectWorkersSearchTerm);
+  const page = useAppSelector(selectWorkersPage);
+  const recordsPerPage = useAppSelector(selectWorkersRecordsPerPage);
+  const sort = useAppSelector(selectWorkersSort);
+  const queryArgs = useAppSelector(selectWorkersQueryArgs);
+
+  const workersQueryResult = useGetWorkersQuery(queryArgs, {
+    pollingInterval: autoRefreshMs > 0 ? autoRefreshMs : undefined,
+  });
+
+  const workersData = workersQueryResult.data as PaginatedResponse<Worker> | undefined;
+  const { isLoading, isFetching, isError, error, refetch } = workersQueryResult;
+
+  const handleSearchTermChange = useCallback(
+    (value: string) => {
+      dispatch(setWorkersSearchTerm(value));
+    },
+    [dispatch],
+  );
+
+  const handleSortStatusChange = useCallback(
+    (status: DataTableSortStatus<WorkersTableRecord>) => {
+      dispatch(
+        setWorkersSort({
+          column: status.columnAccessor,
+          direction: status.direction,
+        }),
+      );
+    },
+    [dispatch],
+  );
+
+  const handlePageChange = useCallback(
+    (nextPage: number) => {
+      dispatch(setWorkersPage(nextPage));
+    },
+    [dispatch],
+  );
+
+  const handleRecordsPerPageChange = useCallback(
+    (value: number) => {
+      dispatch(setWorkersRecordsPerPage(value));
+    },
+    [dispatch],
+  );
+
+  const handleResetFilters = useCallback(() => {
+    dispatch(resetWorkersFilters());
+  }, [dispatch]);
+
+  const handleShowWorkerDetails = useCallback(
+    (worker: Worker) => {
+      dispatch(
+        openDrawer({
+          type: 'worker-detail',
+          worker,
+        }),
+      );
+    },
+    [dispatch],
+  );
+
+  const hasWorkers = Array.isArray(workersData?.items) && workersData.items.length > 0;
+  const showSkeleton = isLoading && !hasWorkers;
+
+  useEffect(() => {
+    if (isError) {
+      const descriptor = getErrorDescriptor(error);
+      const suffix = descriptor ? ` (${descriptor})` : '';
+      dispatch(
+        showAlert({
+          id: 'workers-fetch',
+          message: `Unable to refresh workers${suffix}. The table may be out of date until the connection recovers.`,
+          tone: 'error',
+        }),
+      );
+      return;
+    }
+
+    if (!isLoading && !isFetching) {
+      dispatch(hideAlert({ id: 'workers-fetch' }));
+    }
+  }, [dispatch, error, isError, isFetching, isLoading]);
+
+  useEffect(
+    () => () => {
+      dispatch(hideAlert({ id: 'workers-fetch' }));
+    },
+    [dispatch],
+  );
+
+  return (
+    <Stack gap='md'>
+      <Title order={1}>Runners</Title>
+
+      <TextInput
+        placeholder='Search by Runner ID'
+        value={searchTerm}
+        onChange={(event) => handleSearchTermChange(event.currentTarget.value)}
+        leftSection={<IconSearch size={16} />}
+        data-testid='workers-search-input'
+        w='100%'
+        style={{ maxWidth: 360 }}
+      />
+
+      {showSkeleton ? (
+        <Skeleton height={360} radius='md' />
+      ) : (
+        <WorkersTable
+          workers={workersData?.items}
+          totalRecords={workersData?.total ?? 0}
+          isFetching={isFetching}
+          isError={isError}
+          error={error}
+          searchTerm={searchTerm}
+          sort={sort}
+          page={page}
+          recordsPerPage={recordsPerPage}
+          onSortStatusChange={handleSortStatusChange}
+          onPageChange={handlePageChange}
+          onRecordsPerPageChange={handleRecordsPerPageChange}
+          onResetFilters={handleResetFilters}
+          onRefetch={refetch}
+          onShowDetails={handleShowWorkerDetails}
+        />
+      )}
+    </Stack>
+  );
+}

dashboard/src/store/index.ts

@@ -7,6 +7,7 @@ import { rolloutsApi, rolloutsReducer } from '../features/rollouts';
 import { tracesReducer } from '../features/traces';
 import { alertReducer } from '../features/ui/alert';
 import { drawerReducer } from '../features/ui/drawer';
+import { workersReducer } from '../features/workers';
 
 const rootReducer = combineReducers({
   config: configReducer,
@@ -14,6 +15,7 @@ const rootReducer = combineReducers({
   alert: alertReducer,
   rollouts: rolloutsReducer,
   resources: resourcesReducer,
+  workers: workersReducer,
   traces: tracesReducer,
   [rolloutsApi.reducerPath]: rolloutsApi.reducer,
 });

dashboard/src/types.ts

@@ -28,6 +28,24 @@ export type Attempt = {
   metadata: Record<string, any> | null;
 };
 
+export type WorkerStatus = 'idle' | 'busy' | 'unknown';
+
+/**
+ * Synced with agentlightning.types.core.Worker
+ * with camel case and snake case conversions
+ */
+export type Worker = {
+  workerId: string;
+  status: WorkerStatus;
+  heartbeatStats: Record<string, any> | null;
+  lastHeartbeatTime: Timestamp | null;
+  lastDequeueTime: Timestamp | null;
+  lastBusyTime: Timestamp | null;
+  lastIdleTime: Timestamp | null;
+  currentRolloutId: string | null;
+  currentAttemptId: string | null;
+};
+
 /**
  * Synced with agentlightning.types.core.Rollout
  * with camel case and snake case conversions

dashboard/src/utils/mock.test.ts

@@ -8,27 +8,32 @@
  */
 
 import { describe, expect, it } from 'vitest';
-import type { Attempt, Resources, Rollout, Span } from '@/types';
+import type { Attempt, Resources, Rollout, Span, Worker } from '@/types';
 import {
   buildAttemptsResponse,
   buildResourcesResponse,
   buildRolloutsResponse,
   buildSpansResponse,
+  buildWorkersResponse,
   createMockHandlers,
   createResourcesHandlers,
   createRolloutsHandlers,
   createSpansHandlers,
+  createWorkersHandlers,
   filterResourcesForParams,
   filterRolloutsForParams,
   filterSpansForParams,
+  filterWorkersForParams,
   getResourcesSortValue,
   getRolloutSortValue,
   getSpanSortValue,
+  getWorkerSortValue,
   parseNumberParam,
   sortAttemptsForParams,
   sortResourcesForParams,
   sortRolloutsForParams,
   sortSpansForParams,
+  sortWorkersForParams,
 } from './mock';
 
 const now = Math.floor(Date.now() / 1000);
@@ -219,6 +224,53 @@ const sampleResources: Resources[] = [
   },
 ];
 
+const sampleWorkers: Worker[] = [
+  {
+    workerId: 'worker-alpha',
+    status: 'busy',
+    heartbeatStats: { queueDepth: 2 },
+    lastHeartbeatTime: now - 30,
+    lastDequeueTime: now - 300,
+    lastBusyTime: now - 60,
+    lastIdleTime: now - 600,
+    currentRolloutId: 'ro-001',
+    currentAttemptId: 'at-001',
+  },
+  {
+    workerId: 'worker-beta',
+    status: 'idle',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 120,
+    lastDequeueTime: now - 1200,
+    lastBusyTime: now - 3600,
+    lastIdleTime: now - 180,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+  {
+    workerId: 'worker-gamma',
+    status: 'busy',
+    heartbeatStats: null,
+    lastHeartbeatTime: now - 10,
+    lastDequeueTime: now - 60,
+    lastBusyTime: now - 20,
+    lastIdleTime: now - 4000,
+    currentRolloutId: 'ro-003',
+    currentAttemptId: 'at-003',
+  },
+  {
+    workerId: 'worker-delta',
+    status: 'unknown',
+    heartbeatStats: { queueDepth: 0 },
+    lastHeartbeatTime: now - 5,
+    lastDequeueTime: now - 80,
+    lastBusyTime: null,
+    lastIdleTime: null,
+    currentRolloutId: null,
+    currentAttemptId: null,
+  },
+];
+
 describe('parseNumberParam', () => {
   it('returns default value when param is missing', () => {
     const params = new URLSearchParams();
@@ -725,6 +777,115 @@ describe('createResourcesHandlers', () => {
   });
 });
 
+describe('filterWorkersForParams', () => {
+  it('returns all workers without filters', () => {
+    const params = new URLSearchParams();
+    const result = filterWorkersForParams(sampleWorkers, params);
+    expect(result).toHaveLength(4);
+  });
+
+  it('filters by status and worker ID substring using AND logic', () => {
+    const params = new URLSearchParams('status_in=busy&worker_id_contains=gamma');
+    const result = filterWorkersForParams(sampleWorkers, params);
+    expect(result).toHaveLength(1);
+    expect(result[0].workerId).toBe('worker-gamma');
+  });
+
+  it('supports filter_logic=or', () => {
+    const params = new URLSearchParams('status_in=idle&worker_id_contains=gamma&filter_logic=or');
+    const result = filterWorkersForParams(sampleWorkers, params);
+    expect(result).toHaveLength(2);
+  });
+
+  it('filters by unknown status', () => {
+    const params = new URLSearchParams('status_in=unknown');
+    const result = filterWorkersForParams(sampleWorkers, params);
+    expect(result).toHaveLength(1);
+    expect(result[0].workerId).toBe('worker-delta');
+  });
+});
+
+describe('getWorkerSortValue', () => {
+  const worker = sampleWorkers[0];
+
+  it('returns worker_id', () => {
+    expect(getWorkerSortValue(worker, 'worker_id')).toBe('worker-alpha');
+  });
+
+  it('returns status', () => {
+    expect(getWorkerSortValue(worker, 'status')).toBe('busy');
+  });
+
+  it('returns timestamp fields', () => {
+    expect(getWorkerSortValue(worker, 'last_busy_time')).toBe(worker.lastBusyTime);
+    expect(getWorkerSortValue(worker, 'last_idle_time')).toBe(worker.lastIdleTime);
+    expect(getWorkerSortValue(worker, 'last_dequeue_time')).toBe(worker.lastDequeueTime);
+  });
+
+  it('returns rollout and attempt identifiers', () => {
+    expect(getWorkerSortValue(worker, 'current_rollout_id')).toBe(worker.currentRolloutId);
+    expect(getWorkerSortValue(worker, 'current_attempt_id')).toBe(worker.currentAttemptId);
+  });
+
+  it('falls back to last_heartbeat_time', () => {
+    expect(getWorkerSortValue(worker, 'unknown')).toBe(worker.lastHeartbeatTime);
+  });
+});
+
+describe('sortWorkersForParams', () => {
+  it('sorts by last heartbeat ascending by default', () => {
+    const result = sortWorkersForParams(sampleWorkers, null, 'asc');
+    expect(result.map((worker) => worker.workerId)).toEqual([
+      'worker-beta',
+      'worker-alpha',
+      'worker-gamma',
+      'worker-delta',
+    ]);
+  });
+
+  it('sorts descending by worker_id when requested', () => {
+    const result = sortWorkersForParams(sampleWorkers, 'worker_id', 'desc');
+    expect(result.map((worker) => worker.workerId)).toEqual([
+      'worker-gamma',
+      'worker-delta',
+      'worker-beta',
+      'worker-alpha',
+    ]);
+  });
+
+  it('sorts by current_rollout_id', () => {
+    const result = sortWorkersForParams(sampleWorkers, 'current_rollout_id', 'asc');
+    expect(result.map((worker) => worker.currentRolloutId)).toEqual([null, null, 'ro-001', 'ro-003']);
+  });
+});
+
+describe('buildWorkersResponse', () => {
+  it('applies filters before pagination', () => {
+    const request = new Request('http://localhost/v1/agl/workers?worker_id_contains=beta&limit=5');
+    const response = buildWorkersResponse(sampleWorkers, request);
+    expect(response.items).toHaveLength(1);
+    const items = response.items as Array<Record<string, unknown>>;
+    expect(items[0].worker_id).toBe('worker-beta');
+  });
+
+  it('applies sort and pagination parameters', () => {
+    const request = new Request('http://localhost/v1/agl/workers?sort_by=worker_id&limit=2&offset=1');
+    const response = buildWorkersResponse(sampleWorkers, request);
+    expect(response.items).toHaveLength(2);
+    const items = response.items as Array<Record<string, unknown>>;
+    expect(items[0].worker_id).toBe('worker-beta');
+    expect(response.total).toBe(4);
+  });
+});
+
+describe('createWorkersHandlers', () => {
+  it('creates handler for workers endpoint', () => {
+    const handlers = createWorkersHandlers(sampleWorkers);
+    expect(handlers).toHaveLength(1);
+    expect(handlers[0].info.header).toContain('GET');
+  });
+});
+
 describe('createRolloutsHandlers', () => {
   it('creates handlers that return correct rollout data', async () => {
     const attemptsByRollout = { 'ro-001': sampleAttempts };

dashboard/src/utils/mock.ts

@@ -14,7 +14,7 @@
  */
 
 import { delay, http, HttpResponse } from 'msw';
-import type { Attempt, Resources, Rollout, Span } from '@/types';
+import type { Attempt, Resources, Rollout, Span, Worker } from '@/types';
 import { snakeCaseKeys } from './format';
 
 /**
@@ -434,6 +434,117 @@ export function buildResourcesResponse(resources: Resources[], request: Request)
   });
 }
 
+/**
+ * Filter workers based on query parameters.
+ * Supports: status_in, worker_id_contains
+ */
+export function filterWorkersForParams(workers: Worker[], params: URLSearchParams): Worker[] {
+  const statusFilters = params.getAll('status_in');
+  const workerIdContains = params.get('worker_id_contains');
+  const filterLogic = params.get('filter_logic') === 'or' ? 'or' : 'and';
+
+  return workers.filter((worker) => {
+    const checks: boolean[] = [];
+    if (statusFilters.length > 0) {
+      checks.push(statusFilters.includes(worker.status));
+    }
+    if (workerIdContains) {
+      checks.push(worker.workerId.toLowerCase().includes(workerIdContains.toLowerCase()));
+    }
+    if (checks.length === 0) {
+      return true;
+    }
+    return filterLogic === 'or' ? checks.some(Boolean) : checks.every(Boolean);
+  });
+}
+
+/**
+ * Resolve a worker sort value for the given column.
+ */
+export function getWorkerSortValue(worker: Worker, sortBy: string): string | number | null {
+  switch (sortBy) {
+    case 'worker_id':
+      return worker.workerId;
+    case 'status':
+      return worker.status;
+    case 'current_rollout_id':
+      return worker.currentRolloutId ?? '';
+    case 'current_attempt_id':
+      return worker.currentAttemptId ?? '';
+    case 'last_busy_time':
+      return worker.lastBusyTime ?? null;
+    case 'last_idle_time':
+      return worker.lastIdleTime ?? null;
+    case 'last_dequeue_time':
+      return worker.lastDequeueTime ?? null;
+    case 'last_heartbeat_time':
+    default:
+      return worker.lastHeartbeatTime ?? null;
+  }
+}
+
+/**
+ * Sort workers based on query parameters.
+ * Default sort_by is 'last_heartbeat_time'.
+ */
+export function sortWorkersForParams(workers: Worker[], sortBy: string | null, sortOrder: 'asc' | 'desc'): Worker[] {
+  const resolvedSortBy = sortBy ?? 'last_heartbeat_time';
+  const sorted = [...workers].sort((a, b) => {
+    const aValue = getWorkerSortValue(a, resolvedSortBy);
+    const bValue = getWorkerSortValue(b, resolvedSortBy);
+    if (aValue === bValue) {
+      return 0;
+    }
+    if (aValue == null) {
+      return -1;
+    }
+    if (bValue == null) {
+      return 1;
+    }
+    if (typeof aValue === 'number' && typeof bValue === 'number') {
+      return aValue - bValue;
+    }
+    return String(aValue).localeCompare(String(bValue));
+  });
+
+  if (sortOrder === 'desc') {
+    sorted.reverse();
+  }
+
+  return sorted;
+}
+
+/**
+ * Build a paginated workers response matching the Python server's format.
+ */
+export function buildWorkersResponse(workers: Worker[], request: Request): Record<string, unknown> {
+  const url = new URL(request.url);
+  const params = url.searchParams;
+  const filtered = filterWorkersForParams(workers, params);
+  const sortBy = params.get('sort_by');
+  const sortOrder = params.get('sort_order') === 'desc' ? 'desc' : 'asc';
+  const sorted = sortWorkersForParams(filtered, sortBy, sortOrder);
+  const limitParam = parseNumberParam(params, 'limit', sorted.length);
+  const offsetParam = parseNumberParam(params, 'offset', 0);
+  const effectiveLimit = limitParam < 0 ? sorted.length : limitParam;
+  const offset = offsetParam < 0 ? 0 : offsetParam;
+  const paginated = effectiveLimit >= 0 ? sorted.slice(offset, offset + effectiveLimit) : [...sorted];
+
+  return snakeCaseKeys({
+    items: paginated,
+    limit: effectiveLimit,
+    offset,
+    total: filtered.length,
+  });
+}
+
+/**
+ * Create MSW handlers for workers endpoints.
+ */
+export function createWorkersHandlers(workers: Worker[]) {
+  return [http.get('*/v1/agl/workers', ({ request }) => HttpResponse.json(buildWorkersResponse(workers, request)))];
+}
+
 /**
  * Create MSW handlers for resources endpoints.
  *

dashboard/test-utils/python-server.py

@@ -33,6 +33,7 @@ from agentlightning.types import (
     RolloutConfig,
     Span,
     TraceStatus,
+    Worker,
 )
 
 
@@ -634,6 +635,68 @@ def inject_mock_data(store: InMemoryLightningStore, now: float | None = None) ->
     store._resources["rs-story-005"] = resource5
     store._latest_resources_id = "rs-story-005"
 
+    # Register workers with diverse states and activity windows.
+    workers = [
+        Worker(
+            worker_id="worker-east",
+            status="busy",
+            heartbeat_stats={"queue_depth": 2, "gpu_utilization": 0.82},
+            last_heartbeat_time=now - 20,
+            last_dequeue_time=now - 60,
+            last_busy_time=now - 120,
+            last_idle_time=now - 600,
+            current_rollout_id="ro-story-001",
+            current_attempt_id="at-story-010",
+        ),
+        Worker(
+            worker_id="worker-north",
+            status="idle",
+            heartbeat_stats={"queue_depth": 0, "gpu_utilization": 0.15},
+            last_heartbeat_time=now - 90,
+            last_dequeue_time=now - 3600,
+            last_busy_time=now - 5400,
+            last_idle_time=now - 5400,
+            current_rollout_id=None,
+            current_attempt_id=None,
+        ),
+        Worker(
+            worker_id="worker-west",
+            status="busy",
+            heartbeat_stats={"queue_depth": 1, "gpu_utilization": 0.41},
+            last_heartbeat_time=now - 45,
+            last_dequeue_time=now - 300,
+            last_busy_time=now - 200,
+            last_idle_time=now - 4800,
+            current_rollout_id="ro-story-003",
+            current_attempt_id="at-story-033",
+        ),
+        Worker(
+            worker_id="worker-south",
+            status="idle",
+            heartbeat_stats={"queue_depth": 0},
+            last_heartbeat_time=now - 900,
+            last_dequeue_time=now - 7200,
+            last_busy_time=now - 8600,
+            last_idle_time=now - 8600,
+            current_rollout_id=None,
+            current_attempt_id=None,
+        ),
+        Worker(
+            worker_id="worker-observer",
+            status="unknown",
+            heartbeat_stats={"queue_depth": 0},
+            last_heartbeat_time=now - 15,
+            last_dequeue_time=now - 4000,
+            last_busy_time=None,
+            last_idle_time=None,
+            current_rollout_id=None,
+            current_attempt_id=None,
+        ),
+    ]
+
+    for worker in workers:
+        store._workers[worker.worker_id] = worker
+
 
 async def main():
     parser = argparse.ArgumentParser(description="Run a Python server for the LightningStore")

docs/changelog.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## Agent-lightning v0.2.2 (11/12/2025)
+
+Agent-lightning v0.2.2 is a stabilization release for v0.2.1. It introduces several bug fixes.
+
+* Fix compatibility issues with VERL 0.6.0.
+* Fix model name for pre-downloaded models in VERL.
+* Fix preparing status transition on rollout when creating attempts.
+* Fix OpenAI Agents SDK compatibility issues.
+
+**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.2.1...v0.2.2
+
+---
+
 ## Agent-lightning v0.2.1 (10/30/2025)
 
 Agent-lightning v0.2.1 is a stabilization release for v0.2.0. It introduces several bug fixes and new features, plus a number of unlisted CI improvements.

docs/deep-dive/store.md

@@ -72,6 +72,15 @@ Each attempt begins in **preparing**, created either when a rollout is dequeued
 
 This simple model allows the system to distinguish between normal termination, abnormal stalling, and recoverable interruption without additional state flags.
 
+## Worker Telemetry
+
+Workers track runner-level activity timestamps (`last_heartbeat_time`, `last_dequeue_time`, `last_busy_time`, `last_idle_time`) plus their current rollout assignment. Those fields are now derived automatically:
+
+- [`dequeue_rollout(worker_id=...)`][agentlightning.LightningStore.dequeue_rollout] records which worker polled the queue and refreshes `last_dequeue_time`.
+- [`update_attempt(..., worker_id=...)`][agentlightning.LightningStore.update_attempt] drives the worker status machine. Assigning an attempt marks the worker **busy** and stamps `last_busy_time`; finishing with `status in {"succeeded","failed"}` switches to **idle**, while watchdog transitions such as `timeout`/`unresponsive` make the worker **unknown** and clear `current_rollout_id` / `current_attempt_id`.
+- [`update_worker(...)`][agentlightning.LightningStore.update_worker] is reserved for heartbeats. It snapshots optional `heartbeat_stats` and always updates `last_heartbeat_time`.
+
+Because every transition flows through these APIs, worker status is derived automatically from rollout execution and heartbeats. Note, however, that calling `update_worker` with a new `worker_id` will create a new worker record with status "unknown" if one does not exist. Thus, while manual status changes are not allowed, new worker records can be created externally via heartbeats.
 
 ## Rollout Transition Map
 

docs/reference/internal.md

@@ -28,6 +28,12 @@
 
 ::: agentlightning.llm_proxy.AddReturnTokenIds
 
+::: agentlightning.llm_proxy.StreamConversionMiddleware
+
+::: agentlightning.llm_proxy.MessageInspectionMiddleware
+
+::: agentlightning.llm_proxy.RolloutAttemptMiddleware
+
 ::: agentlightning.store.base.UNSET
 
 ::: agentlightning.store.utils.propagate_status

docs/reference/trainer.md

@@ -23,3 +23,11 @@
 ## CLI Builder
 
 ::: agentlightning.lightning_cli
+
+## Logging
+
+::: agentlightning.configure_logger
+
+::: agentlightning.setup_module_logging
+
+::: agentlightning.setup_logging

docs/reference/types.md

@@ -26,6 +26,10 @@
 
 ::: agentlightning.AttemptedRollout
 
+::: agentlightning.Worker
+
+::: agentlightning.WorkerStatus
+
 ::: agentlightning.Hook
 
 ## Resources

examples/apo/apo_custom_algorithm.py

@@ -181,5 +181,5 @@ async def main():
 
 
 if __name__ == "__main__":
-    agl.configure_logger()
+    agl.setup_logging()
     asyncio.run(main())

examples/apo/apo_custom_algorithm_trainer.py

@@ -19,7 +19,7 @@ python apo_custom_algorithm.py runner
 from apo_custom_algorithm import apo_algorithm, apo_rollout
 from rich.console import Console
 
-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.algorithm import algo
 from agentlightning.store import LightningStore
 
@@ -39,6 +39,6 @@ async def apo_algorithm_usable_in_trainer(*, store: LightningStore):
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
     trainer = Trainer(n_workers=1, algorithm=apo_algorithm_usable_in_trainer)
     trainer.fit(apo_rollout)

examples/apo/apo_debug.py

@@ -8,7 +8,7 @@ from typing import cast
 
 from apo_custom_algorithm import apo_rollout
 
-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.litagent import LitAgent
 from agentlightning.runner import LitAgentRunner
 from agentlightning.store import InMemoryLightningStore
@@ -105,7 +105,7 @@ def debug_with_trainer():
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
 
     parser = argparse.ArgumentParser(description="Debug APO with runner or trainer approach.")
     parser.add_argument(

examples/apo/legacy_apo_client.py

@@ -12,7 +12,7 @@ from typing import Any
 import dotenv
 from openai import OpenAI
 
-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.litagent import LitAgent
 from agentlightning.trainer import Trainer
 
@@ -40,7 +40,7 @@ class SimpleAgent(LitAgent[Any]):
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
     dotenv.load_dotenv()
     agent = SimpleAgent()
     # Use 2 workers to simulate multiple clients

examples/apo/room_selector_apo.py

@@ -8,7 +8,7 @@ from typing import Tuple, cast
 from openai import AsyncOpenAI
 from room_selector import RoomSelectionTask, load_room_tasks, prompt_template_baseline, room_selector
 
-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.adapter import TraceToMessages
 from agentlightning.algorithm.apo import APO
 from agentlightning.types import Dataset
@@ -33,7 +33,7 @@ def setup_apo_logger(file_path: str = "apo.log") -> None:
 
 
 def main() -> None:
-    configure_logger()
+    setup_logging()
     setup_apo_logger()
 
     openai_client = AsyncOpenAI()

examples/azure/tests/test_deployment.py

@@ -2,7 +2,7 @@
 
 from aoai_finetune import AzureOpenAIFinetune
 
-from agentlightning import configure_logger
+from agentlightning import setup_logging
 
 finetune_algo = AzureOpenAIFinetune(
     base_deployment_name="gpt-4.1-mini",
@@ -12,7 +12,7 @@ finetune_algo = AzureOpenAIFinetune(
     data_filter_ratio=0.6,
 )
 
-configure_logger()
+setup_logging()
 
 
 def test_deployment():

examples/azure/train_capital_agent.py

@@ -5,13 +5,13 @@ from aoai_finetune import AzureOpenAIFinetune
 from capital_agent import capital_agent
 from rich.console import Console
 
-from agentlightning import TraceToMessages, Trainer, configure_logger
+from agentlightning import TraceToMessages, Trainer, setup_logging
 
 console = Console()
 
 
 def main():
-    configure_logger()
+    setup_logging()
     finetune_algo = AzureOpenAIFinetune(
         base_deployment_name="gpt-4.1-mini",
         finetuned_deployment_name="gpt-4.1-mini-ft",

examples/calc_x/legacy_calc_agent.py

@@ -19,9 +19,9 @@ from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
 from eval_utils import evaluate_v0_1
 
-from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, setup_logging
 
-configure_logger()
+setup_logging()
 
 calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])
 

examples/calc_x/train_calc_agent.py

@@ -162,16 +162,17 @@ def train(
         print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
 
         # Keep it tiny/light without adding new knobs
-        config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+        config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
         config["trainer"]["total_epochs"] = 1
-        config["trainer"]["total_training_steps"] = 6
-        config["trainer"]["test_freq"] = 6
+        config["trainer"]["total_training_steps"] = 20
+        config["trainer"]["test_freq"] = 20
         config["trainer"]["experiment_name"] = EXPERIMENT_NAME
         config["trainer"]["project_name"] = PROJECT_NAME
         config["trainer"].pop("save_freq", None)
 
         if ci_fast:
             # Extra fast CI toggle for testing purposes.
+            config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
             config["trainer"]["total_training_steps"] = 1
             config["trainer"]["test_freq"] = 1
 

examples/rag/rag_agent.py

@@ -15,10 +15,10 @@ from agentlightning import (
     LitAgent,
     NamedResources,
     Trainer,
-    configure_logger,
+    setup_logging,
 )
 
-configure_logger()
+setup_logging()
 
 agent_prompt = """You are an assistant who answers questions using Wikipedia retriever. Answer the question using only the retrieved passages. Verify your answer directly against the text.
 

examples/search_r1/search_r1_agent.py

@@ -8,9 +8,9 @@ import requests
 from openai import OpenAI
 from qa_em import compute_score_em
 
-from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger, reward
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, reward, setup_logging
 
-configure_logger()
+setup_logging()
 
 # Copied and adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/data_process/nq_search.py
 INSTRUCTION_FORMAT = """Answer the given question. You must conduct reasoning inside <think> and </think> first every time you get new information. After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. You can search as many times as your want. If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. For example, <answer> Beijing </answer>. Question: """

examples/spider/sql_agent.py

@@ -9,6 +9,7 @@ as well as https://langchain-ai.github.io/langgraph/tutorials/sql-agent/
 
 from __future__ import annotations
 
+import logging
 import os
 import re
 import shutil
@@ -29,9 +30,9 @@ from spider_eval.exec_eval import eval_exec_match
 
 import agentlightning as agl
 
-agl.configure_logger()
+agl.setup_logging(apply_to=[__name__])
 
-logger = agl.configure_logger(name=__name__)
+logger = logging.getLogger(__name__)
 
 
 WRITE_QUERY_PROMPT = ChatPromptTemplate(

examples/tinker/agl_tinker/llm.py

@@ -271,7 +271,7 @@ def create_llm_proxy(
     renderer_name: str,
     port: int = 1899,
     store: Optional[LightningStore] = None,
-    _add_return_token_ids: bool = True,
+    add_return_token_ids: bool = True,
 ) -> LLMProxy:
     """Create an LLMProxy configured for a Tinker-based model.
 
@@ -284,6 +284,7 @@ def create_llm_proxy(
         renderer_name: Renderer type for prompt formatting (e.g., "qwen3", "qwen3_instruct").
         port: Port to expose the LiteLLM proxy. Defaults to 1899.
         store: Optional Lightning store for tracking usage. Defaults to None.
+        add_return_token_ids: Whether to add return token ids to the response. Defaults to True.
 
     Returns:
         Configured LLMProxy instance ready to serve the model.
@@ -305,5 +306,5 @@ def create_llm_proxy(
         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,
+        callbacks=["opentelemetry"] if add_return_token_ids else None,
     )

examples/tinker/hello.py

@@ -192,7 +192,7 @@ def main():
 
     args = parser.parse_args()
 
-    agl.configure_logger()
+    agl.setup_logging()
     if args.mode == "algo":
         run_algo()
     elif args.mode == "runner":

examples/tinker/q20_evaluate.py

@@ -79,7 +79,7 @@ async def evaluate_q20(
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
     if model_name.startswith("Qwen/"):
-        llm_proxy = create_llm_proxy(model_name, "qwen3", port, store, _add_return_token_ids=False)
+        llm_proxy = create_llm_proxy(model_name, "qwen3", port, store, add_return_token_ids=False)
     else:
         console.print(f"Assuming {model_name} is an OpenAI model.")
         llm_proxy = LLMProxy(
@@ -90,7 +90,8 @@ async def evaluate_q20(
             ],
             num_retries=2,
             launch_mode="thread",
-            _add_return_token_ids=False,
+            # Not going to add return_token_ids because we are not using Tinker.
+            callbacks=["opentelemetry"],
         )
 
     answerer_model_name = "gpt-5-mini"

examples/tinker/q20_train.py

@@ -367,7 +367,7 @@ def main() -> None:
     runner_parser.set_defaults(func=_run_runner)
 
     args = parser.parse_args()
-    agl.configure_logger()
+    agl.setup_logging()
     args.func(args)
 
 

examples/tinker/tests/test_tinker_llm.py

@@ -7,7 +7,6 @@ It should be included in CI in future if we decided to maintain this example.
 """
 
 import asyncio
-import logging
 from typing import cast
 
 import openai
@@ -24,13 +23,12 @@ from agentlightning import (
     LLMProxy,
     LlmProxyTraceToTriplet,
     TracerTraceToTriplet,
-    configure_logger,
     emit_reward,
+    setup_logging,
 )
 from agentlightning.store import LightningStoreThreaded
 
-configure_logger(name="agentlightning")
-configure_logger(name="agl_tinker", level=logging.INFO)
+setup_logging(apply_to=["agl_tinker"])
 
 
 async def test_tracer():

examples/unsloth/math_agent.py

@@ -29,7 +29,7 @@ from openai import AsyncOpenAI
 from rich.console import Console
 from trl import SFTConfig, SFTTrainer  # type: ignore
 
-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.litagent import rollout
 from agentlightning.types import LLM, Dataset
 
@@ -173,5 +173,5 @@ def math_agent_dry_run() -> None:
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
     math_agent_dry_run()

examples/unsloth/sft_algorithm.py

@@ -32,7 +32,7 @@ from math_agent import GsmProblem, load_math_dataset
 from rich.console import Console
 from unsloth_helper import unsloth_training
 
-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.adapter import LlmProxyTraceToTriplet, TraceToTripletBase
 from agentlightning.llm_proxy import LLMProxy, ModelConfig
 from agentlightning.store import LightningStore, LightningStoreClient
@@ -380,7 +380,7 @@ async def sft_algorithm(*, store: LightningStore) -> None:
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
 
     store = LightningStoreClient("http://localhost:4747")
 

examples/unsloth/sft_allinone.py

@@ -17,7 +17,7 @@ from math_agent import GsmProblem, load_math_dataset, math_agent
 from rich.console import Console
 from sft_algorithm import sft_one_iter
 
-from agentlightning import Trainer, configure_logger
+from agentlightning import Trainer, setup_logging
 from agentlightning.adapter import TraceToTripletBase
 from agentlightning.algorithm import Algorithm
 from agentlightning.llm_proxy import LLMProxy
@@ -94,7 +94,7 @@ class UnslothSupervisedFinetuning(Algorithm):
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
 
     algo = UnslothSupervisedFinetuning(
         max_iterations=2,

examples/unsloth/sft_rollout_runners.py

@@ -17,7 +17,7 @@ import multiprocessing
 from math_agent import GsmProblem, math_agent
 from rich.console import Console
 
-from agentlightning import configure_logger
+from agentlightning import setup_logging
 from agentlightning.runner import LitAgentRunner
 from agentlightning.store import LightningStore, LightningStoreClient
 from agentlightning.tracer import OtelTracer
@@ -67,6 +67,6 @@ def spawn_runners(*, store: LightningStore, n_runners: int) -> None:
 
 
 if __name__ == "__main__":
-    configure_logger()
+    setup_logging()
     store = LightningStoreClient("http://localhost:4747")
     spawn_runners(store=store, n_runners=4)

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.10"
 dependencies = [
   "graphviz",
   "psutil",
+  "gpustat",
   "setproctitle",
   "flask",
   "agentops>=0.4.13",

scripts/badge_aggregation.js

@@ -65,10 +65,14 @@ module.exports = async function badgeAggregation({ github, context, core, depend
       workflow_id: dep.workflow,
       branch: 'main', // Always check the main branch status no matter what
       status: 'completed', // only completed runs
-      per_page: 1, // latest only
+      per_page: 50, // retrieve latest 50 so we can filter
+      sort: 'created',
+      direction: 'desc',
     });
 
-    const run = runsData?.workflow_runs?.[0];
+    const filteredRuns = runsData?.workflow_runs?.filter(run => ['schedule', 'workflow_dispatch'].includes(run.event));
+
+    const run = filteredRuns?.[0];
     if (!run) {
       failures.push(`No completed run found for ${dep.label} on branch "${branch}"`);
       continue;

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -15,12 +15,14 @@ There are some specific TODOs for each test function.
 import ast
 import asyncio
 import json
-from typing import Any, cast
+from typing import Any, Dict, List, Type, Union, cast
 
 import anthropic
 import openai
 import pytest
+from litellm.integrations.custom_logger import CustomLogger
 
+from agentlightning import LlmProxyTraceToTriplet
 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
@@ -114,6 +116,14 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
         assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
         assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
 
+        # Verify start time and end time
+        # TODO: Remove this when this PR is merged: https://github.com/BerriAI/litellm/pull/16558
+        print(f">>> Span: {span.name}")
+        print(f">>> Start time: {span.start_time}")
+        print(f">>> End time: {span.end_time}")
+        assert span.start_time is not None, f"Span {span.name} has no start time"
+        assert span.end_time is not None, f"Span {span.name} has no end time"
+
     # Find the raw_gen_ai_request span and verify token IDs
     raw_gen_ai_spans = [s for s in spans if s.name == "raw_gen_ai_request"]
     assert len(raw_gen_ai_spans) == 1, f"Expected 1 raw_gen_ai_request span, found {len(raw_gen_ai_spans)}"
@@ -171,7 +181,13 @@ 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):
+async def _make_proxy_and_store(
+    qwen25_model: RemoteOpenAIServer,
+    *,
+    retries: int = 0,
+    gunicorn: bool = False,
+    callbacks: List[Union[Type[CustomLogger], str]] | None = None,
+):
     clear_tracer_provider()
     _reset_litellm_logging_worker()  # type: ignore
     store = InMemoryLightningStore()
@@ -192,6 +208,7 @@ async def _make_proxy_and_store(qwen25_model: RemoteOpenAIServer, *, retries: in
         num_workers=4 if gunicorn else 1,
         store=store_server,
         num_retries=retries,
+        callbacks=callbacks,
     )
     await proxy.start()
     return proxy, store_server
@@ -367,7 +384,6 @@ async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
         await store.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)
@@ -382,15 +398,149 @@ async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
         )
         collected: list[str] = []
         for evt in stream:
+            print(f">>> Event: {evt}")
             for c in evt.choices:
                 if c.delta and getattr(c.delta, "content", None):
                     assert isinstance(c.delta.content, str)
                     collected.append(c.delta.content)
-        assert "apple" in "".join(collected).lower()
+        # Sometimes the model responds with "hello" instead of "apple"
+        assert "apple" in "".join(collected).lower() or "hello" in "".join(collected).lower()
 
         spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
         assert len(spans) > 0
-        # TODO: didn't test the token ids in streaming chunks here
+        for span in spans:
+            print(f">>> Span {span.name}: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+            if span.name == "litellm_request":
+                assert "gen_ai.completion.0.content" in span.attributes
+    finally:
+        await proxy.stop()
+        await store.stop()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer):
+    proxy, store = await _make_proxy_and_store(qwen25_model)
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        adapter = LlmProxyTraceToTriplet()
+        client = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+
+        # non-stream
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+        )
+
+        txt = "".join([b.text for b in response.content if b.type == "text"])
+        assert "banana" in txt.lower(), f"Response does not contain 'banana': {txt}"
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+        assert len(spans) > 0
+
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+        assert len(triplets) == 1
+        assert triplets[0].prompt["token_ids"]
+        assert triplets[0].response["token_ids"]
+
+        # stream
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+            stream=True,
+        )
+        chunk_number: int = 0
+        for chunk in response:
+            print(f">>> Chunk: {chunk}")
+            chunk_number += 1
+        assert chunk_number >= 1
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+            if span.name == "litellm_request":
+                assert "gen_ai.completion.0.content" in span.attributes
+        assert len(spans) > 0
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+            assert triplet.prompt["token_ids"]
+            assert triplet.response["token_ids"]
+        assert len(triplets) == 2
+    finally:
+        await proxy.stop()
+        await store.stop()
+
+
+class LogprobsCallback(CustomLogger):
+
+    async def async_pre_call_hook(self, data: Dict[str, Any], **kwargs: Any) -> Dict[str, Any]:  # type: ignore
+        return {**data, "logprobs": 1}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer):
+    proxy, store = await _make_proxy_and_store(
+        qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"]
+    )
+    try:
+        resource, rollout = await _new_resource(proxy, store)
+        client = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+        adapter = LlmProxyTraceToTriplet()
+
+        # test streaming case only
+        response = client.messages.create(
+            model="gpt-4o-arbitrary",
+            max_tokens=64,
+            messages=[{"role": "user", "content": "Say the word: banana"}],
+            stream=True,
+        )
+
+        chunk_number: int = 0
+        for chunk in response:
+            print(f">>> Chunk: {chunk}")
+            chunk_number += 1
+        assert chunk_number >= 1
+
+        spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
+        for i, span in enumerate(spans):
+            print(f">>> Span {i}: {span.name}, attributes: {span.attributes}")
+            if span.name == "raw_gen_ai_request":
+                assert "llm.hosted_vllm.prompt_token_ids" in span.attributes
+                assert "llm.hosted_vllm.choices" in span.attributes
+                choices: list[dict[str, Any]] = ast.literal_eval(span.attributes["llm.hosted_vllm.choices"])  # type: ignore
+
+                # Check for token IDs and logprobs in the first choice
+                assert len(choices) > 0
+                if VLLM_VERSION >= (0, 10, 2):
+                    assert "token_ids" in choices[0]
+                    assert choices[0]["token_ids"]
+                assert "logprobs" in choices[0]
+                assert "content" in choices[0]["logprobs"]
+                assert len(choices[0]["logprobs"]["content"]) > 0
+                assert isinstance(choices[0]["logprobs"]["content"][0], dict)
+                assert "token" in choices[0]["logprobs"]["content"][0]
+                assert "logprob" in choices[0]["logprobs"]["content"][0]
+                assert isinstance(choices[0]["logprobs"]["content"][0]["logprob"], float)
+
+        assert len(spans) > 0
+
+        triplets = adapter.adapt(spans)
+        for i, triplet in enumerate(triplets):
+            print(f">>> Triplet {i}: {triplet}")
+            assert triplet.prompt["token_ids"]
+            assert triplet.response["token_ids"]
+            # TODO: Check logprobs
     finally:
         await proxy.stop()
         await store.stop()

tests/llm_proxy/test_stream.py

@@ -0,0 +1,714 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from typing import Any, AsyncGenerator, Dict, Iterator, List, Optional, cast
+
+import pytest
+
+from agentlightning.llm_proxy import StreamConversionMiddleware
+
+
+def merge_openai_streaming(chunks: Iterator[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Merge chunks from OpenAI chat completion streaming into a single message dict.
+
+    Returns a dict with keys:
+      - role: "assistant" (or whatever)
+      - content: full concatenated content string
+      - function_call: optional dict with keys name, arguments (string or JSON parsed)
+    """
+    role: Optional[str] = None
+    content_parts: List[str] = []
+    function_name: Optional[str] = None
+    function_args_str: Optional[str] = None
+
+    for chunk in chunks:
+        choice = chunk.get("choices", [])[0]
+        delta = choice.get("delta", {})
+
+        if "role" in delta and delta["role"] is not None:
+            role = delta["role"]
+
+        if "content" in delta and delta["content"] is not None:
+            content_parts.append(delta["content"])
+
+        # existing format: function_call
+        if "function_call" in delta and delta["function_call"] is not None:
+            fn = delta["function_call"]
+            if function_name is None:
+                function_name = fn.get("name")
+                function_args_str = fn.get("arguments", "")
+            else:
+                function_args_str += fn.get("arguments", "")
+
+        # new format: tool_calls array
+        if "tool_calls" in delta and delta["tool_calls"]:
+            for tc in delta["tool_calls"]:
+                func = tc.get("function", {})
+                # set name if first time
+                if function_name is None and func.get("name"):
+                    function_name = func["name"]
+                # accumulate arguments
+                if func.get("arguments") is not None:
+                    if function_args_str is None:
+                        function_args_str = func["arguments"]
+                    else:
+                        function_args_str += func["arguments"]
+
+    full_content = "".join(content_parts)
+    result: Dict[str, Any] = {"role": role or "assistant", "content": full_content}
+    if function_name is not None:
+        try:
+            function_args = json.loads(function_args_str or "")  # type: ignore
+        except Exception:
+            function_args = function_args_str
+        result["function_call"] = {"name": function_name, "arguments": function_args}
+    return result
+
+
+def merge_anthropic_streaming(chunks: Iterator[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Merge chunks from Anthropic streaming into a single message dict.
+
+    Returns a dict with keys:
+      - role: "assistant"
+      - content_text: full content text (concatenated)
+      - tool_calls: list of dicts { name, input } if any
+    """
+    role: Optional[str] = None
+    content_text_parts: List[str] = []
+    tool_calls: List[Dict[str, Any]] = []
+    current_tool: Optional[Dict[str, Any]] = None
+    current_tool_input_str: Optional[str] = None
+
+    for chunk in chunks:
+        # role
+        if role is None and "role" in chunk:
+            role = chunk["role"]
+
+        # handle content_block style (fine-grained)
+        typ = chunk.get("type")
+        if typ == "content_block_start":
+            block = chunk.get("content_block", {})
+            if block.get("type") == "tool_use":
+                # finish previous tool if exists
+                if current_tool is not None:
+                    try:
+                        input_obj = json.loads(current_tool_input_str or "")
+                    except Exception:
+                        input_obj = current_tool_input_str
+                    current_tool["input"] = input_obj
+                    tool_calls.append(current_tool)
+                current_tool = {"name": block.get("name"), "id": block.get("id"), "input": None}
+                current_tool_input_str = ""
+            continue
+
+        if typ == "content_block_delta":
+            delta = chunk.get("delta", {})
+            dtyp = delta.get("type")
+            if dtyp == "input_json_delta":
+                current_tool_input_str = (current_tool_input_str or "") + delta.get("partial_json", "")
+            elif dtyp == "text_delta":
+                content_text_parts.append(delta.get("text", ""))
+            continue
+
+        if typ == "content_block_stop":
+            if current_tool is not None:
+                try:
+                    input_obj = json.loads(current_tool_input_str or "")
+                except Exception:
+                    input_obj = current_tool_input_str
+                current_tool["input"] = input_obj
+                tool_calls.append(current_tool)
+                current_tool = None
+                current_tool_input_str = None
+            continue
+
+        # handle normal content items
+        content_items = chunk.get("content", [])
+        for item in content_items:
+            t = item.get("type")
+            if t == "text":
+                content_text_parts.append(item.get("text", ""))
+            elif t == "tool_use":
+                tool_id = item.get("id")
+                name = item.get("name")
+                inp = item.get("input", {})
+                if current_tool and current_tool.get("id") == tool_id:
+                    # merge into same tool
+                    try:
+                        existing = json.loads(current_tool_input_str or "{}")
+                    except Exception:
+                        existing: Dict[str, Any] = {}
+                    if isinstance(existing, dict):  # type: ignore
+                        existing.update(inp)
+                        current_tool_input_str = json.dumps(existing)
+                    else:
+                        # fallback: treat as string concatenation
+                        current_tool_input_str += json.dumps(inp)
+                else:
+                    # finish previous tool
+                    if current_tool is not None:
+                        try:
+                            input_obj = json.loads(current_tool_input_str or "")
+                        except Exception:
+                            input_obj = current_tool_input_str
+                        current_tool["input"] = input_obj
+                        tool_calls.append(current_tool)
+                    current_tool = {"name": name, "id": tool_id, "input": None}
+                    current_tool_input_str = json.dumps(inp)
+            # else: ignore
+    # end loop
+
+    # finish any open tool
+    if current_tool is not None:
+        try:
+            input_obj = json.loads(current_tool_input_str or "")
+        except Exception:
+            input_obj = current_tool_input_str
+        current_tool["input"] = input_obj
+        tool_calls.append(current_tool)
+
+    full_text = "".join(content_text_parts)
+    result: Dict[str, Any] = {"role": role or "assistant", "content_text": full_text}
+    if tool_calls:
+        result["tool_calls"] = tool_calls
+    return result
+
+
+def test_openai_text_only_short():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {"content": " world!"}, "finish_reason": None}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["role"] == "assistant"
+    assert merged["content"] == "Hello world!"
+    assert "function_call" not in merged
+
+
+def test_openai_text_and_function_call_arguments_split():
+    # Mixed content + function_call arguments spread over multiple deltas
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {"choices": [{"index": 0, "delta": {"content": "Starting… "}}]},
+                {
+                    "choices": [
+                        {"index": 0, "delta": {"function_call": {"name": "get_weather", "arguments": '{"city": "'}}}
+                    ]
+                },
+                {"choices": [{"index": 0, "delta": {"function_call": {"arguments": 'Singapore", "unit": "'}}}]},
+                {"choices": [{"index": 0, "delta": {"function_call": {"arguments": 'celsius"}'}}}]},
+                {"choices": [{"index": 0, "delta": {"content": "done."}}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["content"] == "Starting… done."
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"city": "Singapore", "unit": "celsius"}
+
+
+def test_openai_tool_calls_via_tool_calls_field():
+    # Newer shape: delta.tool_calls with function.name/arguments segments
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {"index": 0, "id": "call_1", "type": "function", "function": {"name": "search"}}
+                                ]
+                            },
+                        }
+                    ]
+                },
+                {
+                    "choices": [
+                        {"index": 0, "delta": {"tool_calls": [{"index": 0, "function": {"arguments": '{"q": "'}}]}}
+                    ]
+                },
+                {
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"tool_calls": [{"index": 0, "function": {"arguments": 'python streaming"}'}}]},
+                        }
+                    ]
+                },
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["function_call"]["name"] == "search"
+    assert merged["function_call"]["arguments"] == {"q": "python streaming"}
+
+
+def test_openai_invalid_json_arguments_falls_back_to_string():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {
+                    "choices": [{"index": 0, "delta": {"function_call": {"name": "do", "arguments": '{"bad": '}}}]
+                },  # truncated
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["function_call"]["name"] == "do"
+    # Should be raw string because JSON parsing fails
+    assert isinstance(merged["function_call"]["arguments"], str)
+    assert merged["function_call"]["arguments"].startswith('{"bad": ')
+
+
+def test_anthropic_text_only_multiple_blocks():
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"role": "assistant", "content": [{"type": "text", "text": "Hello "}]},
+                {"content": [{"type": "text", "text": "world!"}]},
+                {"type": "message_delta", "delta": {"stop_reason": "end_turn"}},
+            ],
+        )
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["role"] == "assistant"
+    assert merged["content_text"] == "Hello world!"
+    assert "tool_calls" not in merged
+
+
+def test_anthropic_tool_use_split_inputs_merge():
+    # Tool input is delivered as multiple content fragments that should be merged
+    chunks = iter(
+        [
+            {"role": "assistant", "content": [{"type": "text", "text": "Working… "}]},
+            {"content": [{"type": "tool_use", "id": "toolu_1", "name": "calculate", "input": {"a": 1}}]},
+            {"content": [{"type": "tool_use", "id": "toolu_1", "name": "calculate", "input": {"b": 2}}]},
+            {"content": [{"type": "text", "text": "done."}]},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["content_text"] == "Working… done."
+    assert merged["tool_calls"][0]["name"] == "calculate"
+    assert merged["tool_calls"][0]["input"] == {"a": 1, "b": 2}
+
+
+def test_anthropic_fine_grained_input_json_delta():
+    # Simulate SSE-style events: content_block_start(tool_use) + multiple input_json_delta pieces
+    chunks = iter(
+        [
+            {
+                "type": "content_block_start",
+                "index": 1,
+                "content_block": {"type": "tool_use", "id": "toolu_x", "name": "fetch"},
+            },
+            {
+                "type": "content_block_delta",
+                "index": 1,
+                "delta": {"type": "input_json_delta", "partial_json": '{"url": "'},
+                "active_tool_id": "toolu_x",
+            },
+            {
+                "type": "content_block_delta",
+                "index": 1,
+                "delta": {"type": "input_json_delta", "partial_json": 'https://example.com"}'},
+                "active_tool_id": "toolu_x",
+            },
+            {"type": "content_block_stop", "index": 1},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    [tool] = merged["tool_calls"]
+    assert tool["id"] == "toolu_x"
+    assert tool["name"] == "fetch"
+    assert tool["input"] == {"url": "https://example.com"}
+
+
+def test_anthropic_text_and_tool_interleaved_with_text_deltas():
+    # Mix text via text_delta and plain text content items
+    chunks = iter(
+        [
+            {"role": "assistant", "content": [{"type": "text", "text": "Start "}]},
+            {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "middle "}},
+            {"content": [{"type": "text", "text": "end."}]},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    assert merged["content_text"] == "Start middle end."
+
+
+def test_anthropic_partial_json_left_as_string_when_invalid():
+    # Provide malformed JSON parts; merger should keep raw string for tool input
+    chunks = iter(
+        [
+            {
+                "type": "content_block_start",
+                "index": 2,
+                "content_block": {"type": "tool_use", "id": "toolu_bad", "name": "ingest"},
+            },
+            {
+                "type": "content_block_delta",
+                "index": 2,
+                "delta": {"type": "input_json_delta", "partial_json": '{"alpha": 1, '},
+                "active_tool_id": "toolu_bad",
+            },
+            {
+                "type": "content_block_delta",
+                "index": 2,
+                "delta": {"type": "input_json_delta", "partial_json": '"beta": 2'},
+                "active_tool_id": "toolu_bad",
+            },
+            # missing closing brace
+            {"type": "content_block_stop", "index": 2},
+            {"type": "message_stop"},
+        ]
+    )
+    merged = merge_anthropic_streaming(chunks)
+    [tool] = merged["tool_calls"]
+    assert tool["id"] == "toolu_bad"
+    assert isinstance(tool["input"], str)
+    assert tool["input"].startswith('{"alpha": 1, ')
+
+
+@pytest.mark.parametrize("text_len", [1, 50, 500])
+def test_openai_long_text_stream_rounds_up(text_len: int):
+    # Create a synthetic long content split into ~20-40 char pieces as the merger would see
+    text = "x" * text_len
+    # Simulate content arriving in three chunks
+    part1, part2, part3 = text[: text_len // 3], text[text_len // 3 : 2 * text_len // 3], text[2 * text_len // 3 :]
+    chunks = iter(
+        cast(
+            List[Dict[str, Any]],
+            [
+                {"choices": [{"index": 0, "delta": {"role": "assistant"}}]},
+                {"choices": [{"index": 0, "delta": {"content": part1}}]},
+                {"choices": [{"index": 0, "delta": {"content": part2}}]},
+                {"choices": [{"index": 0, "delta": {"content": part3}}]},
+                {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+            ],
+        )
+    )
+    merged = merge_openai_streaming(chunks)
+    assert merged["content"] == text
+
+
+async def collect_sse(gen: AsyncGenerator[str, Any]) -> List[str]:
+    """Drain an async generator of SSE strings into a list."""
+    out: List[str] = []
+    async for s in gen:
+        assert isinstance(s, str)
+        out.append(s)
+    return out
+
+
+def parse_openai_sse_to_json_events(sse_chunks: List[str]) -> List[Dict[str, Any]]:
+    """From the OpenAI stream (which uses only 'data:' lines), return JSON events.
+    Filters out the literal DONE sentinel.
+    """
+    events: List[Dict[str, Any]] = []
+    for chunk in sse_chunks:
+        # each chunk looks like 'data: {...}\n\n' OR 'data: [DONE]\n\n'
+        for line in chunk.splitlines():
+            line = line.strip()
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]":
+                continue
+            events.append(json.loads(payload))
+    return events
+
+
+def parse_anthropic_sse_to_json_payloads(sse_chunks: List[str]) -> List[Dict[str, Any]]:
+    """Extract the JSON payload from each Anthropic SSE event (ignore pings)."""
+    out: List[Dict[str, Any]] = []
+    for chunk in sse_chunks:
+        # chunks look like 'event: <name>\ndata: {json}\n\n'
+        if "data:" not in chunk:
+            continue
+        data_line = [ln for ln in chunk.splitlines() if ln.startswith("data:")]
+        if not data_line:
+            continue
+        payload = data_line[0][len("data:") :].strip()
+        obj = json.loads(payload)
+        if obj.get("type") == "ping":
+            continue
+        out.append(obj)
+    return out
+
+
+@pytest.fixture
+def mw() -> StreamConversionMiddleware:
+    # BaseHTTPMiddleware requires an ASGI app; we only need the instance for bound methods.
+    class _DummyApp:
+        async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+            pass
+
+    return StreamConversionMiddleware(_DummyApp())
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "text, finish_reason",
+    [
+        ("Hello world.", "stop"),
+        ("This answer was cut off on purpose.", "length"),
+    ],
+)
+async def test_openai_content_only_stream_roundtrip(mw: StreamConversionMiddleware, text: str, finish_reason: str):
+    response_json = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": text},
+                "finish_reason": finish_reason,
+                # include logprobs to ensure it doesn't interfere with streaming
+                "logprobs": None,
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+
+    # basic shape checks
+    assert any('"delta": {"role": ' in s for s in sse_chunks)
+    assert any("[DONE]" in s for s in sse_chunks)
+
+    events = parse_openai_sse_to_json_events(sse_chunks)
+    assert events, "Expected JSON events from stream"
+
+    # the last JSON event before [DONE] should contain the finish_reason
+    last = events[-1]
+    assert last["choices"][0]["finish_reason"] == finish_reason
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["role"] == "assistant"
+    assert merged["content"] == text
+    assert "function_call" not in merged
+
+
+@pytest.mark.asyncio
+async def test_openai_long_text_chunking_and_reassembly(mw: StreamConversionMiddleware):
+    long_text = """
+        This is a deliberately long sentence that should be broken into multiple streaming deltas by the
+        chunking logic so that we can verify reassembly yields the exact same content without loss. """.strip()
+
+    response_json = {
+        "id": "chatcmpl-long",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": long_text}, "finish_reason": "stop"}],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    # ensure multiple content delta chunks were emitted
+    content_deltas = [ev for ev in events if ev["choices"][0]["delta"].get("content")]
+    assert len(content_deltas) > 1
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["content"] == long_text
+
+
+@pytest.mark.asyncio
+async def test_openai_tool_call_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    response_json = {
+        "id": "chatcmpl-tool",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": json.dumps({"location": "Boston"}),
+                            },
+                        }
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    # expect at least one tool_calls delta with name, followed by deltas with arguments
+    assert any(
+        (tc := ev["choices"][0]["delta"].get("tool_calls")) and tc[0].get("function", {}).get("name") == "get_weather"
+        for ev in events
+    )
+    assert any(
+        (tc := ev["choices"][0]["delta"].get("tool_calls")) and "arguments" in tc[0].get("function", {})
+        for ev in events
+    )
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"location": "Boston"}
+
+
+@pytest.mark.asyncio
+async def test_openai_content_and_tool_call_stream_roundtrip(mw: StreamConversionMiddleware):
+    response_json = {
+        "id": "chatcmpl-mixed",
+        "object": "chat.completion",
+        "model": "gpt-4o-mini",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "I'll call the weather tool now...",
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {
+                                "name": "get_weather",
+                                "arguments": json.dumps({"location": "Singapore", "units": "metric"}),
+                            },
+                        }
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+        ],
+    }
+
+    sse_chunks = await collect_sse(mw.openai_stream_generator(response_json))
+    events = parse_openai_sse_to_json_events(sse_chunks)
+
+    merged = merge_openai_streaming(iter(events))
+    assert merged["content"].startswith("I'll call the weather tool")
+    assert merged["function_call"]["name"] == "get_weather"
+    assert merged["function_call"]["arguments"] == {"location": "Singapore", "units": "metric"}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_text_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    original_response = {
+        "id": "msg_123",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {"type": "text", "text": "Hello there from Claude."},
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 7},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+
+    # sanity: stream contains lifecycle events
+    assert any("event: message_start" in s for s in sse_chunks)
+    assert any("event: message_stop" in s for s in sse_chunks)
+
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+    merged = merge_anthropic_streaming(iter(payloads))
+
+    assert merged["role"] == "assistant"
+    assert merged["content_text"] == "Hello there from Claude."
+    assert "tool_calls" not in merged
+
+
+@pytest.mark.asyncio
+async def test_anthropic_tool_use_only_stream_roundtrip(mw: StreamConversionMiddleware):
+    original_response = {
+        "id": "msg_tool",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {
+                "type": "tool_use",
+                "id": "toolu_1",
+                "name": "get_weather",
+                "input": {"location": "Boston"},
+            }
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+
+    merged = merge_anthropic_streaming(iter(payloads))
+
+    assert merged["tool_calls"][0]["name"] == "get_weather"
+    assert merged["tool_calls"][0]["id"] == "toolu_1"
+    assert merged["tool_calls"][0]["input"] == {"location": "Boston"}
+
+
+@pytest.mark.asyncio
+async def test_anthropic_mixed_text_and_tool_use_roundtrip(mw: StreamConversionMiddleware):
+    # tool input is long to ensure multiple input_json_delta chunks
+    long_input = {
+        "location": "Singapore",
+        "units": "metric",
+        "details": {"hourly": True, "with_forecast": True, "days": 5},
+    }
+    original_response = {
+        "id": "msg_mixed",
+        "model": "claude-3.5-sonnet",
+        "content": [
+            {"type": "text", "text": "I'll check the weather tool for you."},
+            {"type": "tool_use", "id": "toolu_2", "name": "get_weather", "input": long_input},
+        ],
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+        "stop_reason": "end_turn",
+    }
+
+    sse_chunks = await collect_sse(mw.anthropic_stream_generator(original_response))
+    payloads = parse_anthropic_sse_to_json_payloads(sse_chunks)
+
+    # Verify we saw content_block_start/stop and deltas for both text and tool input
+    types = [p.get("type") for p in payloads]
+    assert "content_block_start" in types
+    assert "content_block_delta" in types
+    assert "content_block_stop" in types
+    assert any(p.get("delta", {}).get("type") == "text_delta" for p in payloads)
+    assert any(p.get("delta", {}).get("type") == "input_json_delta" for p in payloads)
+
+    merged = merge_anthropic_streaming(iter(payloads))
+    assert merged["content_text"].startswith("I'll check the weather tool")
+    tool = merged["tool_calls"][0]
+    assert tool["name"] == "get_weather"
+    assert tool["id"] == "toolu_2"
+    assert tool["input"] == long_input

tests/runner/test_agent_runner.py

@@ -3,7 +3,7 @@
 import asyncio
 import random
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Dict, List, Optional, Sequence, cast
+from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Sequence, Tuple, cast
 
 import pytest
 from opentelemetry import trace as trace_api
@@ -17,10 +17,10 @@ from agentlightning.litagent import LitAgent
 from agentlightning.reward import emit_reward, find_final_reward
 from agentlightning.runner import LitAgentRunner
 from agentlightning.runner.base import Runner
-from agentlightning.store.base import LightningStore
+from agentlightning.store.base import UNSET, LightningStore, Unset
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.tracer.base import Tracer
-from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, SpanNames
+from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, SpanNames, Worker
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -114,6 +114,49 @@ class DummyTracer(Tracer):
         return span
 
 
+class RecordingStore(InMemoryLightningStore):
+    """In-memory store that records worker heartbeat updates for inspection in tests."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.worker_updates: List[Tuple[str, Optional[Dict[str, Any]]]] = []
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        payload = None if isinstance(heartbeat_stats, Unset) else heartbeat_stats
+        self.worker_updates.append((worker_id, payload))
+        return await super().update_worker(worker_id, heartbeat_stats=heartbeat_stats)
+
+
+class HeartbeatAgent(LitAgent[Dict[str, Any]]):
+    """Minimal agent used for heartbeat-only runner tests."""
+
+    def validation_rollout(self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any) -> float:
+        return 0.0
+
+
+async def setup_heartbeat_runner(
+    *,
+    heartbeat_interval: float = 0.05,
+    heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
+) -> tuple[LitAgentRunner[Any], RecordingStore]:
+    """Create a runner wired to a RecordingStore for heartbeat tests."""
+
+    store = RecordingStore()
+    runner = LitAgentRunner[Any](
+        tracer=DummyTracer(),
+        heartbeat_interval=heartbeat_interval,
+        heartbeat_launch_mode=heartbeat_launch_mode,
+    )
+    agent = HeartbeatAgent()
+    runner.init(agent)
+    runner.init_worker(worker_id=0, store=store)
+    return runner, store
+
+
 async def setup_runner(
     agent: LitAgent[Any],
     *,
@@ -658,3 +701,46 @@ async def test_step_with_custom_resources_returns_rollout() -> None:
 
     # Verify the rollout has the correct resources_id
     assert result.resources_id is not None
+
+
+@pytest.mark.asyncio
+async def test_emit_heartbeat_updates_worker_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
+    snapshot = {"cpu_pct": 42.0, "mem_pct": 10.5}
+    monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda: snapshot)
+
+    runner, store = await setup_heartbeat_runner(heartbeat_interval=0.1)
+    worker_label = runner.get_worker_id()
+    try:
+        await runner._emit_heartbeat(store)  # pyright: ignore[reportPrivateUsage]
+    finally:
+        teardown_runner(runner)
+
+    assert store.worker_updates == [(worker_label, snapshot)]
+    worker = await store.get_worker_by_id(worker_label)
+    assert worker is not None
+    assert worker.heartbeat_stats == snapshot
+    assert worker.last_heartbeat_time is not None
+
+
+@pytest.mark.asyncio
+async def test_heartbeat_loop_runs_until_stopped(monkeypatch: pytest.MonkeyPatch) -> None:
+    snapshot = {"timestamp": 1234567890}
+    monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda: snapshot)
+
+    runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05)
+    stop_heartbeat = runner._start_heartbeat_loop(store)  # pyright: ignore[reportPrivateUsage]
+    assert stop_heartbeat is not None
+
+    try:
+        await asyncio.sleep(0.12)
+    finally:
+        await stop_heartbeat()
+
+    update_count = len(store.worker_updates)
+    assert update_count >= 1
+    assert all(stats == snapshot for _, stats in store.worker_updates if stats is not None)
+
+    await asyncio.sleep(0.06)
+    assert len(store.worker_updates) == update_count
+
+    teardown_runner(runner)

tests/store/dummy_store.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Literal, Optional, Sequence
 
 from opentelemetry.sdk.trace import ReadableSpan
 
+from agentlightning.store import LightningStoreCapabilities
 from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.types import (
     Attempt,
@@ -16,6 +17,7 @@ from agentlightning.types import (
     RolloutStatus,
     Span,
     TaskInput,
+    Worker,
 )
 
 
@@ -25,6 +27,14 @@ class DummyLightningStore(LightningStore):
         self.calls: List[tuple[str, tuple[Any, ...], Dict[str, Any]]] = []
         self.return_values = return_values
 
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        return LightningStoreCapabilities(
+            async_safe=True,
+            thread_safe=False,
+            zero_copy=False,
+        )
+
     async def start_rollout(
         self,
         input: TaskInput,
@@ -47,8 +57,8 @@ class DummyLightningStore(LightningStore):
         self.calls.append(("enqueue_rollout", (input, mode, resources_id, config, metadata), {}))
         return self.return_values["enqueue_rollout"]
 
-    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
-        self.calls.append(("dequeue_rollout", (), {}))
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
+        self.calls.append(("dequeue_rollout", (worker_id,), {}))
         return self.return_values["dequeue_rollout"]
 
     async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
@@ -156,6 +166,31 @@ class DummyLightningStore(LightningStore):
         )
         return self.return_values["update_attempt"]
 
+    async def query_workers(self) -> List[Worker]:
+        self.calls.append(("query_workers", (), {}))
+        return self.return_values["query_workers"]
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        self.calls.append(("get_worker_by_id", (worker_id,), {}))
+        return self.return_values["get_worker_by_id"]
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Any = UNSET,
+    ) -> Worker:
+        self.calls.append(
+            (
+                "update_worker",
+                (
+                    worker_id,
+                    heartbeat_stats,
+                ),
+                {},
+            )
+        )
+        return self.return_values["update_worker"]
+
 
 def minimal_dummy_store() -> DummyLightningStore:
     # Provide minimal return values
@@ -180,5 +215,8 @@ def minimal_dummy_store() -> DummyLightningStore:
             "query_spans": [],
             "update_rollout": None,
             "update_attempt": None,
+            "query_workers": [],
+            "get_worker_by_id": None,
+            "update_worker": Worker(worker_id="worker-0"),
         }
     )

tests/store/test_client_server.py

@@ -19,6 +19,7 @@ from agentlightning.store.base import UNSET, LightningStore
 from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
 from agentlightning.store.memory import InMemoryLightningStore
 from agentlightning.types import LLM, OtelResource, PromptTemplate, RolloutConfig, Span, TraceStatus
+from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncherArgs
 
 
 def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
@@ -71,7 +72,15 @@ async def server_client(
 
 
 @pytest.mark.asyncio
-async def test_server_start_rejects_port_conflict() -> None:
+async def test_mp_server_does_not_work_with_inmemory_store() -> None:
+    store = InMemoryLightningStore()
+    with pytest.raises(ValueError, match="The store does not support zero-copy."):
+        LightningStoreServer(store, "127.0.0.1", pick_unused_port(), launch_mode="mp")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("launch_mode", ["asyncio", "thread"])
+async def test_server_start_rejects_port_conflict(caplog: pytest.LogCaptureFixture, launch_mode: LaunchMode) -> None:
     """Ensure startup fails loudly when the port is already owned by another store."""
     store_a = InMemoryLightningStore()
     port = pick_unused_port()
@@ -79,31 +88,59 @@ async def test_server_start_rejects_port_conflict() -> None:
     await server_a.start()
 
     store_b = InMemoryLightningStore()
-    server_b = LightningStoreServer(store_b, "127.0.0.1", port)
+    server_b = LightningStoreServer(store_b, "127.0.0.1", port, launch_mode=launch_mode)
 
-    with pytest.raises(RuntimeError, match="Another process may already be using this port"):
+    with pytest.raises(RuntimeError, match="did not start up within"):
         await server_b.start()
+    assert "address already in use" in caplog.text
 
     await server_a.stop()
 
 
 @pytest.mark.asyncio
-async def test_run_forever_rejects_port_conflict() -> None:
+@pytest.mark.parametrize("launch_mode", ["asyncio", "thread"])
+async def test_run_forever_rejects_port_conflict(caplog: pytest.LogCaptureFixture, launch_mode: LaunchMode) -> None:
     """Ensure run_forever also reports port conflicts with the friendly message."""
     store_a = InMemoryLightningStore()
     port = pick_unused_port()
-    server_a = LightningStoreServer(store_a, "127.0.0.1", port)
+    server_a = LightningStoreServer(store_a, "127.0.0.1", port, launch_mode=launch_mode)
     await server_a.start()
 
     store_b = InMemoryLightningStore()
-    server_b = LightningStoreServer(store_b, "127.0.0.1", port)
+    server_b = LightningStoreServer(store_b, "127.0.0.1", port, launch_mode=launch_mode)
 
-    with pytest.raises(RuntimeError, match="Another process may already be using this port"):
+    with pytest.raises(RuntimeError, match="did not start up within"):
         await server_b.run_forever()
+    assert "address already in use" in caplog.text
 
     await server_a.stop()
 
 
+@pytest.mark.asyncio
+async def test_server_accepts_custom_launcher_args(store_fixture: LightningStore) -> None:
+    """Ensure providing launcher_args works end-to-end and is propagated to the launcher."""
+    port = pick_unused_port()
+    launcher_args = PythonServerLauncherArgs(
+        host="127.0.0.1",
+        port=port,
+        launch_mode="asyncio",
+        healthcheck_url="/v1/agl/health",
+    )
+    server = LightningStoreServer(store_fixture, launcher_args=launcher_args)
+    assert server.launcher_args is launcher_args
+    assert server.server_launcher.args is launcher_args
+    assert server.server_launcher.health_url == f"http://127.0.0.1:{port}/v1/agl/health"
+
+    await server.start()
+    client = LightningStoreClient(server.endpoint)
+    try:
+        rollout = await client.start_rollout(input={"source": "launcher-args"})
+        assert rollout.rollout_id
+    finally:
+        await client.close()
+        await server.stop()
+
+
 @pytest.mark.asyncio
 async def test_add_resources_via_server(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
     """Test that add_resources works correctly via server."""
@@ -229,7 +266,13 @@ async def test_client_server_end_to_end(
     server_queue_config = RolloutConfig(unresponsive_seconds=4.2, max_attempts=2)
     queued_rollout = await server.enqueue_rollout(input={"origin": "server-queue"}, config=server_queue_config)
     assert queued_rollout.config.unresponsive_seconds == 4.2
-    dequeued = await server.dequeue_rollout()
+    server_worker_id = "server-worker"
+    dequeued = await server.dequeue_rollout(worker_id=server_worker_id)
+    server_worker_after_dequeue = await server.get_worker_by_id(server_worker_id)
+    assert server_worker_after_dequeue is not None
+    assert server_worker_after_dequeue.status == "idle"
+    assert server_worker_after_dequeue.last_dequeue_time is not None
+    dequeue_time = server_worker_after_dequeue.last_dequeue_time
     started_attempt = await server.start_attempt(queued_rollout.rollout_id)
 
     await server.query_rollouts()
@@ -260,10 +303,25 @@ async def test_client_server_end_to_end(
         queued_rollout.rollout_id,
         started_attempt.attempt.attempt_id,
         status="running",
-        worker_id="server-worker",
+        worker_id=server_worker_id,
         metadata={"phase": "warmup"},
     )
+    server_worker_busy = await server.get_worker_by_id(server_worker_id)
+    assert server_worker_busy is not None
+    assert server_worker_busy.status == "busy"
+    assert server_worker_busy.current_rollout_id == queued_rollout.rollout_id
+    assert server_worker_busy.current_attempt_id == started_attempt.attempt.attempt_id
+    assert server_worker_busy.last_busy_time is not None
+    assert server_worker_busy.last_busy_time >= dequeue_time
+
     await server.update_attempt(queued_rollout.rollout_id, "latest", status="succeeded")
+    server_worker_idle = await server.get_worker_by_id(server_worker_id)
+    assert server_worker_idle is not None
+    assert server_worker_idle.status == "idle"
+    assert server_worker_idle.current_rollout_id is None
+    assert server_worker_idle.current_attempt_id is None
+    assert server_worker_idle.last_idle_time is not None
+    assert server_worker_idle.last_idle_time >= server_worker_busy.last_busy_time
     completed = await server.wait_for_rollouts(rollout_ids=[queued_rollout.rollout_id], timeout=0.1)
     assert completed and completed[0].status in {"succeeded", "failed", "cancelled"}
 
@@ -285,8 +343,14 @@ async def test_client_server_end_to_end(
     client_queue_config = RolloutConfig(unresponsive_seconds=6.0)
     enqueued = await client.enqueue_rollout(input={"origin": "client-queue"}, config=client_queue_config)
     assert enqueued.config.unresponsive_seconds == 6.0
-    dequeued_client = await client.dequeue_rollout()
+    client_worker_id = "client-worker"
+    dequeued_client = await client.dequeue_rollout(worker_id=client_worker_id)
     assert dequeued_client is not None
+    client_worker_after_dequeue = await client.get_worker_by_id(client_worker_id)
+    assert client_worker_after_dequeue is not None
+    assert client_worker_after_dequeue.status == "idle"
+    assert client_worker_after_dequeue.last_dequeue_time is not None
+    client_dequeue_time = client_worker_after_dequeue.last_dequeue_time
     started_client_attempt = await client.start_attempt(dequeued_client.rollout_id)
 
     all_rollouts = await client.query_rollouts()
@@ -324,11 +388,26 @@ async def test_client_server_end_to_end(
     await client.update_attempt(
         dequeued_client.rollout_id,
         started_client_attempt.attempt.attempt_id,
-        worker_id="client-worker",
+        worker_id=client_worker_id,
         metadata={"info": "started"},
     )
+    client_worker_busy = await client.get_worker_by_id(client_worker_id)
+    assert client_worker_busy is not None
+    assert client_worker_busy.status == "busy"
+    assert client_worker_busy.current_rollout_id == dequeued_client.rollout_id
+    assert client_worker_busy.current_attempt_id == started_client_attempt.attempt.attempt_id
+    assert client_worker_busy.last_busy_time is not None
+    assert client_worker_busy.last_busy_time >= client_dequeue_time
+
     await client.update_attempt(dequeued_client.rollout_id, "latest", status="succeeded")
     await client.update_rollout(dequeued_client.rollout_id, status="succeeded")
+    client_worker_idle = await client.get_worker_by_id(client_worker_id)
+    assert client_worker_idle is not None
+    assert client_worker_idle.status == "idle"
+    assert client_worker_idle.current_rollout_id is None
+    assert client_worker_idle.current_attempt_id is None
+    assert client_worker_idle.last_idle_time is not None
+    assert client_worker_idle.last_idle_time >= client_worker_busy.last_busy_time
 
     wait_result = await client.wait_for_rollouts(rollout_ids=[dequeued_client.rollout_id], timeout=0.05)
     assert wait_result and wait_result[0].status == "succeeded"
@@ -415,6 +494,77 @@ async def test_update_attempt_none_vs_unset(server_client: Tuple[LightningStoreS
     assert preserved.status == "running"
 
 
+@pytest.mark.asyncio
+async def test_update_worker_records_heartbeat(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    _, client = server_client
+
+    first = await client.update_worker("runner-1", heartbeat_stats={"cpu": 0.4})
+    assert first.status == "unknown"
+    assert first.heartbeat_stats == {"cpu": 0.4}
+    assert first.last_heartbeat_time is not None
+
+    second = await client.update_worker("runner-1")
+    assert second.last_heartbeat_time is not None
+    assert second.last_heartbeat_time >= first.last_heartbeat_time
+    assert second.heartbeat_stats == {"cpu": 0.4}
+
+
+@pytest.mark.asyncio
+async def test_update_worker_rejects_none_stats(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    _, client = server_client
+    with pytest.raises(ClientResponseError) as exc_info:
+        await client.update_worker("runner-err", heartbeat_stats=cast(Any, None))
+    assert exc_info.value.status == 400
+
+
+@pytest.mark.asyncio
+async def test_worker_status_transitions_via_attempts(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient],
+) -> None:
+    _, client = server_client
+
+    await client.enqueue_rollout(input={"payload": "worker"})
+    claimed = await client.dequeue_rollout(worker_id="runner-auto")
+    assert claimed is not None
+
+    await client.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, worker_id="runner-auto")
+    busy = await client.get_worker_by_id("runner-auto")
+    assert busy is not None
+    assert busy.status == "busy"
+    assert busy.current_rollout_id == claimed.rollout_id
+    assert busy.current_attempt_id == claimed.attempt.attempt_id
+    assert busy.last_dequeue_time is not None
+    assert busy.last_busy_time is not None
+
+    await client.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, status="succeeded")
+    idle = await client.get_worker_by_id("runner-auto")
+    assert idle is not None
+    assert idle.status == "idle"
+    assert idle.current_rollout_id is None
+    assert idle.current_attempt_id is None
+
+
+@pytest.mark.asyncio
+async def test_get_worker_by_id(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
+    server, client = server_client
+
+    await server.update_worker("runner-lookup", heartbeat_stats={"cpu": 0.3})
+
+    server_worker = await server.get_worker_by_id("runner-lookup")
+    assert server_worker is not None
+    assert server_worker.worker_id == "runner-lookup"
+    assert await server.get_worker_by_id("missing") is None
+
+    client_worker = await client.get_worker_by_id("runner-lookup")
+    assert client_worker is not None
+    assert client_worker.worker_id == "runner-lookup"
+    assert await client.get_worker_by_id("missing") is None
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "bad_payload",
@@ -609,7 +759,7 @@ async def test_retry_on_400_application_error(
 
     # Force app-side exception so server returns 400 via exception handler.
     call_count = {"n": 0}
-    original = server.store.enqueue_rollout
+    original = server.store.enqueue_rollout  # type: ignore
 
     async def boom(*args: Any, **kwargs: Any) -> Any:
         call_count["n"] += 1
@@ -900,3 +1050,124 @@ async def test_get_next_span_sequence_id_returns_proper_int(
 
     # Verify monotonic increment
     assert seq_id_2 == seq_id_1 + 1
+
+
+@pytest.mark.asyncio
+async def test_empty_retry_delays_disable_retries(monkeypatch: MonkeyPatch) -> None:
+    """
+    When retry_delays is empty, the client should perform only the initial attempt
+    and not retry on transient network errors.
+    """
+    store = InMemoryLightningStore()
+    port = pick_unused_port()
+    server = LightningStoreServer(
+        store,
+        launcher_args=PythonServerLauncherArgs(
+            port=port,
+            host="127.0.0.1",
+            healthcheck_url=None,
+            launch_mode="thread",
+        ),
+    )
+    await server.start()
+
+    # retry_delays=() disables retries; health checks still enabled
+    client = LightningStoreClient(
+        server.endpoint,
+        retry_delays=(),
+        health_retry_delays=(0.01,),
+    )
+
+    try:
+        original_post = aiohttp.ClientSession.post
+        original_get = aiohttp.ClientSession.get
+
+        calls = {"post": 0, "health": 0}
+
+        def failing_post(self: aiohttp.ClientSession, url: Any, *args: Any, **kwargs: Any) -> MockResponse:
+            if str(url).endswith("/rollouts"):
+                calls["post"] += 1
+                # Always raise a transient error
+                raise ServerDisconnectedError("synthetic disconnect for empty retry_delays")
+            return MockResponse(original_post(self, url, *args, **kwargs))
+
+        def ok_health_get(self: aiohttp.ClientSession, url: Any, *args: Any, **kwargs: Any) -> MockResponse:
+            if str(url).endswith("/health"):
+                calls["health"] += 1
+            # delegate to the real get() and wrap in MockResponse so it stays an async CM
+            return MockResponse(original_get(self, url, *args, **kwargs))
+
+        monkeypatch.setattr(aiohttp.ClientSession, "post", failing_post, raising=True)
+        monkeypatch.setattr(aiohttp.ClientSession, "get", ok_health_get, raising=True)
+
+        with pytest.raises(ServerDisconnectedError):
+            await client.start_rollout(input={"origin": "empty-retry-delays"})
+
+        # Only the initial attempt should be made
+        assert calls["post"] == 1
+        # Health should be probed at least once
+        assert calls["health"] >= 1
+    finally:
+        await client.close()
+        await server.stop()
+
+
+@pytest.mark.asyncio
+async def test_empty_health_retry_delays_skip_health_checks(monkeypatch: MonkeyPatch) -> None:
+    """
+    When health_retry_delays is empty, _wait_until_healthy should not perform any
+    /health probes, but retries governed by retry_delays should still occur.
+    """
+    store = InMemoryLightningStore()
+    port = pick_unused_port()
+    server = LightningStoreServer(
+        store,
+        launcher_args=PythonServerLauncherArgs(
+            port=port,
+            host="127.0.0.1",
+            healthcheck_url=None,
+            launch_mode="thread",
+        ),
+    )
+    await server.start()
+
+    # health_retry_delays=() disables health probes; still allow one retry
+    client = LightningStoreClient(
+        server.endpoint,
+        retry_delays=(0.01,),
+        health_retry_delays=(),
+    )
+
+    try:
+        original_post = aiohttp.ClientSession.post
+        original_get = aiohttp.ClientSession.get
+
+        calls = {"post": 0, "health": 0}
+
+        def flaky_post(self: aiohttp.ClientSession, url: Any, *args: Any, **kwargs: Any) -> MockResponse:
+            if str(url).endswith("/rollouts"):
+                calls["post"] += 1
+                # First call fails, second succeeds
+                if calls["post"] == 1:
+                    raise ServerDisconnectedError("synthetic disconnect for empty health_retry_delays")
+            return MockResponse(original_post(self, url, *args, **kwargs))
+
+        def counting_health_get(self: aiohttp.ClientSession, url: Any, *args: Any, **kwargs: Any) -> MockResponse:
+            if str(url).endswith("/health"):
+                calls["health"] += 1
+            return MockResponse(original_get(self, url, *args, **kwargs))
+
+        monkeypatch.setattr(aiohttp.ClientSession, "post", flaky_post, raising=True)
+        monkeypatch.setattr(aiohttp.ClientSession, "get", counting_health_get, raising=True)
+
+        # Should succeed after one retry, without ever calling /health
+        attempted = await client.start_rollout(input={"origin": "empty-health-delays"})
+        assert attempted.rollout_id
+
+        # One failure + one success
+        assert calls["post"] == 2
+        # No health checks should have been performed
+        assert calls["health"] == 0
+    finally:
+        await client.close()
+        await server.stop()

tests/store/test_implementation.py

@@ -446,6 +446,49 @@ async def test_requeue_mechanism(store_fixture: LightningStore) -> None:
     assert latest_attempt.sequence_id == 2
 
 
+@pytest.mark.asyncio
+async def test_update_and_query_workers(store_fixture: LightningStore) -> None:
+    """Workers can be created, heartbeats recorded, and telemetry auto-updated."""
+    first = await store_fixture.update_worker("worker-1", heartbeat_stats={"cpu": 0.5})
+    assert first.worker_id == "worker-1"
+    assert first.heartbeat_stats == {"cpu": 0.5}
+    assert isinstance(first.last_heartbeat_time, float)
+    assert first.status == "unknown"
+
+    rollout = await store_fixture.enqueue_rollout(input={"task": "work"})
+    claimed = await store_fixture.dequeue_rollout(worker_id="worker-1")
+    assert claimed is not None
+    assert claimed.rollout_id == rollout.rollout_id
+
+    await store_fixture.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, worker_id="worker-1")
+    busy = await store_fixture.get_worker_by_id("worker-1")
+    assert busy is not None
+    assert busy.status == "busy"
+    assert busy.current_rollout_id == claimed.rollout_id
+    assert busy.current_attempt_id == claimed.attempt.attempt_id
+    assert isinstance(busy.last_dequeue_time, float)
+    assert isinstance(busy.last_busy_time, float)
+
+    heartbeat = await store_fixture.update_worker("worker-1")
+    assert heartbeat.last_heartbeat_time is not None
+    assert heartbeat.last_heartbeat_time >= first.last_heartbeat_time
+
+    await store_fixture.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, status="succeeded")
+    idle = await store_fixture.get_worker_by_id("worker-1")
+    assert idle is not None
+    assert idle.status == "idle"
+    assert idle.current_rollout_id is None
+    assert idle.current_attempt_id is None
+    assert isinstance(idle.last_idle_time, float)
+
+    workers = await store_fixture.query_workers()
+    assert any(w.worker_id == "worker-1" for w in workers)
+    assert await store_fixture.get_worker_by_id("missing") is None
+
+    with pytest.raises(TypeError):
+        await store_fixture.update_worker("worker-1", heartbeat_stats=None)  # type: ignore[arg-type]
+
+
 # Resource Management Tests
 
 

tests/store/test_restful.py

@@ -1012,3 +1012,123 @@ async def test_client_query_with_filters(
     rollouts = await client.query_rollouts(rollout_ids=[r2.rollout_id])
     assert len(rollouts) == 1
     assert rollouts[0].rollout_id == r2.rollout_id
+
+
+@pytest.mark.asyncio
+async def test_workers_endpoint_supports_updates(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    _server, _client, session, api_endpoint = server_client
+
+    async with session.post(
+        f"{api_endpoint}/workers/worker-1",
+        json={"heartbeat_stats": {"cpu": 0.7}},
+    ) as resp:
+        assert resp.status == 200
+        created = await resp.json()
+        assert created["worker_id"] == "worker-1"
+        assert created["status"] == "unknown"
+        assert created["heartbeat_stats"] == {"cpu": 0.7}
+        first_heartbeat = created["last_heartbeat_time"]
+
+    async with session.get(f"{api_endpoint}/workers") as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        workers = data["items"]
+        assert len(workers) == 1
+        assert workers[0]["worker_id"] == "worker-1"
+
+    async with session.post(
+        f"{api_endpoint}/workers/worker-1",
+        json={"heartbeat_stats": {"cpu": 0.8}},
+    ) as resp:
+        assert resp.status == 200
+        updated = await resp.json()
+        assert updated["last_heartbeat_time"] >= first_heartbeat
+
+    async with session.get(f"{api_endpoint}/workers") as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        workers = data["items"]
+        assert workers[0]["status"] == "unknown"
+
+
+@pytest.mark.asyncio
+async def test_workers_endpoint_rejects_none_stats(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    _server, _client, session, api_endpoint = server_client
+
+    async with session.post(
+        f"{api_endpoint}/workers/worker-err",
+        json={"heartbeat_stats": None},
+    ) as resp:
+        assert resp.status == 400
+
+
+@pytest.mark.asyncio
+async def test_get_worker_by_id_restful(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    await server.update_worker("worker-fetch", heartbeat_stats={"cpu": 0.4})
+
+    async with session.get(f"{api_endpoint}/workers/worker-fetch") as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        assert data["worker_id"] == "worker-fetch"
+
+    async with session.get(f"{api_endpoint}/workers/missing") as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        assert data is None
+
+
+@pytest.mark.asyncio
+async def test_workers_endpoint_filter_and_sort(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str],
+) -> None:
+    server, _client, session, api_endpoint = server_client
+
+    # Worker A: finishes an attempt and becomes idle.
+    await server.update_worker("worker-a", heartbeat_stats={"cpu": 0.1})
+    await server.enqueue_rollout(input={"worker": "a"})
+    claimed_a = await server.dequeue_rollout(worker_id="worker-a")
+    assert claimed_a is not None
+    await server.update_attempt(
+        claimed_a.rollout_id, claimed_a.attempt.attempt_id, worker_id="worker-a", status="succeeded"
+    )
+
+    # Worker B: currently busy on an attempt.
+    await server.update_worker("worker-b", heartbeat_stats={"cpu": 0.9})
+    await server.enqueue_rollout(input={"worker": "b"})
+    claimed_b = await server.dequeue_rollout(worker_id="worker-b")
+    assert claimed_b is not None
+    await server.update_attempt(claimed_b.rollout_id, claimed_b.attempt.attempt_id, worker_id="worker-b")
+
+    # Worker C: also busy.
+    await server.update_worker("worker-c", heartbeat_stats={"cpu": 0.2})
+    await server.enqueue_rollout(input={"worker": "c"})
+    claimed_c = await server.dequeue_rollout(worker_id="worker-c")
+    assert claimed_c is not None
+    await server.update_attempt(claimed_c.rollout_id, claimed_c.attempt.attempt_id, worker_id="worker-c")
+
+    async with session.get(
+        f"{api_endpoint}/workers",
+        params={"status_in": ["busy"], "worker_id_contains": "worker", "sort_by": "worker_id", "sort_order": "desc"},
+    ) as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        items = data["items"]
+        assert [w["worker_id"] for w in items] == ["worker-c", "worker-b"]
+
+    async with session.get(
+        f"{api_endpoint}/workers",
+        params={"limit": 1, "offset": 1, "sort_by": "worker_id", "sort_order": "asc"},
+    ) as resp:
+        assert resp.status == 200
+        data = await resp.json()
+        assert data["limit"] == 1
+        assert data["offset"] == 1
+        assert len(data["items"]) == 1

tests/store/test_threading.py

@@ -25,6 +25,7 @@ from agentlightning.types import (
     SpanContext,
     TaskInput,
     TraceStatus,
+    Worker,
 )
 
 from .dummy_store import DummyLightningStore
@@ -160,6 +161,8 @@ async def test_threaded_store_delegates_all_methods() -> None:
         last_heartbeat_time=1.5,
         metadata={"idx": 0},
     )
+    worker_list = [Worker(worker_id="worker-1", status="busy")]
+    updated_worker = Worker(worker_id="worker-1", status="idle")
 
     return_values = {
         "start_rollout": attempted_rollout,
@@ -180,6 +183,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "query_spans": [span],
         "update_rollout": updated_rollout,
         "update_attempt": updated_attempt,
+        "query_workers": worker_list,
+        "get_worker_by_id": worker_list[0],
+        "update_worker": updated_worker,
     }
 
     dummy_store = DummyLightningStore(return_values)
@@ -231,6 +237,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
         )
         == updated_attempt
     )
+    assert await threaded_store.query_workers() == worker_list
+    assert await threaded_store.get_worker_by_id("worker-1") == worker_list[0]
+    assert await threaded_store.update_worker("worker-1", heartbeat_stats={"cpu": 0.5}) == updated_worker
 
     expected_order = [
         "start_rollout",
@@ -251,6 +260,9 @@ async def test_threaded_store_delegates_all_methods() -> None:
         "query_spans",
         "update_rollout",
         "update_attempt",
+        "query_workers",
+        "get_worker_by_id",
+        "update_worker",
     ]
     assert [name for name, *_ in dummy_store.calls] == expected_order
 

tests/test_logging.py

@@ -0,0 +1,420 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import io
+import logging
+import multiprocessing as mp
+from multiprocessing.queues import Queue
+from pathlib import Path
+from typing import Any, Dict, List
+
+import pytest
+
+from agentlightning.logging import _to_level_value  # pyright: ignore[reportPrivateUsage]
+from agentlightning.logging import (
+    DATE_FORMAT,
+    DEFAULT_FORMAT,
+)
+
+
+def _logging_worker(case: str, queue: Queue[Dict[str, Any]]) -> None:
+    """
+    Runs in a separate process using spawn. It performs a specific logging
+    configuration scenario and returns a summary dict via the queue.
+    """
+    import logging
+    import warnings
+
+    # Re-import inside the subprocess so everything is picklable & isolated
+    from agentlightning.logging import (
+        setup,
+        setup_module,
+    )
+
+    if case == "setup_module_plain_console":
+        logger = setup_module(
+            level="DEBUG",
+            name="agentlightning.test",
+            console=True,
+            color=False,
+            propagate=False,
+        )
+
+        handlers = logger.handlers
+        handler = handlers[0] if handlers else None
+        fmt = handler.formatter._fmt if handler and handler.formatter else None
+        datefmt = handler.formatter.datefmt if handler and handler.formatter else None
+
+        queue.put(
+            {
+                "logger_name": logger.name,
+                "logger_level": logger.level,
+                "num_handlers": len(handlers),
+                "handler_class": handler.__class__.__name__ if handler else None,
+                "handler_level": handler.level if handler else None,
+                "fmt": fmt,
+                "datefmt": datefmt,
+            }
+        )
+
+    elif case == "setup_module_color_rich":
+        # Rich variant: color=True uses RichHandler
+        logger = setup_module(
+            level="INFO",
+            name="agentlightning.rich",
+            console=True,
+            color=True,
+            propagate=False,
+        )
+        handlers = logger.handlers
+        handler = handlers[0] if handlers else None
+
+        queue.put(
+            {
+                "logger_name": logger.name,
+                "logger_level": logger.level,
+                "num_handlers": len(handlers),
+                "handler_class": handler.__class__.__name__ if handler else None,
+                "handler_has_formatter": handler.formatter is not None if handler else None,
+            }
+        )
+
+    elif case == "setup_with_submodules_apply_to_capture_warnings":
+        # Extra handler to attach via extra_handlers
+        stream = io.StringIO()
+        stream_handler = logging.StreamHandler(stream)
+
+        setup(
+            level="INFO",
+            console=False,
+            color=False,
+            propagate=False,
+            disable_existing_loggers=False,
+            capture_warnings=True,
+            submodule_levels={"agentlightning.io": "DEBUG"},
+            extra_handlers=[stream_handler],
+            apply_to=["external"],
+        )
+
+        base = logging.getLogger("agentlightning")
+        sub = logging.getLogger("agentlightning.io")
+        ext = logging.getLogger("external")
+
+        # Capture warnings via logging after capture_warnings=True
+        class ListHandler(logging.Handler):
+            def __init__(self) -> None:
+                super().__init__()
+                self.records: List[logging.LogRecord] = []
+
+            def emit(self, record: logging.LogRecord) -> None:
+                self.records.append(record)
+
+        lh = ListHandler()
+        wlog = logging.getLogger("py.warnings")
+        wlog.handlers.clear()
+        wlog.addHandler(lh)
+        wlog.setLevel(logging.WARNING)
+        wlog.propagate = False
+
+        warnings.warn("from warnings", UserWarning)
+
+        queue.put(
+            {
+                "base_level": base.level,
+                "base_num_handlers": len(base.handlers),
+                "extra_in_base": stream_handler in base.handlers,
+                "sub_level": sub.level,
+                "ext_level": ext.level,
+                "ext_handlers_same": base.handlers == ext.handlers,
+                "ext_propagate": ext.propagate,
+                "warnings_logged": len(lh.records),
+            }
+        )
+
+    elif case == "setup_with_console_and_extra_handler":
+        # Console + extra handler combination to test handler attachment
+        stream = io.StringIO()
+        extra_handler = logging.StreamHandler(stream)
+
+        setup(
+            level="WARNING",
+            console=True,
+            color=False,
+            propagate=False,
+            extra_handlers=[extra_handler],
+        )
+
+        base = logging.getLogger("agentlightning")
+        handler_classes = [h.__class__.__name__ for h in base.handlers]
+        has_extra = extra_handler in base.handlers
+
+        queue.put(
+            {
+                "base_level": base.level,
+                "num_handlers": len(base.handlers),
+                "handler_classes": handler_classes,
+                "has_extra": has_extra,
+            }
+        )
+
+    else:
+        queue.put({})
+
+
+def _logging_worker_files_string(queue: Queue[Dict[str, Any]], base_dir: str) -> None:
+    """
+    Runs in a separate spawned process and configures logging with a single
+    files=str path. Returns information about the attached FileHandler.
+    """
+    import logging
+    import os
+
+    from agentlightning.logging import setup
+
+    log_path = os.path.join(base_dir, "logs", "agent.log")
+
+    setup(
+        level="INFO",
+        console=False,
+        color=False,
+        propagate=False,
+        files=log_path,
+    )
+
+    base = logging.getLogger("agentlightning")
+    file_handlers = [h for h in base.handlers if isinstance(h, logging.FileHandler)]
+    fh = file_handlers[0] if file_handlers else None
+
+    fmt = fh.formatter._fmt if fh and fh.formatter else None
+    datefmt = fh.formatter.datefmt if fh and fh.formatter else None
+
+    queue.put(
+        {
+            "logger_level": base.level,
+            "num_handlers": len(base.handlers),
+            "num_file_handlers": len(file_handlers),
+            "file_base": fh.baseFilename if fh else None,
+            "file_level": fh.level if fh else None,
+            "fmt": fmt,
+            "datefmt": datefmt,
+        }
+    )
+
+
+def _logging_worker_files_mapping(queue: Queue[Dict[str, Any]], base_dir: str) -> None:
+    """
+    Runs in a separate spawned process and configures logging with a files=dict
+    mapping, then calls setup twice to verify idempotent FileHandler attachment.
+    """
+    import logging
+    import os
+
+    from agentlightning.logging import setup
+
+    base_log = os.path.join(base_dir, "agent.log")
+    external_log = os.path.join(base_dir, "external.log")
+
+    files_mapping: Dict[str, str] = {
+        "agentlightning": base_log,
+        "external": external_log,
+    }
+
+    def file_handlers(logger: logging.Logger) -> list[logging.FileHandler]:
+        return [h for h in logger.handlers if isinstance(h, logging.FileHandler)]
+
+    # First setup call
+    setup(
+        level="DEBUG",
+        console=False,
+        color=False,
+        propagate=False,
+        files=files_mapping,
+    )
+
+    base_logger = logging.getLogger("agentlightning")
+    ext_logger = logging.getLogger("external")
+
+    base_fh_first = file_handlers(base_logger)
+    ext_fh_first = file_handlers(ext_logger)
+
+    # Second setup call with the same mapping should not add duplicate FileHandlers
+    setup(
+        level="DEBUG",
+        console=False,
+        color=False,
+        propagate=False,
+        files=files_mapping,
+    )
+
+    base_fh_second = file_handlers(base_logger)
+    ext_fh_second = file_handlers(ext_logger)
+
+    queue.put(
+        {
+            "base_level": base_logger.level,
+            "ext_level": ext_logger.getEffectiveLevel(),
+            "base_first_count": len(base_fh_first),
+            "ext_first_count": len(ext_fh_first),
+            "base_second_count": len(base_fh_second),
+            "ext_second_count": len(ext_fh_second),
+            "base_file_first": base_fh_first[0].baseFilename if base_fh_first else None,
+            "ext_file_first": ext_fh_first[0].baseFilename if ext_fh_first else None,
+            "base_file_second": base_fh_second[0].baseFilename if base_fh_second else None,
+            "ext_file_second": ext_fh_second[0].baseFilename if ext_fh_second else None,
+            # For sanity: capture handler levels as well
+            "base_handler_level": base_fh_first[0].level if base_fh_first else None,
+            "ext_handler_level": ext_fh_first[0].level if ext_fh_first else None,
+        }
+    )
+
+
+def _run_case(case: str) -> Dict[str, Any]:
+    """Helper to run a scenario in a spawn’ed process and fetch the result."""
+    ctx = mp.get_context("spawn")
+    q: Queue[Dict[str, Any]] = ctx.Queue()
+    p = ctx.Process(target=_logging_worker, args=(case, q))
+    p.start()
+    result = q.get(timeout=10)
+    p.join(timeout=10)
+    assert p.exitcode == 0
+    return result
+
+
+def test_to_level_value_int_and_str() -> None:
+    # direct, no multiprocessing needed
+    assert _to_level_value(logging.DEBUG) == logging.DEBUG
+    assert _to_level_value("info") == logging.INFO
+    assert _to_level_value("WARNING") == logging.WARNING
+
+    with pytest.raises(ValueError):
+        _to_level_value("not-a-level")
+
+
+def test_setup_module_plain_console_spawn() -> None:
+    result = _run_case("setup_module_plain_console")
+
+    assert result["logger_name"] == "agentlightning.test"
+    assert result["logger_level"] == logging.DEBUG
+
+    # Console handler with plain formatter configured
+    assert result["num_handlers"] == 1
+    assert result["handler_class"].endswith("StreamHandler")
+    assert result["handler_level"] == logging.DEBUG
+    assert result["fmt"] == DEFAULT_FORMAT
+    assert result["datefmt"] == DATE_FORMAT
+
+
+def test_setup_module_color_rich_spawn() -> None:
+    # Only run this test if rich is installed
+    pytest.importorskip("rich")
+
+    result = _run_case("setup_module_color_rich")
+
+    assert result["logger_name"] == "agentlightning.rich"
+    assert result["logger_level"] == logging.INFO
+    assert result["num_handlers"] == 1
+    # We can’t rely on full module path, just the class name
+    assert result["handler_class"].endswith("RichHandler")
+
+
+def test_setup_with_submodules_apply_to_and_capture_warnings_spawn() -> None:
+    result = _run_case("setup_with_submodules_apply_to_capture_warnings")
+
+    # Base logger level and handler attachment
+    assert result["base_level"] == logging.INFO
+    assert result["base_num_handlers"] >= 1
+    assert result["extra_in_base"] is True
+
+    # Submodule level overridden
+    assert result["sub_level"] == logging.DEBUG
+
+    # apply_to logger mirrors base handlers & level, propagation disabled
+    assert result["ext_level"] == logging.INFO
+    assert result["ext_handlers_same"] is True
+    assert result["ext_propagate"] is False
+
+    # capture_warnings=True causes warnings.warn to go through logging
+    assert result["warnings_logged"] >= 1
+
+
+def test_setup_with_console_and_extra_handler_spawn() -> None:
+    result = _run_case("setup_with_console_and_extra_handler")
+
+    # Level propagated to base logger
+    assert result["base_level"] == logging.WARNING
+
+    # Both console handler and extra handler should be attached
+    assert result["num_handlers"] >= 2
+    assert any(cls.endswith("StreamHandler") for cls in result["handler_classes"])
+    assert result["has_extra"] is True
+
+
+def test_setup_files_string_spawn(tmp_path: Path) -> None:
+    """
+    Verifies that passing files as a string attaches a single FileHandler with
+    the expected level and default formatter in a spawned process.
+    """
+    ctx = mp.get_context("spawn")
+    q: Queue[Dict[str, Any]] = ctx.Queue()
+    p = ctx.Process(target=_logging_worker_files_string, args=(q, str(tmp_path)))
+    p.start()
+    result = q.get(timeout=10)
+    p.join(timeout=10)
+    assert p.exitcode == 0
+
+    assert result["logger_level"] == logging.INFO
+    # We expect at least one handler and exactly one FileHandler
+    assert result["num_handlers"] >= 1
+    assert result["num_file_handlers"] == 1
+
+    # Filename should be inside the tmp_path tree
+    assert str(tmp_path) in result["file_base"]
+    # FileHandler uses the base logger level
+    assert result["file_level"] == logging.INFO
+
+    # Default formatter applied by _ensure_file_handler
+    assert result["fmt"] == DEFAULT_FORMAT
+    assert result["datefmt"] == DATE_FORMAT
+
+
+def test_setup_files_mapping_spawn(tmp_path: Path) -> None:
+    """
+    Verifies that passing files as a mapping attaches FileHandlers to each
+    logger and that calling setup twice does not create duplicate handlers.
+    """
+    ctx = mp.get_context("spawn")
+    q: Queue[Dict[str, Any]] = ctx.Queue()
+    p = ctx.Process(target=_logging_worker_files_mapping, args=(q, str(tmp_path)))
+    p.start()
+    result = q.get(timeout=10)
+    p.join(timeout=10)
+    assert p.exitcode == 0
+
+    # Base logger level is DEBUG
+    assert result["base_level"] == logging.DEBUG
+
+    # External's effective level is WARNING (inherited from root)
+    assert result["ext_level"] == logging.WARNING
+
+    # First setup: one FileHandler per logger
+    assert result["base_first_count"] == 1
+    assert result["ext_first_count"] == 1
+
+    # Second setup: still one FileHandler per logger (idempotence)
+    assert result["base_second_count"] == 1
+    assert result["ext_second_count"] == 1
+
+    # File paths are stable across calls
+    assert result["base_file_first"] == result["base_file_second"]
+    assert result["ext_file_first"] == result["ext_file_second"]
+
+    # Paths should live under tmp_path
+    assert str(tmp_path) in result["base_file_first"]
+    assert str(tmp_path) in result["ext_file_first"]
+
+    # Handler levels:
+    # - base handler uses the base logger level (DEBUG)
+    # - external handler uses external's effective level at creation (WARNING)
+    assert result["base_handler_level"] == logging.DEBUG
+    assert result["ext_handler_level"] == logging.WARNING

tests/utils/test_server_launcher.py

@@ -791,9 +791,6 @@ async def test_run_gunicorn_reports_health_failure_preload():
     """
     Health endpoint returns 503 -> watchdog posts error and requests graceful shutdown.
     """
-    from agentlightning.logging import configure_logger
-
-    configure_logger(logging.DEBUG)
     host = "127.0.0.1"
     port = portpicker.pick_unused_port()
     ctx = multiprocessing.get_context("fork")

tests/utils/test_system_snapshot.py

@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Optional
+
+import pytest
+
+from agentlightning.utils import system_snapshot
+
+try:
+    import torch  # type: ignore
+
+    GPU_AVAILABLE = torch.cuda.is_available()
+except Exception:
+    GPU_AVAILABLE = False  # type: ignore
+
+
+def _patch_system_snapshot(monkeypatch: pytest.MonkeyPatch, include_gpu: bool = False) -> Optional[SimpleNamespace]:
+    monkeypatch.setattr(system_snapshot.platform, "processor", lambda: "test-cpu")
+    monkeypatch.setattr(system_snapshot.platform, "platform", lambda: "test-platform")
+    monkeypatch.setattr(system_snapshot.socket, "gethostname", lambda: "test-host")
+
+    def fake_cpu_count(logical: bool = True) -> int:
+        return 4 if logical else 2
+
+    monkeypatch.setattr(system_snapshot.psutil, "cpu_count", fake_cpu_count)
+    monkeypatch.setattr(system_snapshot.psutil, "cpu_percent", lambda _: 33.3)  # type: ignore
+
+    vm = SimpleNamespace(used=5 * (2**30), total=10 * (2**30), percent=50.0)
+    monkeypatch.setattr(system_snapshot.psutil, "virtual_memory", lambda: vm)
+
+    du = SimpleNamespace(used=2 * (2**30), total=8 * (2**30), percent=25.0)
+    monkeypatch.setattr(system_snapshot.psutil, "disk_usage", lambda _: du)  # type: ignore
+
+    net = SimpleNamespace(bytes_sent=4 * (2**20), bytes_recv=6 * (2**20))
+    monkeypatch.setattr(system_snapshot.psutil, "net_io_counters", lambda: net)
+
+    if include_gpu:
+        dummy_gpu = SimpleNamespace(
+            name="Test GPU",
+            utilization=37.5,
+            memory_used=1024,
+            memory_total=4096,
+            temperature=68,
+        )
+        monkeypatch.setattr(
+            system_snapshot.GPUStatCollection,
+            "new_query",
+            lambda *args, **kwargs: SimpleNamespace(gpus=[dummy_gpu]),  # type: ignore
+        )
+
+        return dummy_gpu
+    return None
+
+
+def test_system_snapshot_excludes_gpu_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    _patch_system_snapshot(monkeypatch)
+
+    snapshot = system_snapshot.system_snapshot()
+
+    assert snapshot["cpu_name"] == "test-cpu"
+    assert snapshot["cpu_cores"] == 2
+    assert snapshot["cpu_threads"] == 4
+    assert snapshot["cpu_usage_pct"] == 33.3
+    assert snapshot["mem_used_gb"] == 5.0
+    assert snapshot["mem_total_gb"] == 10.0
+    assert snapshot["mem_pct"] == 50.0
+    assert snapshot["disk_used_gb"] == 2.0
+    assert snapshot["disk_total_gb"] == 8.0
+    assert snapshot["disk_pct"] == 25.0
+    assert snapshot["bytes_sent_mb"] == 4.0
+    assert snapshot["bytes_recv_mb"] == 6.0
+    assert snapshot["host"] == "test-host"
+    assert snapshot["os"] == "test-platform"
+    assert "gpus" not in snapshot
+
+
+def test_system_snapshot_includes_gpus_when_requested(monkeypatch: pytest.MonkeyPatch) -> None:
+    dummy_gpu = _patch_system_snapshot(monkeypatch, include_gpu=True)
+    assert dummy_gpu is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+
+    assert snapshot["gpus"] == [
+        {
+            "gpu": dummy_gpu.name,
+            "util_pct": dummy_gpu.utilization,
+            "mem_used_mb": dummy_gpu.memory_used,
+            "mem_total_mb": dummy_gpu.memory_total,
+            "temp_c": dummy_gpu.temperature,
+        }
+    ]
+
+
+def test_sanity_check() -> None:
+    snapshot = system_snapshot.system_snapshot()
+    assert snapshot is not None
+
+    snapshot = system_snapshot.system_snapshot(include_gpu=True)
+    assert snapshot is not None
+
+    if GPU_AVAILABLE:
+        assert snapshot["gpus"] is not None
+        assert len(snapshot["gpus"]) > 0

uv.lock

@@ -130,6 +130,7 @@ dependencies = [
     { 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 = "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 = "gpustat", 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')" },
     { name = "gunicorn", 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 = "httpdbg", 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')" },
@@ -348,6 +349,7 @@ requires-dist = [
     { name = "aiohttp" },
     { name = "fastapi" },
     { name = "flask" },
+    { name = "gpustat" },
     { name = "graphviz" },
     { name = "gunicorn" },
     { name = "httpdbg" },
@@ -1482,6 +1484,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/5c/dbd00727a3dd165d7e0e8af40e630cd7e45d77b525a3218afaff8a87358e/blake3-1.0.8-cp314-cp314t-win_amd64.whl", hash = "sha256:421b99cdf1ff2d1bf703bc56c454f4b286fce68454dd8711abbcb5a0df90c19a", size = 215133, upload-time = "2025-10-14T06:47:16.069Z" },
 ]
 
+[[package]]
+name = "blessed"
+version = "1.23.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "wcwidth", 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/c6/70/057c35a79a2015a6e35e45101d710b7a84b92af0c1104399bb83d33cd0d4/blessed-1.23.0.tar.gz", hash = "sha256:56591a32966f704f6131f1400af4151d9e8f5f4144133a5ca034019763dee77b", size = 6745236, upload-time = "2025-11-03T02:50:12.633Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/bc/f4/668909d1273be078ce5fb9a6d75bcfd3bf0832f238b4b09e0eb020d0eec9/blessed-1.23.0-py3-none-any.whl", hash = "sha256:4c432dcde0d45112372d1d096b2c4c0a6a5db1b94d546124872d0c3e64b5ea26", size = 95330, upload-time = "2025-11-03T02:50:10.64Z" },
+]
+
 [[package]]
 name = "blinker"
 version = "1.9.0"
@@ -3178,6 +3192,17 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/86/f1/62a193f0227cf15a920390abe675f386dec35f7ae3ffe6da582d3ade42c7/googleapis_common_protos-1.70.0-py3-none-any.whl", hash = "sha256:b8bfcca8c25a2bb253e0e0b0adaf8c00773e5e6af6fd92397576680b807e0fd8", size = 294530, upload-time = "2025-04-14T10:17:01.271Z" },
 ]
 
+[[package]]
+name = "gpustat"
+version = "1.1.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "blessed", 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 = "nvidia-ml-py", 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 = "psutil", 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/79/c4/46d005aec3bf911cb030467d91e062a5386ff4a03e51874424cacc0f60c1/gpustat-1.1.1.tar.gz", hash = "sha256:c18d3ed5518fc16300c42d694debc70aebb3be55cae91f1db64d63b5fa8af9d8", size = 98052, upload-time = "2023-08-22T19:39:06.062Z" }
+
 [[package]]
 name = "graphviz"
 version = "0.21"
@@ -6695,6 +6720,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2f/d8/a6b0d0d0c2435e9310f3e2bb0d9c9dd4c33daef86aa5f30b3681defd37ea/nvidia_cusparselt_cu12-0.7.1-py3-none-win_amd64.whl", hash = "sha256:f67fbb5831940ec829c9117b7f33807db9f9678dc2a617fbe781cac17b4e1075", size = 271020911, upload-time = "2025-02-26T00:14:47.204Z" },
 ]
 
+[[package]]
+name = "nvidia-ml-py"
+version = "13.580.82"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/dd/6c/4a533f2c0185027c465adb6063086bc3728301e95f483665bfa9ebafb2d3/nvidia_ml_py-13.580.82.tar.gz", hash = "sha256:0c028805dc53a0e2a6985ea801888197765ac2ef8f1c9e29a7bf0d3616a5efc7", size = 47999, upload-time = "2025-09-11T16:44:56.267Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7f/96/d6d25a4c307d6645f4a9b91d620c0151c544ad38b5e371313a87d2761004/nvidia_ml_py-13.580.82-py3-none-any.whl", hash = "sha256:4361db337b0c551e2d101936dae2e9a60f957af26818e8c0c3a1f32b8db8d0a7", size = 49008, upload-time = "2025-09-11T16:44:54.915Z" },
+]
+
 [[package]]
 name = "nvidia-nccl-cu12"
 version = "2.26.2"