add claude code ci file

.github/workflows/examples-claude-code.yml

@@ -0,0 +1,151 @@
+name: Examples - Claude Code
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 4 AM UTC+8
+    - cron: "0 20 * * *"
+
+  workflow_dispatch:
+
+  repository_dispatch:
+    types: [ci-claude-code, ci-all]
+
+run-name: >-
+  ${{ github.event_name == 'repository_dispatch'
+      && format(
+        'Claude Code - PR #{0} - {1} - {2}',
+        github.event.client_payload.pull_number,
+        github.event.client_payload.ci_label,
+        github.event.client_payload.correlation_id
+      )
+      || format('Claude Code - {0}', github.event_name) }}
+
+jobs:
+  claude-code:
+    if: >
+      github.event_name != 'repository_dispatch' ||
+      github.event.action == 'ci-claude-code' ||
+      github.event.action == 'ci-all'
+    name: Claude Code (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+          - 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
+        run: |
+          uv sync --frozen --no-default-groups \
+            --group dev --group experiment --group agents --group torch-gpu-stable
+      - 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-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Download model
+        run: |
+          source .venv/bin/activate
+          python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
+
+      - name: Launch vLLM server
+        run: |
+          set -euo pipefail
+          source .venv/bin/activate
+          vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
+            --max-model-len 131072 \
+            --enable-auto-tool-choice \
+            --tool-call-parser qwen3_coder \
+            --port 45993 &
+
+          VLLM_READY=0
+          for i in {1..100}; do
+            if curl -sSf http://localhost:45993/v1/models > /dev/null 2>&1; then
+              echo "vLLM server is ready!"
+              VLLM_READY=1
+              break
+            fi
+            echo "Waiting for vLLM server to be ready... (${i})"
+            sleep 5
+          done
+          if [[ "$VLLM_READY" != "1" ]]; then
+            echo "vLLM server failed to start!"
+            exit 1
+          fi
+
+      - name: Claude Code sanity check with vLLM models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py vllm --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct --base-url http://localhost:45993/v1 --debug
+        shell: bash
+
+      - name: Upload sanity check artifacts for vLLM
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error
+
+      - name: Cleanup vLLM
+        run: |
+          set -euo pipefail
+          pkill -f vllm
+          for i in {1..60}; do
+            if ! pgrep -f vllm; then
+              break
+            fi
+            sleep 5
+          done
+          rm -rf examples/claude_code/data/
+          rm -rf examples/claude_code/logs/
+
+      - name: Claude Code sanity check with OpenAI models
+        run: |
+          source .venv/bin/activate
+          cd examples/claude_code
+          python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
+        shell: bash
+        env:
+          OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
+          OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
+
+      - name: Upload sanity check artifacts for OpenAI
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
+          path: |
+            examples/claude_code/data/
+            examples/claude_code/logs/
+          if-no-files-found: error

agentlightning/llm_proxy.py

@@ -175,6 +175,24 @@ class AddReturnTokenIds(CustomLogger):
         return {**data, "return_token_ids": True}
 
 
+class AddLogprobs(CustomLogger):
+    """LiteLLM logger hook to request logprobs from vLLM.
+
+    This mutates the outgoing request payload to include `logprobs=1`
+    for backends that support logprobs return (e.g., vLLM).
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload."""
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure logprobs are requested from the backend when supported.
+        return {**data, "logprobs": 1}
+
+
 class LightningSpanExporter(SpanExporter):
     """Buffered OTEL span exporter with subtree flushing and training-store sink.
 
@@ -981,6 +999,7 @@ _MIDDLEWARE_REGISTRY: Dict[str, Type[BaseHTTPMiddleware]] = {
 
 _CALLBACK_REGISTRY = {
     "return_token_ids": AddReturnTokenIds,
+    "logprobs": AddLogprobs,
     "opentelemetry": LightningOpenTelemetry,
 }
 
@@ -1039,7 +1058,7 @@ class LLMProxy:
             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".
+            Available callback aliases are: "return_token_ids", "opentelemetry", "logprobs".
     """
 
     def __init__(
@@ -1053,8 +1072,8 @@ class LLMProxy:
         num_workers: int = 1,
         launch_mode: LaunchMode = "mp",
         launcher_args: PythonServerLauncherArgs | None = None,
-        middlewares: List[Union[Type[BaseHTTPMiddleware], str]] | None = None,
-        callbacks: List[Union[Type[CustomLogger], str]] | None = None,
+        middlewares: Sequence[Union[Type[BaseHTTPMiddleware], str]] | None = None,
+        callbacks: Sequence[Union[Type[CustomLogger], str]] | None = None,
     ):
         self.store = store
 

agentlightning/tracer/otel.py

@@ -20,6 +20,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 
 from agentlightning.semconv import LightningResourceAttributes
 from agentlightning.store.base import LightningStore
+from agentlightning.utils.otel import get_tracer_provider
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
 from .base import Tracer
@@ -51,7 +52,16 @@ class OtelTracer(Tracer):
         logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")
 
         if self._initialized:
-            logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
+            logger.info(f"[Worker {worker_id}] Tracer provider is already initialized. Skipping initialization.")
+            return
+
+        try:
+            get_tracer_provider()
+            logger.error(
+                f"[Worker {worker_id}] Tracer provider is already initialized but not by OtelTracer. OpenTelemetry may not work as expected."
+            )
+        except RuntimeError:
+            logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
 
         self._tracer_provider = TracerProvider()
         trace_api.set_tracer_provider(self._tracer_provider)
@@ -66,8 +76,7 @@ class OtelTracer(Tracer):
 
     def teardown_worker(self, worker_id: int):
         super().teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
-        self._tracer_provider = None
+        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")
 
     @asynccontextmanager
     async def trace_context(

agentlightning/utils/otel.py

@@ -17,7 +17,6 @@ from pydantic import TypeAdapter
 
 from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
 from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
-from agentlightning.tracer.otel import LightningSpanProcessor
 from agentlightning.types import SpanLike
 from agentlightning.utils.otlp import LightningStoreOTLPExporter
 
@@ -52,6 +51,8 @@ def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
         inspect: Whether to inspect the tracer provider and log its configuration.
             When it's on, make sure you also set the logger level to DEBUG to see the logs.
     """
+    from agentlightning.tracer.otel import LightningSpanProcessor
+
     if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
         raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
     tracer_provider = otel_get_tracer_provider()

examples/.gitignore

@@ -4,6 +4,7 @@ outputs/
 checkpoints/
 calc-x-data.zip
 spider-data.zip
+claude_code/logs/
 agentops.log
 unsloth/models/
 unsloth/unsloth_compiled_cache/

examples/claude_code/README.md

@@ -1,15 +1,22 @@
-# Training ANY LLM to Claude Code
+# Training Claude Code with Agent-lightning
+
+This example demonstrates how to train a Claude Code agent with Agent-lightning. **The example is still under development.**
+
+It wraps Claude Code as the agent to:
 
-This example wraps Claude Code as the agent to:
 1. collect traces from agent execution on coding tasks;
 2. train a hosted LLM with the traces ***🔨 Under development***
 
 ## Requirements
-1. install [agentlightning](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/)
-2. `(uv) pip install swebench` for evaluation
+
+1. Install agentlightning following [installation instructions](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/);
+2. `(uv) pip install swebench` for evaluation.
 
 ## Dataset
-We provide a small dataset `swe_debug.jsonl` which is a subset of [SWE-bench](https://huggingface.co/datasets/SWE-bench/SWE-bench) for sanity check.
+
+We provide a small dataset `swebench_samples.jsonl` which is a subset of [SWE-bench](https://huggingface.co/datasets/SWE-bench/SWE-bench) for sanity check.
+
+The instruction to prepare the full dataset is still underway.
 
 ## Included Files
 
@@ -25,12 +32,16 @@ We provide a small dataset `swe_debug.jsonl` which is a subset of [SWE-bench](ht
 | `swebench_utils/`               | Utility module with helper functions for SWE-bench dataset containerized exeuction and evaluation |
 
 ## Trace collection
+
 We support running Claude Code via two ways:
-- Hosted LLM servers, supporting versatile customizations
-- Official Claude Code
+
+- Hosted LLM servers (i.e., vLLM), useful for fine-tuning the LLM;
+- Official Claude Code (i.e., via Anthropic API), useful for prompt tuning.
 
 ### From Hosted LLM server
+
 1. Prepare an OpenAI-compatible server:
+
 ```bash
 vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
     --max-model-len 131072 \
@@ -39,10 +50,9 @@ vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
 ```
 
 2. Sanity check:
-```bash
-cd examples/cc
 
-# Suppose the vllm server is running at localhost
+```bash
+# Suppose the vllm server is running at localhost:8000
 python cc_agent \
     --model_name_or_path Qwen/Qwen3-Coder-30B-A3B-Instruct \
     --server_address http://localhost:8000/v1 \
@@ -50,6 +60,7 @@ python cc_agent \
     --max_step 32 \
     --output_dir data_debug
 ```
+
 The above commands will generate a `data_debug` dir, which contains two targets: (1) a Huggingface Dataset named `dataset-<instance_id>` and (2) a trace file named `stream_<instance_id>.jsonl`, where `instance_id` is a unique key of the SWE-bench samples.
 The dataset showcases the versatile customization capability of agent-lightning. In particular, we support extracting **prompt/response ids**, **logprobs** from the vllm server.
 The trace file is the conversation logs for claude code to tackle the SWE-bench instance.

examples/claude_code/cc_agent.py

@@ -1,404 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Main module for the Claude Code Agent implementation.
-
-This module provides the core functionality for running Claude Code agent experiments
-on SWE-bench datasets. It includes the CodingAgent class that implements the agent logic,
-functions for loading datasets, and asynchronous execution functions for running experiments.
-
-Key components:
-- CodingAgent: Main agent implementation that handles rollout logic
-- Dataset loading utilities
-- Asynchronous execution functions for dry runs and full datasets
-"""
-
-import asyncio
-import json
-import os
-import platform
-import resource
-import time
-from typing import Any, Dict, List, Literal, Optional, cast
-
-from claude_code_controller import ClaudeController
-from custom_adapter import LlmProxyTraceToAugmentedTriplet
-from custom_callbacks import AddLogprobs
-from datasets import Dataset
-from swebench.harness.utils import load_swebench_dataset  # type: ignore
-from swebench_utils.evaluation import evaluate
-from swebench_utils.logger import logger
-from transformers import AutoTokenizer
-
-from agentlightning import (
-    InMemoryLightningStore,
-    LightningStoreServer,
-    LitAgentRunner,
-    OtelTracer,
-    configure_logger,
-)
-from agentlightning.litagent import LitAgent
-from agentlightning.llm_proxy import LLMProxy, ModelConfig
-from agentlightning.types import LLM, AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span
-
-
-def load_dataset(path: str = "swe_debug.jsonl", epoch: int = 0, limit: Optional[int] = None) -> List[Dict[str, Any]]:
-    instances: List[Dict[str, Any]] = []
-    with open(path) as f:
-        for line in f:
-            instance = json.loads(line)
-            instance["epoch"] = epoch
-            instances.append(instance)
-
-    if limit is not None:
-        instances = instances[:limit]
-    return instances
-
-
-class CodingAgent(LitAgent[Dict[str, Any]]):
-    def __init__(
-        self,
-        namespace: Literal["swebench", "starryzhang"] = "swebench",
-        full_set: Literal["princeton-nlp/SWE-bench", "SWE-bench-Live/SWE-bench-Live"] = "princeton-nlp/SWE-bench",
-        split: str = "test",
-        max_step: int = 5,
-        run_method: Literal["python", "cli"] = "cli",
-        open_file_limit: int = 4096,
-        cache_level: str = "env",  # ["none", "base", "env", "instance"]
-        clean: bool = False,
-        force_rebuild: bool = False,
-        timeout: int = 1_800,  # in sec
-        instance_image_tag: str = "latest",
-        rewrite_reports: bool = False,
-    ) -> None:
-        super().__init__()
-        self.namespace = namespace
-        self.full_set = full_set
-        self.split = split
-        self.max_step = max_step
-        self.run_method = run_method
-
-        self.cache_level = cache_level
-        self.clean = clean
-        self.force_rebuild = force_rebuild
-        self.timeout = timeout
-        self.instance_image_tag = instance_image_tag
-        self.rewrite_reports = rewrite_reports
-
-        full_dataset = load_swebench_dataset(full_set, split)
-        self.dataset = {each["instance_id"]: each for each in full_dataset}
-
-        # run instances locally
-        if platform.system() == "Linux":
-            resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
-
-    async def rollout_async(
-        self, task: Dict[str, Any], resources: NamedResources, rollout: Rollout
-    ) -> RolloutRawResult:
-        run_id = f"epoch_{task.get('epoch', 0)}"
-        image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_")
-
-        llm = cast(ProxyLLM, resources.get("llm"))
-        assert llm is not None, "LLM resource is required for rollout."
-
-        llm = self._strip_proxy_helper(llm, rollout)
-
-        try:
-            # 1. init container
-            controller = ClaudeController(
-                image, task, run_id, llm.endpoint, llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy")
-            )
-            # 2. execute task
-            prediction = controller.run_instance(
-                task, max_step=self.max_step, run_method=cast(Literal["python", "cli"], self.run_method)
-            )
-            del controller
-        except Exception as e:
-            logger(run_id, task["instance_id"], f"Exception during rollout: {e}")
-            return 0.0
-
-        # 3. obtain rewards (evaluation result)
-        reward = 0.0
-        # empty patch
-        if prediction["model_patch"] in ["", None]:
-            return reward
-
-        instance_id = prediction["instance_id"]
-
-        result = evaluate(
-            prediction,
-            self.dataset[instance_id],
-            self.cache_level,
-            self.clean,
-            self.force_rebuild,
-            run_id,
-            self.timeout,
-            namespace=self.namespace,
-            instance_image_tag=self.instance_image_tag,
-            rewrite_reports=self.rewrite_reports,
-        )
-
-        # error patch
-        if result is None:
-            return reward
-
-        report = result[1]
-        # resolved/unresolved patch
-        if report[instance_id]["resolved"]:
-            reward = 1.0
-        return reward
-
-    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
-        """Convert [`ProxyLLM`][agentlightning.ProxyLLM] instances into concrete LLMs.
-
-        It resolves ProxyLLM instances to their concrete LLM implementation
-        by attaching the attempted rollout context. This is only used when the function
-        signature accepts an `llm` parameter and strip_proxy is True.
-
-        Args:
-            proxy_llm: Candidate LLM resource.
-            rollout: Rollout metadata that provides rollout and attempt identifiers.
-
-        Returns:
-            [`LLM`][agentlightning.LLM] with rollout context baked into the endpoint.
-
-        Raises:
-            ValueError: If the rollout is not an
-                [`AttemptedRollout`][agentlightning.AttemptedRollout].
-        """
-
-        if not isinstance(proxy_llm, ProxyLLM):
-            # Not a ProxyLLM, nothing to strip here.
-            return proxy_llm
-
-        # Rollout is still a Rollout here because API is not stabilized yet.
-        # In practice, it must be an AttemptedRollout.
-        if not isinstance(rollout, AttemptedRollout):
-            raise ValueError("Rollout is not an AttemptedRollout.")
-
-        return proxy_llm.with_attempted_rollout(rollout)
-
-
-def flatten_messages(messages: List[Any]) -> List[Dict[str, str]]:
-    flattened: List[Dict[str, str]] = []
-    for msg in messages:
-        if msg["role"] in ["system", "user"] and isinstance(msg["content"], list):
-            msg_content: List[str] = []
-            for content in msg["content"]:
-                msg_content.append(content["text"])
-
-            msg["content"] = "".join(msg_content)
-        elif msg["role"] == "assistant" and "tool_calls" in msg:
-            # NOTE:
-            # Tool calls are list of dict, though in most case only one tool call is made per call
-            # We serialize it as json string here to avoid nested structure
-            msg["tool_calls"] = json.dumps(msg["tool_calls"])
-
-        for k in msg:
-            assert isinstance(msg[k], str), f"\n>>> {msg}"
-        flattened.append(msg)
-    return flattened
-
-
-async def cc_agent_dry_run_sample(
-    model_path: str,
-    server_address: str,
-    dataset_path: str,
-    sonnet_name: str,
-    haiku_name: str,
-    max_step: int,
-    output_dir: Optional[str],
-) -> None:
-    """Run a dry run of the cc agent on a single sample.
-
-    This is a simple test function that runs the math agent on the first 4 problems
-    using a single worker. Useful for testing the setup and configuration.
-    """
-    dataset = load_dataset(dataset_path, limit=4)
-    # from_pretrained has partially unknown typing in some stubs; cast via Any to satisfy type-checkers.
-    tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_path)  # type: ignore
-
-    logging = configure_logger(name="Claude Code Agent")
-
-    tracer = OtelTracer()
-    runner = LitAgentRunner[Dict[str, Any]](tracer)
-    adapter = LlmProxyTraceToAugmentedTriplet()
-    store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
-    llm_proxy = LLMProxy(port=12358, store=store, callbacks=["return_token_ids", "opentelemetry", AddLogprobs])
-
-    await store.start()
-
-    llm_proxy.update_model_list(
-        [
-            ModelConfig(
-                model_name=f"{sonnet_name}",
-                litellm_params={
-                    "model": f"hosted_vllm/{model_path}",
-                    "api_base": server_address,
-                },
-            ),
-            ModelConfig(
-                model_name=f"{haiku_name}",
-                litellm_params={
-                    "model": f"hosted_vllm/{model_path}",
-                    "api_base": server_address,
-                },
-            ),
-        ]
-    )
-    await llm_proxy.restart()
-
-    # Put the LLM proxy address into the store as an address
-    await store.add_resources(
-        {
-            "llm": llm_proxy.as_resource(model="local"),
-        }
-    )
-
-    with runner.run_context(agent=CodingAgent(max_step=max_step), store=store):
-        rollout = await runner.step(
-            dataset[0],
-        )
-
-        spans = await store.query_spans(rollout.rollout_id)
-        triplets = adapter.adapt(cast(List[Span], spans))
-        logging.info(f"dump {len(spans)} spans, extract {len(triplets)} triplets")
-        if output_dir is not None:
-            with open(os.path.join(output_dir, f"stream_{dataset[0]['instance_id']}.json"), "w") as f:
-                for span in spans:
-                    f.write(json.dumps(span.model_dump()) + "\n")
-
-            all_triplets: List[Dict[str, Any]] = []
-            recent_reward: Optional[float] = None
-            for triplet in reversed(triplets):
-                if triplet.reward is not None:
-                    recent_reward = triplet.reward
-
-                prompt = tokenizer.decode(triplet.prompt["token_ids"])  # type: ignore
-                all_triplets.append(
-                    {
-                        "repo": rollout.input["repo"],
-                        "instance_id": rollout.input["instance_id"],
-                        "turn": triplet.metadata["sequence_id"],
-                        "prompt_ids": triplet.prompt["token_ids"],
-                        "gold_completion_ids": triplet.response["token_ids"],
-                        "logprobs": triplet.response["logprobs"],
-                        "reward": recent_reward,
-                        "prompt": prompt,
-                        "messages": flatten_messages(triplet.metadata["messages"]),
-                    }
-                )
-
-            ds = Dataset.from_list(all_triplets)  # type: ignore
-            ds.save_to_disk(os.path.join(output_dir, f"dataset-{dataset[0]['instance_id']}"))  # type: ignore
-            logging.info(f"Saved dataset with {len(ds)} samples to dataset-{dataset[0]['instance_id']}")
-
-
-async def gold_cc_agent_run_dataset(
-    sonnet_name: str,
-    haiku_name: str,
-    max_step: int,
-    dataset_path: str,
-    output_dir: Optional[str],
-):
-    """Run a dry run of the cc agent on a single sample.
-
-    This is a simple test function that runs the math agent on the first 4 problems
-    using a single worker. Useful for testing the setup and configuration.
-    """
-    dataset = load_dataset(dataset_path)
-
-    logging = configure_logger(name="Claude Code Agent")
-
-    tracer = OtelTracer()
-    runner = LitAgentRunner[Dict[str, Any]](tracer)
-    store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
-    llm_proxy = LLMProxy(port=12358, store=store)
-
-    await store.start()
-
-    llm_proxy.update_model_list(
-        [
-            ModelConfig(
-                model_name=f"{sonnet_name}",
-                litellm_params={"model": f"anthropic/{sonnet_name}", "api_key": "os.environ/ANTHROPIC_API_KEY"},
-            ),
-            ModelConfig(
-                model_name=f"{haiku_name}",
-                litellm_params={"model": f"anthropic/{haiku_name}", "api_key": "os.environ/ANTHROPIC_API_KEY"},
-            ),
-        ]
-    )
-    await llm_proxy.restart()
-
-    # Put the LLM proxy address into the store as an address
-    await store.add_resources(
-        {
-            "llm": llm_proxy.as_resource(model="local"),
-        }
-    )
-
-    for each in dataset:
-        with runner.run_context(agent=CodingAgent(max_step=max_step), store=store):
-            rollout = await runner.step(each)
-            spans = await store.query_spans(rollout.rollout_id)
-
-        if output_dir is None:
-            logging.info(f"instance {each['instance_id']} generate {len(spans)} spans")
-        else:
-            logging.info(f"instance {each['instance_id']} dump {len(spans)} spans to {output_dir}")
-            with open(os.path.join(output_dir, f"{each['instance_id']}.json"), "w") as f:
-                for span in spans:
-                    f.write(json.dumps(span.model_dump()) + "\n")
-
-        time.sleep(2 * 60)
-
-
-if __name__ == "__main__":
-    from argparse import ArgumentParser
-
-    parser = ArgumentParser()
-    # extract spans from official Claude Code
-    parser.add_argument("--official", action="store_true", help="Whether to run official claude code.")
-
-    # extract spans from hosted LLM server via litellm proxy
-    parser.add_argument(
-        "--model_name_or_path", type=str, default="Qwen/Qwen3-Coder-30B-A3B-Instruct", help="Model name or path."
-    )
-    parser.add_argument("--server_address", type=str, default="http://localhost:8000/v1", help="LLM server address.")
-
-    # common setup
-    parser.add_argument(
-        "--sonnet_name", type=str, default="claude-sonnet-4-5-20250929", help="Name of the sonnet model."
-    )
-    parser.add_argument("--haiku_name", type=str, default="claude-haiku-4-5-20251001", help="Name of the haiku model.")
-    parser.add_argument("--dataset_path", type=str, default="swe_debug.jsonl", help="Path to the dataset.")
-    parser.add_argument("--max_step", type=int, default=5, help="Maximum steps per instance.")
-    parser.add_argument("--output_dir", type=str, default="data", help="Directory to save output logs.")
-
-    args = parser.parse_args()
-
-    if args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    if not args.official:
-        asyncio.run(
-            cc_agent_dry_run_sample(
-                model_path=args.model_name_or_path,
-                server_address=args.server_address,
-                dataset_path=args.dataset_path,
-                sonnet_name=args.sonnet_name,
-                haiku_name=args.haiku_name,
-                max_step=args.max_step,
-                output_dir=args.output_dir,
-            )
-        )
-    else:
-        asyncio.run(
-            gold_cc_agent_run_dataset(
-                sonnet_name=args.sonnet_name,
-                haiku_name=args.haiku_name,
-                dataset_path=args.dataset_path,
-                output_dir=args.output_dir,
-                max_step=args.max_step,
-            )
-        )

examples/claude_code/claude_code_agent.py

@@ -0,0 +1,502 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Main module for the Claude Code Agent implementation.
+
+This module provides the core functionality for running Claude Code agent experiments
+on SWE-bench datasets. It includes the ClaudeCodeAgent class that implements the agent logic,
+functions for loading datasets, and asynchronous execution functions for running experiments.
+
+Key components:
+
+- Dataset loading utilities
+- ClaudeCodeAgent: Main agent implementation that handles rollout logic
+- Asynchronous execution functions for dry runs and full datasets
+"""
+
+import asyncio
+import json
+import logging
+import os
+import resource
+from argparse import ArgumentParser
+from typing import Any, Dict, List, Literal, Optional, Sequence, cast
+
+from claude_code_controller import ClaudeController
+from datasets import Dataset
+from extended_adapter import ExtendedLlmProxyTraceToTriplet
+from swebench.harness.constants import SWEbenchInstance
+from swebench.harness.utils import load_swebench_dataset  # pyright: ignore[reportUnknownVariableType]
+from swebench_utils.evaluation import evaluate
+from swebench_utils.logging import log_for_evaluation
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from agentlightning import (
+    InMemoryLightningStore,
+    LightningStoreServer,
+    LitAgentRunner,
+    OtelTracer,
+    setup_logging,
+    setup_module_logging,
+)
+from agentlightning.litagent import LitAgent
+from agentlightning.llm_proxy import LLMProxy, ModelConfig
+from agentlightning.store import LightningStore
+from agentlightning.types import AttemptedRollout, NamedResources, ProxyLLM, Rollout, RolloutRawResult, Span
+
+logger = logging.getLogger("claude_code_agent")
+
+
+def _load_dataset(path: str, epoch: int = 0, limit: Optional[int] = None) -> List[SWEbenchInstance]:
+    instances: List[SWEbenchInstance] = []
+    with open(path) as f:
+        for line in f:
+            instance = json.loads(line)
+            instance["epoch"] = epoch
+            instances.append(instance)
+
+    if limit is not None:
+        instances = instances[:limit]
+    return instances
+
+
+def _flatten_messages(messages: List[Any]) -> List[Dict[str, str]]:
+    flattened: List[Dict[str, str]] = []
+    for msg in messages:
+        if msg["role"] in ["system", "user"] and isinstance(msg["content"], list):
+            msg_content: List[str] = []
+            for content in msg["content"]:
+                msg_content.append(content["text"])
+
+            msg["content"] = "".join(msg_content)
+        elif msg["role"] == "assistant" and "tool_calls" in msg:
+            # NOTE:
+            # Tool calls are list of dict, though in most case only one tool call is made per call
+            # We serialize it as json string here to avoid nested structure
+            msg["tool_calls"] = json.dumps(msg["tool_calls"])
+
+        for k in msg:
+            assert isinstance(msg[k], str), f"\n>>> {msg}"
+        flattened.append(msg)
+    return flattened
+
+
+class ClaudeCodeAgent(LitAgent[SWEbenchInstance]):
+    """Claude Code Agent implementation.
+
+    This agent is a wrapper of the Claude Code controller,
+    and it should be used to run the Claude Code agent on SWE-bench datasets.
+    """
+
+    def __init__(
+        self,
+        namespace: Literal["swebench", "starryzhang"] = "swebench",
+        max_turns: int = 5,
+        run_method: Literal["python", "cli"] = "cli",
+        open_file_limit: int = 4096,
+        cache_level: str = "env",  # ["none", "base", "env", "instance"]
+        clean: bool = False,
+        force_rebuild: bool = False,
+        timeout: int = 1_800,  # in sec
+        instance_image_tag: str = "latest",
+        rewrite_reports: bool = False,
+        swebench_full_dataset: Optional[List[SWEbenchInstance]] = None,
+    ) -> None:
+        super().__init__()
+        self.namespace = namespace
+        self.max_turns = max_turns
+        self.run_method = run_method
+
+        self.cache_level = cache_level
+        self.clean = clean
+        self.force_rebuild = force_rebuild
+        self.timeout = timeout
+        self.instance_image_tag = instance_image_tag
+        self.rewrite_reports = rewrite_reports
+
+        self.swebench_full_dataset = (
+            {each["instance_id"]: each for each in swebench_full_dataset} if swebench_full_dataset is not None else {}
+        )
+
+        # Set the maximum number of open files to the specified limit.
+        resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
+
+    async def rollout_async(
+        self, task: SWEbenchInstance, resources: NamedResources, rollout: Rollout
+    ) -> RolloutRawResult:
+        if not isinstance(rollout, AttemptedRollout):
+            # Technically, rollout should be an AttemptedRollout here.
+            # but the API is not stabilized yet.
+            raise ValueError("Rollout is not an AttemptedRollout.")
+
+        run_id = f"epoch_{task.get('epoch', 0)}"
+        image = f"{self.namespace}/sweb.eval.x86_64.{task['instance_id'].lower()}".replace("__", "_1776_")
+
+        llm = cast(ProxyLLM, resources["llm"])
+
+        try:
+            # 1. init container
+            controller = ClaudeController(
+                image,
+                task,
+                run_id,
+                llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+                llm.api_key or os.environ.get("ANTHROPIC_AUTH_TOKEN", "dummy"),
+            )
+            # 2. execute task
+            prediction = controller.run_instance(
+                task, max_turns=self.max_turns, run_method=cast(Literal["python", "cli"], self.run_method)
+            )
+            del controller
+        except Exception as e:
+            log_for_evaluation(run_id, task["instance_id"], f"Exception during rollout: {e}")
+            return 0.0
+
+        # 3. obtain rewards (evaluation result)
+        reward = 0.0
+        # empty patch
+        if prediction["model_patch"] in ["", None]:
+            return reward
+
+        instance_id = prediction["instance_id"]
+
+        result = evaluate(
+            cast(Any, prediction),
+            self.swebench_full_dataset[instance_id],
+            self.cache_level,
+            self.clean,
+            self.force_rebuild,
+            run_id,
+            self.timeout,
+            namespace=self.namespace,
+            instance_image_tag=self.instance_image_tag,
+            rewrite_reports=self.rewrite_reports,
+        )
+
+        # error patch
+        if result is None:
+            return reward
+
+        report = result[1]
+        # resolved/unresolved patch
+        if report[instance_id]["resolved"]:
+            reward = 1.0
+        return reward
+
+
+def sanity_check_spans(spans: Sequence[Span]) -> None:
+    assert len(spans) > 1, f"At least two spans are expected for a valid rollout. Found {len(spans)} spans."
+    assert any(span.name == "raw_gen_ai_request" for span in spans), "raw_gen_ai_request span not found"
+    assert any(span.name == "agentlightning.annotation" for span in spans), "agentlightning.annotation span not found"
+
+
+async def run_instance_async(
+    instance: SWEbenchInstance,
+    agent: ClaudeCodeAgent,
+    runner: LitAgentRunner[SWEbenchInstance],
+    store: LightningStore,
+    output_dir: Optional[str],
+    adapter: Optional[ExtendedLlmProxyTraceToTriplet],
+    tokenizer: Optional[PreTrainedTokenizerBase],
+) -> None:
+    """Runs the agent on a specific SWE-bench instance.
+
+    Running on specific SWE-bench instance and queries the traced spans.
+    It then extracts the triplets and saves the dataset.
+    """
+
+    instance_id = instance["instance_id"]
+    logger.info(f"Starting to run instance: {instance_id}")
+
+    # Run the agent and query the traced spans.
+    with runner.run_context(agent=agent, store=store):
+        rollout = await runner.step(instance)
+
+    logger.info(f"Finished running instance: {instance_id}")
+
+    spans = await store.query_spans(rollout.rollout_id)
+
+    if output_dir is None:
+        logger.info(f"Generated {len(spans)} spans for {instance_id}")
+        return
+
+    # 1. Dump raw spans (Common for both types)
+    raw_path = os.path.join(output_dir, f"stream_{instance_id}.json")
+    with open(raw_path, "w") as f:
+        for span in spans:
+            f.write(json.dumps(span.model_dump()) + "\n")
+    logger.info(f"Dumped {len(spans)} spans to {raw_path}")
+
+    # 2. Extract Triplets and Save Dataset (vLLM specific)
+    if adapter is not None and tokenizer is not None:
+        try:
+            triplets = adapter.adapt(cast(List[Span], spans))
+            logger.info(f"Extracted {len(triplets)} triplets for {instance_id}")
+
+            all_triplets: List[Dict[str, Any]] = []
+            recent_reward: Optional[float] = None
+
+            # Process in reverse to propagate rewards if necessary/logic dictates
+            for triplet in reversed(triplets):
+                if triplet.reward is not None:
+                    recent_reward = triplet.reward
+
+                prompt_text = tokenizer.decode(triplet.prompt["token_ids"])  # type: ignore
+                all_triplets.append(
+                    {
+                        "repo": instance.get("repo", ""),
+                        "instance_id": instance_id,
+                        "turn": triplet.metadata["sequence_id"],
+                        "prompt_ids": triplet.prompt["token_ids"],
+                        "gold_completion_ids": triplet.response["token_ids"],
+                        "logprobs": triplet.response["logprobs"],
+                        "reward": recent_reward,
+                        "prompt": prompt_text,
+                        "messages": _flatten_messages(triplet.metadata["messages"]),
+                    }
+                )
+
+            if all_triplets:
+                ds = Dataset.from_list(all_triplets)  # type: ignore
+                save_path = os.path.join(output_dir, f"dataset-{instance_id}")
+                ds.save_to_disk(save_path)  # type: ignore
+                logger.info(f"Saved HuggingFace dataset to {save_path}")
+
+        except Exception as e:
+            logger.error(f"Failed to extract triplets for {instance_id}: {e}")
+
+    logger.info(f"Finished extracting spans and traces for instance: {instance_id}")
+
+    # Quickly sanity check the spans
+    sanity_check_spans(spans)
+    logger.info(f"Sanity check passed for instance: {instance_id}")
+
+
+async def dry_run_claude_code(
+    *,
+    dataset_path: str,
+    haiku_frontend_name: str,
+    haiku_backend_name: str,
+    sonnet_frontend_name: str,
+    sonnet_backend_name: str,
+    backend_type: Literal["vllm", "anthropic", "openai"],
+    api_base_url: Optional[str],
+    output_dir: Optional[str],
+    max_turns: int,
+    limit: Optional[int],
+    cooldown_seconds: float,
+) -> None:
+    """Executes a dry run of the Claude Code agent on a dataset.
+
+    This function handles both 'official' runs (interacting with Anthropic APIs)
+    and 'hosted' runs (interacting with vLLM or compatible servers). It manages
+    initialization of the Lightning Store, LLM Proxy, and the execution loop.
+
+    If running in 'vllm' mode, it will also attempt to extract triplets using
+    the provided backend name as the tokenizer path and save a HuggingFace Dataset.
+
+    Args:
+        dataset_path: Path to the JSONL dataset file.
+        haiku_frontend_name: The model name used in the code to request the 'fast' model.
+        haiku_backend_name: The actual model name/path on the backend.
+        sonnet_frontend_name: The model name used in the code to request the 'strong' model.
+        sonnet_backend_name: The actual model name/path on the backend.
+        backend_type: The type of backend to configure ("vllm", "anthropic" or "openai").
+        api_base_url: Base URL for the API. Required for "vllm" or "openai".
+        output_dir: Directory to save logs, spans, and datasets.
+        max_turns: Maximum number of steps the agent can take per instance.
+        limit: Optional limit on the number of instances to process.
+    """
+    dataset = _load_dataset(dataset_path, limit=limit)
+
+    # Initialize Infrastructure
+    tracer = OtelTracer()
+    runner = LitAgentRunner[SWEbenchInstance](tracer)
+    store = LightningStoreServer(InMemoryLightningStore(), host="0.0.0.0", port=7654)
+    await store.start()
+
+    # Enable callbacks for training data extraction if using vLLM
+    callbacks = ["return_token_ids", "opentelemetry", "logprobs"] if backend_type == "vllm" else ["opentelemetry"]
+    llm_proxy = LLMProxy(port=12358, store=store, callbacks=callbacks)
+
+    # Configure Models based on backend type
+    model_configs: List[ModelConfig] = []
+    model_params: Dict[str, Any] = {}
+
+    if backend_type == "vllm":
+        model_namespace = "hosted_vllm"
+        if api_base_url:
+            model_params["api_base"] = api_base_url
+        else:
+            raise ValueError("api_base_url is required for vllm backend")
+    elif backend_type == "anthropic":
+        model_namespace = "anthropic"
+        model_params["api_key"] = "os.environ/ANTHROPIC_API_KEY"
+        if api_base_url:
+            model_params["api_base"] = api_base_url
+    elif backend_type == "openai":
+        model_namespace = "openai"
+        model_params["api_key"] = "os.environ/OPENAI_API_KEY"
+        if api_base_url:
+            # Users can still override this via environment variables,
+            # even if they don't pass it in as an argument.
+            model_params["api_base"] = api_base_url
+
+    model_configs.extend(
+        [
+            ModelConfig(
+                model_name=sonnet_frontend_name,
+                litellm_params={
+                    "model": f"{model_namespace}/{sonnet_backend_name}",
+                    **model_params,
+                },
+            ),
+            ModelConfig(
+                model_name=haiku_frontend_name,
+                litellm_params={
+                    "model": f"{model_namespace}/{haiku_backend_name}",
+                    **model_params,
+                },
+            ),
+        ]
+    )
+
+    logger.info(f"Updating model list: {model_configs}")
+
+    llm_proxy.update_model_list(model_configs)
+    await llm_proxy.start()
+
+    try:
+        # Add the LLM proxy as a resource to the store
+        await store.add_resources({"llm": llm_proxy.as_resource(model="local")})
+
+        # Prepare for triplet extraction if vllm
+        adapter = ExtendedLlmProxyTraceToTriplet() if backend_type == "vllm" else None
+        tokenizer = None
+        if backend_type == "vllm":
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(sonnet_backend_name)  # type: ignore
+            except Exception as e:
+                logger.warning(f"Could not load tokenizer for {sonnet_backend_name}: {e}")
+
+        # Load full swebench dataset. Mainly for evaluation purposes.
+        swebench_full_dataset = load_swebench_dataset("princeton-nlp/SWE-bench", split="test")
+
+        # Initialize Claude Code Agent
+        claude_code_agent = ClaudeCodeAgent(swebench_full_dataset=swebench_full_dataset, max_turns=max_turns)
+
+        # Execution Loop
+        for instance in dataset:
+            await run_instance_async(
+                instance,
+                claude_code_agent,
+                runner,
+                store,
+                output_dir,
+                adapter,
+                cast(PreTrainedTokenizerBase, tokenizer),
+            )
+            # Basic sleep to allow resource cleanup or rate limit cooling
+            await asyncio.sleep(cooldown_seconds)
+
+    finally:
+        await llm_proxy.stop()
+        await store.stop()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Run Claude Code Agent experiments.")
+
+    # Backend Selection
+    parser.add_argument(
+        "backend_type",
+        type=str,
+        choices=["vllm", "anthropic", "openai"],
+        help="Backend type: 'vllm' for hosted models, 'anthropic' for official API, 'openai' for OpenAI API.",
+    )
+
+    # Model Configuration
+    parser.add_argument(
+        "--backend-model-high",
+        type=str,
+        default=None,
+        help="Backend model path/name for expensive model usages (used as vLLM model name / OpenAI model name).",
+    )
+    parser.add_argument(
+        "--backend-model-low",
+        type=str,
+        default=None,
+        help="Backend model path/name for low-price model usages (used as vLLM model name / OpenAI model name).",
+    )
+    parser.add_argument(
+        "--base-url", type=str, default="http://localhost:8000/v1", help="LLM server address (required for vllm)."
+    )
+
+    # Frontend/Agent Configuration
+    parser.add_argument(
+        "--frontend-model-high",
+        type=str,
+        default="claude-sonnet-4-5-20250929",
+        help="The frontend high-price model name provided to Claude Code.",
+    )
+    parser.add_argument(
+        "--frontend-model-low",
+        type=str,
+        default="claude-haiku-4-5-20251001",
+        help="The frontend low-price model name provided to Claude Code.",
+    )
+
+    # Execution Configuration
+    parser.add_argument("--dataset-path", type=str, default="swebench_samples.jsonl", help="Path to the dataset.")
+    parser.add_argument("--max-turns", type=int, default=5, help="Maximum turns per instance.")
+    parser.add_argument("--output-dir", type=str, default="data", help="Directory to save output logs.")
+    parser.add_argument("--limit", type=int, default=None, help="Limit the number of instances to run (for debugging).")
+    parser.add_argument("--cooldown-seconds", type=float, default=2.0, help="Cooldown seconds between instances.")
+    parser.add_argument("--debug", action="store_true", help="Enable debug loggings.")
+
+    args = parser.parse_args()
+
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    if args.debug:
+        setup_logging()
+        setup_module_logging("DEBUG", name="claude_code_agent")
+    else:
+        setup_logging(apply_to=[logger.name])
+
+    # Map backend_type to the appropriate args
+    backend_mode = cast(Literal["vllm", "anthropic", "openai"], args.backend_type)
+
+    # If using anthropic, the backend name usually matches the frontend or is specific API string.
+    # Otherwise, the backend name is the model name/path (e.g., Qwen/...) and must be provided.
+    if args.backend_model_high is None:
+        if args.backend_type == "anthropic":
+            backend_model_high = args.frontend_model_high
+        else:
+            raise ValueError("--backend-model-high is required for non-anthropic backends")
+    else:
+        backend_model_high = args.backend_model_high
+
+    if args.backend_model_low is None:
+        if args.backend_type == "anthropic":
+            backend_model_low = args.frontend_model_low
+        else:
+            raise ValueError("--backend-model-low is required for non-anthropic backends")
+    else:
+        backend_model_low = args.backend_model_low
+
+    asyncio.run(
+        dry_run_claude_code(
+            dataset_path=args.dataset_path,
+            haiku_frontend_name=args.frontend_model_low,
+            haiku_backend_name=backend_model_low,
+            sonnet_frontend_name=args.frontend_model_high,
+            sonnet_backend_name=backend_model_high,
+            backend_type=backend_mode,
+            api_base_url=args.base_url if backend_mode == "vllm" else None,
+            output_dir=args.output_dir,
+            max_turns=args.max_turns,
+            limit=args.limit,
+            cooldown_seconds=args.cooldown_seconds,
+        )
+    )

examples/claude_code/claude_code_controller.py

@@ -5,27 +5,22 @@
 This module provides the ClaudeController class that manages the execution of Claude Code
 within Docker containers. It handles container initialization, command execution, and
 patch application for SWE-bench evaluation tasks.
-
-Key features:
-- Container lifecycle management
-- Claude Code execution via CLI or Python SDK
-- Patch application and result retrieval
-- Integration with SWE-bench evaluation framework
 """
 
+import logging
 from functools import partial
-from typing import Any, Dict, Literal
+from typing import Literal, TypedDict
 
 import dotenv
+from swebench.harness.constants import SWEbenchInstance
 from swebench_utils.docker_runtime import Runtime
-from swebench_utils.logger import logger
+from swebench_utils.logging import log_for_evaluation
 
-
-class ClaudeController:
-    system_prompt = """
+SWEBENCH_EXTRA_SYSTEM_PROMPT = """
  You are an expert software engineer solving swebench bug fixing tasks.
 """
-    user_prompt = """
+
+SWEBENCH_USER_PROMPT = """
 You are given a code repository in the current directory (/testbed).
 The bug description is:
 {description}
@@ -39,104 +34,194 @@ You task is to fix the bug with the following steps:
 Please do not commit your edits. We will do it later.
 """
 
-    def __init__(self, image: str, instance: Dict[str, Any], run_id: str, endpoint: str, api_key: str) -> None:
+logger = logging.getLogger("claude_code_agent")
+
+
+class RunInstanceResult(TypedDict):
+    instance_id: str
+    model_patch: str
+    model_name_or_path: str
+
+
+class ClaudeController:
+    """Manages the execution of Claude Code within a Docker runtime.
+
+    This controller handles the lifecycle of a SWE-bench task execution, including
+    environment setup, tool installation, agent execution (via CLI or Python SDK),
+    and result extraction.
+
+    Attributes:
+        container: The active Docker runtime session.
+    """
+
+    def __init__(self, image: str, instance: SWEbenchInstance, run_id: str, endpoint: str, api_key: str) -> None:
+        """Initialize the ClaudeController.
+
+        Args:
+            image: The Docker image tag.
+            instance: The dataset instance containing the problem statement and ID.
+            run_id: The identifier for the evaluation run.
+            endpoint: The API endpoint URL.
+            api_key: The API authentication key.
+        """
         self.image = image
         self.instance = instance
         self.run_id = run_id
         self.endpoint = endpoint
         self.api_key = api_key
         self.container: Runtime = self.init_container(self.image, self.instance)
-        return
 
-    def init_container(self, image: str, instance: Dict[str, Any]) -> Runtime:
+    def init_container(self, image: str, instance: SWEbenchInstance) -> Runtime:
+        """Initializes the Docker container and sets up the Claude Code environment.
+
+        This method starts the container session, installs the Claude CLI,
+        configures environment variables for authentication and sandbox mode.
+
+        Args:
+            image: The Docker image tag to start.
+            instance: The dataset instance to load into the environment.
+
+        Returns:
+            An initialized and configured Docker runtime object.
+        """
         container = Runtime.start_session(
             image,
             instance,
-            log_function=partial(logger, run_id=self.run_id, instance_id=instance["instance_id"]),
+            log_function=partial(log_for_evaluation, run_id=self.run_id, instance_id=instance["instance_id"]),
         )
+        # Install Claude CLI
         container.send_command("curl -fsSL https://claude.ai/install.sh | bash")
         container.send_command('alias claude="$HOME/.local/bin/claude"')
+
+        # Configure Environment
         dotenv.load_dotenv()
-        # anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')
-        # container.send_command(f"export ANTHROPIC_API_KEY={anthropic_api_key}")
-        # if (not os.getenv("ANTHROPIC_BASE_URL")) or (not os.getenv("ANTHROPIC_AUTH_TOKEN")):
-        #     raise RuntimeError("ANTHROPIC_BASE_URL and ANTHROPIC_AUTH_TOKEN not found!")
         container.send_command(f"export ANTHROPIC_BASE_URL={self.endpoint}")
         container.send_command(f"export ANTHROPIC_AUTH_TOKEN={self.api_key}")
         container.send_command("export IS_SANDBOX=1")
+
         return container
 
-    def _run_cli(self, instance: Dict[str, Any], max_step: int, timelimit: int) -> None:
-        # prepare prompt safely: write it to a file inside the container using a single-quoted heredoc
+    def _run_cli(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
+        """Executes Claude Code using the Command Line Interface.
+
+        Constructs a safe heredoc for the prompt to avoid shell interpolation issues
+        and executes the `claude` binary directly.
+
+        Args:
+            instance: The problem instance containing the problem statement.
+            max_turns: The maximum number of interaction turns allowed.
+            time_limit: The execution time limit in minutes.
+        """
+        # Prepare prompt safely: write it to a file inside the container using a single-quoted heredoc
         # directly applying prompt for heredoc may raise error for windows line ending \r\n
-        prompt_text = self.user_prompt.format(description=instance["problem_statement"].replace('"""', "'''"))
-        # choose a simple filename and a heredoc delimiter unlikely to collide
+        prompt_text = SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
+
+        # Choose a simple filename and a heredoc delimiter unlikely to collide
         heredoc_cmd = "cat > /tmp/cc_prompt.txt <<'CC_PROMPT'\n" + prompt_text + "\nCC_PROMPT\n"
         self.container.send_command(heredoc_cmd)
 
-        # self.container.send_command("mkdir -p /testbed/.claude")
-        # with open("utils/settings.template.json") as f:
-        #     setting = f.read()
-        # setting_cmd = "cat > /testbed/.claude/settings.json <<'CC_SETTING'\n" + setting + "\nCC_SETTING\n"
-        # self.container.send_command(setting_cmd)
+        # Run claude reading the prompt from the file
+        claude_cmd = (
+            f'claude -p "$(cat /tmp/cc_prompt.txt)" '
+            f'--append-system-prompt "{SWEBENCH_EXTRA_SYSTEM_PROMPT}" '
+            f"--max-turns {max_turns} "
+            f"--dangerously-skip-permissions "
+            f"--output-format json --verbose"
+        )
+        logger.info(f"Running Claude Code CLI command: {claude_cmd}")
+        self.container.send_command(claude_cmd, time_limit * 60)
+        logger.info(f"Claude Code CLI command completed")
 
-        # with open("utils/handle_hook.template.sh") as f:
-        #     handler = f.read()
-        # handler_cmd = "cat > /tmp/handle_hook.sh <<'CC_HOOK'\n" + handler + "\nCC_HOOK\n"
-        # self.container.send_command(handler_cmd)
-        # self.container.send_command("chmod +x /tmp/handle_hook.sh")
+    def _run_python_sdk(self, instance: SWEbenchInstance, max_turns: int, time_limit: int) -> None:
+        """Executes Claude Code using the Python SDK wrapper.
 
-        # run claude reading the prompt from the file to avoid shell interpolation issues
-        claude_cmd = f'claude -p "$(cat /tmp/cc_prompt.txt)" --append-system-prompt "{self.system_prompt}" --max-turns {max_step} --dangerously-skip-permissions --output-format json --verbose'
-        self.container.send_command(claude_cmd, timelimit * 60)
-        # self.container.send_command("cat /tmp/hook.out")
-        return
+        Installs the Python SDK if necessary, hydrates a template script with the
+        problem prompt, and executes the generated Python script.
 
-    def _run_python_sdk(self, instance: Dict[str, Any], max_step: int, timelimit: int):
+        Note:
+            This path is still under development and not yet stable.
+
+        Args:
+            instance: The problem instance containing the problem statement.
+            max_turns: The maximum number of interaction turns allowed.
+            time_limit: The execution time limit in minutes.
+        """
+        # Ensure Python 3.12 is available
         self.container.send_command(
             f"""
 if ! command -v python3 &> /dev/null; then
     echo "Python is not installed. Installing Python 3.12..."
-    sudo apt-get update && sudo apt-get install -y python3.12
+    sudo apt-get update -qq && sudo apt-get install -y -qq python3.12
 else
     echo "Python is already installed."
 fi
 """
         )
         self.container.send_command("python3 -m pip install claude-code-sdk")
+
+        # Load and fill the execution template
         with open("src/agent/cc/claude_code_main.py.template") as f:
             entrance_template = f.read()
-        entrance_template.replace("SYS_PROMPT", self.system_prompt).replace(
-            "PROMPT", self.user_prompt.format(description=instance["problem_statement"].replace('"""', "'''"))
-        ).replace("MAX_STEP", str(max_step))
-        self.container.send_command(f"cat > /tmp/claude_code_main.py <<'CC_MAIN'\n{entrance_template}\nCC_MAIN\n")
-        self.container.send_command("python3 /tmp/claude_code_main.py", timelimit * 60)
+
+        script_content = (
+            entrance_template.replace("SYS_PROMPT", SWEBENCH_EXTRA_SYSTEM_PROMPT)
+            .replace(
+                "PROMPT", SWEBENCH_USER_PROMPT.format(description=instance["problem_statement"].replace('"""', "'''"))
+            )
+            .replace("MAX_STEP", str(max_turns))
+        )
+
+        # Write the script to the container and execute
+        self.container.send_command(f"cat > /tmp/claude_code_main.py <<'CC_MAIN'\n{script_content}\nCC_MAIN\n")
+        self.container.send_command("python3 /tmp/claude_code_main.py", time_limit * 60)
         return
 
     def run_instance(
         self,
-        instance: Dict[str, Any],
-        max_step: int = 40,
-        timelimit: int = 30,
+        instance: SWEbenchInstance,
+        max_turns: int = 40,
+        time_limit: int = 30,
         run_method: Literal["python", "cli"] = "python",
-    ) -> Dict[str, Any]:
-        """
-        timelimit: in minute
+    ) -> RunInstanceResult:
+        """Runs the agent on a specific SWE-bench instance.
+
+        This method orchestrates the agent execution via the specified method (CLI or Python),
+        and extracts the generated git diff (patch) upon completion.
+
+        Args:
+            instance: The dataset instance dictionary.
+            max_turns: Maximum conversation turns allowed for the agent. Defaults to 40.
+            time_limit: Time limit for the execution in minutes. Defaults to 30.
+            run_method: The execution method, either "python" (SDK) or "cli". Defaults to "python".
+
+        Returns:
+            A dictionary containing the result:
+
+            - instance_id: The ID of the processed instance.
+            - model_patch: The git diff generated by the agent.
+            - model_name_or_path: Hardcoded to "cc" (Claude Code).
+
+        Raises:
+            ValueError: If `run_method` is not "python" or "cli".
         """
         if run_method == "python":
-            self._run_python_sdk(instance, max_step, timelimit)
+            logger.warning("Running Claude Code using Python SDK is still under development and not yet stable.")
+            self._run_python_sdk(instance, max_turns, time_limit)
         elif run_method == "cli":
-            self._run_cli(instance, max_step, timelimit)
+            self._run_cli(instance, max_turns, time_limit)
         else:
-            raise ValueError(f"wrong run_method {run_method}, run_method should be in [python, cli]")
+            raise ValueError(f"Wrong run_method '{run_method}', run_method should be in ['python', 'cli']")
+
         result = self.container.send_command("git --no-pager diff HEAD")
         git_diff = result.output.replace("git --no-pager diff HEAD\n", "")
+
         return {
             "instance_id": instance["instance_id"],
             "model_patch": git_diff,
             "model_name_or_path": "cc",
         }
 
-    def __del__(self):
+    def __del__(self) -> None:
+        """Destructor to ensure container resources are cleaned up."""
         if hasattr(self, "container"):
             self.container.cleanup()

examples/claude_code/custom_callbacks.py

@@ -1,45 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""Custom callback module for LiteLLM logging integration.
-
-This module provides the AddLogprobs callback that integrates with LiteLLM to
-request log probabilities from vLLM backends. It modifies outgoing request payloads
-to include logprobs=1 for backends that support log probability returns.
-
-Key features:
-- LiteLLM pre-call hook implementation
-- Log probability request augmentation
-- Backend compatibility for vLLM and similar systems
-"""
-
-from typing import Any, Dict, Optional, Union
-
-from litellm.integrations.custom_logger import CustomLogger
-
-from agentlightning.llm_proxy import _get_pre_call_data  # type: ignore
-
-
-class AddLogprobs(CustomLogger):
-    """LiteLLM logger hook to request logprobs from vLLM.
-
-    This mutates the outgoing request payload to include `logprobs=1`
-    for backends that support logprobs return (e.g., vLLM).
-    """
-
-    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
-        """Async pre-call hook to adjust request payload.
-
-        Args:
-            args: Positional args from LiteLLM.
-            kwargs: Keyword args from LiteLLM.
-
-        Returns:
-            Either an updated payload dict or an Exception to short-circuit.
-        """
-        try:
-            data = _get_pre_call_data(args, kwargs)
-        except Exception as e:
-            return e
-
-        # Ensure logprobs are requested from the backend when supported.
-        return {**data, "logprobs": 1}

examples/claude_code/extended_adapter.py

@@ -2,16 +2,10 @@
 
 """Custom adapter module for converting LLM proxy traces to augmented trajectories.
 
-This module provides the LlmProxyTraceToAugmentedTriplet adapter that converts
-LLM proxy spans into augmented trajectories for analysis and evaluation. It extends
-the base LlmProxyTraceToTriplet to include additional metadata like chat messages,
+This module provides an augmented LlmProxyTraceToTriplet adapter that converts
+LLM proxy spans into augmented trajectories for analysis and evaluation.
+It extends the base LlmProxyTraceToTriplet to include additional metadata like chat messages,
 log probabilities, and sequence IDs.
-
-Key features:
-- Conversion of LLM proxy spans to triplet format
-- Augmentation with chat messages history
-- Log probability extraction
-- Sequence ID tracking for conversation turns
 """
 
 import logging
@@ -23,13 +17,11 @@ from agentlightning.types import Span, Triplet
 logger = logging.getLogger(__name__)
 
 
-class LlmProxyTraceToAugmentedTriplet(LlmProxyTraceToTriplet):
-    """Convert LLM Proxy spans into augmented trajectories
-
-    !!! warning
-        Only support `raw_gen_ai_request` spans for now.
+class ExtendedLlmProxyTraceToTriplet(LlmProxyTraceToTriplet):
+    """Convert LLM Proxy spans into trajectories with logprobs and customized metadata.
 
     Augmented fields include:
+
     - chat messages history from [`llm.hosted_vllm.messages`], saved to `Triplet.metadata['messages']`
     - logprobs from [`llm.hosted_vllm.choices`], saved to `Triplet.response['logprobs']`
     - sequence_id from [`Span.sequence_id`] to locate the order of the span (conversation turn), saved to `Triplet.metadata['sequence_id']`
@@ -94,9 +86,6 @@ class LlmProxyTraceToAugmentedTriplet(LlmProxyTraceToTriplet):
 
             if s.name == "raw_gen_ai_request":
                 prompt_ids, resp_ids, logprobs = self._extract_tokens_from_raw(attrs)
-                # elif s.name == "litellm_request":
-                #     # Some proxies never include token ids here. Ignore unless present.
-                #     prompt_ids, resp_ids = self._extract_tokens_from_openai(attrs)
 
                 if len(prompt_ids) == 0 or len(resp_ids) == 0:
                     logger.warning(

examples/claude_code/swebench_utils/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

examples/claude_code/swebench_utils/docker_runtime.py

@@ -1,6 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-
 """Docker runtime management for repository setup and command execution.
 
 Provides containerized environment for repository testing with command execution,
@@ -10,6 +9,7 @@ file operations, and state management capabilities.
 from __future__ import annotations
 
 import json
+import logging
 import queue
 import re
 import threading
@@ -20,10 +20,14 @@ from typing import Any, Callable, Dict, List, Optional
 
 from docker.errors import DockerException, ImageNotFound
 from docker.models.containers import Container
+from swebench.harness.constants import SWEbenchInstance
 from typing_extensions import Self
 
 import docker
 
+# This will log to the console for debugging purposes.
+claude_code_logger = logging.getLogger("claude_code_agent.docker_runtime")
+
 CMD_OUTPUT_PS1_BEGIN = "\n###PS1JSON###\n"
 CMD_OUTPUT_PS1_END = "\n###PS1END###"
 CMD_OUTPUT_METADATA_PS1_REGEX = re.compile(
@@ -187,7 +191,7 @@ class Runtime:
         ).replace('"', r"\"")
         ps1 = CMD_OUTPUT_PS1_BEGIN + json_str + CMD_OUTPUT_PS1_END + "\n"
         self.send_command(f'export PROMPT_COMMAND=\'export PS1="{ps1}"\'; export PS2=""')
-        self.send_command("apt update && apt install -y git")
+        self.send_command("apt update -qq && apt install -y -qq git")
         self.stopped = False
 
     def _stream_output(self):
@@ -207,6 +211,7 @@ class Runtime:
     def _start_output_thread(self):
         self.output_thread = threading.Thread(target=self._stream_output, daemon=True)
         self.output_thread.start()
+        # TODO: kill the thread if main thread is stopped
 
     def _clear_initial_prompt(self):
         time.sleep(0.5)
@@ -273,7 +278,25 @@ class Runtime:
 
         raise TypeError(f"Don't know how to write to {type(self.sock).__name__}")
 
+    def _log_command_result(self, result: CommandResult) -> None:
+        claude_code_logger.debug("Docker runtime command finished with metadata: %s", result.metadata)
+        if len(result.output) > 2048:
+            logged_output = result.output[:1024] + "\n(... stripped due to length ...)\n" + result.output[-1024:]
+        else:
+            logged_output = result.output
+        claude_code_logger.debug(
+            "Docker runtime command finished with output (length = %d):\n%s", len(result.output), logged_output
+        )
+        # Output to the evaluation logger simultaneously
+        self.logger(text=logged_output)
+
     def send_command(self, command: str, timeout: float = 20 * 60) -> CommandResult:
+        # Redact sensitive API keys from the command before logging
+        redacted_command = command
+        for sensitive_var in ["ANTHROPIC_AUTH_TOKEN", "API_KEY", "SECRET_KEY"]:
+            pattern = rf"(export (.*?){re.escape(sensitive_var)}(.*?)=)[^\s]+"
+            redacted_command = re.sub(pattern, rf"\1****REDACTED****", redacted_command)
+        claude_code_logger.info("Docker runtime receiving command: %s", redacted_command)
         # Normalize newline semantics for interactive shells
         if not command.endswith("\n"):
             command += "\n"
@@ -284,9 +307,10 @@ class Runtime:
         self._send_bytes(command.encode())
 
         output, metadata = self._read_raw_output(timeout=timeout)
+        # TODO: Check exit code of the command (claude code download fail will not be caught by this)
         if metadata is not None:
             result = CommandResult(output=output, metadata=metadata)
-            self.logger(text=result.output)
+            self._log_command_result(result)
             return result
 
         # handle timeout
@@ -298,22 +322,26 @@ class Runtime:
         output = output + kill_output + "\n**Exited due to timeout**\n"
         if kill_metadata is not None:
             kill_metadata.exit_code = TIMEOUT_EXIT_CODE
-            return CommandResult(output=output, metadata=kill_metadata)
+            result = CommandResult(output=output, metadata=kill_metadata)
+            self._log_command_result(result)
+            return result
 
         fallback_metadata = CmdOutputMetadata(
             exit_code=TIMEOUT_EXIT_CODE,
         )
-
         result = CommandResult(output=output, metadata=fallback_metadata)
-        self.logger(text=result.output)
+        self._log_command_result(result)
         return result
 
     def cleanup(self) -> None:
         if self.stopped:
             return
         try:
+            claude_code_logger.info(f"Stopping container: {self.container.id}")
             self.container.stop()
+            claude_code_logger.info(f"Removing container: {self.container.id}")
             self.container.remove(force=True)
+            claude_code_logger.info(f"Container removed: {self.container.id}")
             self.stopped = True
         except Exception as e:
             print(f"Failed to stop container: {e}")
@@ -343,7 +371,7 @@ class Runtime:
     def start_session(
         cls,
         image_name: str,
-        instance: dict[Any, Any],
+        instance: SWEbenchInstance,
         log_function: Callable[..., None] = lambda: None,
     ) -> Runtime:
         """
@@ -376,6 +404,9 @@ class Runtime:
         shell_command = "/bin/bash"
         working_dir = "/testbed"
 
+        claude_code_logger.info(
+            f"Starting container {container_name} with image {image_name}. Shell command: {shell_command}"
+        )
         container = client.containers.run(
             image_name,
             name=container_name,
@@ -392,6 +423,7 @@ class Runtime:
             cpu_quota=int(CPU_CORES * 100000),
             mem_limit=MEM_LIMIT,
         )
+        claude_code_logger.info(f"Container {container_name} started with ID: {container.id}")
 
         session = cls(container, log_function=log_function)
 

examples/claude_code/swebench_utils/evaluation.py

@@ -7,13 +7,6 @@ instances. It handles containerized execution of test scripts, patch application
 and generation of evaluation reports. The module orchestrates the complete evaluation
 process including container management, patch application, test execution, and result
 grading.
-
-Key components:
-- Instance evaluation with containerized test execution
-- Patch application and validation
-- Test script execution with timeouts
-- Evaluation report generation and logging
-- Container lifecycle management
 """
 
 import json

examples/claude_code/swebench_utils/logging.py

@@ -4,20 +4,22 @@
 
 This module provides a simple logging utility function that writes evaluation
 results and logs to timestamped files organized by run ID and instance ID.
-It supports structured logging for tracking the progress and outcomes of
-SWE-bench evaluation experiments.
-
-Key features:
-- Timestamped logging with datetime formatting
-- Run-specific directory organization
-- Simple text-based logging for evaluation outputs
 """
 
 import datetime
 import os
 
 
-def logger(run_id: str, instance_id: str, text: str) -> None:
+def log_for_evaluation(run_id: str, instance_id: str, text: str) -> None:
+    """Log a message for evaluation purposes of SWE-Bench.
+
+    The format follows the SWE-Bench evaluation framework.
+
+    Args:
+        run_id: The run ID of the evaluation.
+        instance_id: The instance ID of the evaluation.
+        text: The text to log.
+    """
     os.makedirs(f"./logs/{run_id}", exist_ok=True)
     current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     with open(f"./logs/{run_id}/{instance_id}", mode="a") as f: