diff --git a/mkdocs/docs/concepts/presets.md b/mkdocs/docs/concepts/presets.md index 47f19a8ad..f265be85c 100644 --- a/mkdocs/docs/concepts/presets.md +++ b/mkdocs/docs/concepts/presets.md @@ -12,7 +12,7 @@ Presets offer a toolkit that streamlines agent-based model inference optimizatio ??? info "Prerequisites" Before using presets, make sure you’ve [installed](../installation.md) the server and CLI, and created a [fleet](fleets.md). - Creating a preset requires the `claude` CLI to be installed on the machine where you create a preset. + Creating a preset requires the `claude` or `codex` CLI to be installed on the machine where you create a preset. ## Apply a configuration @@ -68,7 +68,7 @@ Create the preset dsv4-flash? [y/n]: y > optimization is done against that hardware. Point `dstack apply` to a fleet configured > correspondingly, via `fleets` inside the preset configuration or via `--fleet` in the CLI. -The command executes entirely locally and uses the locally installed `claude` CLI along with `dstack`'s bundled skills. The agent uses a `dstack` task to find the best serving configuration for the available fleet offers, then submits it as a `dstack` service for a final benchmark. +The command executes entirely locally and uses the locally installed `claude` or `codex` CLI along with `dstack`'s bundled skills (see [Agent](#agent)). The agent uses a `dstack` task to find the best serving configuration for the available fleet offers, then submits it as a `dstack` service for a final benchmark. You can stop watching with `Ctrl`+`C` at any time. The agent keeps running, and `dstack preset logs -f` follows it again. Resume an interrupted creation with `dstack preset resume`: @@ -84,24 +84,6 @@ When resuming, the configuration and constraints are read from the original sess To stop a creation and its runs, use `dstack preset stop`. -??? info "Claude configuration" - By default, preset creation uses the existing `claude` login. To use an Anthropic API key instead, set: - - ```shell - export DSTACK_AGENT_ANTHROPIC_API_KEY=... - ``` - - By default, the agent sets neither a model nor an effort level, so the `claude` CLI's built-in defaults apply. To override them, set: - - ```shell - export DSTACK_AGENT_ANTHROPIC_MODEL=claude-fable-5-1 - export DSTACK_AGENT_CLAUDE_EFFORT=max - ``` - - See the [Models overview](https://platform.claude.com/docs/en/models/overview) for the available models and their IDs. - - Supported effort levels are `low`, `medium`, `high`, `xhigh`, and `max`. - ??? info "Presets directory" The verified presets are saved locally under `~/.dstack/presets`, and `dstack preset` reads them from there. Presets aren't stored on the server. @@ -116,6 +98,50 @@ Alternatively, pass `--fleet` to `dstack apply`. > Profile settings such as `spot_policy`, `max_price`, and `backends` are ignored during preset > creation. Configure them on the fleet instead. +### Agent + +Set `agent` to choose which agent CLI creates the preset. Optionally, pin the model and the reasoning effort; otherwise the CLI's defaults apply. + +=== "Claude" + + ```yaml + agent: + provider: claude + model: claude-fable-5-1 + effort: max + ``` + + Model IDs are listed in the [Models overview](https://platform.claude.com/docs/en/models/overview). Supported effort levels are `low`, `medium`, `high`, `xhigh`, and `max`. + + To set the defaults for all presets, use `DSTACK_AGENT_PROVIDER=claude`, `DSTACK_AGENT_ANTHROPIC_MODEL`, and `DSTACK_AGENT_CLAUDE_EFFORT`. + + ??? info "Authentication" + By default, the agent uses your existing `claude` login. To use an Anthropic API key instead, set: + + ```shell + export DSTACK_AGENT_ANTHROPIC_API_KEY=... + ``` + +=== "Codex" + + ```yaml + agent: + provider: codex + model: gpt-6-astra + effort: xhigh + ``` + + Run `codex debug models` to list model IDs. Supported effort levels are `low`, `medium`, `high`, and `xhigh`. + + To set the defaults for all presets, use `DSTACK_AGENT_PROVIDER=codex`, `DSTACK_AGENT_OPENAI_MODEL`, and `DSTACK_AGENT_CODEX_EFFORT`. + + ??? info "Authentication" + By default, the agent uses your existing `codex` login and configuration, without its MCP servers. To use an OpenAI API key instead, set: + + ```shell + export DSTACK_AGENT_OPENAI_API_KEY=... + ``` + ### Model === "Base" @@ -354,7 +380,7 @@ $ dstack preset delete c83375b4 ## Protips -Under the hood, presets run an agent as a subprocess, using the local `claude` CLI. This process writes a real-time trace to `~/.dstack/presets//trace.jsonl`. The subprocess is launched with a built-in harness: how to run trials, submit runs, benchmark, verify presets, and use `dstack`. +Under the hood, presets run an agent as a subprocess, using the local `claude` or `codex` CLI. This process writes a real-time trace to `~/.dstack/presets//trace.jsonl`. The subprocess is launched with a built-in harness: how to run trials, submit runs, benchmark, verify presets, and use `dstack`. At the same time, it's recommended to create presets using your own agent — either via a CLI such as Claude Code, or inside your IDE. Your agent helps you design the preset configuration, formulate hypotheses, and — most importantly — analyze the session's traces as well as the trial results (stored under `~/.dstack/presets//trials//trial.json`), to decide what the next session can be and what instructions to give it via `prompt`. diff --git a/mkdocs/docs/reference/cli/dstack/preset.md b/mkdocs/docs/reference/cli/dstack/preset.md index 1520295cc..074115a0d 100644 --- a/mkdocs/docs/reference/cli/dstack/preset.md +++ b/mkdocs/docs/reference/cli/dstack/preset.md @@ -43,15 +43,22 @@ $ dstack preset create --help ##### Agent settings -Preset creation uses the existing `claude` login unless -`DSTACK_AGENT_ANTHROPIC_API_KEY` is set. +Presets are created with the `claude` CLI by default, or with `codex` when the +configuration's `agent` or `DSTACK_AGENT_PROVIDER` selects it. The configuration +takes precedence over the variables below. Each CLI uses its existing login unless +its API key variable is set. | Variable | Description | | --- | --- | -| `DSTACK_AGENT_ANTHROPIC_API_KEY` | Anthropic API key used by the agent. | +| `DSTACK_AGENT_PROVIDER` | The agent CLI: `claude` (default) or `codex`. | +| `DSTACK_AGENT_ANTHROPIC_API_KEY` | Anthropic API key used by the `claude` agent. | | `DSTACK_AGENT_CLAUDE_PATH` | `claude` executable name or path. Defaults to `claude` from `PATH`. | | `DSTACK_AGENT_ANTHROPIC_MODEL` | Claude model used by the agent. If unset, the `claude` CLI's built-in default is used. | | `DSTACK_AGENT_CLAUDE_EFFORT` | Claude effort level: `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. | +| `DSTACK_AGENT_OPENAI_API_KEY` | OpenAI API key used by the `codex` agent. | +| `DSTACK_AGENT_CODEX_PATH` | `codex` executable name or path. Defaults to `codex` from `PATH`. | +| `DSTACK_AGENT_OPENAI_MODEL` | Model used by the `codex` agent. If unset, the `codex` CLI's built-in default is used. | +| `DSTACK_AGENT_CODEX_EFFORT` | Codex reasoning effort: `low`, `medium`, `high`, or `xhigh`. If unset, the `codex` CLI default is used. | Agent progress is written to `agent.log` under `~/.dstack/presets//`, alongside the effective configuration (`preset.dstack.yml`), the recorded diff --git a/mkdocs/docs/reference/dstack.yml/preset.md b/mkdocs/docs/reference/dstack.yml/preset.md index c4de1a0f3..c43f16e91 100644 --- a/mkdocs/docs/reference/dstack.yml/preset.md +++ b/mkdocs/docs/reference/dstack.yml/preset.md @@ -11,6 +11,12 @@ used to create or apply a [preset](../../concepts/presets.md). type: required: true +### `agent` + +#SCHEMA# dstack._internal.core.models.configurations.PresetAgentConfig + overrides: + show_root_heading: false + ### `model` === "Base model" diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 64ced37ed..382c9f874 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -217,7 +217,12 @@ $ find ~/.dstack/logs/cli/ - `DSTACK_PROJECT`{ #DSTACK_PROJECT } – Has the same effect as `--project`. Defaults to `None`. -- `DSTACK_AGENT_ANTHROPIC_API_KEY`{ #DSTACK_AGENT_ANTHROPIC_API_KEY } – The Anthropic API key used by the preset agent. If unset, the existing `claude` login is used. +- `DSTACK_AGENT_PROVIDER`{ #DSTACK_AGENT_PROVIDER } – The agent CLI that creates presets: `claude` (default) or `codex`. The preset configuration's `agent` block overrides the `DSTACK_AGENT_*` variables. +- `DSTACK_AGENT_ANTHROPIC_API_KEY`{ #DSTACK_AGENT_ANTHROPIC_API_KEY } – The Anthropic API key used by the `claude` preset agent. If unset, the existing `claude` login is used. - `DSTACK_AGENT_CLAUDE_PATH`{ #DSTACK_AGENT_CLAUDE_PATH } – The `claude` executable name or path used by the preset agent. Defaults to `claude` from `PATH`. - `DSTACK_AGENT_ANTHROPIC_MODEL`{ #DSTACK_AGENT_ANTHROPIC_MODEL } – The Claude model used by the preset agent. If unset, the `claude` CLI's built-in default is used. - `DSTACK_AGENT_CLAUDE_EFFORT`{ #DSTACK_AGENT_CLAUDE_EFFORT } – The Claude effort level used by the preset agent. Can be `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. +- `DSTACK_AGENT_OPENAI_API_KEY`{ #DSTACK_AGENT_OPENAI_API_KEY } – The OpenAI API key used by the `codex` preset agent. If unset, the existing `codex` login is used. +- `DSTACK_AGENT_CODEX_PATH`{ #DSTACK_AGENT_CODEX_PATH } – The `codex` executable name or path used by the preset agent. Defaults to `codex` from `PATH`. +- `DSTACK_AGENT_OPENAI_MODEL`{ #DSTACK_AGENT_OPENAI_MODEL } – The model used by the `codex` preset agent. If unset, the `codex` CLI's built-in default is used. +- `DSTACK_AGENT_CODEX_EFFORT`{ #DSTACK_AGENT_CODEX_EFFORT } – The reasoning effort used by the `codex` preset agent. Can be `low`, `medium`, `high`, or `xhigh`. If unset, the `codex` CLI default is used. diff --git a/pyproject.toml b/pyproject.toml index 19ae87be8..961c9211a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "requests", "requests-unixsocket>=0.4.1", "typing-extensions>=4.0.0", + "tomli>=2.0; python_version < '3.11'", "cryptography", "packaging", "python-dateutil", diff --git a/src/dstack/_internal/cli/models/preset_agent.py b/src/dstack/_internal/cli/models/preset_agent.py index b975a135e..812ddce17 100644 --- a/src/dstack/_internal/cli/models/preset_agent.py +++ b/src/dstack/_internal/cli/models/preset_agent.py @@ -12,7 +12,7 @@ ) from dstack._internal.core.models.common import CoreModel -from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.configurations import PresetAgentProvider, ServiceConfiguration from dstack._internal.core.models.presets import PresetBenchmark @@ -98,12 +98,15 @@ class PresetSessionRun(CoreModel): workspace: PresetSessionWorkspace finalize: PresetSessionFinalize - # Only known when this CLI launched the agent; a follower leaves it as is. - claude_model: Optional[str] - # None between claude process attempts and after a detach outlives them. - agent: Optional[PresetSessionProcess] - # None until the agent's stream reveals it. - claude_session_id: Optional[str] + # A resume launches the same agent CLI the session started with. + agent_provider: PresetAgentProvider + # The model pinned by the user; None leaves the choice to the agent CLI. + agent_model: Optional[str] + # The agent process; None between attempts and after a detach outlives them. + session_process: Optional[PresetSessionProcess] + # The agent CLI's own session, which its resume continues; None until the + # agent's stream reveals it. + session_id: Optional[str] class PresetSessionState(CoreModel): @@ -126,14 +129,15 @@ class PresetSessionState(CoreModel): class PresetAgentInfo(CoreModel): - """`agent.json`: how the claude agent was launched. A debug record.""" + """`agent.json`: how the agent CLI was launched. A debug record.""" + provider: PresetAgentProvider executable: str version: Optional[str] auth_status: str - # None is the claude CLI's default. + # None is the agent CLI's default. effort: Optional[str] - # Reported by claude on its init line; None until then. + # Reported by the agent CLI once it starts; None until then. model: Optional[str] @@ -169,3 +173,30 @@ class ClaudeResultEvent(ClaudeStreamEvent): AnyClaudeStreamEvent = Annotated[ Union[ClaudeResultEvent, ClaudeStreamEvent], Field(union_mode="left_to_right") ] + + +class CodexStreamItem(CoreModel): + # "agent_message", "command_execution", "web_search", "reasoning", and others. + type: str + # The message text of an "agent_message"; the final one is the JSON report. + text: Optional[str] = None + + +class CodexStreamError(CoreModel): + message: str + + +class CodexStreamEvent(CoreModel): + """One line of `codex exec --json`. Not our format: unknown fields are dropped + and omitted fields default.""" + + # "thread.started", "turn.started", "item.started", "item.completed", + # "turn.completed", "turn.failed", "error", and whatever a newer codex adds. + type: str + # On "thread.started": the id `codex exec resume` takes. + thread_id: Optional[str] = None + item: Optional[CodexStreamItem] = None + # On "turn.failed". + error: Optional[CodexStreamError] = None + # On "error": a stream or API failure outside a turn. + message: Optional[str] = None diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py index 70d22be98..ee77c788f 100644 --- a/src/dstack/_internal/cli/services/presets/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -9,19 +9,12 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Any, AsyncIterator, Callable, Literal, Optional, Sequence, get_args +from typing import Any, AsyncIterator, Callable, Literal, Optional, Sequence import psutil -from pydantic import ValidationError - -from dstack._internal.cli.models.preset_agent import ( - AnyClaudeStreamEvent, - ClaudeResultEvent, - PresetAgentFailure, - PresetAgentInfo, - PresetAgentSuccess, - PresetSessionProcess, -) + +from dstack._internal.cli.models.preset_agent import PresetSessionProcess +from dstack._internal.cli.services.presets.agents.base import PresetAgent, PresetAgentSpec from dstack._internal.cli.services.presets.redaction import redact, redact_structure from dstack._internal.cli.services.presets.session import ( PresetSession, @@ -44,12 +37,9 @@ ) from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.common import validate_json_extra_ignore from dstack._internal.core.services.configs import ConfigManager from dstack.api import Client -_CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" -ClaudeEffort = Literal["low", "medium", "high", "xhigh", "max"] _RESUME_DELAYS_SECONDS: tuple[int, ...] = (30, 60, 120) _TERMINATE_GRACE_SECONDS = 3 _AGENT_ERROR_MAX_LENGTH = 200 @@ -87,15 +77,6 @@ ) -@dataclass(frozen=True) -class ClaudeAuth: - api_key: Optional[str] - executable: str - # None means the flag is not passed and claude uses its own default. - effort: Optional[ClaudeEffort] - model: Optional[str] - - @dataclass class PresetAgentProcessOutput: report_data: Optional[dict[str, Any]] = None @@ -104,68 +85,12 @@ class PresetAgentProcessOutput: made_progress: bool = False -def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: - try: - result = subprocess.run( - [auth.executable, "--version"], - capture_output=True, - text=True, - timeout=15, - ) - return result.stdout.strip() or None - except (OSError, subprocess.SubprocessError): - return None - - -def _get_claude_auth_status(auth: "ClaudeAuth") -> str: - if auth.api_key: - return "api-key" - try: - result = subprocess.run( - [auth.executable, "auth", "status", "--json"], - capture_output=True, - text=True, - timeout=15, - ) - return result.stdout.strip() or "unknown" - except (OSError, subprocess.SubprocessError): - return "unknown" - - -def get_agent_info(auth: ClaudeAuth) -> PresetAgentInfo: - return PresetAgentInfo( - executable=auth.executable, - version=_get_claude_version(auth), - auth_status=_get_claude_auth_status(auth), - effort=auth.effort, - model=None, - ) - - -def get_claude_auth() -> ClaudeAuth: - api_key = os.getenv("DSTACK_AGENT_ANTHROPIC_API_KEY") or None - configured_path = os.getenv("DSTACK_AGENT_CLAUDE_PATH") or "claude" - executable = shutil.which(configured_path) - if executable is None: - raise CLIError(f"Claude executable not found: {configured_path}") - effort = os.getenv("DSTACK_AGENT_CLAUDE_EFFORT") or None - if effort is not None and effort not in get_args(ClaudeEffort): - raise CLIError( - f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(get_args(ClaudeEffort))}" - ) - return ClaudeAuth( - api_key=api_key, - executable=executable, - effort=effort, - model=os.getenv("DSTACK_AGENT_ANTHROPIC_MODEL") or None, - ) - - def build_preset_agent_env( *, api: Client, preset_env: dict[str, str], - auth: ClaudeAuth, + agent: PresetAgent, + spec: PresetAgentSpec, workspace: PresetAgentWorkspace, token: str, ) -> dict[str, str]: @@ -190,18 +115,7 @@ def build_preset_agent_env( env[PROGRESS_ENV] = str(workspace.progress_path) for name in ["TMPDIR", "TEMP", "TMP"]: env[name] = str(workspace.temp_path) - # Sandbox the agent's Claude config under the workspace home when we pass our - # own API key; under subscription auth keep the real HOME so it reuses the - # user's existing `claude` login. - if auth.api_key is not None: - env["ANTHROPIC_API_KEY"] = auth.api_key - env["HOME"] = str(workspace.dstack_home) - if IS_WINDOWS: - env["USERPROFILE"] = str(workspace.dstack_home) - else: - env["HOME"] = str(Path.home()) - if IS_WINDOWS: - env["USERPROFILE"] = str(Path.home()) + agent.build_env(spec, workspace, env) return env @@ -210,7 +124,8 @@ async def run_preset_agent( prompt: str, env: dict[str, str], workspace: PresetAgentWorkspace, - auth: ClaudeAuth, + agent: PresetAgent, + spec: PresetAgentSpec, redacted_values: Sequence[str], session: PresetSession, initial_resume_session_id: Optional[str] = None, @@ -227,19 +142,20 @@ async def run_preset_agent( retry_delays = list(_RESUME_DELAYS_SECONDS) while True: command = _prepare_subprocess_command( - _build_claude_command(auth=auth, resume_session_id=resume_session_id) + agent.build_command(spec, workspace, resume_session_id) ) - output, returncode = await _run_claude_process( + output, returncode = await _run_agent_process( command=command, prompt=attempt_prompt, env=env, workspace=workspace, + agent=agent, redacted_values=redacted_values, session=session, offset_store=offset_store, ) if output.report_data is None and returncode != 0: - output.error = output.error or f"Claude exited with return code {returncode}" + output.error = output.error or f"The agent exited with return code {returncode}" error = output.error # Retry any process death without a submitted report; a terminal # failure report from the agent returns immediately. @@ -272,12 +188,13 @@ async def run_preset_agent( await asyncio.sleep(delay) -async def _run_claude_process( +async def _run_agent_process( *, command: list[str], prompt: str, env: dict[str, str], workspace: PresetAgentWorkspace, + agent: PresetAgent, redacted_values: Sequence[str], session: PresetSession, offset_store: OffsetStore, @@ -304,7 +221,7 @@ async def _run_claude_process( # CreateProcess on Windows (WinError 87). close_fds=True, ) - session.record_agent( + session.record_session_process( PresetSessionProcess(pid=proc.pid, started_at=process_started_at(proc.pid)) ) assert proc.stdin is not None @@ -320,6 +237,7 @@ def agent_alive() -> bool: _collect_agent_output( workspace=workspace, session=session, + agent=agent, redacted_values=redacted_values, is_alive=agent_alive, offset_store=offset_store, @@ -347,45 +265,14 @@ def agent_alive() -> bool: return output, returncode -def _build_claude_command(*, auth: ClaudeAuth, resume_session_id: Optional[str]) -> list[str]: - command = [ - auth.executable, - "-p", - "--output-format", - "stream-json", - "--verbose", - "--tools", - _CLAUDE_TOOLS, - "--allowedTools", - _CLAUDE_TOOLS, - "--disallowedTools", - "Task,NotebookEdit", - "--permission-mode", - "bypassPermissions", - "--json-schema", - json.dumps(_get_report_json_schema()), - ] - if auth.api_key is None: - command[2:2] = ["--setting-sources", "project,local"] - else: - command[2:2] = ["--bare"] - if auth.effort is not None: - command[2:2] = ["--effort", auth.effort] - if auth.model is not None: - command[2:2] = ["--model", auth.model] - if resume_session_id is not None: - command += ["--resume", resume_session_id] - return command - - def _prepare_subprocess_command(command: list[str]) -> list[str]: - """On Windows a `.bat`/`.cmd` Claude launcher can't be exec'd directly; wrap + """On Windows a `.bat`/`.cmd` agent launcher can't be exec'd directly; wrap it in `cmd.exe /c`. Every other case is returned unchanged.""" if not IS_WINDOWS or Path(command[0]).suffix.lower() not in {".bat", ".cmd"}: return command comspec = os.getenv("COMSPEC") or shutil.which("cmd.exe") if comspec is None: - raise CLIError("Cannot run the Claude batch launcher because cmd.exe was not found") + raise CLIError("Cannot run the agent batch launcher because cmd.exe was not found") return [comspec, "/d", "/s", "/c", subprocess.list2cmdline(command)] @@ -479,6 +366,7 @@ async def _collect_agent_output( *, workspace: PresetAgentWorkspace, session: PresetSession, + agent: PresetAgent, redacted_values: Sequence[str], is_alive: Callable[[], bool], offset_store: OffsetStore, @@ -493,6 +381,7 @@ async def _collect_agent_output( is_alive=is_alive, ), stream_name="stdout", + agent=agent, redacted_values=redacted_values, session=session, ), @@ -504,6 +393,7 @@ async def _collect_agent_output( is_alive=is_alive, ), stream_name="stderr", + agent=agent, redacted_values=redacted_values, session=session, ), @@ -514,6 +404,7 @@ async def _collect_agent_output( async def attach_preset_agent( *, workspace: PresetAgentWorkspace, + agent: PresetAgent, redacted_values: Sequence[str], session: PresetSession, ) -> PresetAgentProcessOutput: @@ -528,11 +419,16 @@ async def attach_preset_agent( state = session.read_state() def agent_alive() -> bool: - return state is not None and state.run is not None and process_alive(state.run.agent) + return ( + state is not None + and state.run is not None + and process_alive(state.run.session_process) + ) return await _collect_agent_output( workspace=workspace, session=session, + agent=agent, redacted_values=redacted_values, is_alive=agent_alive, offset_store=offset_store, @@ -543,6 +439,7 @@ async def _read_process_stream( *, stream: "FileLineReader", stream_name: Literal["stdout", "stderr"], + agent: PresetAgent, redacted_values: Sequence[str], session: PresetSession, ) -> PresetAgentProcessOutput: @@ -563,33 +460,20 @@ async def _read_process_stream( ) if not parse_result: continue - try: - event = validate_json_extra_ignore(AnyClaudeStreamEvent, text) - except ValidationError: + update = agent.parse_line(text) + if update is None: continue - if output.session_id is None and event.session_id: - output.session_id = event.session_id - session.record_claude_session_id(event.session_id) - if event.model: - session.record_agent_model(event.model) - if event.type == "assistant": + if output.session_id is None and update.session_id: + output.session_id = update.session_id + session.record_session_id(update.session_id) + if update.model: + session.record_agent_model(update.model) + if update.made_progress: output.made_progress = True - if not isinstance(event, ClaudeResultEvent): - continue - if event.is_error: - output.error = redact(str(event.result or "Claude failed"), redacted_values) - if event.structured_output is not None: - output.report_data = event.structured_output - continue - # An agent may print the report as its final text instead of submitting - # it through `StructuredOutput`. - if isinstance(event.result, str): - try: - parsed = json.loads(event.result) - except json.JSONDecodeError: - continue - if isinstance(parsed, dict): - output.report_data = parsed + if update.error is not None: + output.error = redact(update.error, redacted_values) + if update.report_data is not None: + output.report_data = update.report_data async def _terminate_process(proc: asyncio.subprocess.Process) -> None: @@ -650,22 +534,3 @@ def _terminate_windows_process_tree(pid: int) -> None: with suppress(psutil.NoSuchProcess): process.kill() psutil.wait_procs(alive, timeout=3) - - -def _get_report_json_schema() -> dict[str, Any]: - """The one shape the API can enforce: a single object, no union, only - `success` required. `AnyPresetAgentResult` enforces the rest at parse.""" - success = PresetAgentSuccess.model_json_schema() - failure = PresetAgentFailure.model_json_schema() - return { - "type": "object", - "properties": { - **success["properties"], - **failure["properties"], - # Each outcome fixes its own value; only the merged shape offers both. - "success": {"type": "boolean"}, - }, - "required": ["success"], - "additionalProperties": False, - "$defs": {**success.get("$defs", {}), **failure.get("$defs", {})}, - } diff --git a/src/dstack/_internal/cli/services/presets/agents/__init__.py b/src/dstack/_internal/cli/services/presets/agents/__init__.py new file mode 100644 index 000000000..e8e6f9bef --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/agents/__init__.py @@ -0,0 +1,45 @@ +"""The agent CLIs that can run a preset creation, selected by the preset's `agent.provider` +or by `DSTACK_AGENT_PROVIDER`.""" + +import os +from typing import Callable, Optional + +from dstack._internal.cli.services.presets.agents.base import ( + PresetAgent, + PresetAgentSpec, + PresetAgentStreamUpdate, +) +from dstack._internal.cli.services.presets.agents.claude import ClaudePresetAgent +from dstack._internal.cli.services.presets.agents.codex import CodexPresetAgent +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.configurations import PresetAgentProvider + +AGENT_PROVIDER_ENV = "DSTACK_AGENT_PROVIDER" +DEFAULT_AGENT_PROVIDER: PresetAgentProvider = "claude" +_AGENTS: dict[str, Callable[[], PresetAgent]] = { + "claude": ClaudePresetAgent, + "codex": CodexPresetAgent, +} + + +def get_preset_agent(provider: Optional[str] = None) -> PresetAgent: + """The agent for `provider`, or for `DSTACK_AGENT_PROVIDER` (default claude) when None. + One instance per run: an agent may keep what it learns while running.""" + if provider is None: + provider = os.getenv(AGENT_PROVIDER_ENV) or DEFAULT_AGENT_PROVIDER + factory = _AGENTS.get(provider) + if factory is None: + raise CLIError( + f"Unknown preset agent provider {provider!r}; supported: {', '.join(_AGENTS)}" + ) + return factory() + + +__all__ = [ + "AGENT_PROVIDER_ENV", + "DEFAULT_AGENT_PROVIDER", + "PresetAgent", + "PresetAgentSpec", + "PresetAgentStreamUpdate", + "get_preset_agent", +] diff --git a/src/dstack/_internal/cli/services/presets/agents/base.py b/src/dstack/_internal/cli/services/presets/agents/base.py new file mode 100644 index 000000000..b9d89bf14 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/agents/base.py @@ -0,0 +1,118 @@ +"""The contract between the shared creation loop and one agent CLI.""" + +import json +import os +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional, Protocol, Sequence + +from dstack._internal.cli.models.preset_agent import PresetAgentInfo +from dstack._internal.cli.services.presets.workspace import PresetAgentWorkspace +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.configurations import PresetAgentConfig, PresetAgentProvider + + +@dataclass(frozen=True) +class PresetAgentSpec: + """How the agent CLI is run: the preset's `agent` block over the + `DSTACK_AGENT_*` variables over the CLI's own defaults.""" + + executable: str + api_key: Optional[str] + # None means the flag is not passed and the CLI uses its own default. + model: Optional[str] + effort: Optional[str] + + +@dataclass +class PresetAgentStreamUpdate: + """What one line of the agent's stdout contributes to the run.""" + + session_id: Optional[str] = None + model: Optional[str] = None + made_progress: bool = False + error: Optional[str] = None + report_data: Optional[dict[str, Any]] = None + + +class PresetAgent(Protocol): + provider: PresetAgentProvider + # Where the CLI discovers project skills, relative to the agent's working directory. + skills_dir: Path + + def get_spec(self, config: Optional[PresetAgentConfig]) -> PresetAgentSpec: ... + + def get_info(self, spec: PresetAgentSpec) -> PresetAgentInfo: ... + + def build_env( + self, spec: PresetAgentSpec, workspace: PresetAgentWorkspace, env: dict[str, str] + ) -> None: + """Adds the CLI's own variables (auth, home) to the shared agent environment + and puts whatever the run needs into the workspace.""" + ... + + def build_command( + self, + spec: PresetAgentSpec, + workspace: PresetAgentWorkspace, + resume_session_id: Optional[str], + ) -> list[str]: ... + + def parse_line(self, text: str) -> Optional[PresetAgentStreamUpdate]: + """None for a line that is not an event of this CLI.""" + ... + + +def resolve_executable(path_env: str, default: str, display_name: str) -> str: + configured = os.getenv(path_env) or default + executable = shutil.which(configured) + if executable is None: + raise CLIError(f"{display_name} executable not found: {configured}") + return executable + + +def choose_model(config: Optional[PresetAgentConfig], env_name: str) -> Optional[str]: + if config is not None and config.model: + return config.model + return os.getenv(env_name) or None + + +def choose_effort( + config: Optional[PresetAgentConfig], env_name: str, allowed: Sequence[str] +) -> Optional[str]: + """The configuration validates `agent.effort` itself; only the variable is checked here.""" + if config is not None and config.effort is not None: + return config.effort + effort = os.getenv(env_name) or None + if effort is not None and effort not in allowed: + raise CLIError(f"{env_name} must be one of: {', '.join(allowed)}") + return effort + + +def set_home(env: dict[str, str], home: Path) -> None: + env["HOME"] = str(home) + if IS_WINDOWS: + env["USERPROFILE"] = str(home) + + +def probe_cli( + args: Sequence[str], *, env: Optional[dict[str, str]] = None, with_stderr: bool = False +) -> Optional[str]: + """Output of a short CLI probe such as `--version`; None when it cannot run.""" + try: + result = subprocess.run(list(args), capture_output=True, text=True, timeout=15, env=env) + except (OSError, subprocess.SubprocessError): + return None + text = result.stdout + (result.stderr if with_stderr else "") + return text.strip() or None + + +def parse_json_object(text: str) -> Optional[dict[str, Any]]: + try: + parsed = json.loads(text) + except json.JSONDecodeError: + return None + return parsed if isinstance(parsed, dict) else None diff --git a/src/dstack/_internal/cli/services/presets/agents/claude.py b/src/dstack/_internal/cli/services/presets/agents/claude.py new file mode 100644 index 000000000..c99a6d518 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/agents/claude.py @@ -0,0 +1,126 @@ +"""Claude Code as the preset agent.""" + +import json +import os +from pathlib import Path +from typing import Literal, Optional, get_args + +from pydantic import ValidationError + +from dstack._internal.cli.models.preset_agent import ( + AnyClaudeStreamEvent, + ClaudeResultEvent, + PresetAgentInfo, +) +from dstack._internal.cli.services.presets.agents.base import ( + PresetAgentSpec, + PresetAgentStreamUpdate, + choose_effort, + choose_model, + parse_json_object, + probe_cli, + resolve_executable, + set_home, +) +from dstack._internal.cli.services.presets.report_schema import get_report_json_schema +from dstack._internal.cli.services.presets.workspace import PresetAgentWorkspace +from dstack._internal.core.models.common import validate_json_extra_ignore +from dstack._internal.core.models.configurations import PresetAgentConfig, PresetAgentProvider + +_CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" +ClaudeEffort = Literal["low", "medium", "high", "xhigh", "max"] + + +class ClaudePresetAgent: + provider: PresetAgentProvider = "claude" + skills_dir = Path(".claude") / "skills" + + def get_spec(self, config: Optional[PresetAgentConfig]) -> PresetAgentSpec: + return PresetAgentSpec( + executable=resolve_executable("DSTACK_AGENT_CLAUDE_PATH", "claude", "Claude"), + api_key=os.getenv("DSTACK_AGENT_ANTHROPIC_API_KEY") or None, + model=choose_model(config, "DSTACK_AGENT_ANTHROPIC_MODEL"), + effort=choose_effort(config, "DSTACK_AGENT_CLAUDE_EFFORT", get_args(ClaudeEffort)), + ) + + def get_info(self, spec: PresetAgentSpec) -> PresetAgentInfo: + if spec.api_key: + auth_status = "api-key" + else: + auth_status = probe_cli([spec.executable, "auth", "status", "--json"]) or "unknown" + return PresetAgentInfo( + provider=self.provider, + executable=spec.executable, + version=probe_cli([spec.executable, "--version"]), + auth_status=auth_status, + effort=spec.effort, + model=None, + ) + + def build_env( + self, spec: PresetAgentSpec, workspace: PresetAgentWorkspace, env: dict[str, str] + ) -> None: + # With our own API key, a blank home keeps claude out of the user's `~/.claude`; + # otherwise the user's home, so claude reuses their login. + if spec.api_key is not None: + env["ANTHROPIC_API_KEY"] = spec.api_key + set_home(env, workspace.dstack_home) + else: + set_home(env, Path.home()) + + def build_command( + self, + spec: PresetAgentSpec, + workspace: PresetAgentWorkspace, + resume_session_id: Optional[str], + ) -> list[str]: + command = [ + spec.executable, + "-p", + "--output-format", + "stream-json", + "--verbose", + "--tools", + _CLAUDE_TOOLS, + "--allowedTools", + _CLAUDE_TOOLS, + "--disallowedTools", + "Task,NotebookEdit", + "--permission-mode", + "bypassPermissions", + "--json-schema", + json.dumps(get_report_json_schema()), + ] + if spec.api_key is None: + command[2:2] = ["--setting-sources", "project,local"] + else: + command[2:2] = ["--bare"] + if spec.effort is not None: + command[2:2] = ["--effort", spec.effort] + if spec.model is not None: + command[2:2] = ["--model", spec.model] + if resume_session_id is not None: + command += ["--resume", resume_session_id] + return command + + def parse_line(self, text: str) -> Optional[PresetAgentStreamUpdate]: + try: + event = validate_json_extra_ignore(AnyClaudeStreamEvent, text) + except ValidationError: + return None + update = PresetAgentStreamUpdate( + session_id=event.session_id or None, model=event.model or None + ) + if event.type == "assistant": + update.made_progress = True + if not isinstance(event, ClaudeResultEvent): + return update + if event.is_error: + update.error = str(event.result or "Claude failed") + if event.structured_output is not None: + update.report_data = event.structured_output + elif isinstance(event.result, str): + # An agent may print the report as its final text instead of submitting + # it through `StructuredOutput`. + update.report_data = parse_json_object(event.result) + return update diff --git a/src/dstack/_internal/cli/services/presets/agents/codex.py b/src/dstack/_internal/cli/services/presets/agents/codex.py new file mode 100644 index 000000000..ae176a00c --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/agents/codex.py @@ -0,0 +1,209 @@ +"""OpenAI Codex CLI as the preset agent.""" + +import json +import os +import re +from pathlib import Path +from typing import Literal, Optional, get_args + +from pydantic import ValidationError + +try: + import tomllib +except ModuleNotFoundError: # Python 3.10 + import tomli as tomllib # type: ignore[import-not-found, no-redef] + +from dstack._internal.cli.models.preset_agent import CodexStreamEvent, PresetAgentInfo +from dstack._internal.cli.services.presets.agents.base import ( + PresetAgentSpec, + PresetAgentStreamUpdate, + choose_effort, + choose_model, + parse_json_object, + probe_cli, + resolve_executable, + set_home, +) +from dstack._internal.cli.services.presets.report_schema import get_strict_report_json_schema +from dstack._internal.cli.services.presets.workspace import PresetAgentWorkspace +from dstack._internal.core.models.common import validate_json_extra_ignore +from dstack._internal.core.models.configurations import PresetAgentConfig, PresetAgentProvider + +CodexEffort = Literal["low", "medium", "high", "xhigh"] +_REPORT_SCHEMA_FILENAME = ".report-schema.json" +_LAST_MESSAGE_FILENAME = ".agent-last-message.json" +# A `-c` dotted path cannot address any other kind of table key. +_PLAIN_TOML_KEY = re.compile(r"[A-Za-z0-9_-]+") + + +class CodexPresetAgent: + provider: PresetAgentProvider = "codex" + skills_dir = Path(".codex") / "skills" + + def __init__(self) -> None: + # Set by `build_env`; a follower that never launched codex leaves it None. + self._codex_home: Optional[Path] = None + self._thread_id: Optional[str] = None + self._rollout_path: Optional[Path] = None + # The stream never names the model; the session rollout does once the first + # turn has started. Read on the first item and again at the end of the turn. + self._model_reads_left = 0 + self._last_message: Optional[str] = None + + def get_spec(self, config: Optional[PresetAgentConfig]) -> PresetAgentSpec: + return PresetAgentSpec( + executable=resolve_executable("DSTACK_AGENT_CODEX_PATH", "codex", "Codex"), + api_key=os.getenv("DSTACK_AGENT_OPENAI_API_KEY") or None, + model=choose_model(config, "DSTACK_AGENT_OPENAI_MODEL"), + effort=choose_effort(config, "DSTACK_AGENT_CODEX_EFFORT", get_args(CodexEffort)), + ) + + def get_info(self, spec: PresetAgentSpec) -> PresetAgentInfo: + if spec.api_key: + auth_status = "api-key" + else: + auth_status = ( + probe_cli([spec.executable, "login", "status"], with_stderr=True) or "unknown" + ) + return PresetAgentInfo( + provider=self.provider, + executable=spec.executable, + version=probe_cli([spec.executable, "--version"]), + auth_status=auth_status, + effort=spec.effort, + model=None, + ) + + def build_env( + self, spec: PresetAgentSpec, workspace: PresetAgentWorkspace, env: dict[str, str] + ) -> None: + # Like claude: with our own API key, codex runs from the workspace home with a + # private CODEX_HOME; otherwise exactly as the user runs it. + if spec.api_key is not None: + env["CODEX_API_KEY"] = spec.api_key + set_home(env, workspace.dstack_home) + codex_home = workspace.dstack_home / ".codex" + codex_home.mkdir(mode=0o700, exist_ok=True) + else: + set_home(env, Path.home()) + codex_home = get_user_codex_home() + env["CODEX_HOME"] = str(codex_home) + self._codex_home = codex_home + (workspace.path / _REPORT_SCHEMA_FILENAME).write_text( + json.dumps(get_strict_report_json_schema()), encoding="utf-8" + ) + + def build_command( + self, + spec: PresetAgentSpec, + workspace: PresetAgentWorkspace, + resume_session_id: Optional[str], + ) -> list[str]: + command = [spec.executable, "exec"] + if resume_session_id is not None: + command += ["resume", resume_session_id] + command += [ + "--json", + "--skip-git-repo-check", + "--dangerously-bypass-approvals-and-sandbox", + "--disable", + "multi_agent", + "--disable", + "apps", + "-c", + 'web_search="live"', + "--output-schema", + str(workspace.path / _REPORT_SCHEMA_FILENAME), + "--output-last-message", + str(workspace.path / _LAST_MESSAGE_FILENAME), + ] + if resume_session_id is None: + # `exec resume` has no `-C`; the thread remembers its working directory. + command += ["-C", str(workspace.path)] + if spec.api_key is None: + # The user's codex as they use it (auth, model provider, hooks), minus + # their MCP servers: the agent gets no tools beyond the shell. + command += _mcp_server_overrides(get_user_codex_home() / "config.toml") + if spec.model is not None: + command += ["-m", spec.model] + if spec.effort is not None: + command += ["-c", f'model_reasoning_effort="{spec.effort}"'] + command.append("-") + return command + + def parse_line(self, text: str) -> Optional[PresetAgentStreamUpdate]: + try: + event = validate_json_extra_ignore(CodexStreamEvent, text) + except ValidationError: + return None + update = PresetAgentStreamUpdate() + if event.type == "thread.started": + update.session_id = event.thread_id + self._thread_id = event.thread_id + self._model_reads_left = 2 + elif event.type == "item.completed" and event.item is not None: + update.made_progress = True + if event.item.type == "agent_message": + self._last_message = event.item.text + elif event.type == "turn.completed": + # Only the message that ends the turn is the schema-checked report; an + # earlier message that happens to be JSON is not. + if self._last_message: + update.report_data = parse_json_object(self._last_message) + elif event.type == "turn.failed" and event.error is not None: + update.error = event.error.message + elif event.type == "error" and event.message: + update.error = event.message + if self._model_reads_left and ( + event.type.startswith("item.") or event.type == "turn.completed" + ): + self._model_reads_left -= 1 + update.model = self._read_rollout_model() + if update.model is not None: + self._model_reads_left = 0 + return update + + def _read_rollout_model(self) -> Optional[str]: + """The `model` of the rollout's `turn_context` record, if it has been written.""" + if self._codex_home is None or self._thread_id is None: + return None + if self._rollout_path is None: + pattern = f"sessions/*/*/*/rollout-*-{self._thread_id}.jsonl" + self._rollout_path = next(iter(sorted(self._codex_home.glob(pattern))), None) + if self._rollout_path is None: + return None + try: + with self._rollout_path.open(encoding="utf-8") as rollout: + for line in rollout: + record = parse_json_object(line) + if record is None or record.get("type") != "turn_context": + continue + payload = record.get("payload") + if isinstance(payload, dict) and isinstance(payload.get("model"), str): + return payload["model"] + except OSError: + return None + return None + + +def get_user_codex_home() -> Path: + return Path(os.getenv("CODEX_HOME") or Path.home() / ".codex") + + +def _mcp_server_overrides(config_path: Path) -> list[str]: + """`-c` overrides that disable every MCP server in the user's config. A `-c` + dotted path replaces the whole server table, so each gets a stub table that + still passes codex's transport validation.""" + try: + config = tomllib.loads(config_path.read_text(encoding="utf-8")) + except (OSError, tomllib.TOMLDecodeError): + return [] + servers = config.get("mcp_servers") + if not isinstance(servers, dict): + return [] + return [ + arg + for name in servers + if _PLAIN_TOML_KEY.fullmatch(name) + for arg in ("-c", f'mcp_servers.{name}={{command="disabled", enabled=false}}') + ] diff --git a/src/dstack/_internal/cli/services/presets/create.py b/src/dstack/_internal/cli/services/presets/create.py index e2617098b..da097d86e 100644 --- a/src/dstack/_internal/cli/services/presets/create.py +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -18,21 +18,25 @@ PresetAgentFailure, PresetAgentSuccess, PresetSessionFinalize, + PresetSessionRun, PresetSessionState, PresetSessionStatus, PresetSessionWorkspace, ) from dstack._internal.cli.models.presets import VerifiedPreset from dstack._internal.cli.services.presets.agent import ( - ClaudeAuth, PresetAgentProcessOutput, attach_preset_agent, build_preset_agent_env, - get_agent_info, - get_claude_auth, run_preset_agent, terminate_agent_process, ) +from dstack._internal.cli.services.presets.agents import ( + AGENT_PROVIDER_ENV, + PresetAgent, + PresetAgentSpec, + get_preset_agent, +) from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt from dstack._internal.cli.services.presets.redaction import ( contains_redacted_value, @@ -319,7 +323,7 @@ def stop_preset_session(api: Client, preset_id: str) -> None: # Stop wins, like `dstack stop`: record the intent first so a live owner's # retry loop exits instead of resurrecting the agent, then terminate. _finish_agent_session(session, "interrupted") - terminate_agent_process(state.run.agent if state.run else None) + terminate_agent_process(state.run.session_process if state.run else None) _stop_active_session_runs(api, session) _suspend_agent_session(session) @@ -453,7 +457,9 @@ class _CreationSetup: """Per-mode inputs to the shared creation path, built by one of `_fresh_setup`, `_resume_setup`, or `_attach_setup`.""" - auth: Optional[ClaudeAuth] + agent: PresetAgent + # None when this CLI only follows an agent it did not start. + spec: Optional[PresetAgentSpec] workspace: PresetAgentWorkspace workspace_record: PresetSessionWorkspace build_name: str @@ -477,8 +483,9 @@ def _fresh_setup( allowed_fleets = _get_allowed_fleets(api, configuration) if not allowed_fleets: raise CLIError(_NO_FLEETS_ERROR) - auth = get_claude_auth() - workspace, workspace_record = create_agent_workspace(session) + agent = get_preset_agent(configuration.agent.provider if configuration.agent else None) + spec = agent.get_spec(configuration.agent) + workspace, workspace_record = create_agent_workspace(session, agent.skills_dir) previous_ids = tuple(session.preset_id for session in previous) if previous_ids: install_previous_records(workspace, previous) @@ -486,7 +493,8 @@ def _fresh_setup( configuration.name, configuration.model.api_model_name, session.preset_id ) return _CreationSetup( - auth=auth, + agent=agent, + spec=spec, workspace=workspace, workspace_record=workspace_record, build_name=build_name, @@ -498,7 +506,15 @@ def _fresh_setup( ) +def _read_session_run(session: PresetSession) -> PresetSessionRun: + state = session.read_state() + if state is None or state.run is None: + raise CLIError(f"Preset {session.preset_id} session state is unreadable") + return state.run + + def _resume_setup( + configuration: PresetConfiguration, session: PresetSession, build_name: Optional[str], user_prompt: Optional[str], @@ -510,19 +526,34 @@ def _resume_setup( "The configuration prompt is ignored when resuming: the preset keeps its original prompt" ) user_prompt = pinned_prompt - auth = get_claude_auth() workspace, workspace_record = attach_agent_workspace(session) - state = session.read_state() - if state is None or state.run is None: - raise CLIError(f"Preset {session.preset_id} session state is unreadable") - if state.run.claude_model: - auth = dataclasses.replace(auth, model=state.run.claude_model) - initial_resume_session_id = state.run.claude_session_id - previous_ids = tuple(state.previous) + run = _read_session_run(session) + # The session keeps the agent it started with; neither the environment nor an + # edited `agent` block can switch it. + agent = get_preset_agent(run.agent_provider) + agent_config = configuration.agent + if agent_config is not None and agent_config.provider != agent.provider: + warn( + f"agent.provider={agent_config.provider} is ignored when resuming:" + f" the preset keeps its original agent ({agent.provider})" + ) + agent_config = None + requested_provider = os.getenv(AGENT_PROVIDER_ENV) + if requested_provider and requested_provider != agent.provider: + warn( + f"{AGENT_PROVIDER_ENV}={requested_provider} is ignored when resuming:" + f" the preset keeps its original agent ({agent.provider})" + ) + spec = agent.get_spec(agent_config) + if run.agent_model: + spec = dataclasses.replace(spec, model=run.agent_model) + initial_resume_session_id = run.session_id + previous_ids = _read_previous_ids(session) if previous_ids: install_previous_records(workspace, _load_pinned_previous_sessions(previous_ids)) return _CreationSetup( - auth=auth, + agent=agent, + spec=spec, workspace=workspace, workspace_record=workspace_record, build_name=build_name or _load_build_name(workspace), @@ -540,7 +571,8 @@ def _attach_setup( ) -> _CreationSetup: workspace, workspace_record = attach_agent_workspace(session) return _CreationSetup( - auth=None, + agent=get_preset_agent(_read_session_run(session).agent_provider), + spec=None, workspace=workspace, workspace_record=workspace_record, build_name=build_name or _load_build_name(workspace), @@ -570,7 +602,7 @@ async def _create_preset( if mode == "attach": setup = _attach_setup(session, build_name) elif mode == "resume": - setup = _resume_setup(session, build_name, user_prompt) + setup = _resume_setup(configuration, session, build_name, user_prompt) else: setup = _fresh_setup( api, configuration, session, build_name, allowed_fleets, user_prompt, previous @@ -580,7 +612,8 @@ async def _create_preset( session.begin_run( workspace=setup.workspace_record, finalize=PresetSessionFinalize(project=api.project, keep_service=keep_service), - claude_model=setup.auth.model if setup.auth is not None else None, + agent_provider=setup.agent.provider, + agent_model=setup.spec.model if setup.spec is not None else None, ) preset_env = configuration.env.as_dict() @@ -590,7 +623,7 @@ async def _create_preset( redacted_values = get_redacted_values( [ token, - (setup.auth.api_key if setup.auth is not None else None) or "", + (setup.spec.api_key if setup.spec is not None else None) or "", # Passthrough values are resolved from the caller's environment and # are secrets; literal values are the user's own configuration text. # The passthrough keys come from the source configuration, since @@ -610,11 +643,12 @@ async def _create_preset( creation_succeeded = False interrupted = False cleanup_error: Optional[str] = None - if setup.auth is not None: + if setup.spec is not None: env = build_preset_agent_env( api=api, preset_env=preset_env, - auth=setup.auth, + agent=setup.agent, + spec=setup.spec, workspace=setup.workspace, token=token, ) @@ -623,6 +657,7 @@ async def _create_preset( baseline=configuration.effective_baseline, previous=setup.previous, custom_dataset=configuration.dataset is not None, + provider=setup.agent.provider, ) if setup.write_constraints: if setup.user_prompt: @@ -637,12 +672,13 @@ async def _create_preset( # while the listing and `--previous` read constraints from the session dir. session.write_constraints(constraints_text) session.write_prompt(prompt) - if setup.auth is not None: - session.write_agent_info(get_agent_info(setup.auth)) + if setup.spec is not None: + session.write_agent_info(setup.agent.get_info(setup.spec)) try: if mode == "attach": process_output = await attach_preset_agent( workspace=setup.workspace, + agent=setup.agent, redacted_values=redacted_values, session=session, ) @@ -653,12 +689,13 @@ async def _create_preset( ): raise AgentExitedWithoutReport(process_output.error) else: - assert setup.auth is not None + assert setup.spec is not None process_output = await run_preset_agent( prompt=prompt, env=env, workspace=setup.workspace, - auth=setup.auth, + agent=setup.agent, + spec=setup.spec, redacted_values=redacted_values, session=session, initial_resume_session_id=setup.initial_resume_session_id, @@ -670,7 +707,7 @@ async def _create_preset( redacted_values=redacted_values, ) if isinstance(result, PresetAgentFailure): - raise CLIError(result.failure_summary or "Claude did not create a preset") + raise CLIError(result.failure_summary or "The agent did not create a preset") report = result run = api.client.runs.get(api.project, report.run_name) preset = build_verified_preset( @@ -773,7 +810,9 @@ def _stop_or_detach_agent_session(session: PresetSession, api: Client) -> None: """`create` interrupt: stop the session, or detach and leave the agent working as a running session in `dstack preset`.""" state = session.read_state() - agent_alive = state is not None and state.run is not None and process_alive(state.run.agent) + agent_alive = ( + state is not None and state.run is not None and process_alive(state.run.session_process) + ) stop = True if agent_alive: try: @@ -788,7 +827,7 @@ def _stop_or_detach_agent_session(session: PresetSession, api: Client) -> None: ) return if state is not None: - terminate_agent_process(state.run.agent if state.run else None) + terminate_agent_process(state.run.session_process if state.run else None) _stop_active_session_runs(api, session) _suspend_agent_session(session) diff --git a/src/dstack/_internal/cli/services/presets/prompt.py b/src/dstack/_internal/cli/services/presets/prompt.py index ca19b03c1..0910acdf5 100644 --- a/src/dstack/_internal/cli/services/presets/prompt.py +++ b/src/dstack/_internal/cli/services/presets/prompt.py @@ -4,6 +4,7 @@ from typing import Optional, Sequence, Union from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.configurations import PresetAgentProvider _SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" @@ -175,6 +176,7 @@ def get_preset_agent_system_prompt( baseline: bool, previous: Sequence[str], custom_dataset: bool, + provider: PresetAgentProvider = "claude", ) -> str: text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() variables = { @@ -185,6 +187,9 @@ def get_preset_agent_system_prompt( "previous": ", ".join(previous) if previous else None, # Rendered for its presence only; the dataset itself is in constraints.json. "dataset": "on" if custom_dataset else None, + # Rendered for its presence only: the codex CLI loads skills and takes the + # final report differently from claude, whose text is the `else` branch. + "codex": "on" if provider == "codex" else None, } applied: set[str] = set() rendered = _render_branch(_parse_directives(text, variables), variables, applied, dedent=False) diff --git a/src/dstack/_internal/cli/services/presets/report_schema.py b/src/dstack/_internal/cli/services/presets/report_schema.py new file mode 100644 index 000000000..c96de1346 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/report_schema.py @@ -0,0 +1,110 @@ +"""The final report schema the agent CLIs enforce on the agent's last message.""" + +from typing import Any + +from dstack._internal.cli.models.preset_agent import PresetAgentFailure, PresetAgentSuccess + +# Keywords the strict structured-output mode of the OpenAI API rejects. +_STRICT_UNSUPPORTED_KEYWORDS = frozenset( + { + "title", + "default", + "examples", + "format", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "minLength", + "maxLength", + "pattern", + "minItems", + "maxItems", + "uniqueItems", + "multipleOf", + "minProperties", + "maxProperties", + "patternProperties", + "readOnly", + "deprecated", + "$comment", + "contentMediaType", + } +) + + +def get_report_json_schema() -> dict[str, Any]: + """The one shape the API can enforce: a single object, no union, only + `success` required. `AnyPresetAgentResult` enforces the rest at parse.""" + success = PresetAgentSuccess.model_json_schema() + failure = PresetAgentFailure.model_json_schema() + return { + "type": "object", + "properties": { + **success["properties"], + **failure["properties"], + # Each outcome fixes its own value; only the merged shape offers both. + "success": {"type": "boolean"}, + }, + "required": ["success"], + "additionalProperties": False, + "$defs": {**success.get("$defs", {}), **failure.get("$defs", {})}, + } + + +def get_strict_report_json_schema() -> dict[str, Any]: + """`get_report_json_schema` in the strict form the OpenAI API enforces: every + property required, optional ones nullable, no unsupported keywords. A field + the agent has no value for comes back as `null`, which the report parser drops.""" + return _strict(get_report_json_schema()) + + +def _strict(node: Any) -> Any: + if isinstance(node, list): + return [_strict(item) for item in node] + if not isinstance(node, dict): + return node + if "$ref" in node: + # A `$ref` may carry no sibling keywords. + return {"$ref": node["$ref"]} + if node.get("type") == "object" and "properties" in node: + required = set(node.get("required", [])) + properties = {} + for name, schema in node["properties"].items(): + # An optional field with a concrete default stays non-null and becomes + # required, so the model writes the default; one without becomes nullable. + has_default = schema.get("default") is not None + schema = _strict(schema) + if name not in required and not has_default: + schema = _nullable(schema) + properties[name] = schema + rest = { + key: _strict(value) + for key, value in node.items() + if key not in _STRICT_UNSUPPORTED_KEYWORDS and key not in ("properties", "required") + } + return { + **rest, + "properties": properties, + "required": list(properties), + "additionalProperties": False, + } + return { + key: _strict(value) + for key, value in node.items() + if key not in _STRICT_UNSUPPORTED_KEYWORDS + } + + +def _nullable(schema: dict[str, Any]) -> dict[str, Any]: + if "anyOf" in schema: + if not any(option.get("type") == "null" for option in schema["anyOf"]): + schema["anyOf"].append({"type": "null"}) + return schema + if "$ref" in schema: + return {"anyOf": [schema, {"type": "null"}]} + if isinstance(schema.get("type"), str): + schema["type"] = [schema["type"], "null"] + elif isinstance(schema.get("type"), list) and "null" not in schema["type"]: + schema["type"].append("null") + return schema diff --git a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md index 0e439a0d0..101191b94 100644 --- a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md @@ -77,11 +77,11 @@ submitting `dstack` runs (e.g. tasks for trials and services for the final verification). To do this, use the real `dstack` CLI and shell commands in this workspace. -Load and follow `/dstack` for `dstack` CLI/YAML syntax. Load and follow -`/dstack-prototyping` for how to test a model-serving configuration with +Load and follow the `$dstack` skill`/dstack` for `dstack` CLI/YAML syntax. Load and follow +the `$dstack-prototyping` skill`/dstack-prototyping` for how to test a model-serving configuration with `dstack` tasks before verifying it as a `dstack` service. If a skill cannot -be loaded with its -slash command, read it from `.claude/skills//SKILL.md` and +be loaded by namewith its +slash command, read it from `.codex.claude/skills//SKILL.md` and follow it the same way. # Workspace Files @@ -131,7 +131,7 @@ id with `dstack run get --json` and append one JSON line: Never edit or delete existing lines. Stop runs you no longer need unless they are still needed for attach/SSH debugging, logs, or backend diagnosis. -After stopping a `dstack` task or service, follow `/dstack` structured status +After stopping a `dstack` task or service, follow the `$dstack` skill's`/dstack` structured status guidance and confirm that the run reached a terminal status before continuing. # Progress @@ -471,7 +471,7 @@ Trials are done entirely using `dstack` tasks. For maximum efficiency, it is a requirement that you always set the task `commands` to `sleep infinity` (for a task with `groups`, in each group's `commands`) and run commands inside the task interactively, via SSH. It is important that -you follow the `/dstack-prototyping` skill when working with tasks. +you follow the `$/dstack-prototyping` skill when working with tasks. # Hardware @@ -480,7 +480,7 @@ hardware from one trial to another if this can help the outcome. The available hardware is defined by the allowed `dstack` fleets and their offers (coming from the configured backends). Follow the -`/dstack-prototyping` skill on how to efficiently select offers among the +`$/dstack-prototyping` skill on how to efficiently select offers among the available backends or SSH fleets. Pick the offer whose hardware best fits the trial's idea. Only when several offers fit comparably, prefer backends or SSH fleets that support idle instances/instance volumes: later runs reuse the instance and cached model @@ -626,8 +626,8 @@ On failure (no trial verification was successful), include exactly: - `failure_summary`: the reason a preset could not be created and any change required from the user or administrator -Write the report to `final_report.json`, then submit the identical JSON object -through `StructuredOutput`. +Write the report to `final_report.json`, then end your turn with a final message that is exactly that JSON object and nothing else; it is checked against the report schemasubmit the identical JSON object +through `StructuredOutput`. Verify that `final_report.json` is correct and matches the required schema. diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index 8e054cf11..f92b70b82 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -30,7 +30,7 @@ from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError from dstack._internal.core.models.common import validate_extra_ignore, validate_json_extra_ignore -from dstack._internal.core.models.configurations import PresetConfiguration +from dstack._internal.core.models.configurations import PresetAgentProvider, PresetConfiguration from dstack._internal.utils.common import get_dstack_dir _PROGRESS_FILENAME = "progress.jsonl" @@ -152,6 +152,8 @@ def read_state(self) -> Optional[PresetSessionState]: return None if isinstance(data, dict) and "run" not in data and ("pid" in data or "workspace" in data): data = _upgrade_pre_0_21_2_state(data) + if isinstance(data, dict) and isinstance(data.get("run"), dict): + data["run"] = _upgrade_pre_0_22_run(data["run"]) try: return validate_extra_ignore(PresetSessionState, data) except ValidationError: @@ -168,10 +170,11 @@ def begin_run( *, workspace: PresetSessionWorkspace, finalize: PresetSessionFinalize, - claude_model: Optional[str], + agent_provider: PresetAgentProvider, + agent_model: Optional[str], ) -> None: """This CLI takes ownership and starts (or joins) the agent run. Everything - an earlier run established survives: the claude session id and model pin so + an earlier run established survives: the agent session id and model pin so a resume finds them, and the agent process reference so following a live detached agent keeps it alive instead of reading it as dead.""" state = self.read_state() @@ -183,26 +186,27 @@ def begin_run( state.run = PresetSessionRun( workspace=workspace, finalize=finalize, - claude_model=claude_model or (earlier.claude_model if earlier else None), - agent=earlier.agent if earlier else None, - claude_session_id=earlier.claude_session_id if earlier else None, + agent_provider=agent_provider, + agent_model=agent_model or (earlier.agent_model if earlier else None), + session_process=earlier.session_process if earlier else None, + session_id=earlier.session_id if earlier else None, ) self.write_state(state) - def record_agent(self, agent: PresetSessionProcess) -> None: + def record_session_process(self, process: PresetSessionProcess) -> None: state = self.read_state() if state is None or state.run is None: return - state.run.agent = agent + state.run.session_process = process self.write_state(state) - def record_claude_session_id(self, session_id: str) -> None: + def record_session_id(self, session_id: str) -> None: # An unreadable state stays as it is: rewriting it would fabricate a # session record out of one field. state = self.read_state() if state is None or state.run is None: return - state.run.claude_session_id = session_id + state.run.session_id = session_id self.write_state(state) def detach(self) -> None: @@ -247,6 +251,7 @@ def _upgrade_pre_0_21_2_state(data: dict[str, Any]) -> dict[str, Any]: # Without the finalize context there is no run to reconcile or resume. data["run"] = None else: + # The 0.21.2 shape; `_upgrade_pre_0_22_run` renames its fields next. data["run"] = { "workspace": {"path": workspace, "alias": alias or workspace}, "finalize": {"project": project, "keep_service": bool(keep_service)}, @@ -262,6 +267,20 @@ def _upgrade_pre_0_21_2_state(data: dict[str, Any]) -> dict[str, Any]: return data +# TODO: Remove in 0.23 +def _upgrade_pre_0_22_run(run: dict[str, Any]) -> dict[str, Any]: + """A run written before 0.22 named its fields after Claude, the only agent + then. Pure renaming for backward compatibility.""" + if "claude_model" not in run and "claude_session_id" not in run: + return run + run = dict(run) + run["agent_model"] = run.pop("claude_model", None) + run["session_id"] = run.pop("claude_session_id", None) + run["session_process"] = run.pop("agent", None) + run.setdefault("agent_provider", "claude") + return run + + def get_presets_dir() -> Path: return get_dstack_dir() / "presets" @@ -333,7 +352,7 @@ def load_resumable_session(preset_id: str) -> PresetSession: f"Preset {preset_id} is still being created;" f" follow it with dstack preset logs -f {preset_id}" ) - if state.run is None or state.run.claude_session_id is None: + if state.run is None or state.run.session_id is None: raise CLIError(f"Preset {preset_id} creation stopped before it started; create a new one") return session @@ -356,7 +375,7 @@ def process_alive(process: Optional[PresetSessionProcess]) -> bool: def session_process_alive(state: PresetSessionState) -> bool: """True if either a live agent (possibly detached) or a live CLI (possibly between agent retries) still owns the session.""" - if state.run is not None and process_alive(state.run.agent): + if state.run is not None and process_alive(state.run.session_process): return True if state.owner is None or state.owner.pid == os.getpid(): return False diff --git a/src/dstack/_internal/cli/services/presets/verify.py b/src/dstack/_internal/cli/services/presets/verify.py index 36fcda5fd..643522c50 100644 --- a/src/dstack/_internal/cli/services/presets/verify.py +++ b/src/dstack/_internal/cli/services/presets/verify.py @@ -67,7 +67,7 @@ def load_preset_agent_report( if report_data is None: raise CLIError( redact( - output.error or "Claude exited without a final report", + output.error or "The agent exited without a final report", redacted_values, ) ) @@ -76,7 +76,7 @@ def load_preset_agent_report( report_data, context={"redacted_values": tuple(redacted_values)} ) except ValidationError as e: - raise CLIError(f"Claude returned an invalid final report: {e}") from e + raise CLIError(f"The agent returned an invalid final report: {e}") from e def build_verified_preset( @@ -98,7 +98,7 @@ def build_verified_preset( service = _verified_run_service(run, report) _check_report_answers_request(report, preset_configuration) if service.model is None or service.model.name != preset_configuration.model.api_model_name: - raise CLIError("Claude final service model name does not match the requested model") + raise CLIError("The agent's final service model name does not match the requested model") return build_preset( name=name, service=_portable_service( @@ -126,12 +126,12 @@ def _verified_run_service(run: Run, report: PresetAgentSuccess) -> ServiceConfig """The service the server actually runs, after proving the report talks about this run and the run is a live model service.""" if run.id != report.run_id or run.run_spec.run_name != report.run_name: - raise CLIError("Claude final report identifies a different service run") + raise CLIError("The agent's final report identifies a different service run") if run.status != RunStatus.RUNNING or run.service is None: - raise CLIError("Claude final service is not running") + raise CLIError("The agent's final service is not running") service = run.run_spec.configuration if not isinstance(service, ServiceConfiguration): - raise CLIError("Claude final run is not a model service") + raise CLIError("The agent's final run is not a model service") return service @@ -144,9 +144,9 @@ def _check_report_answers_request( _check_workload_answers_request(report.benchmark.workload, configuration) if configuration.model.allows_variant_selection: if report.base != configuration.model.api_model_name: - raise CLIError("Claude final report base does not match the requested model") + raise CLIError("The agent's final report base does not match the requested model") elif report.model != configuration.model.exact_repo: - raise CLIError("Claude changed an exact model request") + raise CLIError("The agent changed an exact model request") def _check_workload_answers_request( @@ -162,19 +162,19 @@ def _check_workload_answers_request( if configuration.dataset is not None: if workload.dataset != configuration.dataset: raise CLIError( - f"Claude final benchmark dataset {workload.dataset!r} does not match the" + f"The agent's final benchmark dataset {workload.dataset!r} does not match the" f" requested dataset {configuration.dataset!r}" ) else: shared_prefix_tokens = configuration.shared_prefix_tokens or 0 if workload.shared_prefix_tokens != shared_prefix_tokens: raise CLIError( - f"Claude final benchmark shared prefix of {workload.shared_prefix_tokens}" + f"The agent's final benchmark shared prefix of {workload.shared_prefix_tokens}" f" tokens does not match the requested {shared_prefix_tokens}" ) if configuration.concurrency is not None and workload.concurrency != configuration.concurrency: raise CLIError( - f"Claude final benchmark concurrency of {workload.concurrency} does not match the" + f"The agent's final benchmark concurrency of {workload.concurrency} does not match the" f" requested concurrency of {configuration.concurrency}" ) @@ -222,13 +222,15 @@ def _mirrored_file_path(local_path: str, *, workspace_path: Path, session_path: try: relative = Path(local_path).resolve().relative_to(workspace_path.resolve()) except ValueError: - raise CLIError(f"Claude final service file '{local_path}' is outside the agent workspace") + raise CLIError( + f"The agent's final service file '{local_path}' is outside the agent workspace" + ) if ( relative.parts[:1] not in (("trials",), ("service",)) or not (session_path / relative).exists() ): raise CLIError( - f"Claude final service file '{local_path}' has no mirrored copy" + f"The agent's final service file '{local_path}' has no mirrored copy" f" at '{session_path / relative}'" ) return relative.as_posix() diff --git a/src/dstack/_internal/cli/services/presets/workspace.py b/src/dstack/_internal/cli/services/presets/workspace.py index c10613175..97ade98b3 100644 --- a/src/dstack/_internal/cli/services/presets/workspace.py +++ b/src/dstack/_internal/cli/services/presets/workspace.py @@ -80,7 +80,10 @@ def agent_stderr_path(self) -> Path: def create_agent_workspace( session: PresetSession, + skills_dir: Path, ) -> tuple[PresetAgentWorkspace, PresetSessionWorkspace]: + """`skills_dir` is where the agent CLI discovers project skills, relative to + the working directory.""" real = session.path / "workspace" try: real.mkdir(mode=0o700) @@ -92,7 +95,7 @@ def create_agent_workspace( alias = _create_workspace_alias(real) _validate_control_socket_path(alias) workspace = PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") - _prepare_workspace(workspace) + _prepare_workspace(workspace, skills_dir) except OSError as e: raise CLIError(f"Could not create the agent workspace under {real}: {e}") from e return workspace, PresetSessionWorkspace(path=str(real), alias=str(alias)) @@ -183,7 +186,7 @@ def _validate_control_socket_path(build_root: Path) -> None: raise CLIError(f"Temporary path is too long for an SSH control socket: {build_root}") -def _prepare_workspace(workspace: PresetAgentWorkspace) -> None: +def _prepare_workspace(workspace: PresetAgentWorkspace, skills_dir: Path) -> None: workspace.path.mkdir(mode=0o700, parents=True, exist_ok=False) workspace.dstack_home.mkdir(mode=0o700) workspace.temp_path.mkdir(mode=0o700) @@ -196,7 +199,7 @@ def _prepare_workspace(workspace: PresetAgentWorkspace) -> None: (workspace.dstack_home / ".ssh").mkdir(mode=0o700) _install_dstack_wrapper(workspace.bin_path, workspace.dstack_home) _install_home_wrapper(workspace.bin_path, "ssh", workspace.dstack_home) - _install_skills(workspace.path) + _install_skills(workspace.path / skills_dir) def _install_dstack_wrapper(bin_dir: Path, home: Path) -> None: @@ -322,10 +325,8 @@ def _copy_session_records(source_root: Path, target_root: Path) -> bool: return copied -def _install_skills(workspace: Path) -> None: +def _install_skills(target_dir: Path) -> None: source_dir = _get_skills_dir() - target_dir = workspace / ".claude" / "skills" - target_dir.mkdir(parents=True) for skill_name in _SKILL_NAMES: source = source_dir / skill_name if not (source / "SKILL.md").is_file(): diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 02860daf3..81d5d6226 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -1637,6 +1637,41 @@ def validate_base(cls, value: str) -> str: PresetModelSpec = Union[PresetModelRepo, PresetModelBase] +PresetAgentProvider = Literal["claude", "codex"] +PresetAgentEffort = Literal["low", "medium", "high", "xhigh", "max"] + + +class PresetAgentConfig(CoreModel): + provider: Annotated[ + PresetAgentProvider, + Field(description="The agent CLI that creates the preset: `claude` or `codex`"), + ] + model: Annotated[ + Optional[str], + Field( + description=( + "The model the agent runs with, e.g. `claude-opus-5` or `gpt-6-astra`." + " Defaults to the agent CLI's own default" + ) + ), + ] = None + effort: Annotated[ + Optional[PresetAgentEffort], + Field( + description=( + "The reasoning effort. `max` is supported by `claude` only." + " Defaults to the agent CLI's own default" + ) + ), + ] = None + + @model_validator(mode="after") + def validate_effort(self) -> Self: + if self.provider == "codex" and self.effort == "max": + raise ValueError("effort `max` is not supported by `codex`") + return self + + MAX_PROMPT_LENGTH = 10_000 @@ -1805,6 +1840,15 @@ class PresetConfiguration( ), ), ] = None + agent: Annotated[ + Optional[PresetAgentConfig], + Field( + description=( + "The agent that creates the preset. Overrides the `DSTACK_AGENT_*` environment" + " variables. Defaults to `claude`" + ) + ), + ] = None env: Annotated[Env, Field(description="The mapping or the list of environment variables")] = ( Env() ) diff --git a/src/tests/_internal/cli/commands/test_preset.py b/src/tests/_internal/cli/commands/test_preset.py index 4bc098c80..d8a52097f 100644 --- a/src/tests/_internal/cli/commands/test_preset.py +++ b/src/tests/_internal/cli/commands/test_preset.py @@ -85,7 +85,7 @@ def test_resume_uses_the_configuration_and_prompt_pinned_at_creation(self, tmp_p (session_dir / "session.json").write_text( json.dumps( get_session_state( - status="interrupted", run=get_session_run(claude_session_id="sid-1") + status="interrupted", run=get_session_run(session_id="sid-1") ).model_dump(mode="json") ) ) diff --git a/src/tests/_internal/cli/common.py b/src/tests/_internal/cli/common.py index 35ae8bb78..233274625 100644 --- a/src/tests/_internal/cli/common.py +++ b/src/tests/_internal/cli/common.py @@ -18,6 +18,7 @@ PresetSessionWorkspace, ) from dstack._internal.cli.models.presets import VerifiedPreset +from dstack._internal.cli.services.presets.agents.base import PresetAgentSpec from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.models.configurations import ( DEFAULT_REPLICA_GROUP_NAME, @@ -232,13 +233,25 @@ def get_session_state(**overrides: Any) -> PresetSessionState: return PresetSessionState(**fields) +def get_agent_spec(**overrides: Any) -> PresetAgentSpec: + fields: dict[str, Any] = { + "executable": "claude", + "api_key": "anthropic-secret", + "model": "claude-test", + "effort": None, + } + fields.update(overrides) + return PresetAgentSpec(**fields) + + def get_session_run(**overrides: Any) -> PresetSessionRun: fields: dict[str, Any] = { "workspace": PresetSessionWorkspace(path="/tmp/preset-ws", alias="/tmp/preset-ws"), "finalize": PresetSessionFinalize(project="main", keep_service=False), - "claude_model": None, - "agent": None, - "claude_session_id": None, + "agent_provider": "claude", + "agent_model": None, + "session_process": None, + "session_id": None, } fields.update(overrides) return PresetSessionRun(**fields) diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py index f36940e87..6276a58c1 100644 --- a/src/tests/_internal/cli/services/presets/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -15,19 +15,28 @@ from dstack._internal.cli.models.preset_agent import PresetAgentInfo, PresetSessionProcess from dstack._internal.cli.services.presets.agent import ( - ClaudeAuth, - _build_claude_command, _prepare_subprocess_command, _terminate_process, build_preset_agent_env, - get_agent_info, - get_claude_auth, run_preset_agent, ) +from dstack._internal.cli.services.presets.agents.base import ( + PresetAgentSpec, + PresetAgentStreamUpdate, +) +from dstack._internal.cli.services.presets.agents.claude import ClaudePresetAgent +from dstack._internal.cli.services.presets.agents.codex import ( + CodexPresetAgent, + _mcp_server_overrides, +) from dstack._internal.cli.services.presets.redaction import ( contains_redacted_value, redact, ) +from dstack._internal.cli.services.presets.report_schema import ( + get_report_json_schema, + get_strict_report_json_schema, +) from dstack._internal.cli.services.presets.session import ( PresetSession, _read_last_session_verification, @@ -49,9 +58,9 @@ ) from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.configurations import PresetConfiguration +from dstack._internal.core.models.configurations import PresetAgentConfig, PresetConfiguration from dstack._internal.core.services.configs import ConfigManager -from tests._internal.cli.common import get_session_run, get_session_state +from tests._internal.cli.common import get_agent_spec, get_session_run, get_session_state def _record_run(session, workspace_record): @@ -64,18 +73,10 @@ def _record_run(session, workspace_record): pytestmark = pytest.mark.windows -def _claude_auth( - *, api_key: str | None = "anthropic-secret", effort=None, model: str | None = "claude-test" -) -> ClaudeAuth: - return ClaudeAuth( - api_key=api_key, - executable="claude", - effort=effort, - model=model, - ) +_WORKSPACE = PresetAgentWorkspace(path=Path("/w"), dstack_home=Path("/w/h")) -class TestClaudeAuth: +class TestClaudePresetAgent: @pytest.mark.parametrize("api_key_env", ["key", None]) def test_uses_api_key_only_when_env_is_set(self, monkeypatch, api_key_env): if api_key_env is None: @@ -84,9 +85,9 @@ def test_uses_api_key_only_when_env_is_set(self, monkeypatch, api_key_env): monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_API_KEY", api_key_env) monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") - auth = get_claude_auth() + spec = ClaudePresetAgent().get_spec(None) - assert auth.api_key == api_key_env + assert spec.api_key == api_key_env @pytest.mark.parametrize("model_env", ["claude-pinned", "", None]) def test_leaves_model_to_claude_unless_env_is_set(self, monkeypatch, model_env): @@ -96,23 +97,84 @@ def test_leaves_model_to_claude_unless_env_is_set(self, monkeypatch, model_env): monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_MODEL", model_env) monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") - auth = get_claude_auth() + spec = ClaudePresetAgent().get_spec(None) + + assert spec.model == (model_env or None) + + def test_the_agent_block_overrides_the_environment(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") + monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_MODEL", "claude-from-env") + monkeypatch.setenv("DSTACK_AGENT_CLAUDE_EFFORT", "low") + + spec = ClaudePresetAgent().get_spec( + PresetAgentConfig(provider="claude", model="claude-from-config", effort="max") + ) - assert auth.model == (model_env or None) + assert spec.model == "claude-from-config" + assert spec.effort == "max" + + def test_an_empty_agent_block_falls_back_to_the_environment(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") + monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_MODEL", "claude-from-env") + monkeypatch.delenv("DSTACK_AGENT_CLAUDE_EFFORT", raising=False) + + spec = ClaudePresetAgent().get_spec(PresetAgentConfig(provider="claude")) + + assert spec.model == "claude-from-env" + assert spec.effort is None + + def test_rejects_an_effort_the_agent_does_not_have(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") + monkeypatch.setenv("DSTACK_AGENT_CLAUDE_EFFORT", "ultra") + + with pytest.raises(CLIError, match="DSTACK_AGENT_CLAUDE_EFFORT must be one of"): + ClaudePresetAgent().get_spec(None) @pytest.mark.parametrize("api_key", ["key", None]) def test_builds_command_for_selected_auth_mode(self, api_key): - command = _build_claude_command( - auth=_claude_auth(api_key=api_key, effort="high"), resume_session_id=None + command = ClaudePresetAgent().build_command( + get_agent_spec(api_key=api_key, effort="high"), _WORKSPACE, None ) assert ("--bare" in command) is (api_key is not None) assert ("--setting-sources" in command) is (api_key is None) assert command[command.index("--effort") + 1] == "high" + def test_builds_the_exact_command(self): + """The full argv, pinned so an agent-neutral refactor cannot change what + claude receives: login mode, no model or effort, then a resume.""" + tools = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" + expected = [ + "claude", + "-p", + "--setting-sources", + "project,local", + "--output-format", + "stream-json", + "--verbose", + "--tools", + tools, + "--allowedTools", + tools, + "--disallowedTools", + "Task,NotebookEdit", + "--permission-mode", + "bypassPermissions", + "--json-schema", + json.dumps(get_report_json_schema()), + ] + spec = get_agent_spec(api_key=None, model=None) + + assert ClaudePresetAgent().build_command(spec, _WORKSPACE, None) == expected + assert ClaudePresetAgent().build_command(spec, _WORKSPACE, "sid-1") == [ + *expected, + "--resume", + "sid-1", + ] + @pytest.mark.parametrize("model", ["claude-pinned", None]) def test_passes_model_only_when_pinned(self, model): - command = _build_claude_command(auth=_claude_auth(model=model), resume_session_id=None) + command = ClaudePresetAgent().build_command(get_agent_spec(model=model), _WORKSPACE, None) if model is None: assert "--model" not in command @@ -135,6 +197,281 @@ def test_runs_windows_batch_launcher(self, tmp_path): assert result.stdout.strip() == "batch-ok" +def _codex_workspace(tmp_path) -> PresetAgentWorkspace: + workspace = PresetAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace.path.mkdir() + workspace.dstack_home.mkdir() + return workspace + + +class TestCodexPresetAgent: + def test_reads_the_spec_from_the_environment(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/codex") + monkeypatch.setenv("DSTACK_AGENT_OPENAI_API_KEY", "sk-test") + monkeypatch.setenv("DSTACK_AGENT_OPENAI_MODEL", "gpt-5.5") + monkeypatch.setenv("DSTACK_AGENT_CODEX_EFFORT", "high") + + assert CodexPresetAgent().get_spec(None) == PresetAgentSpec( + executable="/usr/bin/codex", + api_key="sk-test", + model="gpt-5.5", + effort="high", + ) + + def test_the_agent_block_overrides_the_environment(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/codex") + monkeypatch.setenv("DSTACK_AGENT_OPENAI_MODEL", "gpt-from-env") + + spec = CodexPresetAgent().get_spec( + PresetAgentConfig(provider="codex", model="gpt-6-astra", effort="xhigh") + ) + + assert spec.model == "gpt-6-astra" + assert spec.effort == "xhigh" + + def test_rejects_an_unknown_effort(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/codex") + monkeypatch.setenv("DSTACK_AGENT_CODEX_EFFORT", "max") + + with pytest.raises(CLIError, match="DSTACK_AGENT_CODEX_EFFORT"): + CodexPresetAgent().get_spec(None) + + def test_builds_the_exact_command_with_an_api_key(self, tmp_path): + workspace = _codex_workspace(tmp_path) + schema_path = workspace.path / ".report-schema.json" + agent = CodexPresetAgent() + agent.build_env(get_agent_spec(executable="codex", api_key="sk-test"), workspace, {}) + + command = agent.build_command( + get_agent_spec(executable="codex", api_key="sk-test", model="gpt-5.5", effort="high"), + workspace, + None, + ) + + assert command == [ + "codex", + "exec", + "--json", + "--skip-git-repo-check", + "--dangerously-bypass-approvals-and-sandbox", + "--disable", + "multi_agent", + "--disable", + "apps", + "-c", + 'web_search="live"', + "--output-schema", + str(schema_path), + "--output-last-message", + str(workspace.path / ".agent-last-message.json"), + "-C", + str(workspace.path), + "-m", + "gpt-5.5", + "-c", + 'model_reasoning_effort="high"', + "-", + ] + assert json.loads(schema_path.read_text()) == get_strict_report_json_schema() + + def test_resume_keeps_the_thread_and_disables_the_user_mcp_servers( + self, tmp_path, monkeypatch + ): + codex_home = tmp_path / "codex-home" + codex_home.mkdir() + (codex_home / "config.toml").write_text( + 'model_provider = "proxy"\n' + "[mcp_servers.computer-use]\n" + 'command = "/usr/bin/cua"\n' + "[mcp_servers.playwright]\n" + 'command = "npx"\n' + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agents.codex.get_user_codex_home", + lambda: codex_home, + ) + workspace = _codex_workspace(tmp_path) + + command = CodexPresetAgent().build_command( + get_agent_spec(executable="codex", api_key=None, model=None), workspace, "thread-1" + ) + + assert command[:4] == ["codex", "exec", "resume", "thread-1"] + # `exec resume` has no `-C`; the thread remembers its working directory. + assert "-C" not in command + # The user's config, provider, and hooks stay; only the MCP servers go. + assert "--ignore-user-config" not in command + assert command[-5:] == [ + "-c", + 'mcp_servers.computer-use={command="disabled", enabled=false}', + "-c", + 'mcp_servers.playwright={command="disabled", enabled=false}', + "-", + ] + + def test_env_isolates_the_home_and_uses_our_key(self, tmp_path): + workspace = _codex_workspace(tmp_path) + env: dict[str, str] = {} + + CodexPresetAgent().build_env( + get_agent_spec(executable="codex", api_key="sk-test", model=None), workspace, env + ) + + assert env["HOME"] == str(workspace.dstack_home) + assert env["CODEX_API_KEY"] == "sk-test" + assert env["CODEX_HOME"] == str(workspace.dstack_home / ".codex") + assert (workspace.dstack_home / ".codex").is_dir() + + def test_env_keeps_the_user_codex_home_without_a_key(self, tmp_path, monkeypatch): + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agents.codex.get_user_codex_home", + lambda: tmp_path / "user-codex", + ) + workspace = _codex_workspace(tmp_path) + env: dict[str, str] = {} + + CodexPresetAgent().build_env( + get_agent_spec(executable="codex", api_key=None, model=None), workspace, env + ) + + # The user's codex as they run it: their home, their CODEX_HOME. + assert env["HOME"] == str(Path.home()) + assert env["CODEX_HOME"] == str(tmp_path / "user-codex") + assert "CODEX_API_KEY" not in env + + +class TestMcpServerOverrides: + def test_disables_every_server_in_the_user_config(self, tmp_path): + config = tmp_path / "config.toml" + config.write_text( + 'model = "gpt-5.5"\n' + "[mcp_servers.computer-use]\n" + 'command = "/usr/bin/cua"\n' + 'args = ["mcp"]\n' + "[mcp_servers.docs]\n" + 'url = "https://example.com/mcp"\n' + '[mcp_servers."acme.docs"]\n' + 'url = "https://acme.example.com/mcp"\n' + ) + + # A quoted key has no `-c` dotted path, so it stays as the user has it. + assert _mcp_server_overrides(config) == [ + "-c", + 'mcp_servers.computer-use={command="disabled", enabled=false}', + "-c", + 'mcp_servers.docs={command="disabled", enabled=false}', + ] + + @pytest.mark.parametrize("content", ['model = "gpt-5.5"\n', "not = valid = toml\n", None]) + def test_nothing_to_disable_adds_nothing(self, tmp_path, content): + config = tmp_path / "config.toml" + if content is not None: + config.write_text(content) + + assert _mcp_server_overrides(config) == [] + + +class TestCodexStream: + def test_maps_events_onto_stream_updates(self): + agent = CodexPresetAgent() + + assert agent.parse_line("not json") is None + assert agent.parse_line( + '{"type":"thread.started","thread_id":"t-1"}' + ) == PresetAgentStreamUpdate(session_id="t-1") + assert agent.parse_line('{"type":"turn.started"}') == PresetAgentStreamUpdate() + assert agent.parse_line( + '{"type":"item.completed","item":{"id":"i1","type":"command_execution","command":"ls","exit_code":0}}' + ) == PresetAgentStreamUpdate(made_progress=True) + assert agent.parse_line( + '{"type":"item.completed","item":{"id":"i2","type":"agent_message","text":"{\\"note\\": \\"looks like json\\"}"}}' + ) == PresetAgentStreamUpdate(made_progress=True) + assert agent.parse_line( + '{"type":"item.completed","item":{"id":"i3","type":"agent_message","text":"{\\"success\\": false, \\"failure_summary\\": \\"no\\"}"}}' + ) == PresetAgentStreamUpdate(made_progress=True) + assert agent.parse_line('{"type":"turn.completed","usage":{}}') == PresetAgentStreamUpdate( + report_data={"success": False, "failure_summary": "no"} + ) + assert agent.parse_line( + '{"type":"error","message":"stream disconnected"}' + ) == PresetAgentStreamUpdate(error="stream disconnected") + assert agent.parse_line( + '{"type":"turn.failed","error":{"message":"boom"}}' + ) == PresetAgentStreamUpdate(error="boom") + + def test_a_turn_ending_in_prose_has_no_report(self): + agent = CodexPresetAgent() + + agent.parse_line( + '{"type":"item.completed","item":{"id":"i1","type":"agent_message","text":"{\\"success\\": true}"}}' + ) + agent.parse_line( + '{"type":"item.completed","item":{"id":"i2","type":"agent_message","text":"done"}}' + ) + + assert ( + agent.parse_line('{"type":"turn.completed","usage":{}}') == PresetAgentStreamUpdate() + ) + + def test_reads_the_model_from_the_session_rollout_once(self, tmp_path, monkeypatch): + codex_home = tmp_path / "codex" + day = codex_home / "sessions" / "2026" / "09" / "05" + day.mkdir(parents=True) + (day / "rollout-2026-09-05T18-00-00-t-1.jsonl").write_text( + '{"type":"session_meta","payload":{"id":"t-1"}}\n' + '{"type":"turn_context","payload":{"model":"gpt-5.5"}}\n' + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agents.codex.get_user_codex_home", + lambda: codex_home, + ) + agent = CodexPresetAgent() + agent.build_env( + get_agent_spec(executable="codex", api_key=None, model=None), + _codex_workspace(tmp_path), + {}, + ) + + agent.parse_line('{"type":"thread.started","thread_id":"t-1"}') + first = agent.parse_line( + '{"type":"item.started","item":{"id":"i1","type":"command_execution"}}' + ) + second = agent.parse_line( + '{"type":"item.completed","item":{"id":"i1","type":"command_execution"}}' + ) + + assert first is not None and first.model == "gpt-5.5" + assert second is not None and second.model is None + + def test_stops_looking_for_the_model_after_the_turn(self, tmp_path, monkeypatch): + codex_home = tmp_path / "codex" + day = codex_home / "sessions" / "2026" / "09" / "05" + day.mkdir(parents=True) + rollout = day / "rollout-2026-09-05T18-00-00-t-1.jsonl" + rollout.write_text('{"type":"session_meta","payload":{"id":"t-1"}}\n') + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agents.codex.get_user_codex_home", + lambda: codex_home, + ) + agent = CodexPresetAgent() + agent.build_env( + get_agent_spec(executable="codex", api_key=None, model=None), + _codex_workspace(tmp_path), + {}, + ) + agent.parse_line('{"type":"thread.started","thread_id":"t-1"}') + agent.parse_line('{"type":"item.started","item":{"id":"i1","type":"command_execution"}}') + agent.parse_line('{"type":"turn.completed","usage":{}}') + + # Written too late: both reads are spent. + rollout.write_text('{"type":"turn_context","payload":{"model":"gpt-5.5"}}\n') + late = agent.parse_line( + '{"type":"item.completed","item":{"id":"i1","type":"command_execution"}}' + ) + + assert late is not None and late.model is None + + class TestAgentIsolation: def test_inherits_only_required_environment(self, tmp_path, monkeypatch): monkeypatch.setenv("PATH", "/usr/bin") @@ -148,7 +485,8 @@ def test_inherits_only_required_environment(self, tmp_path, monkeypatch): env = build_preset_agent_env( api=api, preset_env={"HF_TOKEN": "hf-secret"}, - auth=_claude_auth(), + agent=ClaudePresetAgent(), + spec=get_agent_spec(), workspace=PresetAgentWorkspace( path=tmp_path, dstack_home=tmp_path / "home", @@ -203,7 +541,7 @@ def _session_workspace(tmp_path): session_dir.mkdir() session = PresetSession(path=session_dir, preset_id="abcd1234") session.write_state(get_session_state(id="abcd1234")) - workspace, _ = create_agent_workspace(session) + workspace, _ = create_agent_workspace(session, ClaudePresetAgent.skills_dir) return workspace @@ -338,8 +676,9 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, ) (tmp_path / "progress.jsonl").touch() monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._build_claude_command", - lambda **_: [sys.executable, str(script)], + ClaudePresetAgent, + "build_command", + lambda self, spec, workspace, resume_session_id: [sys.executable, str(script)], ) workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") @@ -355,7 +694,8 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, prompt="full preset prompt", env=os.environ.copy(), workspace=workspace, - auth=_claude_auth(), + agent=ClaudePresetAgent(), + spec=get_agent_spec(), redacted_values=("secret-token",), session=session, ) @@ -394,8 +734,9 @@ async def test_mirrors_trial_and_service_records_into_the_session(self, tmp_path (tmp_path / "trials").mkdir() (tmp_path / "service").mkdir() monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._build_claude_command", - lambda **_: [sys.executable, str(script)], + ClaudePresetAgent, + "build_command", + lambda self, spec, workspace, resume_session_id: [sys.executable, str(script)], ) workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") session_path = tmp_path / "session" @@ -407,7 +748,8 @@ async def test_mirrors_trial_and_service_records_into_the_session(self, tmp_path prompt="p", env=os.environ.copy(), workspace=workspace, - auth=_claude_auth(), + agent=ClaudePresetAgent(), + spec=get_agent_spec(), redacted_values=("secret-token",), session=session, ) @@ -438,15 +780,17 @@ async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypat session_path.mkdir() (session_path / "agent.log").touch() monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._build_claude_command", - lambda **_: [sys.executable, str(script)], + ClaudePresetAgent, + "build_command", + lambda self, spec, workspace, resume_session_id: [sys.executable, str(script)], ) output = await run_preset_agent( prompt="prompt", env=os.environ.copy(), workspace=PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), - auth=_claude_auth(), + agent=ClaudePresetAgent(), + spec=get_agent_spec(), redacted_values=(), session=PresetSession( path=session_path, @@ -640,14 +984,14 @@ def test_skips_files_above_the_size_limit(self, tmp_path, monkeypatch): class TestGetAgentInfo: - def test_describes_the_launch_without_a_model(self, tmp_path, monkeypatch): - monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._get_claude_version", - lambda auth: "2.1.0 (Claude Code)", - ) + def test_describes_the_spec_without_a_model(self, tmp_path, monkeypatch): + probes = { + "--version": "2.1.0 (Claude Code)", + "--json": '{"authMethod": "claude.ai", "loggedIn": true}', + } monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._get_claude_auth_status", - lambda auth: '{"authMethod": "claude.ai", "loggedIn": true}', + "dstack._internal.cli.services.presets.agents.claude.probe_cli", + lambda args, **_: probes[args[-1]], ) session_dir = tmp_path / "session" session_dir.mkdir() @@ -656,12 +1000,13 @@ def test_describes_the_launch_without_a_model(self, tmp_path, monkeypatch): # Even a pinned model is not recorded here: claude reports the model it # runs on its init line, and that is what gets recorded. session.write_agent_info( - get_agent_info( - ClaudeAuth(api_key=None, executable="claude", effort="high", model="claude-pinned") + ClaudePresetAgent().get_info( + get_agent_spec(api_key=None, effort="high", model="claude-pinned") ) ) assert json.loads((session_dir / "agent.json").read_text()) == { + "provider": "claude", "executable": "claude", "version": "2.1.0 (Claude Code)", "auth_status": '{"authMethod": "claude.ai", "loggedIn": true}', @@ -677,7 +1022,12 @@ def test_records_the_model_claude_reports(self, tmp_path): session = PresetSession(path=session_dir, preset_id="ab12cd34") session.write_agent_info( PresetAgentInfo( - executable="claude", version=None, auth_status="{}", effort=None, model=None + provider="claude", + executable="claude", + version=None, + auth_status="{}", + effort=None, + model=None, ) ) @@ -733,12 +1083,76 @@ def _agent_setup(tmp_path): def _patch_claude_command(monkeypatch, script): monkeypatch.setattr( - "dstack._internal.cli.services.presets.agent._build_claude_command", - lambda **kwargs: [sys.executable, str(script)] - + (["--resume", kwargs["resume_session_id"]] if kwargs.get("resume_session_id") else []), + ClaudePresetAgent, + "build_command", + lambda self, spec, workspace, resume_session_id: [sys.executable, str(script)] + + (["--resume", resume_session_id] if resume_session_id else []), ) +class TestCodexResume: + @pytest.mark.asyncio + async def test_resumes_the_thread_after_a_failed_turn(self, tmp_path, monkeypatch): + script = _write_fake_claude( + tmp_path, + """import json +import sys +from pathlib import Path + +args = sys.argv[1:] +with Path("calls.jsonl").open("a") as f: + f.write(json.dumps({"args": args, "prompt": sys.stdin.read()}) + "\\n") +print(json.dumps({"type": "thread.started", "thread_id": "t-1"})) +if "resume" in args: + report = {"success": False, "failure_summary": "resumed"} + print(json.dumps({"type": "item.completed", "item": {"id": "i1", "type": "agent_message", "text": json.dumps(report)}})) + print(json.dumps({"type": "turn.completed", "usage": {}})) +else: + print(json.dumps({"type": "turn.started"})) + print(json.dumps({"type": "turn.failed", "error": {"message": "unexpected status 401"}})) + sys.exit(1) +""", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) + ) + real_build_command = CodexPresetAgent.build_command + + def fake_build_command(self, spec, workspace, resume_session_id): + # The real argv with the fake script as the executable. + command = real_build_command(self, spec, workspace, resume_session_id) + return [sys.executable, str(script), *command[1:]] + + monkeypatch.setattr(CodexPresetAgent, "build_command", fake_build_command) + workspace, session = _agent_setup(tmp_path) + workspace.dstack_home.mkdir(exist_ok=True) + session.write_state(get_session_state(run=get_session_run(agent_provider="codex"))) + agent = CodexPresetAgent() + spec = get_agent_spec(executable="codex", api_key="sk-test", model=None) + agent.build_env(spec, workspace, {}) + + output = await run_preset_agent( + prompt="system prompt", + env=_subprocess_env(), + workspace=workspace, + agent=agent, + spec=spec, + redacted_values=(), + session=session, + ) + + assert output.report_data == {"success": False, "failure_summary": "resumed"} + assert output.session_id == "t-1" + calls = [json.loads(line) for line in (tmp_path / "calls.jsonl").read_text().splitlines()] + assert calls[0]["args"][:2] == ["exec", "--json"] + assert calls[0]["prompt"] == "system prompt" + assert calls[1]["args"][:3] == ["exec", "resume", "t-1"] + assert calls[1]["prompt"].startswith("The previous agent process was interrupted") + state = session.read_state() + assert state is not None and state.run is not None + assert state.run.session_id == "t-1" + + class TestConnectionResume: @pytest.mark.asyncio async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsys): @@ -774,7 +1188,12 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy workspace, session = _agent_setup(tmp_path) session.write_agent_info( PresetAgentInfo( - executable="claude", version=None, auth_status="{}", effort=None, model=None + provider="claude", + executable="claude", + version=None, + auth_status="{}", + effort=None, + model=None, ) ) @@ -782,7 +1201,8 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy prompt="system prompt", env=_subprocess_env(), workspace=workspace, - auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model=None), + agent=ClaudePresetAgent(), + spec=get_agent_spec(api_key=None, model=None), redacted_values=(), session=session, ) @@ -828,7 +1248,8 @@ async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): prompt="system prompt", env=_subprocess_env(), workspace=workspace, - auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + agent=ClaudePresetAgent(), + spec=get_agent_spec(api_key=None, model="m"), redacted_values=(), session=session, ) @@ -867,7 +1288,8 @@ async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monk prompt="system prompt", env=_subprocess_env(), workspace=workspace, - auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + agent=ClaudePresetAgent(), + spec=get_agent_spec(api_key=None, model="m"), redacted_values=(), session=session, ) @@ -911,7 +1333,8 @@ async def test_does_not_resurrect_an_externally_stopped_agent(self, tmp_path, mo prompt="system prompt", env=_subprocess_env(), workspace=workspace, - auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + agent=ClaudePresetAgent(), + spec=get_agent_spec(api_key=None, model="m"), redacted_values=(), session=session, ) @@ -932,7 +1355,7 @@ def _session(self, tmp_path): def test_create_attach_and_remove(self, tmp_path): session = self._session(tmp_path) - workspace, workspace_record = create_agent_workspace(session) + workspace, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) _record_run(session, workspace_record) alias = Path(workspace_record.alias) assert Path(workspace_record.path) == session.path / "workspace" @@ -951,7 +1374,7 @@ def test_create_attach_and_remove(self, tmp_path): @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") def test_attach_refuses_occupied_alias(self, tmp_path): session = self._session(tmp_path) - _, workspace_record = create_agent_workspace(session) + _, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) _record_run(session, workspace_record) alias = Path(workspace_record.alias) os.unlink(alias) @@ -964,7 +1387,7 @@ def test_attach_refuses_occupied_alias(self, tmp_path): def test_attach_fails_when_workspace_is_gone(self, tmp_path): session = self._session(tmp_path) - _, workspace_record = create_agent_workspace(session) + _, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) _record_run(session, workspace_record) remove_agent_workspace(session) with pytest.raises(CLIError, match="no longer exists"): @@ -1042,10 +1465,39 @@ def test_reads_a_pre_0_22_flat_session_file(self, tmp_path): assert state.run.workspace.alias == "/tmp/dpe-1" assert state.run.finalize.project == "main" assert state.run.finalize.keep_service is True - assert state.run.agent == PresetSessionProcess(pid=70521, started_at=1755116374.0) - assert state.run.claude_session_id == "71b025f9-fba0-42b9-8734-e357deca5281" + assert state.run.session_process == PresetSessionProcess( + pid=70521, started_at=1755116374.0 + ) + assert state.run.agent_provider == "claude" + assert state.run.agent_model == "claude-opus-5" + assert state.run.session_id == "71b025f9-fba0-42b9-8734-e357deca5281" assert state.previous == [] + def test_reads_a_pre_0_22_run_named_after_claude(self, tmp_path): + # Verbatim `run` shape written by the 0.21.5 CLI, before the agent-neutral names. + state = get_session_state(id="30a012bf").model_dump(mode="json") + state["run"] = { + "workspace": {"path": "/tmp/w", "alias": "/tmp/dpe-1"}, + "finalize": {"project": "main", "keep_service": True}, + "claude_model": "claude-opus-5", + "agent": {"pid": 70521, "started_at": 1755116374.0}, + "claude_session_id": "71b025f9-fba0-42b9-8734-e357deca5281", + } + session_dir = tmp_path / "30a012bf" + session_dir.mkdir() + (session_dir / "session.json").write_text(json.dumps(state)) + session = PresetSession(path=session_dir, preset_id="30a012bf") + + loaded = session.read_state() + + assert loaded is not None and loaded.run is not None + assert loaded.run.agent_provider == "claude" + assert loaded.run.agent_model == "claude-opus-5" + assert loaded.run.session_id == "71b025f9-fba0-42b9-8734-e357deca5281" + assert loaded.run.session_process == PresetSessionProcess( + pid=70521, started_at=1755116374.0 + ) + class TestLoadResumableSession: def _write_session(self, tmp_path, monkeypatch, state): @@ -1065,7 +1517,7 @@ def test_loads_interrupted_session(self, tmp_path, monkeypatch): { "id": "ab12cd34", "status": "interrupted", - "run": get_session_run(claude_session_id="sid-1"), + "run": get_session_run(session_id="sid-1"), "created_at": "2026-07-20T10:00:00Z", }, ) @@ -1082,7 +1534,7 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): "id": "ab12cd34", "status": "running", "owner": {"pid": 4242, "started_at": None}, - "run": get_session_run(claude_session_id="sid-1"), + "run": get_session_run(session_id="sid-1"), }, ) monkeypatch.setattr( @@ -1105,7 +1557,7 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): "id": "aa000003", "status": "running", "owner": {"pid": 4242, "started_at": None}, - "run": get_session_run(claude_session_id="sid-1"), + "run": get_session_run(session_id="sid-1"), }, "still being created", id="still-running", @@ -1307,7 +1759,9 @@ def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypat state = session.read_state() assert state is not None state.status = "running" - state.run = get_session_run(agent=PresetSessionProcess(pid=agent.pid, started_at=None)) + state.run = get_session_run( + session_process=PresetSessionProcess(pid=agent.pid, started_at=None) + ) session.write_state(state) monkeypatch.setattr(create_module, "confirm_ask", lambda *_: False) diff --git a/src/tests/_internal/cli/services/presets/test_create.py b/src/tests/_internal/cli/services/presets/test_create.py index e3ae148d8..58e90014b 100644 --- a/src/tests/_internal/cli/services/presets/test_create.py +++ b/src/tests/_internal/cli/services/presets/test_create.py @@ -8,21 +8,24 @@ from pydantic import ValidationError from dstack._internal.cli.models.preset_agent import ( + PresetAgentInfo, PresetSessionFinalize, PresetSessionProcess, PresetSessionState, PresetSessionWorkspace, ) from dstack._internal.cli.services.presets.agent import ( - ClaudeAuth, PresetAgentProcessOutput, ) +from dstack._internal.cli.services.presets.agents.base import PresetAgentSpec +from dstack._internal.cli.services.presets.agents.claude import ClaudePresetAgent from dstack._internal.cli.services.presets.create import ( PresetCreateResult, SessionBusyError, _build_constraints, _cleanup_runs, _create_preset, + _fresh_setup, _get_build_name, _print_fleet_offers, _save_final_report_copy, @@ -50,10 +53,11 @@ remove_agent_workspace, ) from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.configurations import PresetConfiguration +from dstack._internal.core.models.configurations import PresetAgentConfig, PresetConfiguration from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.runs import Run, RunStatus from tests._internal.cli.common import ( + get_agent_spec, get_preset, get_running_service_run, get_session_run, @@ -64,13 +68,85 @@ pytestmark = pytest.mark.windows -def _claude_auth() -> ClaudeAuth: - return ClaudeAuth( - api_key="anthropic-secret", - executable="claude", - effort=None, - model="claude-test", - ) +class _TestClaudePresetAgent(ClaudePresetAgent): + """Claude without the environment lookup and the `claude` subprocess probes.""" + + def get_spec(self, config) -> PresetAgentSpec: + return get_agent_spec() + + def get_info(self, spec: PresetAgentSpec) -> PresetAgentInfo: + return PresetAgentInfo( + provider="claude", + executable=spec.executable, + version=None, + auth_status="api-key", + effort=spec.effort, + model=None, + ) + + +class TestFreshSetup: + def test_the_agent_block_picks_the_agent_and_reaches_the_spec(self, tmp_path, monkeypatch): + seen: dict = {} + + class _PresetAgent(_TestClaudePresetAgent): + provider = "codex" + + def get_spec(self, config): + seen["config"] = config + return get_agent_spec() + + def get_preset_agent(provider=None): + seen["provider"] = provider + return _PresetAgent() + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.get_preset_agent", get_preset_agent + ) + configuration = PresetConfiguration( + name="qwen-build", + base="Qwen/Qwen3.5-27B", + agent={"provider": "codex", "model": "gpt-6-astra", "effort": "xhigh"}, + ) + + setup = _fresh_setup( + api=SimpleNamespace(project="main"), + configuration=configuration, + session=_agent_session(tmp_path), + build_name="qwen-build", + allowed_fleets=("gpu-fleet",), + user_prompt=None, + previous=(), + ) + + assert seen["provider"] == "codex" + assert seen["config"] == PresetAgentConfig( + provider="codex", model="gpt-6-astra", effort="xhigh" + ) + assert setup.agent.provider == "codex" + + def test_without_an_agent_block_the_environment_decides(self, tmp_path, monkeypatch): + seen: dict = {} + + def get_preset_agent(provider=None): + seen["provider"] = provider + return _TestClaudePresetAgent() + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.get_preset_agent", get_preset_agent + ) + + _fresh_setup( + api=SimpleNamespace(project="main"), + configuration=PresetConfiguration(name="qwen-build", base="Qwen/Qwen3.5-27B"), + session=_agent_session(tmp_path), + build_name="qwen-build", + allowed_fleets=("gpu-fleet",), + user_prompt=None, + previous=(), + ) + + assert seen["provider"] is None def _session_dirs(tmp_path): @@ -116,8 +192,8 @@ def creation_context(tmp_path, monkeypatch): env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], ) monkeypatch.setattr( - "dstack._internal.cli.services.presets.create.get_claude_auth", - _claude_auth, + "dstack._internal.cli.services.presets.create.get_preset_agent", + lambda provider=None: _TestClaudePresetAgent(), ) monkeypatch.setattr( "dstack._internal.cli.services.presets.create._get_build_name", @@ -256,14 +332,16 @@ def fail_finish(self, preset_id=None): ) @pytest.mark.asyncio - async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypatch): + async def test_checks_active_fleets_before_claude_launch(self, tmp_path, monkeypatch): api = SimpleNamespace( project="main", client=SimpleNamespace(fleets=SimpleNamespace(list=lambda *args, **kwargs: [])), ) monkeypatch.setattr( - "dstack._internal.cli.services.presets.create.get_claude_auth", - lambda: pytest.fail("Claude auth must not be checked without an active fleet"), + "dstack._internal.cli.services.presets.create.get_preset_agent", + lambda provider=None: pytest.fail( + "The agent must not be set up without an active fleet" + ), ) with pytest.raises(CLIError, match="no fleets"): @@ -920,7 +998,7 @@ async def test_resume_uses_saved_claude_session(self, creation_context, monkeypa (session_dir / "agent.log").touch() session = PresetSession(path=session_dir, preset_id="fe98dc76") session.write_state(get_session_state(id="fe98dc76")) - workspace, workspace_record = create_agent_workspace(session) + workspace, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) workspace.constraints_path.write_text( '{"run_name_prefix": "qwen-build"}', encoding="utf-8" ) @@ -928,8 +1006,8 @@ async def test_resume_uses_saved_claude_session(self, creation_context, monkeypa assert state is not None state.run = get_session_run( workspace=workspace_record, - claude_model="claude-pinned", - claude_session_id="sid-xyz", + agent_model="claude-pinned", + session_id="sid-xyz", ) session.write_state(state) captured = {} @@ -957,7 +1035,7 @@ async def run_agent(**kwargs): ) assert captured["initial_resume_session_id"] == "sid-xyz" - assert captured["auth"].model == "claude-pinned" + assert captured["spec"].model == "claude-pinned" assert result.preset.id == "fe98dc76" assert (session_dir / "workspace").is_dir() remove_agent_workspace(session) @@ -1007,13 +1085,13 @@ async def test_resume_keeps_the_pinned_user_prompt( (session_dir / "agent.log").touch() session = PresetSession(path=session_dir, preset_id="ab34ef12") session.write_state(get_session_state(id="ab34ef12")) - workspace, workspace_record = create_agent_workspace(session) + workspace, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) workspace.constraints_path.write_text( '{"run_name_prefix": "qwen-build"}', encoding="utf-8" ) state = session.read_state() assert state is not None - state.run = get_session_run(workspace=workspace_record, claude_session_id="sid-abc") + state.run = get_session_run(workspace=workspace_record, session_id="sid-abc") session.write_state(state) session.write_user_prompt("Optimize for RAG traffic.") captured = {} @@ -1124,13 +1202,13 @@ def _detached_session(self, tmp_path, configuration_yaml: str) -> PresetSession: (session_dir / "agent.log").touch() (session_dir / "preset.dstack.yml").write_text(configuration_yaml) session = PresetSession(path=session_dir, preset_id="ab12cd34") - workspace, workspace_record = create_agent_workspace(session) + workspace, workspace_record = create_agent_workspace(session, ClaudePresetAgent.skills_dir) workspace.constraints_path.write_text('{"run_name_prefix": "qwen-build"}') session.write_state( get_session_state( run=get_session_run( workspace=workspace_record, - agent=PresetSessionProcess(pid=987654321, started_at=None), + session_process=PresetSessionProcess(pid=987654321, started_at=None), ) ) ) @@ -1288,7 +1366,9 @@ def test_stop_wins_over_a_live_owner(self, tmp_path, monkeypatch, capsys): tmp_path, # A live owner: this very process. get_session_state( - run=get_session_run(agent=PresetSessionProcess(pid=os.getpid(), started_at=None)) + run=get_session_run( + session_process=PresetSessionProcess(pid=os.getpid(), started_at=None) + ) ), ) order = [] @@ -1404,7 +1484,7 @@ def _session_dir( ) if owner_alive and state.run is not None: # A live pid with no recorded start time reads as an active owner. - state.run.agent = PresetSessionProcess(pid=os.getpid(), started_at=None) + state.run.session_process = PresetSessionProcess(pid=os.getpid(), started_at=None) (session_dir / "session.json").write_text(state.model_dump_json()) if with_report: (session_dir / "workspace" / "w" / "final_report.json").write_text("{}") @@ -1510,7 +1590,7 @@ def test_recycled_pid_with_stale_start_time_is_not_alive(self): session_process_alive( get_session_state( run=get_session_run( - agent=PresetSessionProcess(pid=os.getpid(), started_at=0.0) + session_process=PresetSessionProcess(pid=os.getpid(), started_at=0.0) ) ) ) @@ -1538,9 +1618,9 @@ def test_records_the_run_whole_and_keeps_claude_state(self, tmp_path): run=get_session_run( workspace=workspace, finalize=PresetSessionFinalize(project="old", keep_service=False), - claude_model="claude-pinned", - agent=PresetSessionProcess(pid=1, started_at=None), - claude_session_id="sid-1", + agent_model="claude-pinned", + session_process=PresetSessionProcess(pid=1, started_at=None), + session_id="sid-1", ), ) ) @@ -1548,7 +1628,8 @@ def test_records_the_run_whole_and_keeps_claude_state(self, tmp_path): session.begin_run( workspace=workspace, finalize=PresetSessionFinalize(project="main", keep_service=True), - claude_model=None, + agent_provider="claude", + agent_model=None, ) state = session.read_state() @@ -1561,6 +1642,7 @@ def test_records_the_run_whole_and_keeps_claude_state(self, tmp_path): # Everything the earlier run established survives: the claude state so a # resume finds it, and the agent reference so following a live detached # agent does not read it as dead (and kill it). - assert state.run.claude_model == "claude-pinned" - assert state.run.claude_session_id == "sid-1" - assert state.run.agent == PresetSessionProcess(pid=1, started_at=None) + assert state.run.agent_provider == "claude" + assert state.run.agent_model == "claude-pinned" + assert state.run.session_id == "sid-1" + assert state.run.session_process == PresetSessionProcess(pid=1, started_at=None) diff --git a/src/tests/_internal/cli/services/presets/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py index 160d6a642..a66092f02 100644 --- a/src/tests/_internal/cli/services/presets/test_prompt.py +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -1,3 +1,5 @@ +import difflib + import pytest from dstack._internal.cli.services.presets import prompt as prompt_module @@ -8,6 +10,27 @@ class TestSystemPrompt: + def test_codex_prompt_swaps_only_the_agent_specific_parts(self): + claude = get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=(), custom_dataset=False, provider="claude" + ) + codex = get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=(), custom_dataset=False, provider="codex" + ) + + assert "`$dstack`" in codex and "`/dstack`" not in codex + assert "$dstack-prototyping" in codex and "/dstack-prototyping" not in codex + assert ".codex/skills//SKILL.md" in codex and ".claude/skills" not in codex + assert "StructuredOutput" not in codex and "exactly that JSON object" in codex + assert "$dstack" not in claude and "StructuredOutput" in claude + changed = [ + line + for line in difflib.unified_diff(claude.splitlines(), codex.splitlines(), n=0) + if line[:1] in "+-" and line[:3] not in ("+++", "---") + ] + # The skills paragraph, three skill mentions, and the report sentence. + assert 0 < len(changed) < 24, changed + def test_stays_byte_identical_without_user_prompt(self): text = get_preset_agent_system_prompt( user_prompt=None, baseline=False, previous=(), custom_dataset=False diff --git a/src/tests/_internal/cli/services/presets/test_report_schema.py b/src/tests/_internal/cli/services/presets/test_report_schema.py new file mode 100644 index 000000000..9298e06d2 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_report_schema.py @@ -0,0 +1,150 @@ +import pytest + +from dstack._internal.cli.models.preset_agent import PresetAgentFailure, PresetAgentSuccess +from dstack._internal.cli.services.presets.report_schema import ( + get_report_json_schema, + get_strict_report_json_schema, +) +from dstack._internal.cli.services.presets.verify import _AGENT_RESULT_ADAPTER +from dstack._internal.core.models.configurations import ServiceConfiguration + +pytestmark = pytest.mark.windows + +# Every report key with no value, as the strict schema makes codex write them. +_NO_VALUES = { + "run_id": None, + "run_name": None, + "service_yaml": None, + "trial": None, + "base": None, + "model": None, + "context_length": None, + "benchmark": None, + "failure_summary": None, +} +# Keywords the OpenAI strict structured-output mode rejects. +_FORBIDDEN = { + "title", + "default", + "examples", + "format", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "minLength", + "maxLength", + "pattern", + "minItems", + "maxItems", +} + + +def _walk(node): + if isinstance(node, dict): + yield node + for value in node.values(): + yield from _walk(value) + elif isinstance(node, list): + for item in node: + yield from _walk(item) + + +class TestGetStrictReportJsonSchema: + def test_every_object_requires_all_of_its_properties(self): + objects = [ + node + for node in _walk(get_strict_report_json_schema()) + if node.get("type") == "object" and "properties" in node + ] + + assert objects + for node in objects: + assert node["required"] == list(node["properties"]) + assert node["additionalProperties"] is False + + def test_drops_the_keywords_strict_mode_rejects(self): + for node in _walk(get_strict_report_json_schema()): + assert not (set(node) & _FORBIDDEN), node + if "$ref" in node: + assert set(node) == {"$ref"} + + def test_keeps_the_agent_facing_shape(self): + strict = get_strict_report_json_schema() + + assert set(strict["properties"]) == set(get_report_json_schema()["properties"]) + # The service stays a YAML string; an optional field is nullable; a field + # with a default stays non-null so the model writes the default. + assert strict["properties"]["service_yaml"] == {"type": ["string", "null"]} + workload = strict["$defs"]["PresetWorkload"]["properties"] + assert workload["dataset"]["anyOf"][-1] == {"type": "null"} + assert workload["shared_prefix_tokens"]["type"] == "integer" + + @pytest.mark.parametrize( + "data, expected", + [ + ( + { + **_NO_VALUES, + "success": False, + "failure_summary": "codex smoke test: no trials were run", + }, + PresetAgentFailure, + ), + ( + { + **_NO_VALUES, + "success": True, + "run_id": "2b0e3b0c-6c1e-4b7e-9c3a-1f2a3b4c5d6e", + "run_name": "qwen-preset-1", + "service_yaml": ( + "type: service\nname: qwen-preset-1\nimage: vllm/vllm-openai:latest\n" + "commands:\n - vllm serve Qwen/Qwen2.5-0.5B-Instruct\nport: 8000\n" + "model: Qwen/Qwen2.5-0.5B-Instruct\n" + ), + "trial": 1, + "base": "Qwen/Qwen2.5-0.5B-Instruct", + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "context_length": 2048, + "benchmark": { + "tool": "vllm.bench serve", + "tool_version": "0.27.0", + "command": "vllm bench serve --model Qwen/Qwen2.5-0.5B-Instruct", + "workload": { + "api": "chat_completions", + "dataset": None, + "num_requests": 32, + "input_tokens": 128, + "output_tokens": 64, + "concurrency": 1, + "shared_prefix_tokens": 0, + }, + "metrics": { + "successful_requests": 32, + "failed_requests": 0, + "duration_seconds": 40.5, + "total_input_tokens": 4096, + "total_output_tokens": 2048, + "output_tok_per_s": 50.6, + "per_user_tok_per_s": 50.6, + "ttft_ms": {"mean": 210.0, "p50": 200.0, "p99": 300.0}, + "tpot_ms": {"mean": 19.8, "p50": 19.5, "p99": 25.0}, + }, + }, + }, + PresetAgentSuccess, + ), + ], + ) + def test_reports_codex_produced_against_the_schema_parse(self, data, expected): + # The shape `codex exec --output-schema` writes: every key present, the + # ones the agent had no value for as `null`. + assert set(data) == set(get_strict_report_json_schema()["properties"]) + + report = _AGENT_RESULT_ADAPTER.validate_python(data, context={"redacted_values": []}) + + assert isinstance(report, expected) + if isinstance(report, PresetAgentSuccess): + assert isinstance(report.service_yaml, ServiceConfiguration) + assert report.benchmark.workload.dataset is None + assert report.benchmark.workload.shared_prefix_tokens == 0 diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index d72b5ba9d..d1949a6ed 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -8,6 +8,7 @@ from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.configurations import ( DevEnvironmentConfigurationParams, + PresetAgentConfig, PresetConfiguration, PresetModelBase, PresetModelRepo, @@ -1208,6 +1209,7 @@ def test_schema_documents_supported_input(self): assert all(field.description for field in PresetConfiguration.model_fields.values()) assert all(field.description for field in PresetModelBase.model_fields.values()) assert all(field.description for field in PresetModelRepo.model_fields.values()) + assert all(field.description for field in PresetAgentConfig.model_fields.values()) assert {"type": "string"} in PresetConfiguration.model_json_schema()["properties"][ "model" ]["anyOf"] @@ -1228,6 +1230,34 @@ def test_parses_base_model(self): assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" assert configuration.model.allows_variant_selection + def test_agent_is_optional_and_carries_provider_model_and_effort(self): + assert PresetConfiguration(model="Qwen/Qwen3.5-27B").agent is None + + configuration = PresetConfiguration( + model="Qwen/Qwen3.5-27B", + agent={"provider": "codex", "model": "gpt-6-astra", "effort": "xhigh"}, + ) + + assert configuration.agent == PresetAgentConfig( + provider="codex", model="gpt-6-astra", effort="xhigh" + ) + assert PresetConfiguration( + model="Qwen/Qwen3.5-27B", agent={"provider": "claude"} + ).agent == (PresetAgentConfig(provider="claude", model=None, effort=None)) + + @pytest.mark.parametrize("agent", [{}, {"model": "gpt-6-astra"}, {"provider": "gemini"}]) + def test_agent_requires_a_known_provider(self, agent): + with pytest.raises(ValidationError): + PresetConfiguration(model="Qwen/Qwen3.5-27B", agent=agent) + + @pytest.mark.parametrize( + "agent", + [{"provider": "claude", "effort": "ultra"}, {"provider": "codex", "effort": "max"}], + ) + def test_agent_effort_must_exist_for_the_provider(self, agent): + with pytest.raises(ValidationError): + PresetConfiguration(model="Qwen/Qwen3.5-27B", agent=agent) + def test_parses_exact_repo_with_client_facing_name(self): configuration = PresetConfiguration( model={