diff --git a/.gitignore b/.gitignore index 7c5950d8..e8179f86 100644 --- a/.gitignore +++ b/.gitignore @@ -43,6 +43,9 @@ configs/**/*.local.yaml .secrets/ .codex_azure*/ +# Local MCP server config — references the machine's GITHUB_PAT, never commit +.mcp.json + # Internal docs (not for open-source release) docs/ablation_plan.md docs/ablation_paper_tables.md diff --git a/configs/_base_/default.yaml b/configs/_base_/default.yaml index ccd12590..175739f7 100644 --- a/configs/_base_/default.yaml +++ b/configs/_base_/default.yaml @@ -35,6 +35,7 @@ model: copilot_chat_target_model: "" copilot_chat_timeout: null # preserves COPILOT_CHAT_TIMEOUT or the built-in default codex_trace_to_optimizer: true + claude_trace_to_optimizer: true azure_openai_endpoint: "" # e.g. "https://your-resource.openai.azure.com/" azure_openai_api_version: "2024-12-01-preview" azure_openai_api_key: "" # Fill locally if you do not export AZURE_OPENAI_API_KEY diff --git a/docs/reference/config.md b/docs/reference/config.md index c1e8ae53..e8dea064 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -67,6 +67,8 @@ defaults to `claude` and can be overridden with `CLAUDE_CLI_BIN`. | `model.minimax_*` | MiniMax `base_url`, `api_key`, shared `minimax_model`, `temperature`, `max_tokens`, and `enable_thinking`; `minimax_model` applies when MiniMax is the target | | `model.codex_exec_*` | Codex path, sandbox, profile, SDK mode, reasoning, network/search, and approval policy; see compatibility notes below | | `model.claude_code_exec_*` | Claude path, profile, SDK mode, effort, and thinking-token cap | +| `model.codex_trace_to_optimizer` | When `true` (default) and target is `codex_exec`, inject the agent's codex trace steps into the reflection prompt | +| `model.claude_trace_to_optimizer` | When `true` (default) and target is `claude_code_exec`, inject the agent's claude trace steps into the reflection prompt | | `model.cursor_exec_path` | Cursor Agent executable path; default `cursor-agent` | | `model.cursor_exec_sandbox` | Cursor sandbox mode: `enabled` (default) or `disabled`; file-edit rollouts require `enabled` | | `model.copilot_exec_path` | GitHub Copilot CLI executable path; default `copilot` | diff --git a/scripts/eval_only.py b/scripts/eval_only.py index 85b71c7e..844e4258 100644 --- a/scripts/eval_only.py +++ b/scripts/eval_only.py @@ -382,7 +382,11 @@ def _set_role(key: str, value: str) -> None: _set_role("optimizer_backend", "codex_exec") _set_role("target_backend", "codex_exec") elif backend == "claude_code_exec": - _set_role("optimizer_backend", "openai_chat") + # Both roles default to Claude Code so reflection sees the full + # trajectory. A role pinned to a non-default value (e.g. minimax_chat) + # still overrides; an explicit --optimizer_backend openai_chat does + # not, because openai_chat is one of the base-config defaults. + _set_role("optimizer_backend", "claude_code_exec") _set_role("target_backend", "claude_code_exec") elif backend == "cursor_exec": _set_role("optimizer_backend", "openai_chat") diff --git a/scripts/train.py b/scripts/train.py index 977974a4..affd03f3 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -648,7 +648,11 @@ def _set_role(key: str, value: str) -> None: _set_role("optimizer_backend", "codex_exec") _set_role("target_backend", "codex_exec") elif backend == "claude_code_exec": - _set_role("optimizer_backend", "openai_chat") + # Both roles default to Claude Code so reflection sees the full + # trajectory. A role pinned to a non-default value (e.g. minimax_chat) + # still overrides; an explicit --optimizer_backend openai_chat does + # not, because openai_chat is one of the base-config defaults. + _set_role("optimizer_backend", "claude_code_exec") _set_role("target_backend", "claude_code_exec") elif backend == "cursor_exec": _set_role("optimizer_backend", "openai_chat") diff --git a/skillopt/config.py b/skillopt/config.py index 4de2cbed..be8fdfac 100644 --- a/skillopt/config.py +++ b/skillopt/config.py @@ -73,6 +73,7 @@ "model.copilot_chat_target_model": "copilot_chat_target_model", "model.copilot_chat_timeout": "copilot_chat_timeout", "model.codex_trace_to_optimizer": "codex_trace_to_optimizer", + "model.claude_trace_to_optimizer": "claude_trace_to_optimizer", "model.azure_endpoint": "azure_endpoint", "model.azure_api_version": "azure_api_version", "model.azure_api_key": "azure_api_key", diff --git a/skillopt/engine/trainer.py b/skillopt/engine/trainer.py index a648f0c6..169aee3c 100644 --- a/skillopt/engine/trainer.py +++ b/skillopt/engine/trainer.py @@ -452,6 +452,27 @@ def _resolve_train_size(cfg: dict, dataloader) -> int: _ROLE_BACKEND_DEFAULTS = (None, "", "openai_chat") +def _configure_trace_to_optimizer_gates(target_backend: str, cfg: dict) -> None: + """Turn on trace-to-optimizer gates for the exec target's trace artifact. + + Sets ``REFLACT_CODEX_TRACE_TO_OPTIMIZER`` (codex) and + ``REFLACT_CLAUDE_TRACE_TO_OPTIMIZER`` (claude) to ``"1"`` only when the + target actually runs on that exec backend and the matching config knob is + on. ``skillopt.gradient.reflect.fmt_minibatch_trajectories`` reads these + env vars, so a non-exec target never pays the injection. + """ + os.environ["REFLACT_CODEX_TRACE_TO_OPTIMIZER"] = ( + "1" + if target_backend == "codex_exec" and cfg.get("codex_trace_to_optimizer", False) + else "0" + ) + os.environ["REFLACT_CLAUDE_TRACE_TO_OPTIMIZER"] = ( + "1" + if target_backend == "claude_code_exec" and cfg.get("claude_trace_to_optimizer", False) + else "0" + ) + + def _resolve_role_backends( backend: str, optimizer_backend: str | None, target_backend: str | None ) -> tuple[str, str]: @@ -478,7 +499,12 @@ def _resolve_role_backends( if target_backend in _ROLE_BACKEND_DEFAULTS: target_backend = "codex_exec" elif backend == "claude_code_exec": - optimizer_backend = optimizer_backend or "openai_chat" + # Both roles default to Claude Code so reflection sees the full + # trajectory. A role pinned to a non-default value (e.g. minimax_chat) + # still overrides; an explicit --optimizer_backend openai_chat does not, + # because openai_chat is one of the base-config defaults. + if optimizer_backend in _ROLE_BACKEND_DEFAULTS: + optimizer_backend = "claude_code_exec" if target_backend in _ROLE_BACKEND_DEFAULTS: target_backend = "claude_code_exec" elif backend == "cursor_exec": @@ -801,11 +827,7 @@ def _build_eval_env(split: str, env_num: int, seed: int): minimax_model_cfg = cfg.get("minimax_model") if minimax_model_cfg and cfg.get("target_backend") == "minimax_chat": set_target_deployment(str(minimax_model_cfg)) - os.environ["REFLACT_CODEX_TRACE_TO_OPTIMIZER"] = ( - "1" - if target_backend == "codex_exec" and cfg.get("codex_trace_to_optimizer", False) - else "0" - ) + _configure_trace_to_optimizer_gates(target_backend, cfg) reasoning = cfg.get("reasoning_effort", "") or None set_reasoning_effort(reasoning) print( diff --git a/skillopt/gradient/reflect.py b/skillopt/gradient/reflect.py index 8078f852..df68c9b7 100644 --- a/skillopt/gradient/reflect.py +++ b/skillopt/gradient/reflect.py @@ -209,6 +209,19 @@ def fmt_minibatch_trajectories( f"{codex_probe_trace_steps}\n" ) + # Claude Code exec backend (issue #233): the SDK's full session trace is + # persisted as claude_trace_steps.txt; surface it so the analyst sees the + # agent's actual tool activity instead of only the collapsed final answer. + # Gated like the codex summary above: only the trainer turns it on, and + # only when the target actually runs on claude_code_exec. + if os.environ.get("REFLACT_CLAUDE_TRACE_TO_OPTIMIZER", "0") == "1": + claude_steps_path = os.path.join(prediction_dir, tid, "claude_trace_steps.txt") + if os.path.exists(claude_steps_path): + with open(claude_steps_path, encoding="utf-8") as f: + claude_steps = f.read().strip() + if claude_steps: + header += f"\n#### Claude Trace Steps\n{claude_steps}\n" + preview = item.get("spreadsheet_preview", "") if not preview: preview_path = os.path.join(prediction_dir, tid, "spreadsheet_preview.txt") diff --git a/skillopt/model/__init__.py b/skillopt/model/__init__.py index bcde7219..a5912529 100644 --- a/skillopt/model/__init__.py +++ b/skillopt/model/__init__.py @@ -6,6 +6,7 @@ from skillopt.model import azure_openai as _openai from skillopt.model import claude_backend as _claude +from skillopt.model import claude_code_backend as _claude_code from skillopt.model import codex_backend as _codex from skillopt.model import copilot_backend as _copilot from skillopt.model import minimax_backend as _minimax @@ -55,7 +56,11 @@ def set_backend(name: str | None) -> str: set_target_backend("codex_exec") return normalized if normalized == "claude_code_exec": - set_optimizer_backend("openai_chat") + # Both roles default to Claude Code so reflection sees the full + # trajectory. A role pinned to a non-default value (e.g. minimax_chat) + # still overrides; an explicit --optimizer_backend openai_chat does not, + # because openai_chat is one of the base-config defaults. + set_optimizer_backend("claude_code_exec") set_target_backend(normalized) return normalized if normalized == "cursor_exec": @@ -181,6 +186,16 @@ def chat_optimizer( stage=stage, timeout=timeout, ) + if get_optimizer_backend() == "claude_code_exec": + return _claude_code.chat_optimizer( + system=system, + user=user, + max_completion_tokens=max_completion_tokens, + retries=retries, + stage=stage, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) return _openai.chat_optimizer( system=system, user=user, @@ -346,6 +361,18 @@ def chat_optimizer_messages( return_message=return_message, timeout=timeout, ) + if get_optimizer_backend() == "claude_code_exec": + return _claude_code.chat_optimizer_messages( + messages=messages, + max_completion_tokens=max_completion_tokens, + retries=retries, + stage=stage, + tools=tools, + tool_choice=tool_choice, + return_message=return_message, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) return _openai.chat_optimizer_messages( messages=messages, max_completion_tokens=max_completion_tokens, @@ -509,6 +536,17 @@ def get_token_summary() -> dict: summary[stage]["prompt_tokens"] += values["prompt_tokens"] summary[stage]["completion_tokens"] += values["completion_tokens"] summary[stage]["total_tokens"] += values["total_tokens"] + claude_code_summary = _claude_code.get_token_summary() + for stage, values in claude_code_summary.items(): + if stage == "_total": + continue + if stage not in summary: + summary[stage] = values + continue + summary[stage]["calls"] += values["calls"] + summary[stage]["prompt_tokens"] += values["prompt_tokens"] + summary[stage]["completion_tokens"] += values["completion_tokens"] + summary[stage]["total_tokens"] += values["total_tokens"] qwen_summary = _qwen.get_token_summary() for stage, values in qwen_summary.items(): if stage == "_total": @@ -584,6 +622,7 @@ def get_token_summary() -> dict: def reset_token_tracker() -> None: _openai.reset_token_tracker() _claude.reset_token_tracker() + _claude_code.reset_token_tracker() _qwen.reset_token_tracker() _minimax.reset_token_tracker() _openai_compat.reset_token_tracker() @@ -736,6 +775,7 @@ def configure_openai_compatible( def set_reasoning_effort(effort: str | None) -> None: _openai.set_reasoning_effort(effort) _claude.set_reasoning_effort(effort) + _claude_code.set_reasoning_effort(effort) _qwen.set_reasoning_effort(effort) _minimax.set_reasoning_effort(effort) _openai_compat.set_reasoning_effort(effort) @@ -745,6 +785,7 @@ def set_reasoning_effort(effort: str | None) -> None: def set_target_deployment(deployment: str) -> None: _openai.set_target_deployment(deployment) _claude.set_target_deployment(deployment) + _claude_code.set_target_deployment(deployment) _qwen.set_target_deployment(deployment) _minimax.set_target_deployment(deployment) _openai_compat.set_target_deployment(deployment) @@ -754,6 +795,7 @@ def set_target_deployment(deployment: str) -> None: def set_optimizer_deployment(deployment: str) -> None: _openai.set_optimizer_deployment(deployment) _claude.set_optimizer_deployment(deployment) + _claude_code.set_optimizer_deployment(deployment) _qwen.set_optimizer_deployment(deployment) _openai_compat.set_optimizer_deployment(deployment) _codex.set_optimizer_deployment(deployment) diff --git a/skillopt/model/backend_config.py b/skillopt/model/backend_config.py index 5db7d9e8..8b84d213 100644 --- a/skillopt/model/backend_config.py +++ b/skillopt/model/backend_config.py @@ -133,11 +133,12 @@ def set_optimizer_backend(backend: str) -> None: "openai_compatible", "copilot_chat", "codex_exec", + "claude_code_exec", }: raise ValueError( f"Unsupported optimizer backend: {OPTIMIZER_BACKEND!r}. " "Supported values are 'openai_chat', 'claude_chat', 'qwen_chat', 'minimax_chat', " - "'openai_compatible', 'copilot_chat', and 'codex_exec'." + "'openai_compatible', 'copilot_chat', 'codex_exec', and 'claude_code_exec'." ) os.environ["OPTIMIZER_BACKEND"] = OPTIMIZER_BACKEND @@ -176,6 +177,7 @@ def is_optimizer_chat_backend() -> bool: "openai_compatible", "copilot_chat", "codex_exec", + "claude_code_exec", } diff --git a/skillopt/model/claude_code_backend.py b/skillopt/model/claude_code_backend.py new file mode 100644 index 00000000..7cc626d2 --- /dev/null +++ b/skillopt/model/claude_code_backend.py @@ -0,0 +1,313 @@ +"""Claude Code CLI/SDK chat backend for ReflACT (optimizer role). + +Runs Claude Code (the same CLI/SDK that powers the ``claude_code_exec`` target +backend) as a plain chat model for reflection. This gives the optimizer access +to Claude's full context window, so minibatch trajectories are not truncated by +a narrower chat backend. +""" +from __future__ import annotations + +import json +import os +import time +from typing import Any + +from skillopt.model import codex_backend as _codex +from skillopt.model.claude_backend import _build_prompt_from_messages +from skillopt.model.codex_harness import run_claude_code_chat +from skillopt.model.common import TokenTracker + +OPTIMIZER_DEPLOYMENT = os.environ.get("OPTIMIZER_DEPLOYMENT", "claude-sonnet-4-6") +TARGET_DEPLOYMENT = os.environ.get("TARGET_DEPLOYMENT", "claude-sonnet-4-6") +REASONING_EFFORT: str | None = None +tracker = TokenTracker() + + +def _assistant_message_schema() -> dict[str, Any]: + return _codex._assistant_message_schema() + + +def _compat_message_from_payload( + payload: dict[str, Any], + *, + tool_choice: str | dict[str, Any] | None = None, +): + return _codex._compat_message_from_payload(payload, tool_choice=tool_choice) + + +def _chat_messages_impl( + model: str, + messages: list[dict[str, Any]], + max_completion_tokens: int, + retries: int, + stage: str, + *, + tools: list[dict[str, Any]] | None = None, + tool_choice: str | dict[str, Any] | None = None, + return_message: bool = False, + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[Any, dict[str, int]]: + del max_completion_tokens # Claude Code does not expose a completion-token cap + last_err = None + structured_output = bool(tools) or return_message + schema = _assistant_message_schema() if structured_output else None + # An explicit per-call effort wins; otherwise fall back to the value set via + # set_reasoning_effort (the `--reasoning_effort` / config path). + effort = reasoning_effort if reasoning_effort is not None else REASONING_EFFORT + + for attempt in range(retries): + try: + system, prompt, attachments = _build_prompt_from_messages( + messages, + tools=tools, + tool_choice=tool_choice, + structured_output=structured_output, + ) + if attachments: + raise RuntimeError( + "claude_code_exec backend does not support image attachments" + ) + raw_text, usage_info = run_claude_code_chat( + system=system, + prompt=prompt, + model=model, + timeout=timeout, + schema=schema, + effort=effort, + ) + tracker.record( + stage, + usage_info["prompt_tokens"], + usage_info["completion_tokens"], + ) + if not structured_output: + return raw_text, usage_info + payload = json.loads(raw_text) + compat = _compat_message_from_payload(payload, tool_choice=tool_choice) + return (compat if return_message else compat.content), usage_info + except Exception as exc: # noqa: BLE001 + last_err = exc + time.sleep(min(2 ** attempt, 30)) + + raise RuntimeError(f"Claude Code call failed after {retries} retries: {last_err}") + + +def chat_with_model( + model: str, + system: str, + user: str, + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "custom", + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[str, dict[str, int]]: + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ] + return _chat_messages_impl( + model, + messages, + max_completion_tokens, + retries, + stage, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_messages_with_model( + model: str, + messages: list[dict[str, Any]], + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "custom", + *, + tools: list[dict[str, Any]] | None = None, + tool_choice: str | dict[str, Any] | None = None, + return_message: bool = False, + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[Any, dict[str, int]]: + return _chat_messages_impl( + model, + messages, + max_completion_tokens, + retries, + stage, + tools=tools, + tool_choice=tool_choice, + return_message=return_message, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_optimizer( + system: str, + user: str, + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "optimizer", + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[str, dict[str, int]]: + return chat_with_model( + model=OPTIMIZER_DEPLOYMENT, + system=system, + user=user, + max_completion_tokens=max_completion_tokens, + retries=retries, + stage=stage, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_target( + system: str, + user: str, + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "target", + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[str, dict[str, int]]: + return chat_with_model( + model=TARGET_DEPLOYMENT, + system=system, + user=user, + max_completion_tokens=max_completion_tokens, + retries=retries, + stage=stage, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_with_deployment( + deployment: str, + system: str, + user: str, + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "custom", + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[str, dict[str, int]]: + return chat_with_model( + model=deployment, + system=system, + user=user, + max_completion_tokens=max_completion_tokens, + retries=retries, + stage=stage, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_optimizer_messages( + messages: list[dict[str, Any]], + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "optimizer", + *, + tools: list[dict[str, Any]] | None = None, + tool_choice: str | dict[str, Any] | None = None, + return_message: bool = False, + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[Any, dict[str, int]]: + return _chat_messages_impl( + OPTIMIZER_DEPLOYMENT, + messages, + max_completion_tokens, + retries, + stage, + tools=tools, + tool_choice=tool_choice, + return_message=return_message, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_target_messages( + messages: list[dict[str, Any]], + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "target", + *, + tools: list[dict[str, Any]] | None = None, + tool_choice: str | dict[str, Any] | None = None, + return_message: bool = False, + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[Any, dict[str, int]]: + return _chat_messages_impl( + TARGET_DEPLOYMENT, + messages, + max_completion_tokens, + retries, + stage, + tools=tools, + tool_choice=tool_choice, + return_message=return_message, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def chat_messages_with_deployment( + deployment: str, + messages: list[dict[str, Any]], + max_completion_tokens: int = 16384, + retries: int = 5, + stage: str = "custom", + *, + tools: list[dict[str, Any]] | None = None, + tool_choice: str | dict[str, Any] | None = None, + return_message: bool = False, + timeout: int | None = None, + reasoning_effort: str | None = None, +) -> tuple[Any, dict[str, int]]: + return _chat_messages_impl( + deployment, + messages, + max_completion_tokens, + retries, + stage, + tools=tools, + tool_choice=tool_choice, + return_message=return_message, + timeout=timeout, + reasoning_effort=reasoning_effort, + ) + + +def get_token_summary() -> dict[str, dict[str, int]]: + return tracker.summary() + + +def reset_token_tracker() -> None: + tracker.reset() + + +def set_reasoning_effort(effort: str | None) -> None: + global REASONING_EFFORT + REASONING_EFFORT = effort if effort else None + + +def set_target_deployment(deployment: str) -> None: + global TARGET_DEPLOYMENT + TARGET_DEPLOYMENT = deployment + os.environ["TARGET_DEPLOYMENT"] = deployment + + +def set_optimizer_deployment(deployment: str) -> None: + global OPTIMIZER_DEPLOYMENT + OPTIMIZER_DEPLOYMENT = deployment + os.environ["OPTIMIZER_DEPLOYMENT"] = deployment diff --git a/skillopt/model/codex_harness.py b/skillopt/model/codex_harness.py index 056539bc..47a93069 100644 --- a/skillopt/model/codex_harness.py +++ b/skillopt/model/codex_harness.py @@ -7,6 +7,7 @@ import re import shutil import subprocess +import tempfile import threading import traceback import warnings @@ -322,6 +323,14 @@ def _persist_claude_artifacts(work_dir: str, raw: str, response: str) -> None: prefix="claude", summary_builder=_build_claude_trace_summary, ) + # Structured trace steps for the reflector (issue #233): expose what the + # agent actually did, not just the collapsed final answer. + steps_text = format_claude_trace_steps(raw) + if steps_text: + pred_dir = os.path.dirname(work_dir.rstrip(os.sep)) + steps_path = os.path.join(pred_dir, "claude_trace_steps.txt") + with open(steps_path, "w", encoding="utf-8") as f: + f.write(steps_text) def _persist_cursor_artifacts(work_dir: str, raw: str, response: str) -> None: @@ -431,6 +440,125 @@ def extract_codex_trace_prefix(raw: str, *, after_step: int) -> str: return "\n".join(lines[:end_line]).strip() +# ── Claude Code trace steps (SDK messages → compact steps) ────────────────── +# The Claude Code SDK serializes its full session into ``messages``: most +# entries are bookkeeping (init / thinking_tokens), the rest are assistant text, +# tool calls, and tool results. Flatten those into numbered steps so the +# reflector can see what the agent actually did without paying the full raw +# payload size. + + +def _claude_step_truncate(text: str, limit: int) -> str: + text = str(text or "").strip() + if len(text) <= limit: + return text + return text[:limit] + f"...[+{len(text) - limit} chars]" + + +def _summarize_claude_tool_call(name: str, input_data: Any) -> str: + name = str(name or "") + if isinstance(input_data, dict): + if name == "Read": + return f"Read {input_data.get('file_path', '')}" + if name == "Glob": + return f"Glob {input_data.get('pattern', '')}" + if name == "Grep": + return f"Grep {input_data.get('pattern', '')}" + if name == "Bash": + return f"Bash {input_data.get('command', '')}" + return _claude_step_truncate(f"{name} {json.dumps(input_data, ensure_ascii=False)}", 500) + + +def _iter_claude_json_blocks(raw: str): + """Yield each JSON object embedded in ``raw``. + + ``run_claude_code_exec`` prefixes every attempt with a + ``===== CLAUDE ... ATTEMPT n =====`` header, so the persisted payload is not + a single JSON document. Split on those headers and parse each block. + """ + for chunk in re.split(r"(?m)^={5,}.*={5,}\s*$", raw or ""): + chunk = chunk.strip() + if not chunk: + continue + try: + yield json.loads(chunk) + except json.JSONDecodeError: + continue + + +def parse_claude_trace_steps(raw: str) -> list[dict]: + """Parse serialized Claude Code SDK messages into ordered, compact steps. + + Returns a list of ``{"index", "type", "summary"}`` dicts where ``type`` is + one of ``text`` / ``tool_call`` / ``tool_result``. System bookkeeping + events (init, thinking tokens) are dropped; tool results are truncated to + keep the trace small. + """ + steps: list[dict] = [] + for block in _iter_claude_json_blocks(raw): + messages = block.get("messages") if isinstance(block, dict) else None + if not isinstance(messages, list): + continue + for message in messages: + if not isinstance(message, dict): + continue + if message.get("subtype") in {"init", "thinking_tokens"}: + continue + data = message.get("data") + if isinstance(data, dict) and data.get("type") == "system": + continue + content = message.get("content") + if not isinstance(content, list): + # Terminal result message carries the final text. + text = str(message.get("result") or "").strip() + if text: + steps.append({"type": "text", "summary": _claude_step_truncate(text, 500)}) + continue + for item in content: + if not isinstance(item, dict): + continue + if "name" in item and "input" in item: + steps.append({ + "type": "tool_call", + "summary": _summarize_claude_tool_call(item.get("name"), item.get("input")), + }) + elif "tool_use_id" in item: + body = item.get("content") + if isinstance(body, list): + text_parts: list[str] = [] + for part in body: + if isinstance(part, dict): + part_text = part.get("content") + if isinstance(part_text, str): + text_parts.append(part_text) + elif isinstance(part, str): + text_parts.append(part) + body = "\n".join(text_parts) + summary = _claude_step_truncate(body, 200) + if item.get("is_error"): + summary = f"[error] {summary}" + steps.append({"type": "tool_result", "summary": summary}) + else: + text = item.get("text") + if isinstance(text, str) and text.strip(): + steps.append({"type": "text", "summary": _claude_step_truncate(text, 500)}) + for index, step in enumerate(steps, 1): + step["index"] = index + return steps + + +def format_claude_trace_steps(raw: str, *, max_chars: int = 4000) -> str: + """Render parsed Claude Code SDK trace into numbered compact steps.""" + steps = parse_claude_trace_steps(raw) + if not steps: + return "" + rendered = [f"[{step['index']}] {step['type']}: {step['summary']}" for step in steps] + text = "\n".join(rendered) + if len(text) > max_chars: + text = text[:max_chars] + "\n...[claude trace steps truncated]..." + return text + + _DENIED_DATA_DIR_NAMES = {"officeqa_split", "sealqa_split"} @@ -859,6 +987,243 @@ def run_claude_code_exec( return last_response, combined +# ── Claude Code *chat* mode (optimizer role) ──────────────────────────────── +# The functions above run Claude Code as the *target* exec backend: they embed +# the target preamble, force the ANSWER_SCHEMA structured output, and read from +# a prepared workspace. When Claude Code is instead selected as the optimizer +# backend (claude_code_exec), reflection calls need a plain-text model call with +# the analyst's own system prompt and no tooling — mirroring claude_backend's +# chat path but driven through the same Claude Code CLI/SDK as the target. + + +def _claude_chat_text_from_messages(messages: list[Any]) -> str: + """Extract the final assistant text from SDK chat-mode messages. + + With ``output_format={"type": "text"}`` the SDK ends with a result message + whose ``result`` holds the final text; fall back to the last text block of + the final assistant message. + """ + for msg in reversed(messages): + result = getattr(msg, "result", None) + if isinstance(result, str) and result.strip(): + return result + content = getattr(msg, "content", None) + if content is None and isinstance(msg, dict): + content = msg.get("content") + if not isinstance(content, list): + continue + for item in content: + text = item.get("text") if isinstance(item, dict) else getattr(item, "text", None) + if isinstance(text, str) and text.strip(): + return text + return "" + + +def _claude_chat_usage_from_event(event: Any) -> dict[str, int]: + """Convert an SDK/CLI result usage payload into the shared usage shape.""" + usage = getattr(event, "usage", {}) if not isinstance(event, dict) else (event or {}).get("usage", {}) + if isinstance(usage, dict): + input_tokens = int(usage.get("input_tokens", 0) or 0) + output_tokens = int(usage.get("output_tokens", 0) or 0) + else: + input_tokens = int(getattr(usage, "input_tokens", 0) or 0) + output_tokens = int(getattr(usage, "output_tokens", 0) or 0) + return { + "prompt_tokens": input_tokens, + "completion_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + + +def _run_claude_code_sdk_chat_exec( + *, + system: str, + prompt: str, + model: str, + timeout: int, + schema: dict[str, Any] | None = None, + effort: str | None = None, +) -> tuple[str, dict]: + from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient + + async def _query() -> tuple[str, dict]: + # Optimizer chat call: no workspace, no tools, optional schema. + with tempfile.TemporaryDirectory(prefix="skillopt_claude_code_chat_") as tmp: + system_prompt: dict[str, Any] = { + "type": "preset", + "preset": "claude_code", + "append": system or "", + } + kwargs: dict[str, Any] = { + "system_prompt": system_prompt, + "output_format": ( + {"type": "json_schema", "schema": schema} + if schema is not None + else {"type": "text"} + ), + "allowed_tools": [], + "cwd": tmp, + "permission_mode": "bypassPermissions", + } + config = get_claude_code_exec_config() + effort_value = _claude_effort(effort if effort is not None else config.get("effort")) + if effort_value: + kwargs["effort"] = effort_value + max_thinking_tokens = int(config.get("max_thinking_tokens", 0) or 0) + if max_thinking_tokens > 0: + kwargs["max_thinking_tokens"] = max_thinking_tokens + options = ClaudeAgentOptions(**kwargs) + if model: + options.model = model.split("/", 1)[1] if model.startswith("anthropic/") else model + + messages = [] + async with ClaudeSDKClient(options) as client: + await client.query(prompt) + messages = [msg async for msg in client.receive_response()] + last = messages[-1] if messages else None + if schema is not None: + payload = _extract_claude_structured_output(messages) + text = _json_dumps(payload) if isinstance(payload, dict) else "" + if not text: + result = getattr(last, "result", None) + if isinstance(result, str) and result.strip(): + text = result + else: + text = _claude_chat_text_from_messages(messages) + usage_info = _claude_chat_usage_from_event(last) + return text, usage_info + + return _run_async(asyncio.wait_for(_query(), timeout=timeout)) + + +def _run_claude_code_cli_chat_exec( + *, + system: str, + prompt: str, + model: str, + timeout: int, + schema: dict[str, Any] | None = None, + effort: str | None = None, +) -> tuple[str, dict]: + config = get_claude_code_exec_config() + cmd = [ + str(config["path"]), + "-p", + "--output-format", + "json", + "--permission-mode", + "dontAsk", + ] + if model: + cmd.extend(["--model", model]) + if schema is not None: + cmd.extend(["--schema", json.dumps(schema, ensure_ascii=False)]) + if config.get("profile"): + cmd.extend(["--settings", '{"env":{"CLAUDE_CODE_USE_BEDROCK":"0"}}']) + cmd.extend(["--append-system-prompt", f"Profile: {config['profile']}"]) + effort_value = _claude_effort(effort if effort is not None else config.get("effort")) + if effort_value: + cmd.extend(["--effort", effort_value]) + + with tempfile.TemporaryDirectory(prefix="skillopt_claude_code_chat_") as tmp: + # System prompt via file, not argv, to avoid the Windows argv cap. + system_path = os.path.join(tmp, "system_prompt.txt") + with open(system_path, "w", encoding="utf-8") as system_fh: + system_fh.write(system or "") + cmd.extend(["--append-system-prompt-file", system_path]) + proc = subprocess.run( + cmd, + input=prompt, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=timeout or 300, + cwd=tmp, + ) + + stderr_text = (proc.stderr or "").strip() + if proc.returncode != 0: + raise RuntimeError(stderr_text or f"Claude Code CLI exited with code {proc.returncode}") + stream = [] + for raw_line in (proc.stdout or "").splitlines(): + raw_line = raw_line.strip() + if not raw_line: + continue + try: + stream.append(json.loads(raw_line)) + except json.JSONDecodeError: + continue + result_event = None + for event in reversed(stream): + if event.get("type") == "result": + result_event = event + break + if result_event is None: + raise RuntimeError("Claude Code CLI did not return a result event.") + text = str(result_event.get("result") or result_event.get("content") or "") + usage_info = _claude_chat_usage_from_event(result_event) + return text, usage_info + + +def run_claude_code_chat( + *, + system: str, + prompt: str, + model: str, + timeout: int, + schema: dict[str, Any] | None = None, + effort: str | None = None, +) -> tuple[str, dict]: + """Run Claude Code as a plain chat model (optimizer role). + + ``effort`` overrides the configured ``claude_code_exec_effort``; when ``None`` + the config value (default "medium") is used, matching the target-exec path. + """ + config = get_claude_code_exec_config() + mode = _sdk_mode(config.get("use_sdk")) + retries = int(config.get("empty_response_retries", 0) or 0) + last_text = "" + last_usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + for _attempt in range(retries + 1): + if mode != "cli": + try: + text, usage_info = _run_claude_code_sdk_chat_exec( + system=system, + prompt=prompt, + model=model, + timeout=timeout, + schema=schema, + effort=effort, + ) + last_text = text + last_usage = usage_info + if text.strip(): + return text, usage_info + except (ImportError, ModuleNotFoundError): + if mode == "sdk": + raise + except Exception: # noqa: BLE001 + if mode == "sdk": + raise + if mode != "sdk": + text, usage_info = _run_claude_code_cli_chat_exec( + system=system, + prompt=prompt, + model=model, + timeout=timeout, + schema=schema, + effort=effort, + ) + last_text = text + last_usage = usage_info + if text.strip(): + return text, usage_info + + return last_text, last_usage + + def _run_codex_sdk_exec( *, work_dir: str, diff --git a/tests/test_claude_code_backend.py b/tests/test_claude_code_backend.py new file mode 100644 index 00000000..a473cacf --- /dev/null +++ b/tests/test_claude_code_backend.py @@ -0,0 +1,351 @@ +"""claude_code_exec optimizer backend: trace parsing, dispatch, persistence, retries. + +Covers the four highest-risk, previously-untested points introduced with +``claude_code_backend`` (issue #233): + +- ``parse_claude_trace_steps`` extracts text / tool_call / tool_result and drops + init / thinking_tokens bookkeeping (``skillopt/model/codex_harness.py``). +- the dispatcher routes ``chat_optimizer`` to the claude_code branch when the + optimizer backend is ``claude_code_exec``. +- ``_persist_claude_artifacts`` writes ``claude_trace_steps.txt`` for the reflector. +- a non-JSON structured reply is retried and then surfaces as ``RuntimeError``. + +Plus a gating regression: ``fmt_minibatch_trajectories`` only injects +``#### Claude Trace Steps`` when ``REFLACT_CLAUDE_TRACE_TO_OPTIMIZER == "1"``. +""" +from __future__ import annotations + +import importlib.util +import json +import os +import sys +import types +from collections.abc import Iterator +from typing import Any + +import pytest + +from skillopt.gradient.reflect import fmt_minibatch_trajectories +from skillopt.model import codex_harness +from skillopt.model.codex_harness import ( + _json_dumps, + _persist_claude_artifacts, + format_claude_trace_steps, + parse_claude_trace_steps, +) + + +class _OpenAIClientStub: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.args = args + self.kwargs = kwargs + + +def _install_openai_stub() -> None: + if "openai" in sys.modules or importlib.util.find_spec("openai") is not None: + return + openai_stub = types.ModuleType("openai") + openai_stub.AzureOpenAI = _OpenAIClientStub + openai_stub.OpenAI = _OpenAIClientStub + sys.modules["openai"] = openai_stub + + +@pytest.fixture(autouse=True) +def isolate_backend_state() -> Iterator[None]: + _install_openai_stub() + from skillopt.model import backend_config + + optimizer_backend = backend_config.get_optimizer_backend() + target_backend = backend_config.get_target_backend() + env = { + key: os.environ.get(key) + for key in ( + "OPTIMIZER_BACKEND", + "TARGET_BACKEND", + "OPTIMIZER_DEPLOYMENT", + "TARGET_DEPLOYMENT", + ) + } + yield + backend_config.set_optimizer_backend(optimizer_backend) + backend_config.set_target_backend(target_backend) + for key, value in env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _sdk_payload() -> str: + """Two SDK attempt blocks: bookkeeping noise + real steps + tool result.""" + block1 = { + "messages": [ + {"subtype": "init", "content": []}, + {"subtype": "thinking_tokens", "content": []}, + {"data": {"type": "system", "text": "system banner"}, "content": []}, + { + "content": [ + {"type": "text", "text": "Let me read the task."}, + {"type": "tool_use", "id": "tu_1", "name": "Read", "input": {"file_path": "task.md"}}, + ], + "data": {"type": "assistant"}, + }, + { + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [{"type": "text", "content": "X" * 300}], + "is_error": False, + } + ], + "data": {"type": "user"}, + }, + {"result": "THE ANSWER", "data": {"type": "result"}}, + ] + } + block2 = { + "messages": [ + { + "content": [{"type": "text", "text": "Final answer body"}], + "data": {"type": "assistant"}, + } + ] + } + return ( + _json_dumps(block1) + + "\n===== CLAUDE SDK ATTEMPT 2 =====\n" + + _json_dumps(block2) + ) + + +def test_parse_claude_trace_steps_extracts_and_filters() -> None: + steps = parse_claude_trace_steps(_sdk_payload()) + + types_seen = [step["type"] for step in steps] + # init / thinking_tokens / system bookkeeping are dropped. + assert types_seen == ["text", "tool_call", "tool_result", "text", "text"] + + # Indices are renumbered sequentially across attempt blocks. + assert [step["index"] for step in steps] == [1, 2, 3, 4, 5] + + assert steps[0]["summary"] == "Let me read the task." + assert steps[1]["summary"] == "Read task.md" + # tool_result is truncated to 200 chars + a [+N chars] trailer. + assert steps[2]["summary"].startswith("X" * 200) + assert "[+100 chars]" in steps[2]["summary"] + assert steps[3]["summary"] == "THE ANSWER" + assert steps[4]["summary"] == "Final answer body" + + +def test_parse_claude_trace_steps_marks_errors() -> None: + block = { + "messages": [ + { + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_9", + "content": [{"type": "text", "content": "boom"}], + "is_error": True, + } + ] + } + ] + } + steps = parse_claude_trace_steps(_json_dumps(block)) + assert len(steps) == 1 + assert steps[0]["type"] == "tool_result" + assert steps[0]["summary"] == "[error] boom" + + +def test_format_claude_trace_steps_truncates_total() -> None: + text = format_claude_trace_steps(_sdk_payload(), max_chars=40) + trailer = "\n...[claude trace steps truncated]..." + assert text.endswith(trailer) + assert text == text[:40] + trailer + + +def test_persist_claude_artifacts_writes_trace_steps(tmp_path) -> None: + work_dir = tmp_path / "pred" / "work" + work_dir.mkdir(parents=True) + + _persist_claude_artifacts(str(work_dir), _sdk_payload(), "response") + + steps_path = tmp_path / "pred" / "claude_trace_steps.txt" + assert steps_path.exists() + content = steps_path.read_text(encoding="utf-8") + assert content.strip() + # text step is index 1, so the tool_call is index 2. + assert "[2] tool_call: Read task.md" in content + + +def test_chat_optimizer_routes_to_claude_code_backend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from skillopt.model import backend_config, claude_code_backend + from skillopt.model import azure_openai + + claude_calls: list[dict[str, Any]] = [] + + def fake_claude_optimizer(**kwargs: Any) -> tuple[str, dict[str, int]]: + claude_calls.append(kwargs) + return "claude result", { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + } + + def fail_openai_optimizer(**_kwargs: Any) -> tuple[str, dict[str, int]]: + raise AssertionError("openai optimizer should not be called for claude_code_exec") + + monkeypatch.setattr(claude_code_backend, "chat_optimizer", fake_claude_optimizer) + monkeypatch.setattr(azure_openai, "chat_optimizer", fail_openai_optimizer) + backend_config.set_optimizer_backend("claude_code_exec") + + from skillopt.model import chat_optimizer + + text, usage = chat_optimizer("system", "user", retries=1, timeout=5) + + assert text == "claude result" + assert usage["total_tokens"] == 3 + assert claude_calls[0]["system"] == "system" + assert claude_calls[0]["user"] == "user" + assert claude_calls[0]["timeout"] == 5 + + +def test_reasoning_effort_forwarded_to_run_claude_code_chat( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from skillopt.model import claude_code_backend + + calls: list[dict[str, Any]] = [] + + def fake_chat(**kwargs: Any) -> tuple[str, dict[str, int]]: + calls.append(kwargs) + return "plain reply", {"prompt_tokens": 1, "completion_tokens": 1} + + monkeypatch.setattr(claude_code_backend, "run_claude_code_chat", fake_chat) + claude_code_backend.set_reasoning_effort("high") + try: + text, _usage = claude_code_backend.chat_optimizer("s", "u", retries=1) + finally: + claude_code_backend.set_reasoning_effort(None) + + assert text == "plain reply" + assert calls[0]["effort"] == "high" + + +def test_reasoning_effort_param_beats_module_global( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from skillopt.model import claude_code_backend + + calls: list[dict[str, Any]] = [] + + def fake_chat(**kwargs: Any) -> tuple[str, dict[str, int]]: + calls.append(kwargs) + return "plain reply", {"prompt_tokens": 1, "completion_tokens": 1} + + monkeypatch.setattr(claude_code_backend, "run_claude_code_chat", fake_chat) + claude_code_backend.set_reasoning_effort("low") + try: + claude_code_backend.chat_optimizer( + "s", "u", retries=1, reasoning_effort="max" + ) + finally: + claude_code_backend.set_reasoning_effort(None) + + assert calls[0]["effort"] == "max" + + +def test_claude_code_backend_retry_on_bad_json( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from skillopt.model import claude_code_backend + + monkeypatch.setattr( + claude_code_backend, + "run_claude_code_chat", + lambda **kwargs: ("this is not json", {"prompt_tokens": 3, "completion_tokens": 4}), + ) + + claude_code_backend.reset_token_tracker() + try: + # structured output (tools set) forces a json.loads on the reply; a + # non-JSON reply must be retried, then surfaced as RuntimeError. + with pytest.raises(RuntimeError, match="failed after 2 retries"): + claude_code_backend.chat_optimizer_messages( + [{"role": "user", "content": "hi"}], + retries=2, + tools=[{"name": "lookup"}], + ) + + summary = claude_code_backend.get_token_summary() + optimizer = summary["optimizer"] + assert optimizer["calls"] == 2 + assert optimizer["prompt_tokens"] == 6 + assert optimizer["completion_tokens"] == 8 + finally: + claude_code_backend.reset_token_tracker() + + +@pytest.mark.parametrize( + ("target_backend", "config_on", "expect_codex", "expect_claude"), + [ + ("claude_code_exec", True, "0", "1"), + ("claude_code_exec", False, "0", "0"), + ("codex_exec", True, "1", "0"), + ("codex_exec", False, "0", "0"), + ("openai_chat", True, "0", "0"), + ], +) +def test_trainer_configures_trace_gates( + monkeypatch: pytest.MonkeyPatch, + target_backend: str, + config_on: bool, + expect_codex: str, + expect_claude: str, +) -> None: + from skillopt.engine.trainer import _configure_trace_to_optimizer_gates + + monkeypatch.delenv("REFLACT_CODEX_TRACE_TO_OPTIMIZER", raising=False) + monkeypatch.delenv("REFLACT_CLAUDE_TRACE_TO_OPTIMIZER", raising=False) + _configure_trace_to_optimizer_gates( + target_backend, + {"codex_trace_to_optimizer": config_on, "claude_trace_to_optimizer": config_on}, + ) + assert os.environ["REFLACT_CODEX_TRACE_TO_OPTIMIZER"] == expect_codex + assert os.environ["REFLACT_CLAUDE_TRACE_TO_OPTIMIZER"] == expect_claude + + +@pytest.mark.parametrize( + ("gate_value", "expect_injected"), + [("0", False), ("1", True)], +) +def test_claude_trace_steps_gated_in_fmt_minibatch( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + gate_value: str, + expect_injected: bool, +) -> None: + tid = "tid0" + pred_dir = tmp_path / "predictions" + (pred_dir / tid).mkdir(parents=True) + (pred_dir / tid / "conversation.json").write_text( + json.dumps([{"role": "assistant", "content": "hi"}]), + encoding="utf-8", + ) + (pred_dir / tid / "claude_trace_steps.txt").write_text( + "[1] tool_call: Read task.md", + encoding="utf-8", + ) + + monkeypatch.setenv("REFLACT_CLAUDE_TRACE_TO_OPTIMIZER", gate_value) + + formatted = fmt_minibatch_trajectories( + [{"id": tid, "task_description": "t", "task_type": "q"}], + str(pred_dir), + ) + + assert ("#### Claude Trace Steps" in formatted) is expect_injected diff --git a/tests/test_role_backend_resolution.py b/tests/test_role_backend_resolution.py index afb92e5a..aa58f30d 100644 --- a/tests/test_role_backend_resolution.py +++ b/tests/test_role_backend_resolution.py @@ -52,7 +52,7 @@ def test_eval_only_backend_label_preserves_explicit_cli_role_override() -> None: ("cursor_exec", ("openai_chat", "cursor_exec")), ("claude", ("claude_chat", "claude_chat")), ("claude_chat", ("claude_chat", "claude_chat")), - ("claude_code_exec", ("openai_chat", "claude_code_exec")), + ("claude_code_exec", ("claude_code_exec", "claude_code_exec")), ("codex", ("codex_exec", "codex_exec")), ("codex_exec", ("codex_exec", "codex_exec")), ("qwen", ("openai_chat", "qwen_chat")), @@ -125,7 +125,7 @@ def test_explicit_optimizer_is_preserved_while_default_target_is_resolved( [ ("claude", "claude_chat"), ("codex", "codex_exec"), - ("claude_code_exec", "openai_chat"), + ("claude_code_exec", "claude_code_exec"), ("cursor", "openai_chat"), ("copilot", "copilot_chat"), ("copilot_exec", "openai_chat"),