diff --git a/agent-flow/agent_flow/agent_runtime.py b/agent-flow/agent_flow/agent_runtime.py new file mode 100644 index 000000000000..59ffc08c4230 --- /dev/null +++ b/agent-flow/agent_flow/agent_runtime.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-role model and backend selection shared by performance workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +from .config import CLAUDE_CODE_DEFAULT_MODEL, CODEX_DEFAULT_MODEL, BackendKind + +AGENTS_FIELD = "agents" +_FIELDS = frozenset({"backend", "model", "reasoning_effort", "extra_mcp_servers"}) +_BACKENDS = frozenset({"claude-code", "codex"}) + + +@dataclass(frozen=True) +class AgentConfig: + """Resolved backend configuration for one workflow role.""" + + backend: BackendKind + model: str + reasoning_effort: str | None = None + extra_mcp_servers: dict[str, Any] | None = None + + +def validate_agents(data: Mapping[str, Any], roles: tuple[str, ...]) -> list[str]: + """Validate the optional ``agents`` block for a workflow's roles.""" + if AGENTS_FIELD not in data: + return [] + + agents = data[AGENTS_FIELD] + if not isinstance(agents, Mapping): + return ["'agents' must be a mapping"] + + errors: list[str] = [] + unknown = set(agents) - {"defaults", "roles"} + if unknown: + errors.append(f"'agents' has unknown field(s) {sorted(unknown)}") + + _validate_block(agents.get("defaults", {}), "agents.defaults", errors) + role_blocks = agents.get("roles", {}) + if not isinstance(role_blocks, Mapping): + errors.append("'agents.roles' must be a mapping") + return errors + + unknown_roles = set(role_blocks) - set(roles) + if unknown_roles: + errors.append(f"'agents.roles' has unknown role(s) {sorted(unknown_roles)}") + for role, block in role_blocks.items(): + _validate_block(block, f"agents.roles.{role}", errors) + return errors + + +def _validate_block(value: Any, path: str, errors: list[str]) -> None: + if not isinstance(value, Mapping): + errors.append(f"'{path}' must be a mapping") + return + + unknown = set(value) - _FIELDS + if unknown: + errors.append(f"'{path}' has unknown field(s) {sorted(unknown)}") + backend = value.get("backend") + if backend is not None and backend not in _BACKENDS: + errors.append(f"'{path}.backend' must be 'claude-code' or 'codex'") + for field in ("model", "reasoning_effort"): + item = value.get(field) + if item is not None and (not isinstance(item, str) or not item.strip()): + errors.append(f"'{path}.{field}' must be a non-empty string") + servers = value.get("extra_mcp_servers") + if servers is not None and not isinstance(servers, Mapping): + errors.append(f"'{path}.extra_mcp_servers' must be a mapping") + + +def resolve_agent_config( + data: Mapping[str, Any], + role: str, + *, + default_backend: BackendKind = "claude-code", + default_model: str = CLAUDE_CODE_DEFAULT_MODEL, +) -> AgentConfig: + """Resolve one role from workflow defaults and its optional override.""" + agents = data.get(AGENTS_FIELD, {}) + if not isinstance(agents, Mapping): + return AgentConfig(default_backend, default_model) + defaults = agents.get("defaults", {}) + defaults = defaults if isinstance(defaults, Mapping) else {} + role_blocks = agents.get("roles", {}) + role_blocks = role_blocks if isinstance(role_blocks, Mapping) else {} + override = role_blocks.get(role, {}) + override = override if isinstance(override, Mapping) else {} + + backend = override.get("backend", defaults.get("backend", default_backend)) + model = override.get("model") + if model is None: + inherited_model = defaults.get("model") + inherited_backend = defaults.get("backend", default_backend) + if inherited_model is not None and backend == inherited_backend: + model = inherited_model + elif backend == default_backend and "backend" not in defaults and "backend" not in override: + model = default_model + else: + model = CODEX_DEFAULT_MODEL if backend == "codex" else CLAUDE_CODE_DEFAULT_MODEL + + servers = override.get("extra_mcp_servers", defaults.get("extra_mcp_servers")) + return AgentConfig( + backend=backend, + model=model, + reasoning_effort=override.get("reasoning_effort", defaults.get("reasoning_effort")), + extra_mcp_servers=dict(servers) if isinstance(servers, Mapping) else None, + ) diff --git a/agent-flow/agent_flow/backends/__init__.py b/agent-flow/agent_flow/backends/__init__.py index 58727067f927..d439fe45d2e5 100644 --- a/agent-flow/agent_flow/backends/__init__.py +++ b/agent-flow/agent_flow/backends/__init__.py @@ -14,11 +14,19 @@ def create_backend(config: BackendConfig | str) -> Backend: kind = config.kind if isinstance(config, BackendConfig) else config + effort = config.reasoning_effort if isinstance(config, BackendConfig) else None + disabled_skills = config.disabled_skills if isinstance(config, BackendConfig) else () if kind == "claude-code": - return claude_code.ClaudeCodeBackend() + return claude_code.ClaudeCodeBackend( + reasoning_effort=effort, + disabled_skills=disabled_skills, + ) if kind == "codex": - return codex.CodexBackend() + return codex.CodexBackend( + reasoning_effort=effort, + disabled_skills=disabled_skills, + ) raise ValueError(f"Unknown backend: {kind!r}") diff --git a/agent-flow/agent_flow/backends/claude_code.py b/agent-flow/agent_flow/backends/claude_code.py index 4063384b5410..abfcb9a49aa7 100644 --- a/agent-flow/agent_flow/backends/claude_code.py +++ b/agent-flow/agent_flow/backends/claude_code.py @@ -446,11 +446,19 @@ async def handler(arguments, definition=definition): class ClaudeCodeBackend(Backend): + def __init__( + self, + reasoning_effort: str | None = None, + disabled_skills: tuple[str, ...] = (), + ) -> None: + self._reasoning_effort = reasoning_effort or _REASONING_EFFORT + self._disabled_skills = disabled_skills + def version(self) -> str: return _claude_backend_version() def reasoning_effort(self) -> str: - return _REASONING_EFFORT + return self._reasoning_effort @asynccontextmanager async def create_client( @@ -490,12 +498,15 @@ async def create_client( ), mcp_servers=mcp_servers, model=model, - effort=_REASONING_EFFORT, + effort=self._reasoning_effort, cwd=cwd or Path.cwd(), sandbox={"enabled": False}, permission_mode="bypassPermissions", hooks=hooks, - disallowed_tools=list(disallowed_tools or []), + disallowed_tools=[ + *(disallowed_tools or []), + *(f"Skill({name})" for name in self._disabled_skills), + ], ) async with ClaudeSDKClient(options=options) as sdk_client: diff --git a/agent-flow/agent_flow/backends/codex.py b/agent-flow/agent_flow/backends/codex.py index 14a7f1e44a64..3f105967bb08 100644 --- a/agent-flow/agent_flow/backends/codex.py +++ b/agent-flow/agent_flow/backends/codex.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import json import os import shutil import subprocess @@ -141,19 +142,36 @@ def _codex_backend_version() -> str: _REASONING_EFFORT = "max" +def _disabled_skills_override(names: tuple[str, ...]) -> tuple[str, ...]: + if not names: + return () + entries = ", ".join(f"{{name={json.dumps(name)}, enabled=false}}" for name in names) + return (f"skills.config=[{entries}]",) + + class CodexBackend(Backend): - def __init__(self) -> None: + def __init__( + self, + reasoning_effort: str | None = None, + disabled_skills: tuple[str, ...] = (), + ) -> None: self._transport: CodexTransport | None = None + self._reasoning_effort = reasoning_effort or _REASONING_EFFORT + self._disabled_skills = disabled_skills def version(self) -> str: return _codex_backend_version() def reasoning_effort(self) -> str: - return _REASONING_EFFORT + return self._reasoning_effort async def __aenter__(self) -> "CodexBackend": self._transport = CodexTransport( - CodexConfig(codex_bin=_resolve_codex_bin(), experimental_api=True) + CodexConfig( + codex_bin=_resolve_codex_bin(), + experimental_api=True, + config_overrides=_disabled_skills_override(self._disabled_skills), + ) ) await self._transport.start() return self @@ -222,7 +240,10 @@ async def create_client( server["disabled_tools"] = list( dict.fromkeys([*inherited, *server["disabled_tools"]]) ) - config.update(model_reasoning_effort=_REASONING_EFFORT, model_context_window=1000000) + config.update( + model_reasoning_effort=self._reasoning_effort, + model_context_window=1000000, + ) params = ThreadStartParams( model=model, developer_instructions=system_prompt or None, diff --git a/agent-flow/agent_flow/config.py b/agent-flow/agent_flow/config.py index 76d81617bddc..dedf6290605e 100644 --- a/agent-flow/agent_flow/config.py +++ b/agent-flow/agent_flow/config.py @@ -16,6 +16,10 @@ class BackendConfig: kind: BackendKind model: str + # Provider reasoning tier. ``None`` retains the backend's historical + # maximum-effort default. + reasoning_effort: str | None = None + disabled_skills: tuple[str, ...] = () tools: list[Any] | None = None # Native SDK hook configuration. Claude accepts HookMatcher callbacks. # Codex requires hooks to be configured and trusted in its native config; diff --git a/agent-flow/agent_flow/workflows/perf_analyze/README.md b/agent-flow/agent_flow/workflows/perf_analyze/README.md index 6b3a34ef2893..f2ba17887563 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/README.md +++ b/agent-flow/agent_flow/workflows/perf_analyze/README.md @@ -5,7 +5,7 @@ benchmarks and profiles it with TensorRT-LLM's `tensorrt_llm/serve/scripts/benchmark_serving.py`, and writes a report whose headline is the **main performance bottleneck**. -All roles run on the **Claude Code** backend: +By default, all roles run on the **Claude Code** backend: ``` benchmarker ──▶ projector ──▶ analyzer ──▶ reporter @@ -81,6 +81,20 @@ Copy [`task.example.yaml`](./task.example.yaml) and fill it in. ## `task.yaml` +An optional `agents` block selects the backend, model, reasoning effort, +and external MCP servers per role. Unspecified roles retain the historical +Claude defaults: + +```yaml +agents: + roles: + projector: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} + analyzer: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} +``` + +Set `casebook.enabled: false` for a control run that hides and blocks the +`perf-optimization-casebook` skill. It is enabled by default. + | Field | Required | Notes | | --- | --- | --- | | `checkpoint_path` | ✅ | Model checkpoint dir to serve. Remote when `cluster_ssh` is set; otherwise local. | diff --git a/agent-flow/agent_flow/workflows/perf_analyze/cli.py b/agent-flow/agent_flow/workflows/perf_analyze/cli.py index cbf04ac39151..8265b74d26fe 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/cli.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/cli.py @@ -4,11 +4,14 @@ import sys from pathlib import Path +from agent_flow.agent_runtime import resolve_agent_config + from .prompts import build_perf_analyze_prompts from .sol_methodology import resolve_sol_methodology from .state import STATE_FILENAME from .task_schema import ( TaskSchemaError, + casebook_enabled, has_slurm_environment, load_and_validate_task_yaml, sol_enabled, @@ -61,15 +64,19 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> None: args = _parse_args(argv) + task_path = args.task + if not args.clean and (args.workspace / STATE_FILENAME).is_file(): + task_path = args.workspace / "task.yaml" try: - task_data = load_and_validate_task_yaml(args.task) + task_data = load_and_validate_task_yaml(task_path) except TaskSchemaError as exc: print(f"error: {exc}", file=sys.stderr) sys.exit(2) # Resolve the projector's methodology skill once, before the run, so # it is told to load a skill this session actually has. Skipped (free) # when the stage is off. - methodology = resolve_sol_methodology(sol_enabled(task_data)) + projector_backend = resolve_agent_config(task_data, "projector").backend + methodology = resolve_sol_methodology(sol_enabled(task_data), backend_kind=projector_backend) note = methodology.console_note() if note: print(note, file=sys.stderr) @@ -79,6 +86,7 @@ def main(argv: list[str] | None = None) -> None: sol_methodology=methodology.name, remote_execution=task_data, campaign_name=args.workspace.resolve().name, + include_casebook=casebook_enabled(task_data), ) with PerfAnalyzeWorkflow( workspace=args.workspace, diff --git a/agent-flow/agent_flow/workflows/perf_analyze/prompts/__init__.py b/agent-flow/agent_flow/workflows/perf_analyze/prompts/__init__.py index 8138882e832e..3aeff04b061c 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/prompts/__init__.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/prompts/__init__.py @@ -4,6 +4,7 @@ from ..task_schema import cluster_ssh, remote_run_root from ._common import ( + CASEBOOK_DISABLED, EXECUTION_SLURM_BOOTSTRAP, REMOTE_SLURM_EXECUTION, SOL_ANALYZER_CONTEXT, @@ -94,6 +95,7 @@ def build_perf_analyze_prompts( sol_methodology: str = "full", remote_execution: Mapping[str, Any] | None = None, campaign_name: str = "perf-analyze", + include_casebook: bool = True, ) -> PromptBundle: """Return the workflow's prompt bundle, optionally augmented. @@ -120,6 +122,11 @@ def build_perf_analyze_prompts( appended to the three roles that may inspect or produce runtime data. """ bundle = DEFAULT_PROMPTS + if not include_casebook: + bundle = bundle.with_extensions( + benchmarker=CASEBOOK_DISABLED, + analyzer=CASEBOOK_DISABLED, + ) if sol_methodology != "full": bundle = dataclasses.replace(bundle, projector=build_projector_prompt(sol_methodology)) if include_slurm_environment: diff --git a/agent-flow/agent_flow/workflows/perf_analyze/prompts/_common.py b/agent-flow/agent_flow/workflows/perf_analyze/prompts/_common.py index 4637754099cd..4eb5218009f4 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/prompts/_common.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/prompts/_common.py @@ -1085,6 +1085,15 @@ line and proceed — never block the run on it. """ +CASEBOOK_DISABLED = """\ +## Casebook experiment + +This run intentionally disables `perf-optimization-casebook`. Do not invoke or +use that skill, even if an earlier section asks for it. Ground conclusions only +in the task, source, measurements, and other enabled skills; `casebook_ref` may +be omitted. +""" + # --------------------------------------------------------------------------- # # Remote execution boundary (appended with task-specific values by the CLI) diff --git a/agent-flow/agent_flow/workflows/perf_analyze/roles.py b/agent-flow/agent_flow/workflows/perf_analyze/roles.py new file mode 100644 index 000000000000..8e379cfca7be --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_analyze/roles.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +ROLES = ("benchmarker", "projector", "analyzer", "reporter") diff --git a/agent-flow/agent_flow/workflows/perf_analyze/sol_methodology.py b/agent-flow/agent_flow/workflows/perf_analyze/sol_methodology.py index cd9292bdd15b..e4de392db553 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/sol_methodology.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/sol_methodology.py @@ -88,7 +88,9 @@ def console_note(self) -> str | None: return None -def resolve_sol_methodology(enabled: bool = True) -> SolMethodology: +def resolve_sol_methodology( + enabled: bool = True, backend_kind: str = "claude-code" +) -> SolMethodology: """Probe the live skill list and return the methodology the projector has. Costs one backend connection and no model call (~1 s), and is @@ -101,7 +103,9 @@ def resolve_sol_methodology(enabled: bool = True) -> SolMethodology: from agent_flow.utils import resolve_first_available_skill try: - loaded, probe_ok = resolve_first_available_skill(SOL_SKILL_CANDIDATES) + loaded, probe_ok = resolve_first_available_skill( + SOL_SKILL_CANDIDATES, backend_kinds=(backend_kind,) + ) except Exception: # noqa: BLE001 - a probe failure must never fail the run return SolMethodology(probed=False) if not probe_ok: diff --git a/agent-flow/agent_flow/workflows/perf_analyze/task.example.yaml b/agent-flow/agent_flow/workflows/perf_analyze/task.example.yaml index c15c271711dc..ee4d5646619a 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/task.example.yaml +++ b/agent-flow/agent_flow/workflows/perf_analyze/task.example.yaml @@ -12,6 +12,17 @@ checkpoint_path: /path/to/checkpoints/- # remote when usin trtllm_repo_path: /path/to/TensorRT-LLM # always local # extra_llm_api_options: /local/path/to/extra_llm_api_options.yaml +# Optional per-role backend routing. Roles omitted here keep the historical +# Claude backend and model. +# agents: +# roles: +# projector: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} +# analyzer: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} + +# Optional A/B switch. The optimization casebook is enabled by default. +# casebook: +# enabled: false + # Optional: the benchmark operating point(s) driven by # benchmark_serving.py. benchmark: diff --git a/agent-flow/agent_flow/workflows/perf_analyze/task_schema.py b/agent-flow/agent_flow/workflows/perf_analyze/task_schema.py index 6e3e8a746f6b..2acee4c4039c 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/task_schema.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/task_schema.py @@ -45,6 +45,10 @@ import yaml +from agent_flow.agent_runtime import AGENTS_FIELD, validate_agents + +from .roles import ROLES + REQUIRED_PATH_FIELDS: tuple[str, ...] = ( "checkpoint_path", "trtllm_repo_path", @@ -99,6 +103,14 @@ SOL_FIELDS: tuple[str, ...] = (SOL_ENABLED_FIELD, *SOL_OPTIONAL_STR_FIELDS) SOL_DEFAULTS: dict[str, Any] = {SOL_ENABLED_FIELD: True} +CASEBOOK_FIELD = "casebook" +CASEBOOK_ENABLED_FIELD = "enabled" +CASEBOOK_DEFAULTS: dict[str, Any] = {CASEBOOK_ENABLED_FIELD: True} +CASEBOOK_SKILL_NAMES = ( + "perf-optimization-casebook", + "trtllm-agent-toolkit:perf-optimization-casebook", +) + # The pre-rename spelling of the ``sol`` block (when the projector still # cross-checked a dlsim checkout) — rejected with an actionable error so # a stale task.yaml fails loudly instead of having its projector settings @@ -188,6 +200,8 @@ "profile", SLURM_ENVIRONMENT_FIELD, SOL_FIELD, + CASEBOOK_FIELD, + AGENTS_FIELD, ) ) @@ -420,7 +434,11 @@ def profile_ranks(data: Mapping[str, Any]) -> tuple[int, ...]: return (0,) -def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: +def load_and_validate_task_yaml( + path: str | Path, + *, + agent_roles: tuple[str, ...] = ROLES, +) -> dict[str, Any]: """Parse ``path`` as YAML and validate the perf-analyze schema. Returns the parsed mapping with the optional ``benchmark`` / ``profile`` @@ -460,7 +478,7 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: f"{task_path} must be a YAML mapping at the top level, got {type(data).__name__}" ) - errors: list[str] = [] + errors: list[str] = validate_agents(data, agent_roles) # Decided once, from the spec, before any path is looked at. In remote mode # the checkpoint is remote, but the checkout and optional tuning YAML remain @@ -549,6 +567,14 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: errors.append("'profile.nsys_iter_range' must be a non-empty string (e.g. \"100-150\")") _validate_profile_ranks(profile, methods, errors) + casebook = _validate_mapping_block(data, CASEBOOK_FIELD, errors) + unknown = set(casebook) - {CASEBOOK_ENABLED_FIELD} + if unknown: + errors.append(f"'{CASEBOOK_FIELD}' has unknown field(s) {sorted(unknown)}") + enabled = casebook.get(CASEBOOK_ENABLED_FIELD) + if CASEBOOK_ENABLED_FIELD in casebook and not isinstance(enabled, bool): + errors.append(f"'{CASEBOOK_FIELD}.{CASEBOOK_ENABLED_FIELD}' must be a boolean") + if SLURM_ENVIRONMENT_FIELD in data and data[SLURM_ENVIRONMENT_FIELD] is not None: slurm_environment = data[SLURM_ENVIRONMENT_FIELD] if not isinstance(slurm_environment, dict): @@ -670,6 +696,7 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: benchmark["concurrency"] = sorted(set(benchmark["concurrency"])) data["benchmark"] = {**BENCHMARK_DEFAULTS, **benchmark} data["profile"] = {**PROFILE_DEFAULTS, **profile} + data[CASEBOOK_FIELD] = {**CASEBOOK_DEFAULTS, **casebook} # ``sol`` is materialized even when the user never wrote the block — # the projector is on by default, so the resolved spec has to state # the gate rather than leave it to a reader's assumption. @@ -749,6 +776,14 @@ def sol_enabled(data: Mapping[str, Any]) -> bool: return True +def casebook_enabled(data: Mapping[str, Any]) -> bool: + """Return whether agents may consult the optimization casebook.""" + casebook = data.get(CASEBOOK_FIELD) + if isinstance(casebook, Mapping): + return casebook.get(CASEBOOK_ENABLED_FIELD, True) is not False + return True + + def is_curve_mode(data: Mapping[str, Any]) -> bool: """True iff ``benchmark.concurrency`` is a list (Pareto-curve mode). @@ -823,6 +858,10 @@ def dump_task_yaml(data: Mapping[str, Any]) -> str: __all__ = [ "BENCHMARK_DEFAULTS", + "CASEBOOK_DEFAULTS", + "CASEBOOK_ENABLED_FIELD", + "CASEBOOK_FIELD", + "CASEBOOK_SKILL_NAMES", "EXTRA_LLM_API_OPTIONS_FIELD", "PROFILE_DEFAULTS", "PROFILE_RANKS_FIELD", @@ -841,6 +880,7 @@ def dump_task_yaml(data: Mapping[str, Any]) -> str: "SOL_OPTIONAL_STR_FIELDS", "VALID_PROFILE_METHODS", "TaskSchemaError", + "casebook_enabled", "concurrency_points", "dump_task_yaml", "cluster_ssh", diff --git a/agent-flow/agent_flow/workflows/perf_analyze/workflow.py b/agent-flow/agent_flow/workflows/perf_analyze/workflow.py index ef3cec87c943..ca62bde35ef7 100644 --- a/agent-flow/agent_flow/workflows/perf_analyze/workflow.py +++ b/agent-flow/agent_flow/workflows/perf_analyze/workflow.py @@ -6,13 +6,8 @@ import yaml -from agent_flow import ( - CLAUDE_CODE_DEFAULT_MODEL, - AgentLayer, - AgentLayerConfig, - BackendConfig, - SessionConfig, -) +from agent_flow import AgentLayer, AgentLayerConfig, BackendConfig, SessionConfig +from agent_flow.agent_runtime import AgentConfig, resolve_agent_config from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger @@ -25,6 +20,7 @@ ) from .prompts import DEFAULT_PROMPTS, PromptBundle from .prompts._common import profile_ranks_note +from .roles import ROLES from .sol_methodology import SolMethodology, output_instruction, projector_instruction from .state import ( STAGE_ANALYZER, @@ -37,6 +33,8 @@ save_state, ) from .task_schema import ( + CASEBOOK_SKILL_NAMES, + casebook_enabled, concurrency_points, dump_task_yaml, is_curve_mode, @@ -103,20 +101,23 @@ def _progress_has_entries(path: Path) -> bool: def _make_agent( name: str, system_prompt: str, + agent_config: AgentConfig, tools: list | None = None, required_tools: list[str] | None = None, - backend_kind: str = "claude-code", - model: str = CLAUDE_CODE_DEFAULT_MODEL, session_mode: str = "persistent", + disabled_skills: tuple[str, ...] = (), ) -> AgentLayer: return AgentLayer( AgentLayerConfig( name=name, system_prompt=system_prompt, backend=BackendConfig( - kind=backend_kind, - model=model, + kind=agent_config.backend, + model=agent_config.model, + reasoning_effort=agent_config.reasoning_effort, + disabled_skills=disabled_skills, tools=tools, + extra_mcp_servers=agent_config.extra_mcp_servers, ), session=SessionConfig(mode=session_mode), required_tools=tuple(required_tools or ()), @@ -234,33 +235,11 @@ def __init__( self._progress_ctx = ProgressContext(path=self.progress_path) progress_tools = build_progress_tools(self._progress_ctx) - self.benchmarker = _make_agent( - "benchmarker", - self.prompts.benchmarker, - progress_tools["benchmarker"], - required_tools=["append_benchmarker_progress"], - ) - # Constructed unconditionally (the stage gate lives in ``run``); - # the backend client is lazy, so a skipped projector costs nothing. - self.projector = _make_agent( - "projector", - self.prompts.projector, - progress_tools["projector"], - required_tools=["append_projector_progress"], - ) - self.analyzer = _make_agent( - "analyzer", - self.prompts.analyzer, - progress_tools["analyzer"], - required_tools=["append_analyzer_progress"], - ) - self.reporter = _make_agent( - "reporter", - self.prompts.reporter, - progress_tools["reporter"], - required_tools=["append_reporter_progress"], - ) self._progress_tools = progress_tools + self._agent_configs: dict[str, AgentConfig] = {} + self._disabled_skills: tuple[str, ...] = () + for role in ROLES: + setattr(self, role, None) def __enter__(self) -> "PerfAnalyzeWorkflow": return self @@ -269,8 +248,28 @@ def __exit__(self, exc_type, exc, tb) -> None: self.close() def close(self) -> None: - for layer in (self.benchmarker, self.projector, self.analyzer, self.reporter): - layer.__exit__(None, None, None) + for role in ROLES: + layer = getattr(self, role) + if layer is not None: + layer.__exit__(None, None, None) + + def _configure_agents(self) -> None: + task_data = self._task_data() + self._disabled_skills = () if casebook_enabled(task_data) else CASEBOOK_SKILL_NAMES + self._agent_configs = {role: resolve_agent_config(task_data, role) for role in ROLES} + for role in ROLES: + setattr( + self, + role, + _make_agent( + role, + getattr(self.prompts, role), + self._agent_configs[role], + self._progress_tools[role], + required_tools=[f"append_{role}_progress"], + disabled_skills=self._disabled_skills, + ), + ) # ------------------------------------------------------------- orchestration @@ -280,6 +279,7 @@ def run(self, task: str) -> None: state = self._init_state(task, log) if state is None: return + self._configure_agents() try: # Each stage checkpoints before advancing, so a crash / Ctrl-C @@ -424,11 +424,13 @@ def _run_benchmarker(self) -> None: f"Read `{self.task_path}` for the spec — resolve `checkpoint_path`, " f"`trtllm_repo_path`, the optional `extra_llm_api_options` path, " f"and the `benchmark` block.\n\n" - f"Then **load the `perf-optimization-casebook` skill** (via the " - f"`Skill` tool) as read-only reference, as your system prompt " - f"directs, so your Configuration/Notes are grounded in known " - f"TRT-LLM performance precedents.\n\n" - f"Launch `trtllm-serve` (passing `--extra_llm_api_options` when " + + self._casebook_instruction( + "Then **load the `perf-optimization-casebook` skill** (via the " + "`Skill` tool) as read-only reference, as your system prompt " + "directs, so your Configuration/Notes are grounded in known " + "TRT-LLM performance precedents.\n\n" + ) + + f"Launch `trtllm-serve` (passing `--extra_llm_api_options` when " f"set), poll it to " f"readiness, {load_instruction}. Use the " f"**canonical `benchmark_serving.py` command in your system " @@ -532,12 +534,14 @@ def _run_analyzer(self) -> None: f"baseline.\n\n" + curve_context + projection_context - + f"Early on, **load the `perf-optimization-casebook` skill** (via " - f"the `Skill` tool) as read-only reference, as your system prompt " - f"directs, and match each ranked bottleneck hypothesis against its " - f"*bottleneck signal → candidate pattern* index so the Reporter " - f"inherits a known precedent.\n\n" - f"First **verify this checkout's profiling knobs** with " + + self._casebook_instruction( + "Early on, **load the `perf-optimization-casebook` skill** (via " + "the `Skill` tool) as read-only reference, as your system prompt " + "directs, and match each ranked bottleneck hypothesis against its " + "*bottleneck signal → candidate pattern* index so the Reporter " + "inherits a known precedent.\n\n" + ) + + f"First **verify this checkout's profiling knobs** with " f"`grep -rn`/`rg` via `Bash` under " f"`{self._trtllm_hint()}` — `py_executor.py` for " f"`TLLM_PROFILE_START_STOP` (the iteration-window gate), and " @@ -687,6 +691,9 @@ def _sol_enabled(self) -> bool: """ return sol_enabled(self._task_data()) + def _casebook_instruction(self, text: str) -> str: + return "" if self._disabled_skills else text + def _profile_ranks(self) -> tuple[int, ...]: """The rank ids nsys must capture, from the resolved spec.""" return profile_ranks(self._task_data()) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/README.md b/agent-flow/agent_flow/workflows/perf_optimize/README.md index 8d68a7c0b14f..70842c238cee 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/README.md +++ b/agent-flow/agent_flow/workflows/perf_optimize/README.md @@ -204,6 +204,27 @@ exactly when the round produced an `items.json`. A replan-only round runs no profiler and writes none; a round whose skill was unavailable or whose pipeline errored writes none either and records the reason under *Caveats*. Neither owes the block anything. +## Agent backend and model routing + +An optional top-level `agents` block in `task.yaml` selects `codex` or +`claude-code`, a model slug, reasoning effort, and portable external MCP +servers. `defaults` applies to every role and `roles.` overrides +individual fields. For example, this runs projector/analyzer on Astra ultra +and leaves every other role on the historical Claude default: + +```yaml +agents: + roles: + projector: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} + analyzer: {backend: codex, model: gpt-6-astra, reasoning_effort: ultra} +``` + +For an A/B run, set `casebook.enabled: false` in the control task. The +optimization casebook remains enabled when the block is omitted. + +Omitting `agents` preserves the historical assignment. On resume, the +checkpointed workspace's `task.yaml` remains authoritative, so a different +new `--task` cannot change models midway through a campaign. ## The acceptance gate diff --git a/agent-flow/agent_flow/workflows/perf_optimize/cli.py b/agent-flow/agent_flow/workflows/perf_optimize/cli.py index 5fd791e7578f..047cca485055 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/cli.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/cli.py @@ -4,7 +4,9 @@ import sys from pathlib import Path +from agent_flow.agent_runtime import resolve_agent_config from agent_flow.workflows.perf_analyze.sol_methodology import resolve_sol_methodology +from agent_flow.workflows.perf_analyze.task_schema import casebook_enabled from .disagg import has_disagg from .prompts import build_perf_optimize_prompts @@ -46,7 +48,8 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Path to the task.yaml spec. Requires `checkpoint_path` and " "`trtllm_repo_path`; optional top-level `extra_llm_api_options` " "path, optional `benchmark` / `profile` / `optimize` / `accuracy` " - "blocks, an optional `slurm-environment` block, and an optional " + "blocks, an optional `agents` block for per-role backend/model routing, " + "an optional `slurm-environment` block, and an optional " "`sol` block (all fields optional: `enabled` gates the one-shot " "SOL projector stage — on by default — and `gpu` names the GPU " "part for the SOL skill's peaks calculator). " @@ -109,10 +112,12 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> None: args = _parse_args(argv) + resume = not args.clean and (args.workspace / STATE_FILENAME).is_file() + task_path = args.workspace / "task.yaml" if resume else args.task try: task_data = load_and_validate_task_yaml( - args.task, - max_rounds_override=args.max_rounds, + task_path, + max_rounds_override=None if resume else args.max_rounds, ) except TaskSchemaError as exc: print(f"error: {exc}", file=sys.stderr) @@ -126,7 +131,8 @@ def main(argv: list[str] | None = None) -> None: # Resolve the projector's methodology skill once, before the run, so # it is told to load a skill this session actually has. Skipped (free) # when the stage is off. - methodology = resolve_sol_methodology(sol_enabled(task_data)) + projector_backend = resolve_agent_config(task_data, "projector").backend + methodology = resolve_sol_methodology(sol_enabled(task_data), backend_kind=projector_backend) note = methodology.console_note() if note: print(note, file=sys.stderr) @@ -139,6 +145,7 @@ def main(argv: list[str] | None = None) -> None: kernel_coverage=kernel_coverage(task_data), sol_methodology=methodology.name, include_disagg=has_disagg(task_data), + include_casebook=casebook_enabled(task_data), ) with PerfOptimizeWorkflow( workspace=args.workspace, diff --git a/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py b/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py index 6491eea8be70..cf71a59bb245 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py @@ -3,6 +3,7 @@ from typing import Any, Mapping, Sequence from agent_flow.workflows.perf_analyze.prompts import build_remote_execution_context +from agent_flow.workflows.perf_analyze.prompts._common import CASEBOOK_DISABLED from ._common import ( DISAGG_CAMPAIGN, @@ -101,6 +102,7 @@ def build_perf_optimize_prompts( kernel_coverage: Mapping[str, Any] | None = None, sol_methodology: str = "full", include_disagg: bool = False, + include_casebook: bool = True, ) -> PromptBundle: """Return the workflow's prompt bundle, augmented per the task spec. @@ -176,6 +178,12 @@ def build_perf_optimize_prompts( wrong. """ bundle = DEFAULT_PROMPTS + if not include_casebook: + bundle = bundle.with_extensions( + benchmarker=CASEBOOK_DISABLED, + analyzer=CASEBOOK_DISABLED, + optimizer=CASEBOOK_DISABLED, + ) if sol_methodology != "full": bundle = dataclasses.replace(bundle, projector=build_projector_prompt(sol_methodology)) restriction = approach_restriction_note(approaches) if approaches is not None else "" diff --git a/agent-flow/agent_flow/workflows/perf_optimize/roles.py b/agent-flow/agent_flow/workflows/perf_optimize/roles.py new file mode 100644 index 000000000000..d822d365d799 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/roles.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +ROLES = ( + "benchmarker", + "projector", + "analyzer", + "optimizer", + "evaluator", + "integrator", + "qa", + "reporter", +) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml b/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml index 593464a5d000..f95bf9ae356e 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml +++ b/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml @@ -23,6 +23,32 @@ trtllm_repo_path: /path/to/TensorRT-LLM # always local # config. # extra_llm_api_options: /path/to/extra_llm_api_options.yaml +# Optional: select the backend/model/reasoning tier globally and override +# individual workflow roles. Omit this block to preserve the workflow's +# historical backend assignment. A role inherits fields from defaults. +# agents: +# roles: +# projector: +# backend: codex # codex | claude-code +# model: gpt-6-astra +# reasoning_effort: ultra +# analyzer: +# backend: codex +# model: gpt-6-astra +# reasoning_effort: ultra +# # Optional portable external MCP servers (STDIO or streamable HTTP). +# # extra_mcp_servers belongs under defaults or an individual role. +# # defaults: +# # extra_mcp_servers: +# # knowledge-base: +# # type: http +# # url: https://example.test/mcp +# # headers: {X-Environment: staging} + +# Optional A/B switch. The optimization casebook is enabled by default. +# casebook: +# enabled: false + # Optional: the benchmark operating point(s) driven by # benchmark_serving.py. Every measurement in the campaign (baseline, # evaluator, final verification) replays the same point(s). diff --git a/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py b/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py index 0c51c4aea75e..061330b34b40 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py @@ -87,6 +87,7 @@ user_set_benchmark_keys, ) from agent_flow.workflows.perf_optimize.roadmap_schema import APPROACHES +from agent_flow.workflows.perf_optimize.roles import ROLES # Defaults merged under the user's values. ``target_improvement_pct`` is # deliberately absent: when the user does not set it, there is no @@ -421,7 +422,7 @@ def load_and_validate_task_yaml( agents read on disk is fully explicit; ``max_rounds_override`` (the CLI ``--max-rounds`` flag) is applied last, over the user's value. """ - data = _base_load_and_validate(path) + data = _base_load_and_validate(path, agent_roles=ROLES) errors: list[str] = [] # Disagg first: the harness config is the source of truth for the diff --git a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py index 1c3baca92f43..a4ef967aeeb8 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py @@ -13,13 +13,8 @@ import yaml from rich.markup import escape -from agent_flow import ( - CLAUDE_CODE_DEFAULT_MODEL, - AgentLayer, - AgentLayerConfig, - BackendConfig, - SessionConfig, -) +from agent_flow import AgentLayer, AgentLayerConfig, BackendConfig, SessionConfig +from agent_flow.agent_runtime import AgentConfig, resolve_agent_config from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger from agent_flow.workflows.perf_analyze.prompts._common import profile_ranks_note @@ -28,6 +23,7 @@ output_instruction, projector_instruction, ) +from agent_flow.workflows.perf_analyze.task_schema import CASEBOOK_SKILL_NAMES, casebook_enabled from agent_flow.workflows.perf_analyze.workflow import clear_stale_benchmark_results from . import gitops, kernel_ledger, nsys_items, reuse, roadmap_schema @@ -45,6 +41,7 @@ ) from .prompts import DEFAULT_PROMPTS, PromptBundle from .roadmap_schema import RoadmapError +from .roles import ROLES from .state import ( ROUND_STAGES, STAGE_ANALYZER, @@ -97,22 +94,25 @@ def _progress_has_entries(path: Path) -> bool: def _make_agent( name: str, system_prompt: str, + agent_config: AgentConfig, tools: list | None = None, required_tools: list[str] | None = None, - backend_kind: str = "claude-code", - model: str = CLAUDE_CODE_DEFAULT_MODEL, session_mode: str = "persistent", cwd: Path | None = None, + disabled_skills: tuple[str, ...] = (), ) -> AgentLayer: return AgentLayer( AgentLayerConfig( name=name, system_prompt=system_prompt, backend=BackendConfig( - kind=backend_kind, - model=model, + kind=agent_config.backend, + model=agent_config.model, + reasoning_effort=agent_config.reasoning_effort, + disabled_skills=disabled_skills, tools=tools, cwd=cwd, + extra_mcp_servers=agent_config.extra_mcp_servers, ), session=SessionConfig(mode=session_mode), required_tools=tuple(required_tools or ()), @@ -120,16 +120,7 @@ def _make_agent( ) -_ROLES = ( - "benchmarker", - "projector", - "analyzer", - "optimizer", - "evaluator", - "integrator", - "qa", - "reporter", -) +_ROLES = ROLES class PerfOptimizeWorkflow: @@ -344,26 +335,11 @@ def __init__( ) progress_tools = build_progress_tools(self._progress_ctx) - for role in _ROLES: - setattr( - self, - role, - _make_agent( - role, - getattr(self.prompts, role), - progress_tools[role], - required_tools=[f"append_{role}_progress"], - # Sessions are scoped to each role's unit of work: the - # judges (evaluator, qa) are stateless so every verdict - # gets fresh eyes, uninfluenced by earlier attempts' / - # rounds' conclusions; the analyzer keeps campaign-long - # memory of the roadmap it authored. - session_mode=( - "stateless" if role in ("qa", "evaluator", "integrator") else "persistent" - ), - ), - ) self._progress_tools = progress_tools + self._agent_configs: dict[str, AgentConfig] = {} + self._disabled_skills: tuple[str, ...] = () + for role in _ROLES: + setattr(self, role, None) def __enter__(self) -> "PerfOptimizeWorkflow": return self @@ -377,6 +353,37 @@ def close(self) -> None: if hasattr(layer, "__exit__"): layer.__exit__(None, None, None) + def _configure_agents(self) -> None: + task_data = self._task_data() + self._disabled_skills = () if casebook_enabled(task_data) else CASEBOOK_SKILL_NAMES + self._agent_configs = { + role: resolve_agent_config( + task_data, + role, + ) + for role in _ROLES + } + for role in _ROLES: + if getattr(self, role) is not None: + continue + setattr( + self, + role, + _make_agent( + role, + getattr(self.prompts, role), + self._agent_configs[role], + self._progress_tools[role], + required_tools=[f"append_{role}_progress"], + # Judges are stateless; other roles retain their existing + # unit-of-work session scope. + session_mode=( + "stateless" if role in ("qa", "evaluator", "integrator") else "persistent" + ), + disabled_skills=self._disabled_skills, + ), + ) + # ------------------------------------------------------------- orchestration def run(self, task: str) -> None: @@ -385,6 +392,7 @@ def run(self, task: str) -> None: state = self._init_state(task, log) if state is None: return + self._configure_agents() try: self._ensure_optimization_branch(state, log) @@ -1066,17 +1074,21 @@ def _run_opt_item(self, state: WorkflowState, entry: dict[str, Any], log) -> Non optimizer = _make_agent( f"optimizer-{item_id}", self.prompts.optimizer, + self._agent_configs["optimizer"], tools["optimizer"], required_tools=["append_optimizer_progress"], cwd=Path(item_state.item_worktree_path), + disabled_skills=self._disabled_skills, ) evaluator = _make_agent( f"evaluator-{item_id}", self.prompts.evaluator, + self._agent_configs["evaluator"], tools["evaluator"], required_tools=["append_evaluator_progress"], session_mode="stateless", cwd=Path(item_state.item_worktree_path), + disabled_skills=self._disabled_skills, ) repo = item_state.item_worktree_path live_config, accepted_config = self._state_tuning_paths(item_state) @@ -1786,6 +1798,9 @@ def _sol_enabled(self) -> bool: """ return sol_enabled(self._task_data()) + def _casebook_instruction(self, text: str) -> str: + return "" if self._disabled_skills else text + def _focus_points(self) -> list[int] | None: """``optimize.focus_concurrencies`` when set, else ``None``. @@ -2146,11 +2161,13 @@ def _run_benchmarker(self, state: WorkflowState) -> None: self._disagg_directive() + f"Workspace: {self.workspace}\n\n" f"Read `{self.task_path}` for the spec — resolve `checkpoint_path`, " f"`trtllm_repo_path`, and the `benchmark` / `optimize` blocks.\n\n" - f"Then **load the `perf-optimization-casebook` skill** (via the " - f"`Skill` tool) as read-only reference, as your system prompt " - f"directs, so your Configuration/Notes are grounded in known " - f"TRT-LLM performance precedents.\n\n" - f"Launch `trtllm-serve` with " + + self._casebook_instruction( + "Then **load the `perf-optimization-casebook` skill** (via the " + "`Skill` tool) as read-only reference, as your system prompt " + "directs, so your Configuration/Notes are grounded in known " + "TRT-LLM performance precedents.\n\n" + ) + + f"Launch `trtllm-serve` with " f"`--extra_llm_api_options {self.tuning_config_path}` (the live " f"tuning config — always passed in this workflow), poll it to " f"readiness, {load_instruction}, and tear the server down " @@ -2285,11 +2302,13 @@ def _run_reused_analyzer(self, state: WorkflowState) -> None: f"block).\n\n" + projection_context + prior_roadmap_context - + f"Then **load the `perf-optimization-casebook` skill** (via the " - f"`Skill` tool) as your system prompt directs, and tag each " - f"roadmap item's `casebook_ref` with the matching *bottleneck " - f"signal → candidate pattern* row.\n\n" - f"Two checks you still owe — both read-only, neither needs a " + + self._casebook_instruction( + "Then **load the `perf-optimization-casebook` skill** (via the " + "`Skill` tool) as your system prompt directs, and tag each " + "roadmap item's `casebook_ref` with the matching *bottleneck " + "signal → candidate pattern* row.\n\n" + ) + + f"Two checks you still owe — both read-only, neither needs a " f"GPU: verify the imported analysis actually describes **this** " f"task (same model/checkpoint, parallel mapping in " f"`{self.tuning_config_path}`, and operating point as " @@ -2434,11 +2453,13 @@ def _run_analyzer(self, state: WorkflowState) -> None: f"recover the serve + benchmark commands and operating point.\n\n" f"{round_context}\n\n" + projection_context - + f"Early on, **load the `perf-optimization-casebook` skill** (via " - f"the `Skill` tool) as read-only reference, as your system prompt " - f"directs — tag each roadmap item's `casebook_ref` with the " - f"matching *bottleneck signal → candidate pattern* row.\n\n" - f"First **verify this checkout's profiling knobs** with " + + self._casebook_instruction( + "Early on, **load the `perf-optimization-casebook` skill** (via " + "the `Skill` tool) as read-only reference, as your system prompt " + "directs — tag each roadmap item's `casebook_ref` with the " + "matching *bottleneck signal → candidate pattern* row.\n\n" + ) + + f"First **verify this checkout's profiling knobs** with " f"`grep -rn`/`rg` via `Bash` under `{self._trtllm_hint()}` as your " f"system prompt directs, then profile the current build under the " f"methods in `profile.methods`: relaunch `trtllm-serve` with " @@ -2583,11 +2604,14 @@ def _run_replan_analyzer(self, state: WorkflowState) -> None: f"`expected_gain_pct` / `evidence` of pending items the " f"measurements bound, re-order what survives, and add items the " f"failures themselves imply (a REJECT often names the real " - f"constraint) — **load the `perf-optimization-casebook` skill** " - f"(via the `Skill` tool) as your system prompt directs before " - f"authoring any, and tag each new item's `casebook_ref` with the " - f"matching *bottleneck signal → candidate pattern* row. Never " - f"rewrite `accepted` / `failed` history, " + f"constraint). " + + self._casebook_instruction( + "Before authoring any, **load the `perf-optimization-casebook` skill** " + "(via the `Skill` tool) as your system prompt directs, and tag each " + "new item's `casebook_ref` with the matching *bottleneck signal → " + "candidate pattern* row. " + ) + + f"Never rewrite `accepted` / `failed` history, " f"`baseline`, `current_best`, or existing ids; new items get " f"fresh ids continuing the sequence.\n\n" f"**If the evidence leaves nothing actionable, leave the roadmap " @@ -2688,11 +2712,13 @@ def _run_optimizer( f"Inside the Slurm job script, before any Python command or " f"`trtllm-serve` launch:\n\n" f'`export PYTHONPATH="{repo}${{PYTHONPATH:+:$PYTHONPATH}}"`\n\n' - f"Read `{self.task_path}` and the roadmap item, then **load the " - f"`perf-optimization-casebook` skill** (via the `Skill` tool) as " - f"your system prompt directs and implement **exactly this one " - f"item** following its `how_to_apply` and the matched casebook " - f"case: `approach: config` → edit `{tuning_config}`; " + f"Read `{self.task_path}` and the roadmap item. " + + self._casebook_instruction( + "Then **load the `perf-optimization-casebook` skill** (via the `Skill` " + "tool) as your system prompt directs and follow the matched casebook case. " + ) + + f"Implement **exactly this one item** following its `how_to_apply`: " + f"`approach: config` → edit `{tuning_config}`; " f"`approach: code` → edit the source under `{repo}` under " f"the git discipline in your system prompt (active-runtime " f"check first; locate code paths with shell `grep -rn`/`rg` via " diff --git a/agent-flow/tests/test_agent_runtime.py b/agent-flow/tests/test_agent_runtime.py new file mode 100644 index 000000000000..41aed9c042f3 --- /dev/null +++ b/agent-flow/tests/test_agent_runtime.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from agent_flow.agent_runtime import resolve_agent_config, validate_agents +from agent_flow.config import CLAUDE_CODE_DEFAULT_MODEL, CODEX_DEFAULT_MODEL + +ROLES = ("projector", "analyzer", "reporter") + + +def test_defaults_and_role_overrides_are_resolved_per_role(): + task = { + "agents": { + "defaults": {"backend": "claude-code"}, + "roles": { + "projector": { + "backend": "codex", + "model": "gpt-6-astra", + "reasoning_effort": "ultra", + } + }, + } + } + + projector = resolve_agent_config(task, "projector") + reporter = resolve_agent_config(task, "reporter") + + assert (projector.backend, projector.model, projector.reasoning_effort) == ( + "codex", + "gpt-6-astra", + "ultra", + ) + assert (reporter.backend, reporter.model) == ( + "claude-code", + CLAUDE_CODE_DEFAULT_MODEL, + ) + + +def test_switching_backend_without_model_uses_that_backends_default(): + config = resolve_agent_config( + { + "agents": { + "defaults": {"backend": "claude-code", "model": "claude-test"}, + "roles": {"analyzer": {"backend": "codex"}}, + } + }, + "analyzer", + ) + assert (config.backend, config.model) == ("codex", CODEX_DEFAULT_MODEL) + + +def test_validate_agents_checks_shape_fields_and_roles(): + errors = validate_agents( + { + "agents": { + "defaults": {"backend": "other", "extra_mcp_servers": []}, + "roles": {"optimizer": {"model": ""}}, + } + }, + ROLES, + ) + assert len(errors) == 4 + assert any("backend" in error for error in errors) + assert any("extra_mcp_servers" in error for error in errors) + assert any("optimizer" in error for error in errors) diff --git a/agent-flow/tests/test_backends.py b/agent-flow/tests/test_backends.py index fa237f557d01..1e51008339a1 100644 --- a/agent-flow/tests/test_backends.py +++ b/agent-flow/tests/test_backends.py @@ -27,6 +27,7 @@ from agent_flow.backends.base import ResultEvent from agent_flow.backends.claude_code import ClaudeCodeBackend, ClaudeCodeClient from agent_flow.backends.codex import CodexBackend +from agent_flow.config import BackendConfig from agent_flow.types import ( AgentTextEvent, CompactBoundaryEvent, @@ -83,6 +84,14 @@ def test_factory_rejects_unknown_backends(self): with pytest.raises(ValueError, match="Unknown backend"): create_backend("unknown") + def test_factory_passes_disabled_skills(self): + config = BackendConfig( + kind="codex", + model="gpt-6-astra", + disabled_skills=("perf-optimization-casebook",), + ) + assert create_backend(config)._disabled_skills == ("perf-optimization-casebook",) + class TestClaudeBackend: async def test_client_maps_tool_use_and_result_messages(self): @@ -742,6 +751,9 @@ async def receive_response(): class TestClaudeBackendCreateClient: + def test_reasoning_effort_override(self): + assert ClaudeCodeBackend(reasoning_effort="medium").reasoning_effort() == "medium" + async def test_framework_tools_keep_annotations_and_independent_handlers(self): from agent_flow.tools import tool @@ -797,7 +809,7 @@ async def explicit(args): with pytest.raises(jsonschema.ValidationError): jsonschema.validate(instance=rejected, schema=advertised) - async def _capture_options(self, monkeypatch, **kwargs): + async def _capture_options(self, monkeypatch, backend=None, **kwargs): # Stand-in for ``ClaudeSDKClient`` that just records the options # ``create_client`` would have launched the real SDK with. @@ -818,11 +830,20 @@ async def __aexit__(self, *args): # but we patch defensively in case that changes. monkeypatch.setattr(cc_mod, "create_sdk_mcp_server", lambda **_: object()) - backend = ClaudeCodeBackend() + backend = backend or ClaudeCodeBackend() async with backend.create_client(system_prompt="hi", model="claude-test", **kwargs): pass return captured["options"] + async def test_create_client_denies_disabled_skills(self, monkeypatch): + options = await self._capture_options( + monkeypatch, + backend=ClaudeCodeBackend( + disabled_skills=("perf-optimization-casebook",), + ), + ) + assert options.disallowed_tools == ["Skill(perf-optimization-casebook)"] + async def test_create_client_disables_bash_sandbox(self, monkeypatch): # ``danger_full_access`` analogue for Claude Code: the bash # sandbox is fully off so commands can touch anything on disk diff --git a/agent-flow/tests/test_codex_backend.py b/agent-flow/tests/test_codex_backend.py index b39603756063..28e71a7b35bc 100644 --- a/agent-flow/tests/test_codex_backend.py +++ b/agent-flow/tests/test_codex_backend.py @@ -106,6 +106,16 @@ async def test_backend_start_and_close_use_configured_runtime(monkeypatch): assert instance._transport is None +async def test_backend_disables_skills_in_codex_config(monkeypatch): + monkeypatch.setattr(codex_module, "CodexTransport", FakeTransport) + monkeypatch.setattr(codex_module, "_resolve_codex_bin", lambda: "/configured/codex") + instance = CodexBackend(disabled_skills=("perf-optimization-casebook",)) + async with instance: + assert instance._transport.config.config_overrides == ( + 'skills.config=[{name="perf-optimization-casebook", enabled=false}]', + ) + + async def test_session_options_use_developer_instructions_and_preserve_native_defaults( backend, tmp_path ): @@ -130,6 +140,10 @@ async def test_session_options_use_developer_instructions_and_preserve_native_de assert ("thread/unsubscribe", {"threadId": "thread-1"}) in transport.requests +def test_reasoning_effort_override(): + assert CodexBackend(reasoning_effort="ultra").reasoning_effort() == "ultra" + + async def test_empty_prompt_does_not_replace_runtime_instructions(backend, tmp_path): async with backend.create_client("", "gpt-5.4", cwd=tmp_path): payload = next( diff --git a/agent-flow/tests/workflows/perf_analyze/test_prompts.py b/agent-flow/tests/workflows/perf_analyze/test_prompts.py index fe51c61adcb2..41d15bed88e2 100644 --- a/agent-flow/tests/workflows/perf_analyze/test_prompts.py +++ b/agent-flow/tests/workflows/perf_analyze/test_prompts.py @@ -1102,3 +1102,9 @@ def test_profile_ranks_note_states_the_duty_for_each_shape(): # Degrades where the topology cannot deliver per-rank traces. assert "spawn-launched `trtllm-serve` cannot" in several assert "make no imbalance claim" in several + + +def test_casebook_control_adds_an_explicit_override(): + bundle = build_perf_analyze_prompts(include_casebook=False) + assert "This run intentionally disables `perf-optimization-casebook`" in bundle.benchmarker + assert "This run intentionally disables `perf-optimization-casebook`" in bundle.analyzer diff --git a/agent-flow/tests/workflows/perf_analyze/test_sol_methodology.py b/agent-flow/tests/workflows/perf_analyze/test_sol_methodology.py index a10a525d1618..94ac549d2a54 100644 --- a/agent-flow/tests/workflows/perf_analyze/test_sol_methodology.py +++ b/agent-flow/tests/workflows/perf_analyze/test_sol_methodology.py @@ -54,6 +54,18 @@ def test_resolution_prefers_the_internal_skill(monkeypatch): assert methodology.console_note() is None +def test_resolution_probes_the_projector_backend(monkeypatch): + seen = [] + + def resolve(candidates, backend_kinds=("claude-code",)): + seen.append(tuple(backend_kinds)) + return _QUALIFIED_SOL, True + + monkeypatch.setattr("agent_flow.utils.resolve_first_available_skill", resolve) + resolve_sol_methodology(backend_kind="codex") + assert seen == [("codex",)] + + def test_resolution_falls_back_to_perf_analysis(monkeypatch): """The open-source-toolkit case: internal skills stripped, the rest kept.""" _stub_probes( diff --git a/agent-flow/tests/workflows/perf_analyze/test_task_schema.py b/agent-flow/tests/workflows/perf_analyze/test_task_schema.py index 209b2f14aaeb..f877a0995fdf 100644 --- a/agent-flow/tests/workflows/perf_analyze/test_task_schema.py +++ b/agent-flow/tests/workflows/perf_analyze/test_task_schema.py @@ -9,6 +9,7 @@ from agent_flow.workflows.perf_analyze.task_schema import ( TaskSchemaError, + casebook_enabled, concurrency_points, dump_task_yaml, has_slurm_environment, @@ -58,6 +59,7 @@ def test_valid_minimal_applies_defaults(tmp_path): # spawn-launched server can give. "profile_ranks": [0], } + assert data["casebook"] == {"enabled": True} assert has_slurm_environment(data) is False # The projector is on by default, and the block is materialized so # the resolved spec states the gate the agents read. @@ -984,3 +986,51 @@ def test_profile_ranks_requires_nsys_in_methods(tmp_path): ) with pytest.raises(TaskSchemaError, match="requires 'nsys' in 'profile.methods'"): load_and_validate_task_yaml(path) + + +def test_agents_accept_only_perf_analyze_roles(tmp_path): + ckpt, repo = _paths(tmp_path) + valid = _write( + tmp_path, + { + "checkpoint_path": ckpt, + "trtllm_repo_path": repo, + "agents": {"roles": {"analyzer": {"backend": "codex"}}}, + }, + ) + assert load_and_validate_task_yaml(valid)["agents"]["roles"]["analyzer"] == {"backend": "codex"} + + invalid = _write( + tmp_path, + { + "checkpoint_path": ckpt, + "trtllm_repo_path": repo, + "agents": {"roles": {"optimizer": {"backend": "codex"}}}, + }, + ) + with pytest.raises(TaskSchemaError, match="optimizer"): + load_and_validate_task_yaml(invalid) + + +def test_casebook_can_be_disabled(tmp_path): + ckpt, repo = _paths(tmp_path) + path = _write( + tmp_path, + { + "checkpoint_path": ckpt, + "trtllm_repo_path": repo, + "casebook": {"enabled": False}, + }, + ) + assert casebook_enabled(load_and_validate_task_yaml(path)) is False + + +@pytest.mark.parametrize("block", [False, {"enabled": "no"}, {"enable": False}]) +def test_casebook_rejects_invalid_config(tmp_path, block): + ckpt, repo = _paths(tmp_path) + path = _write( + tmp_path, + {"checkpoint_path": ckpt, "trtllm_repo_path": repo, "casebook": block}, + ) + with pytest.raises(TaskSchemaError, match="casebook"): + load_and_validate_task_yaml(path) diff --git a/agent-flow/tests/workflows/perf_analyze/test_workflow.py b/agent-flow/tests/workflows/perf_analyze/test_workflow.py index aaa4e47035b2..c51de894a913 100644 --- a/agent-flow/tests/workflows/perf_analyze/test_workflow.py +++ b/agent-flow/tests/workflows/perf_analyze/test_workflow.py @@ -9,6 +9,7 @@ import yaml from agent_flow import CLAUDE_CODE_DEFAULT_MODEL +from agent_flow.workflows.perf_analyze import cli as cli_module from agent_flow.workflows.perf_analyze import progress as progress_module from agent_flow.workflows.perf_analyze import state as state_module from agent_flow.workflows.perf_analyze import workflow as workflow_module @@ -471,6 +472,8 @@ def test_clean_overwrites_stale_managed_files(tmp_path): def test_all_agents_use_claude_code_backend(tmp_path): workflow = Workflow(workspace=tmp_path / "ws") try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for layer in ( workflow.benchmarker, workflow.projector, @@ -486,6 +489,92 @@ def test_all_agents_use_claude_code_backend(tmp_path): workflow.close() +def test_projector_and_analyzer_can_use_codex(tmp_path): + workflow = Workflow(workspace=tmp_path / "ws") + workflow.task_path.write_text( + yaml.safe_dump( + { + "agents": { + "roles": { + role: { + "backend": "codex", + "model": "gpt-6-astra", + "reasoning_effort": "ultra", + } + for role in ("projector", "analyzer") + } + } + } + ), + encoding="utf-8", + ) + try: + workflow._configure_agents() + for role in ("projector", "analyzer"): + backend = getattr(workflow, role).config.backend + assert (backend.kind, backend.model, backend.reasoning_effort) == ( + "codex", + "gpt-6-astra", + "ultra", + ) + assert workflow.benchmarker.config.backend.kind == "claude-code" + assert workflow.reporter.config.backend.kind == "claude-code" + finally: + workflow.close() + + +def test_casebook_disable_reaches_every_backend(tmp_path): + workflow = Workflow(workspace=tmp_path / "ws") + workflow.task_path.write_text("casebook: {enabled: false}\n", encoding="utf-8") + try: + workflow._configure_agents() + for role in ("benchmarker", "projector", "analyzer", "reporter"): + assert getattr(workflow, role).config.backend.disabled_skills + finally: + workflow.close() + + +@pytest.mark.parametrize( + ("checkpoint_casebook", "input_casebook"), + [(True, False), (False, True)], +) +def test_cli_resume_builds_prompts_from_checkpointed_task( + tmp_path, monkeypatch, checkpoint_casebook, input_casebook +): + task = _write_task(tmp_path) + input_data = yaml.safe_load(task.read_text(encoding="utf-8")) + input_data["casebook"] = {"enabled": input_casebook} + input_data["agents"] = {"roles": {"projector": {"backend": "claude-code"}}} + task.write_text(yaml.safe_dump(input_data), encoding="utf-8") + + workspace = tmp_path / "workspace" + workspace.mkdir() + checkpoint_data = dict(input_data) + checkpoint_data["casebook"] = {"enabled": checkpoint_casebook} + checkpoint_data["agents"] = {"roles": {"projector": {"backend": "codex"}}} + (workspace / "task.yaml").write_text(yaml.safe_dump(checkpoint_data), encoding="utf-8") + (workspace / state_module.STATE_FILENAME).write_text("{}", encoding="utf-8") + + captured = {} + + def resolve_methodology(enabled, backend_kind="claude-code"): + captured["backend_kind"] = backend_kind + return SolMethodology() + + def build_prompts(**kwargs): + captured.update(kwargs) + raise RuntimeError("stop after prompt construction") + + monkeypatch.setattr(cli_module, "resolve_sol_methodology", resolve_methodology) + monkeypatch.setattr(cli_module, "build_perf_analyze_prompts", build_prompts) + + with pytest.raises(RuntimeError, match="stop after prompt construction"): + cli_module.main(["--task", str(task), "--workspace", str(workspace)]) + + assert captured["include_casebook"] is checkpoint_casebook + assert captured["backend_kind"] == "codex" + + def test_no_role_wires_an_external_mcp_server(tmp_path): """No role ships a hosted endpoint. @@ -496,6 +585,8 @@ def test_no_role_wires_an_external_mcp_server(tmp_path): """ workflow = Workflow(workspace=tmp_path / "ws") try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for layer in ( workflow.benchmarker, workflow.projector, @@ -614,6 +705,8 @@ def test_each_agent_has_its_progress_tools(tmp_path): "reporter": "append_reporter_progress", } try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for role, append_name in expected.items(): layer = getattr(workflow, role) tool_names = [t.name for t in layer.config.backend.tools] diff --git a/agent-flow/tests/workflows/perf_optimize/test_prompts.py b/agent-flow/tests/workflows/perf_optimize/test_prompts.py index e6dcb1cb512d..d615659c7a09 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_prompts.py +++ b/agent-flow/tests/workflows/perf_optimize/test_prompts.py @@ -1443,3 +1443,11 @@ def test_analyzer_categorizes_imbalance_by_the_work_not_the_collective(): assert "often not fixable in-campaign" in prompt # Bounded by the measured share, not the raw spread. assert "bound `expected_gain_pct` by `pct_of_iter`, never by the whole spread" in prompt + + +def test_casebook_control_adds_an_explicit_override(): + bundle = build_perf_optimize_prompts(include_casebook=False) + for role in ("benchmarker", "analyzer", "optimizer"): + assert "This run intentionally disables `perf-optimization-casebook`" in getattr( + bundle, role + ) diff --git a/agent-flow/tests/workflows/perf_optimize/test_task_schema.py b/agent-flow/tests/workflows/perf_optimize/test_task_schema.py index 1b0e6286cd15..5707df5f7bd7 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_task_schema.py +++ b/agent-flow/tests/workflows/perf_optimize/test_task_schema.py @@ -26,6 +26,7 @@ def test_minimal_task_gets_all_defaults(tmp_path): # perf-analyze base defaults still merge. assert data["benchmark"]["random_input_len"] == 1024 assert data["profile"]["methods"] == ["nsys", "ncu"] + assert data["casebook"] == {"enabled": True} # perf-optimize defaults merge. assert data["optimize"] == { "max_rounds": 5, @@ -639,3 +640,26 @@ def test_the_census_matches_a_fully_populated_spec(tmp_path): ) assert unknown == [] + + +def test_optimize_roles_survive_the_reused_analyze_validation_pass(tmp_path): + task = _write_task( + tmp_path, + { + "agents": { + "roles": { + "optimizer": {"backend": "codex"}, + "evaluator": {"model": "gpt-5.6-sol"}, + "integrator": {"reasoning_effort": "medium"}, + "qa": {"reasoning_effort": "medium"}, + } + } + }, + ) + resolved = task_schema.load_and_validate_task_yaml(task) + assert set(resolved["agents"]["roles"]) == { + "optimizer", + "evaluator", + "integrator", + "qa", + } diff --git a/agent-flow/tests/workflows/perf_optimize/test_workflow.py b/agent-flow/tests/workflows/perf_optimize/test_workflow.py index 9987f9a800c6..d629be23cd93 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_workflow.py +++ b/agent-flow/tests/workflows/perf_optimize/test_workflow.py @@ -13,6 +13,7 @@ from agent_flow import CLAUDE_CODE_DEFAULT_MODEL from agent_flow.workflows.perf_analyze.sol_methodology import SolMethodology +from agent_flow.workflows.perf_optimize import cli as cli_module from agent_flow.workflows.perf_optimize import ( kernel_ledger, nsys_items, @@ -567,6 +568,133 @@ def test_git_stays_local_when_slurm_is_remote(tmp_path, fake_git): assert fake_git.count("create_branch") == 1 +def test_resume_routes_agents_from_checkpointed_task_and_ignores_new_task(tmp_path, fake_git): + ws = tmp_path / "ws" + ws.mkdir() + checkpointed = { + "checkpoint_path": "/checkpoint", + "trtllm_repo_path": "/repo", + "agents": { + "defaults": { + "backend": "codex", + "model": "gpt-5.6-sol", + "reasoning_effort": "medium", + }, + "roles": { + "projector": {"model": "gpt-6-astra", "reasoning_effort": "ultra"}, + "analyzer": {"model": "gpt-6-astra", "reasoning_effort": "ultra"}, + }, + }, + } + (ws / "task.yaml").write_text(yaml.safe_dump(checkpointed), encoding="utf-8") + state_module.save_state( + ws / state_module.STATE_FILENAME, + state_module.WorkflowState( + task_path=str(ws / "task.yaml"), stage=state_module.STAGE_BENCHMARKER + ), + ) + new_task = _write_task( + tmp_path, + {"agents": {"defaults": {"backend": "claude-code", "model": "changed"}}}, + ) + workflow = Workflow(workspace=ws) + + class StopAfterRouting(RuntimeError): + pass + + workflow._ensure_optimization_branch = lambda *_: (_ for _ in ()).throw(StopAfterRouting()) + try: + with pytest.raises(StopAfterRouting): + workflow.run(str(new_task)) + assert workflow.projector.config.backend.model == "gpt-6-astra" + assert workflow.projector.config.backend.reasoning_effort == "ultra" + assert workflow.analyzer.config.backend.model == "gpt-6-astra" + assert workflow.reporter.config.backend.model == "gpt-5.6-sol" + before = workflow.optimizer.config.backend + workflow.optimizer.reset_session() + assert workflow.optimizer.config.backend == before + finally: + workflow.close() + + +def test_casebook_disable_reaches_every_backend(tmp_path): + workflow = Workflow(workspace=tmp_path / "ws") + workflow.task_path.write_text("casebook: {enabled: false}\n", encoding="utf-8") + try: + workflow._configure_agents() + for role in _AGENT_ROLES: + assert getattr(workflow, role).config.backend.disabled_skills + finally: + workflow.close() + + +@pytest.mark.parametrize( + ("checkpoint_casebook", "input_casebook"), + [(True, False), (False, True)], +) +def test_cli_resume_builds_prompts_from_checkpointed_task( + tmp_path, monkeypatch, checkpoint_casebook, input_casebook +): + task = _write_task( + tmp_path, + { + "casebook": {"enabled": input_casebook}, + "agents": {"roles": {"projector": {"backend": "claude-code"}}}, + }, + ) + input_data = yaml.safe_load(task.read_text(encoding="utf-8")) + + workspace = tmp_path / "workspace" + workspace.mkdir() + checkpoint_data = dict(input_data) + checkpoint_data["casebook"] = {"enabled": checkpoint_casebook} + checkpoint_data["agents"] = {"roles": {"projector": {"backend": "codex"}}} + (workspace / "task.yaml").write_text(yaml.safe_dump(checkpoint_data), encoding="utf-8") + (workspace / state_module.STATE_FILENAME).write_text("{}", encoding="utf-8") + + captured = {} + + def resolve_methodology(enabled, backend_kind="claude-code"): + captured["backend_kind"] = backend_kind + return SolMethodology() + + def build_prompts(**kwargs): + captured.update(kwargs) + raise RuntimeError("stop after prompt construction") + + monkeypatch.setattr(cli_module, "resolve_sol_methodology", resolve_methodology) + monkeypatch.setattr(cli_module, "build_perf_optimize_prompts", build_prompts) + + with pytest.raises(RuntimeError, match="stop after prompt construction"): + cli_module.main(["--task", str(task), "--workspace", str(workspace)]) + + assert captured["include_casebook"] is checkpoint_casebook + assert captured["backend_kind"] == "codex" + + +def test_completed_resume_constructs_no_agents_and_close_is_safe(tmp_path, monkeypatch): + ws = tmp_path / "ws" + ws.mkdir() + (ws / "task.yaml").write_text("agents: {defaults: {backend: codex}}\n", encoding="utf-8") + state_module.save_state( + ws / state_module.STATE_FILENAME, + state_module.WorkflowState( + task_path=str(ws / "task.yaml"), + stage=state_module.STAGE_REPORTER, + done=True, + ), + ) + monkeypatch.setattr( + workflow_module, + "_make_agent", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("agent constructed")), + ) + workflow = Workflow(workspace=ws) + workflow.run(str(tmp_path / "ignored.yaml")) + workflow.close() + assert all(getattr(workflow, role) is None for role in _ROLES) + + def test_resume_parked_at_projector_with_block_runs_it(tmp_path, fake_git): """A checkpoint parked at the projector resumes into it (sol set).""" task = _write_task(tmp_path, _sol_extra(tmp_path)) @@ -1457,6 +1585,8 @@ def test_each_parallel_item_uses_its_own_optimizer_session(tmp_path, fake_git): task = _write_task(tmp_path) ws = tmp_path / "ws" workflow = Workflow(workspace=ws) + workflow.task_path.write_text(task.read_text(encoding="utf-8"), encoding="utf-8") + workflow._configure_agents() trace = _stub_agents( workflow, analyzer_items=[[_item("opt-001", gain=10.0), _item("opt-002", gain=5.0)]], @@ -2721,6 +2851,8 @@ def test_clean_wipes_managed_files_and_dirs(tmp_path): def test_all_agents_use_claude_code_backend_with_scoped_sessions(tmp_path): workflow = Workflow(workspace=tmp_path / "ws") try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for role in _AGENT_ROLES: layer = getattr(workflow, role) assert layer.config.backend.kind == "claude-code", role @@ -2742,6 +2874,8 @@ def test_all_agents_use_claude_code_backend_with_scoped_sessions(tmp_path): def test_each_agent_has_its_progress_tools(tmp_path): workflow = Workflow(workspace=tmp_path / "ws") try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for role in _AGENT_ROLES: layer = getattr(workflow, role) tool_names = [t.name for t in layer.config.backend.tools] @@ -2761,6 +2895,8 @@ def test_no_role_wires_an_external_mcp_server(tmp_path): """ workflow = Workflow(workspace=tmp_path / "ws") try: + workflow.task_path.write_text("{}\n", encoding="utf-8") + workflow._configure_agents() for role in _AGENT_ROLES: assert getattr(workflow, role).config.backend.extra_mcp_servers is None, role finally: