From 387547980305a5ae66c4f86c40cc43d62b9345bb Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 24 Jul 2026 00:48:30 -0300 Subject: [PATCH] Complete local subagent delegation for Codex, Claude plan mode, and Pi (#7329) Add session-scoped MCP bridges so Codex, Claude plan mode, and Pi subagents run on the loaded local model, with cloud credentials and Codex state isolated per session and a process-wide Pi agent cap. --- unsloth_cli/claude_subagent_mcp.py | 139 +++++-- unsloth_cli/codex_subagent_mcp.py | 171 ++++++++ unsloth_cli/commands/start.py | 373 ++++++++++++++---- unsloth_cli/pi_subagent.ts | 353 ++++++++++++----- unsloth_cli/tests/test_claude_subagent_mcp.py | 67 ++++ unsloth_cli/tests/test_codex_subagent_mcp.py | 228 +++++++++++ unsloth_cli/tests/test_pi_subagent.py | 301 +++++++++++++- unsloth_cli/tests/test_start.py | 274 +++++++++++-- 8 files changed, 1645 insertions(+), 261 deletions(-) create mode 100644 unsloth_cli/codex_subagent_mcp.py create mode 100644 unsloth_cli/tests/test_codex_subagent_mcp.py diff --git a/unsloth_cli/claude_subagent_mcp.py b/unsloth_cli/claude_subagent_mcp.py index b86368515b..e044d78705 100644 --- a/unsloth_cli/claude_subagent_mcp.py +++ b/unsloth_cli/claude_subagent_mcp.py @@ -19,6 +19,8 @@ from unsloth_cli.commands.start import ( _CLAUDE_ENV_UNSET, _SUBAGENT_DESCRIPTION, _SUBAGENT_INSTRUCTIONS, + _SUBAGENT_PLAN_DESCRIPTION, + _SUBAGENT_PLAN_INSTRUCTIONS, _claude_flags, _claude_local_env, _wsl_shim_env, @@ -113,7 +115,11 @@ def _stop_child(process: subprocess.Popen) -> None: pass -def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str: +def run_local_agent( + task: str, + cancel_event: threading.Event | None = None, + read_only: bool = False, +) -> str: base = _required_env("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL") key = _required_env("UNSLOTH_CLAUDE_SUBAGENT_API_KEY") model = _required_env("UNSLOTH_CLAUDE_SUBAGENT_MODEL") @@ -135,16 +141,20 @@ def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> s *_claude_flags(model), "--permission-mode", ( - "bypassPermissions" - if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1" - else "acceptEdits" + "plan" + if read_only + else ( + "bypassPermissions" + if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1" + else "acceptEdits" + ) ), "--print", "--output-format", "json", "--no-session-persistence", "--append-system-prompt", - _SUBAGENT_INSTRUCTIONS, + _SUBAGENT_PLAN_INSTRUCTIONS if read_only else _SUBAGENT_INSTRUCTIONS, f"Task: {task}", ] bridged, wsl_names = _wsl_shim_env(command, local_env, _CLAUDE_ENV_UNSET) @@ -196,7 +206,15 @@ def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> s return _result_text(stdout) -def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) -> dict | None: +def _response( + request: dict, + run_agent: Callable[[str], str] = run_local_agent, + tool_name: str = "unsloth_agent", + tool_description: str | None = None, + run_read_only_agent: Callable[[str], str] | None = None, + read_only_tool_name: str | None = None, + instructions: str | None = None, +) -> dict | None: request_id = request.get("id") method = request.get("method") if request_id is None: @@ -208,40 +226,55 @@ def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) "capabilities": {"tools": {"listChanged": False}}, "serverInfo": {"name": "unsloth-local-agent", "version": "1.0.0"}, } + if instructions: + result["instructions"] = instructions elif method == "ping": result = {} elif method == "tools/list": - result = { - "tools": [ - { - "name": "unsloth_agent", - "title": "Unsloth local agent", - "description": _SUBAGENT_DESCRIPTION, - "inputSchema": { - "type": "object", - "properties": { - "task": { - "type": "string", - "description": "The complete task for the local Unsloth agent.", - } - }, - "required": ["task"], - "additionalProperties": False, + + def tool_definition(name: str, description: str, read_only: bool) -> dict: + return { + "name": name, + "title": "Unsloth local plan agent" if read_only else "Unsloth local agent", + "description": description, + "inputSchema": { + "type": "object", + "properties": { + "task": { + "type": "string", + "description": "The complete task for the local Unsloth agent.", + } }, - "annotations": { - "readOnlyHint": False, - "destructiveHint": True, - "idempotentHint": False, - "openWorldHint": True, - }, - "_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS}, - } - ] - } + "required": ["task"], + "additionalProperties": False, + }, + "annotations": { + "readOnlyHint": read_only, + "destructiveHint": not read_only, + "idempotentHint": read_only, + "openWorldHint": True, + }, + "_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS}, + } + + tools = [tool_definition(tool_name, tool_description or _SUBAGENT_DESCRIPTION, False)] + if read_only_tool_name and run_read_only_agent: + tools.append(tool_definition(read_only_tool_name, _SUBAGENT_PLAN_DESCRIPTION, True)) + result = {"tools": tools} elif method == "tools/call": params = request.get("params") or {} arguments = params.get("arguments") or {} - task = arguments.get("task") if params.get("name") == "unsloth_agent" else None + requested_tool = params.get("name") + selected_agent = ( + run_agent + if requested_tool == tool_name + else ( + run_read_only_agent + if requested_tool == read_only_tool_name and run_read_only_agent + else None + ) + ) + task = arguments.get("task") if selected_agent else None if not isinstance(task, str) or not task.strip(): result = { "content": [{"type": "text", "text": "A non-empty task is required."}], @@ -249,7 +282,7 @@ def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) } else: try: - text = run_agent(task.strip()) + text = selected_agent(task.strip()) result = {"content": [{"type": "text", "text": text}], "isError": False} except Exception as exc: result = { @@ -269,6 +302,11 @@ def serve( stdin: Any = sys.stdin, stdout: Any = sys.stdout, run_agent: Callable[[str, threading.Event], str] = run_local_agent, + tool_name: str = "unsloth_agent", + tool_description: str | None = None, + run_read_only_agent: Callable[[str, threading.Event], str] | None = None, + read_only_tool_name: str | None = None, + instructions: str | None = None, ) -> None: active: dict[object, threading.Event] = {} workers: list[threading.Thread] = [] @@ -308,6 +346,15 @@ def serve( response = _response( request, run_agent = lambda task: run_agent(task, cancel_event), + tool_name = tool_name, + tool_description = tool_description, + run_read_only_agent = ( + (lambda task: run_read_only_agent(task, cancel_event)) + if run_read_only_agent + else None + ), + read_only_tool_name = read_only_tool_name, + instructions = instructions, ) if not cancel_event.is_set(): send(response) @@ -343,7 +390,18 @@ def serve( worker.start() response = None else: - response = _response(request) + response = _response( + request, + tool_name = tool_name, + tool_description = tool_description, + run_read_only_agent = ( + (lambda task: run_read_only_agent(task, threading.Event())) + if run_read_only_agent + else None + ), + read_only_tool_name = read_only_tool_name, + instructions = instructions, + ) except Exception as exc: response = { "jsonrpc": "2.0", @@ -362,5 +420,14 @@ def serve( signal.signal(signum, handler) +def main() -> None: + serve( + run_read_only_agent = lambda task, cancel_event: run_local_agent( + task, cancel_event, read_only = True + ), + read_only_tool_name = "unsloth_plan_agent", + ) + + if __name__ == "__main__": - serve() + main() diff --git a/unsloth_cli/codex_subagent_mcp.py b/unsloth_cli/codex_subagent_mcp.py new file mode 100644 index 0000000000..f075e66404 --- /dev/null +++ b/unsloth_cli/codex_subagent_mcp.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small stdio MCP bridge from cloud Codex to an explicit local Codex child.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import threading +from pathlib import Path +from typing import Any + +from unsloth_cli.claude_subagent_mcp import _bounded, _stop_child, serve +from unsloth_cli.commands.start import ( + _CODEX_ENV_KEY, + _CODEX_ENV_UNSET, + _CODEX_PROFILE, + _CODEX_SUBAGENT_CONFIG_ENV, + _CODEX_SUBAGENT_MCP_TOOL, + _CODEX_SUBAGENT_TOOL_DESCRIPTION, + _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS, + _SUBAGENT_INSTRUCTIONS, + _merge_wslenv, + _wsl_shim_env, +) + +_CANCEL_POLL_SECONDS = 0.1 +_SERVER_INSTRUCTIONS = _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS + + +def _config() -> dict: + path = os.environ.get(_CODEX_SUBAGENT_CONFIG_ENV, "").strip() + if not path: + raise RuntimeError(f"Missing {_CODEX_SUBAGENT_CONFIG_ENV}.") + try: + config = json.loads(Path(path).read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise RuntimeError("Could not read the local Codex agent configuration.") from exc + if not isinstance(config, dict): + raise RuntimeError("The local Codex agent configuration must be an object.") + for name in ("api_key", "codex_home"): + if not isinstance(config.get(name), str) or not config[name].strip(): + raise RuntimeError(f"The local Codex agent configuration is missing {name}.") + return config + + +def _result_text(stdout: str) -> str: + messages = [] + errors = [] + for line in stdout.splitlines(): + try: + event = json.loads(line) + except ValueError: + continue + if not isinstance(event, dict): + continue + item = event.get("item") + if ( + event.get("type") == "item.completed" + and isinstance(item, dict) + and item.get("type") == "agent_message" + and isinstance(item.get("text"), str) + and item["text"].strip() + ): + messages.append(item["text"].strip()) + if event.get("type") in ("error", "turn.failed"): + detail = event.get("message") or event.get("error") + if isinstance(detail, dict): + detail = detail.get("message") or json.dumps(detail) + if detail: + errors.append(str(detail)) + if errors: + raise RuntimeError(_bounded(errors[-1])) + if messages: + return _bounded(messages[-1]) + raise RuntimeError("The local Codex agent returned no readable result.") + + +def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str: + config = _config() + executable = shutil.which("codex") + if executable is None: + raise RuntimeError("`codex` is not installed or is not on PATH.") + cancel_event = cancel_event or threading.Event() + if cancel_event.is_set(): + raise RuntimeError("The local Codex agent was cancelled.") + + permissions = ( + ["--dangerously-bypass-approvals-and-sandbox"] + if config.get("bypass_permissions") is True + else ["--sandbox", "workspace-write", "--ask-for-approval", "never"] + ) + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *permissions, + "exec", + "--ephemeral", + "--json", + "--skip-git-repo-check", + f"{_SUBAGENT_INSTRUCTIONS}\n\nTask: {task}", + ] + local_env = { + _CODEX_ENV_KEY: config["api_key"], + "CODEX_HOME": config["codex_home"], + "CODEX_SQLITE_HOME": config["codex_home"], + } + bridged, wsl_names = _wsl_shim_env(command, local_env, _CODEX_ENV_UNSET) + child_env = dict(os.environ) + if wsl_names: + bridged = {**bridged, "PWD": os.getcwd()} + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_names) + for name in _CODEX_ENV_UNSET: + child_env[name] = "" + else: + for name in _CODEX_ENV_UNSET: + child_env.pop(name, None) + child_env.update(bridged) + popen_kwargs: dict[str, Any] = { + "cwd": os.getcwd(), + "env": child_env, + "stdin": subprocess.DEVNULL, + "stdout": subprocess.PIPE, + "stderr": subprocess.PIPE, + "text": True, + } + if os.name == "nt": + popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + popen_kwargs["start_new_session"] = True + process = subprocess.Popen([executable, *command[1:]], **popen_kwargs) + try: + while True: + try: + stdout, stderr = process.communicate(timeout = _CANCEL_POLL_SECONDS) + break + except subprocess.TimeoutExpired: + if cancel_event.is_set(): + _stop_child(process) + raise RuntimeError("The local Codex agent was cancelled.") + except BaseException: + if process.poll() is None: + _stop_child(process) + raise + if process.returncode != 0: + detail = stderr.strip() or stdout.strip() + raise RuntimeError( + _bounded(detail) or f"Local Codex exited with code {process.returncode}." + ) + return _result_text(stdout) + + +def main() -> None: + if len(sys.argv) > 1: + os.environ[_CODEX_SUBAGENT_CONFIG_ENV] = sys.argv[1] + serve( + run_agent = run_local_agent, + tool_name = _CODEX_SUBAGENT_MCP_TOOL, + tool_description = _CODEX_SUBAGENT_TOOL_DESCRIPTION, + instructions = _SERVER_INSTRUCTIONS, + ) + + +if __name__ == "__main__": + main() diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 3d407576f8..434c4a0ad5 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -84,8 +84,33 @@ _SUBAGENT_INSTRUCTIONS = ( "use the available tools when useful, verify your work, and return a concise result to the " "parent agent." ) +_SUBAGENT_PLAN_DESCRIPTION = ( + "Read-only local coding subagent powered by Unsloth for planning and codebase research. " + "Use this local agent when Claude is in plan mode." +) +_SUBAGENT_PLAN_INSTRUCTIONS = ( + "You are a read-only local coding subagent powered by Unsloth. Investigate the assigned " + "task with read-only tools, produce a concrete plan or answer, and return a concise result " + "to the parent agent. Do not modify files." +) _CLAUDE_SUBAGENT_MCP_MODULE = "unsloth_cli.claude_subagent_mcp" _CLAUDE_SUBAGENT_TOOL = "mcp__plugin_unsloth-local-agent_unsloth__unsloth_agent" +_CLAUDE_SUBAGENT_PLAN_TOOL = "mcp__plugin_unsloth-local-agent_unsloth__unsloth_plan_agent" +_CODEX_SUBAGENT_MCP_MODULE = "unsloth_cli.codex_subagent_mcp" +_CODEX_SUBAGENT_MCP_SERVER = "unsloth_local_agent" +_CODEX_SUBAGENT_MCP_TOOL = "spawn_local_agent" +_CODEX_SUBAGENT_CONFIG_ENV = "UNSLOTH_CODEX_SUBAGENT_CONFIG" +_CODEX_PARENT_OVERLAY_MANIFEST = ".unsloth-parent-overlay.json" +_CODEX_SUBAGENT_TOOL_DESCRIPTION = ( + f"{_SUBAGENT_DESCRIPTION} Use this tool instead of the built-in spawn_agent tool for those " + "requests. Other subagent requests may use the built-in tools normally." +) +_CODEX_SUBAGENT_ROUTING_INSTRUCTIONS = ( + "When the user asks to spawn an Unsloth agent or local agent, you must call the " + "spawn_local_agent MCP tool once with the complete task. Do not answer, simulate the " + "result, call wait, or use a built-in subagent before calling the tool. Use built-in " + "subagents for other delegation requests." +) _PI_SUBAGENT_EXTENSION = Path(__file__).parent.parent / "pi_subagent.ts" # OpenCode selects a model by "/". Use a dedicated id to avoid # colliding with a user's providers; provider filters are set in the launch-time overlay. @@ -113,6 +138,7 @@ class _PassthroughCommand(TyperCommand): _CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") +_CODEX_ENV_UNSET = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN") # Shared by every agent command; only the config/env/command differ. # Help is grouped into rich panels so `--help` reads as Model / Server / Session @@ -1093,6 +1119,13 @@ def _write_private_json(path: Path, data: dict) -> None: handle.write(json.dumps(data, indent = 2) + "\n") +def _write_private_text(path: Path, text: str) -> None: + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding = "utf-8") as handle: + handle.write(text) + + def _read_json_object(path: Path) -> Optional[dict]: # {} when missing, None when it can't be parsed as an object (so the caller # leaves a user-managed file untouched rather than clobbering it). @@ -1599,62 +1632,214 @@ def write_codex_config(base: str, model: dict, home: Path) -> None: typer.echo(f"Updated {profile}") -def write_codex_subagent_config(base: str, key: str, model: dict, home: Path) -> Path: - """Write a session-scoped Codex custom agent without replacing the main model.""" - home.mkdir(parents = True, exist_ok = True) - model_id = model["id"] - window = model.get("context_length") or model.get("max_context_length") - catalog_name = "unsloth-model-catalog.json" - text = ( - f"name = {json.dumps(_SUBAGENT_NAME)}\n" - f"description = {json.dumps(_SUBAGENT_DESCRIPTION)}\n" - f"developer_instructions = {json.dumps(_SUBAGENT_INSTRUCTIONS)}\n" - f"model_provider = {json.dumps(_CODEX_PROFILE)}\n" - f"model = {json.dumps(model_id)}\n" +def write_codex_subagent_bridge( + base: str, key: str, model: dict, home: Path, *, yolo: bool +) -> Path: + """Write private config for an explicit local Codex child launched through MCP.""" + child_home = home / "child" + write_codex_config(base, model, child_home) + path = home / "subagent.json" + _write_private_json( + path, + { + "api_key": key, + "codex_home": str(child_home), + "bypass_permissions": yolo, + }, ) - if _codex_supports_model_catalog() and _CODEX_FALLBACK_PROMPT.is_file(): - catalog = home / catalog_name - catalog_text = json.dumps(_codex_model_catalog(model), indent = 2) + "\n" - if not catalog.exists() or catalog.read_text(encoding = "utf-8") != catalog_text: - catalog.write_text(catalog_text, encoding = "utf-8") - typer.echo(f"Updated {catalog}") - text += f"model_catalog_json = {json.dumps(catalog_name)}\n" - if window: - text += f"model_context_window = {int(window)}\n" - credential = home / "unsloth-auth.json" - _write_private_json(credential, {"token": key}) - auth_command = sys.executable - auth_args = [ - "-c", - "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])", - str(credential), - ] - if _wsl_windows_executable(["codex"]): - auth_command = "wsl.exe" - auth_args = [ - "-d", - os.environ["WSL_DISTRO_NAME"], - "--", - sys.executable, - *auth_args, - ] - text += ( - f"\n{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - 'wire_api = "responses"\n' - f"\n{_PROVIDER_HEADER[:-1]}.auth]\n" - f"command = {json.dumps(auth_command)}\n" - f"args = {json.dumps(auth_args)}\n" - "timeout_ms = 5000\n" - ) - path = home / f"{_SUBAGENT_NAME}.toml" - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") return path +def _wsl_windows_user_profile(executable: str) -> Path: + """Return the Windows user profile as a path accessible from WSL.""" + profile = os.environ.get("USERPROFILE", "").strip() + if not profile: + try: + profile = subprocess.check_output( + ["cmd.exe", "/d", "/c", "echo %USERPROFILE%"], + text = True, + stderr = subprocess.DEVNULL, + cwd = str(Path(executable).parent), + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not find the Windows user profile for Codex: {exc}") + if not profile or profile == "%USERPROFILE%": + _fail("Could not find the Windows user profile for Codex.") + if profile.startswith("/"): + return Path(profile) + try: + translated = subprocess.check_output( + ["wslpath", "-u", profile], + text = True, + stderr = subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not translate Windows user profile {profile}: {exc}") + if not translated: + _fail(f"Could not translate Windows user profile {profile}.") + return Path(translated) + + +def _codex_source_home(*, ignore_configured: bool = False) -> Path: + configured = None if ignore_configured else os.environ.get("CODEX_HOME") + if configured: + if _wsl_windows_executable(["codex"]) and _looks_like_path(configured): + if not configured.startswith("/"): + try: + configured = subprocess.check_output( + ["wslpath", "-u", configured], + text = True, + stderr = subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not translate Windows CODEX_HOME {configured}: {exc}") + if not configured: + _fail("Could not translate Windows CODEX_HOME.") + return Path(configured).expanduser() + executable = _wsl_windows_executable(["codex"]) + if executable: + return _wsl_windows_user_profile(executable) / ".codex" + return Path.home() / ".codex" + + +def _remove_overlay_entry(path: Path) -> None: + is_junction = getattr(path, "is_junction", None) + if is_junction and is_junction(): + path.rmdir() + elif path.is_symlink() or path.is_file(): + path.unlink() + elif path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + +def _create_directory_junction(source: Path, target: Path) -> bool: + if os.name != "nt": + return False + try: + result = subprocess.run( + ["cmd.exe", "/d", "/c", "mklink", "/J", str(target), str(source)], + capture_output = True, + text = True, + timeout = 30, + check = False, + ) + except (OSError, subprocess.SubprocessError): + return False + return result.returncode == 0 + + +def write_codex_parent_overlay(overlay: Path) -> Path: + """Add local-agent routing without replacing the cloud parent's configuration.""" + overlay.mkdir(parents = True, exist_ok = True, mode = 0o700) + + manifest_path = overlay / _CODEX_PARENT_OVERLAY_MANIFEST + try: + manifest = json.loads(manifest_path.read_text(encoding = "utf-8")) + except (FileNotFoundError, OSError, json.JSONDecodeError): + manifest = None + source_home = _codex_source_home() + overlay_key = str(overlay.resolve(strict = False)) + source_key = str(source_home.resolve(strict = False)) + if source_key == overlay_key: + previous_source = manifest.get("source_home") if isinstance(manifest, dict) else None + if isinstance(previous_source, str) and previous_source: + candidate = Path(previous_source).expanduser() + if str(candidate.resolve(strict = False)) != overlay_key: + source_home = candidate + else: + source_home = _codex_source_home(ignore_configured = True) + else: + source_home = _codex_source_home(ignore_configured = True) + source_key = str(source_home.resolve(strict = False)) + same_source = isinstance(manifest, dict) and manifest.get("source_home") == source_key + if same_source: + managed_entries = manifest.get("entries", []) + if not isinstance(managed_entries, list): + managed_entries = [] + for name in managed_entries: + if isinstance(name, str) and name not in {"", ".", ".."} and Path(name).name == name: + _remove_overlay_entry(overlay / name) + else: + # A reused overlay must never mix credentials, config, or plugins from two + # different Codex homes. Legacy overlays have no manifest, so rebuild them once. + for target in list(overlay.iterdir()): + _remove_overlay_entry(target) + + # Keep the user's auth, config, plugins, agents, skills, rules, and session state visible. + # Symlinks make this an overlay rather than a stale copy. If Windows denies them, + # use directory junctions so large runtime state remains shared without a bulk copy. + # Copy the configuration surfaces and sessions only if both link forms are unavailable. + fallback_dirs = {"agents", "skills", "rules", "plugins", "marketplaces", "sessions"} + entries = [] + if source_home.is_dir(): + for source in source_home.iterdir(): + if source.name in { + "AGENTS.md", + "AGENTS.override.md", + _CODEX_PARENT_OVERLAY_MANIFEST, + }: + continue + target = overlay / source.name + _remove_overlay_entry(target) + try: + target.symlink_to(source, target_is_directory = source.is_dir()) + entries.append(source.name) + except OSError: + if source.is_file(): + shutil.copy2(source, target) + entries.append(source.name) + elif source.is_dir(): + if _create_directory_junction(source, target): + entries.append(source.name) + elif source.name in fallback_dirs: + shutil.copytree(source, target) + entries.append(source.name) + + _write_private_json( + manifest_path, + {"source_home": source_key, "entries": sorted(entries)}, + ) + + inherited = "" + instruction_name = "AGENTS.md" + for candidate in (source_home / "AGENTS.override.md", source_home / "AGENTS.md"): + try: + text = candidate.read_text(encoding = "utf-8") + except FileNotFoundError: + continue + except OSError as exc: + _fail(f"Could not preserve Codex instructions from {candidate}: {exc}") + if text.strip(): + inherited = text.rstrip() + instruction_name = candidate.name + break + + other_name = "AGENTS.md" if instruction_name == "AGENTS.override.md" else "AGENTS.override.md" + other = overlay / other_name + if other.is_file() or other.is_symlink(): + other.unlink() + routing = _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS + combined = f"{inherited}\n\n{routing}\n" if inherited else f"{routing}\n" + _write_private_text(overlay / instruction_name, combined) + return overlay + + +@contextlib.contextmanager +def _codex_parent_overlay(session_home: Path, *, launch: bool, persist: bool): + if launch and not persist: + temp_root = _agents_config_root() / ".tmp" + temp_root.mkdir(parents = True, exist_ok = True, mode = 0o700) + overlay = Path(tempfile.mkdtemp(prefix = "codex-parent-", dir = temp_root)) + try: + yield write_codex_parent_overlay(overlay) + finally: + shutil.rmtree(overlay, ignore_errors = True) + else: + yield write_codex_parent_overlay(session_home / "parent") + + def _agent_config_path(path: Path, command: list) -> str: """Translate a generated config path when a Windows agent runs through WSL.""" return _wsl_windows_path(path) if _wsl_windows_executable(command) else str(path) @@ -1778,25 +1963,42 @@ def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path: "description: Delegate a task to the local agent powered by Unsloth. Use when the " "user asks to spawn an Unsloth agent or local agent.\n" "---\n\n" - "Call the Unsloth local agent tool once with the complete task. Return its result " - "to the user without claiming that the cloud parent completed the local work.\n", + "Call the Unsloth local agent tool once with the complete task. In plan mode, call " + "the read-only Unsloth plan agent instead. Return its result to the user without " + "claiming that the cloud parent completed the local work.\n", encoding = "utf-8", ) return plugin def _codex_subagent_flags(path: Path) -> list[str]: - config_path = _agent_config_path(path, ["codex"]) - return [ - "--enable", - "multi_agent", - "-c", - "agents.max_depth=1", - "-c", - f"agents.{_SUBAGENT_NAME}.description={json.dumps(_SUBAGENT_DESCRIPTION)}", - "-c", - f"agents.{_SUBAGENT_NAME}.config_file={json.dumps(config_path)}", - ] + command = sys.executable + package_root = str(Path(__file__).resolve().parents[2]) + bootstrap = ( + f"import sys;sys.path.insert(0,{json.dumps(package_root)});" + f"from {_CODEX_SUBAGENT_MCP_MODULE} import main;main()" + ) + args = ["-c", bootstrap, str(path)] + if _wsl_windows_executable(["codex"]): + command = "wsl.exe" + args = [ + "-d", + os.environ["WSL_DISTRO_NAME"], + "--", + sys.executable, + "-c", + bootstrap, + str(path), + ] + server = ( + "{ " + f"command = {json.dumps(command)}, " + f"args = {json.dumps(args)}, " + f"required = true, enabled_tools = [{json.dumps(_CODEX_SUBAGENT_MCP_TOOL)}], " + 'default_tools_approval_mode = "approve", ' + "startup_timeout_sec = 15, tool_timeout_sec = 3600 }" + ) + return ["-c", f"mcp_servers.{_CODEX_SUBAGENT_MCP_SERVER}={server}"] def _wsl_windows_executable(command: list) -> Optional[str]: @@ -2647,7 +2849,7 @@ def claude( _agent_config_path(plugin, ["claude"]), # Before ctx.args: a forwarded `--` would turn later flags positional. "--allowedTools", - _CLAUDE_SUBAGENT_TOOL, + f"{_CLAUDE_SUBAGENT_TOOL},{_CLAUDE_SUBAGENT_PLAN_TOOL}", *_yolo_command_flags("claude", yolo), *ctx.args, ] @@ -2734,25 +2936,32 @@ def codex( subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant) subagent_model = {**entry, "id": subagent_id} with _session_config("codex-subagent", launch, persist = persist) as home: - agent_config = write_codex_subagent_config(base, key, subagent_model, home) - command = [ - "codex", - *_codex_subagent_flags(agent_config), - *_yolo_command_flags("codex", yolo), - *ctx.args, - ] - typer.echo( - "Unsloth is available as the `unsloth` local agent. " - "Ask Codex to spawn an Unsloth or local agent." - ) - _run( + bridge_config = write_codex_subagent_bridge( base, + key, subagent_model, - {}, - command, - launch = launch, - install_hint = "npm install -g @openai/codex", + home, + yolo = yolo, ) + with _codex_parent_overlay(home, launch = launch, persist = persist) as parent_home: + command = [ + "codex", + *_codex_subagent_flags(bridge_config), + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + typer.echo( + "Unsloth is available as a local agent. " + "Ask Codex to spawn an Unsloth or local agent." + ) + _run( + base, + subagent_model, + {"CODEX_HOME": str(parent_home)}, + command, + launch = launch, + install_hint = "npm install -g @openai/codex", + ) return command = [ "codex", diff --git a/unsloth_cli/pi_subagent.ts b/unsloth_cli/pi_subagent.ts index d712fc89ae..f4ef0c7d9e 100644 --- a/unsloth_cli/pi_subagent.ts +++ b/unsloth_cli/pi_subagent.ts @@ -7,6 +7,7 @@ import { Type } from "typebox"; const provider = "unsloth"; const maxResultCharacters = 100_000; +const maxParallelAgents = 4; const cancelGraceMilliseconds = 2_000; const configPath = process.env.UNSLOTH_PI_SUBAGENT_CONFIG || ""; delete process.env.UNSLOTH_PI_SUBAGENT_CONFIG; @@ -27,6 +28,8 @@ const baseUrl = typeof config.baseUrl === "string" ? config.baseUrl : ""; const apiKey = typeof config.apiKey === "string" ? config.apiKey : ""; const contextWindow = positiveInt(config.contextWindow, 32768); const maxTokens = positiveInt(config.maxTokens, Math.min(Math.floor(contextWindow / 4), 8192)); +let activeAgents = 0; +const waitingAgents: Array<() => boolean> = []; function positiveInt(value: unknown, fallback: number): number { const parsed = Number.parseInt(typeof value === "string" ? value : String(value || ""), 10); @@ -47,6 +50,45 @@ function boundedResult(text: string): string { return `${text.slice(0, maxResultCharacters)}\n\n[Local agent output truncated]`; } +function agentSlotRelease(): () => void { + let released = false; + return () => { + if (released) return; + released = true; + while (waitingAgents.length) { + if (waitingAgents.shift()!()) return; + } + activeAgents -= 1; + }; +} + +function acquireAgentSlot(signal: AbortSignal | undefined): Promise<() => void> { + if (signal?.aborted) return Promise.reject(new Error("The local Unsloth agent was cancelled.")); + if (activeAgents < maxParallelAgents) { + activeAgents += 1; + return Promise.resolve(agentSlotRelease()); + } + return new Promise((resolve, reject) => { + let waiting = true; + const grant = () => { + if (!waiting) return false; + waiting = false; + signal?.removeEventListener("abort", cancel); + resolve(agentSlotRelease()); + return true; + }; + const cancel = () => { + if (!waiting) return; + waiting = false; + const index = waitingAgents.indexOf(grant); + if (index >= 0) waitingAgents.splice(index, 1); + reject(new Error("The local Unsloth agent was cancelled.")); + }; + waitingAgents.push(grant); + signal?.addEventListener("abort", cancel, { once: true }); + }); +} + function piInvocation(args: string[]): { command: string; args: string[] } { const currentScript = process.argv[1]; const bunVirtualScript = currentScript?.startsWith("/$bunfs/root/"); @@ -107,6 +149,136 @@ async function stopChildTree(child: ChildProcess): Promise { signalProcessGroup(child, "SIGKILL"); } +interface LocalAgentResult { + task: string; + response: string; + transcript: any[]; + error?: string; +} + +async function runLocalAgent( + task: string, + cwd: string, + signal: AbortSignal | undefined, + onProgress: (result: LocalAgentResult) => void, +): Promise { + const extension = fileURLToPath(import.meta.url); + const args = [ + "--mode", + "json", + "--print", + "--no-session", + "--provider", + provider, + "--model", + model, + "--no-extensions", + "--extension", + extension, + `Task: ${task}`, + ]; + const invocation = piInvocation(args); + let output = ""; + let stderr = ""; + let childError = ""; + let aborted = false; + const result: LocalAgentResult = { task, response: "", transcript: [] }; + const transcriptEntries = new Set(); + const appendTranscript = (messages: any[]): boolean => { + let changed = false; + for (const message of messages) { + const entry = JSON.stringify(message); + if (transcriptEntries.has(entry)) continue; + transcriptEntries.add(entry); + result.transcript.push(message); + changed = true; + } + return changed; + }; + const processLine = (line: string) => { + try { + const event = JSON.parse(line); + if (event.type === "message_end" && event.message && appendTranscript([event.message])) { + onProgress(result); + } + if ( + event.type === "turn_end" && + Array.isArray(event.toolResults) && + event.toolResults.length && + appendTranscript(event.toolResults) + ) { + onProgress(result); + } + if (event.type !== "message_end") return; + const message = event.message; + // Pi reports model/API failures as message_end events while still + // exiting 0, so the exit status alone cannot surface them. + if (message?.stopReason === "error" || message?.stopReason === "aborted") { + childError = + (typeof message.errorMessage === "string" && message.errorMessage) || + `The local Unsloth agent stopped: ${message.stopReason}.`; + return; + } + const response = finalText(message); + if (response) { + result.response = boundedResult(response); + childError = ""; + } + } catch { + // Ignore non-JSON diagnostic lines. The exit status still reports failures. + } + }; + + const exitCode = await new Promise((resolve, reject) => { + const child = spawn(invocation.command, invocation.args, { + cwd, + detached: process.platform !== "win32", + shell: false, + stdio: ["ignore", "pipe", "pipe"], + env: { + ...process.env, + UNSLOTH_PI_SUBAGENT_CHILD: "1", + UNSLOTH_PI_SUBAGENT_CONFIG: configPath, + }, + }); + let cleanup: Promise | undefined; + const cancel = () => { + if (aborted) return; + aborted = true; + cleanup = stopChildTree(child); + }; + child.on("error", (error) => { + signal?.removeEventListener("abort", cancel); + reject(error); + }); + child.stdout.on("data", (chunk) => { + output += chunk.toString(); + const lines = output.split("\n"); + output = lines.pop() || ""; + for (const line of lines) processLine(line); + }); + child.stderr.on("data", (chunk) => { + stderr = (stderr + chunk.toString()).slice(-100_000); + }); + child.on("close", async (code) => { + signal?.removeEventListener("abort", cancel); + await cleanup; + if (output.trim()) processLine(output); + resolve(code ?? 1); + }); + signal?.addEventListener("abort", cancel, { once: true }); + if (signal?.aborted) cancel(); + }); + + if (aborted) throw new Error("The local Unsloth agent was cancelled."); + if (exitCode !== 0) { + result.error = stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`; + } + if (childError) result.error = boundedResult(childError); + if (!result.response && !result.error) result.response = "The local agent returned no text."; + return result; +} + export default function unslothSubagent(pi: ExtensionAPI): void { if (!model || !baseUrl || !apiKey || !configPath) { throw new Error("Unsloth subagent configuration is incomplete."); @@ -137,104 +309,97 @@ export default function unslothSubagent(pi: ExtensionAPI): void { name: "unsloth_agent", label: "Unsloth agent", description: - "Local coding subagent powered by Unsloth for debugging, implementation, and codebase research. Use when the user asks to spawn an Unsloth or local agent.", + "Run local coding agents powered by Unsloth for debugging, implementation, and codebase research. Use task for one agent. To run multiple independent agents, use tasks; up to four run concurrently. The tool returns only after every requested agent finishes.", parameters: Type.Object({ - task: Type.String({ description: "The complete task for the local Unsloth agent." }), + task: Type.Optional( + Type.String({ description: "The complete task for one local Unsloth agent." }), + ), + tasks: Type.Optional( + Type.Array(Type.String({ description: "A complete task for one local Unsloth agent." }), { + description: "Independent tasks to run concurrently, one local agent per task.", + minItems: 2, + maxItems: maxParallelAgents, + }), + ), }), - async execute(_toolCallId, params, signal, _onUpdate, ctx) { - const extension = fileURLToPath(import.meta.url); - const args = [ - "--mode", - "json", - "--print", - "--no-session", - "--provider", - provider, - "--model", - model, - "--no-extensions", - "--extension", - extension, - `Task: ${params.task}`, - ]; - const invocation = piInvocation(args); - let output = ""; - let stderr = ""; - let lastResponse = ""; - let childError = ""; - let aborted = false; - const processLine = (line: string) => { - try { - const event = JSON.parse(line); - if (event.type !== "message_end") return; - const message = event.message; - // Pi reports model/API failures as message_end events while still - // exiting 0, so the exit status alone cannot surface them. - if (message?.stopReason === "error" || message?.stopReason === "aborted") { - childError = - (typeof message.errorMessage === "string" && message.errorMessage) || - `The local Unsloth agent stopped: ${message.stopReason}.`; - return; - } - const response = finalText(message); - if (response) { - lastResponse = boundedResult(response); - childError = ""; - } - } catch { - // Ignore non-JSON diagnostic lines. The exit status still reports failures. - } - }; - - const exitCode = await new Promise((resolve, reject) => { - const child = spawn(invocation.command, invocation.args, { - cwd: ctx.cwd, - detached: process.platform !== "win32", - shell: false, - stdio: ["ignore", "pipe", "pipe"], - env: { - ...process.env, - UNSLOTH_PI_SUBAGENT_CHILD: "1", - UNSLOTH_PI_SUBAGENT_CONFIG: configPath, - }, - }); - let cleanup: Promise | undefined; - const cancel = () => { - if (aborted) return; - aborted = true; - cleanup = stopChildTree(child); - }; - child.on("error", (error) => { - signal?.removeEventListener("abort", cancel); - reject(error); - }); - child.stdout.on("data", (chunk) => { - output += chunk.toString(); - const lines = output.split("\n"); - output = lines.pop() || ""; - for (const line of lines) processLine(line); - }); - child.stderr.on("data", (chunk) => { - stderr = (stderr + chunk.toString()).slice(-100_000); - }); - child.on("close", async (code) => { - signal?.removeEventListener("abort", cancel); - await cleanup; - if (output.trim()) processLine(output); - resolve(code ?? 1); - }); - signal?.addEventListener("abort", cancel, { once: true }); - if (signal?.aborted) cancel(); - }); - - if (aborted) throw new Error("The local Unsloth agent was cancelled."); - if (exitCode !== 0) { - throw new Error(stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`); + executionMode: "parallel", + async execute(_toolCallId, params, signal, onUpdate, ctx) { + const singleTask = typeof params.task === "string" && params.task.trim() ? params.task.trim() : ""; + const parallelTasks = Array.isArray(params.tasks) + ? params.tasks.map((task) => task.trim()).filter(Boolean) + : []; + if (Boolean(singleTask) === Boolean(parallelTasks.length)) { + throw new Error("Provide exactly one of task or tasks."); } - if (childError) throw new Error(boundedResult(childError)); + if (parallelTasks.length > maxParallelAgents) { + throw new Error(`At most ${maxParallelAgents} local agents can run concurrently.`); + } + if (parallelTasks.length === 1) { + throw new Error("Use task for one local agent, or tasks for two to four agents."); + } + + const tasks = singleTask ? [singleTask] : parallelTasks; + const results: Array = new Array(tasks.length); + let completed = 0; + const details = () => ({ + provider, + model, + mode: tasks.length === 1 ? "single" : "parallel", + results: results.filter((result): result is LocalAgentResult => Boolean(result)), + }); + const emitUpdate = () => { + onUpdate?.({ + content: [ + { + type: "text", + text: `Local agents: ${completed}/${tasks.length} completed`, + }, + ], + details: details(), + }); + }; + await Promise.all( + tasks.map(async (task, index) => { + let releaseAgentSlot: (() => void) | undefined; + try { + releaseAgentSlot = await acquireAgentSlot(signal); + results[index] = await runLocalAgent(task, ctx.cwd, signal, (partial) => { + results[index] = partial; + emitUpdate(); + }); + } catch (error) { + results[index] = { + task, + response: "", + transcript: results[index]?.transcript || [], + error: String(error), + }; + } finally { + releaseAgentSlot?.(); + completed += 1; + emitUpdate(); + } + }), + ); + if (signal?.aborted) throw new Error("The local Unsloth agent was cancelled."); + const completedResults = results.filter( + (result): result is LocalAgentResult => Boolean(result), + ); + const succeeded = completedResults.filter((result) => !result.error).length; + const response = + completedResults.length === 1 + ? completedResults[0].error || completedResults[0].response + : [ + `Parallel: ${succeeded}/${tasks.length} local agents succeeded`, + ...completedResults.map( + (result, index) => + `\n### Agent ${index + 1}${result.error ? " failed" : ""}\n\n${result.error || result.response}`, + ), + ].join("\n"); + if (succeeded !== completedResults.length) throw new Error(response); return { - content: [{ type: "text", text: lastResponse || "The local agent returned no text." }], - details: { provider, model }, + content: [{ type: "text", text: response }], + details: details(), }; }, }); diff --git a/unsloth_cli/tests/test_claude_subagent_mcp.py b/unsloth_cli/tests/test_claude_subagent_mcp.py index 13a9bd6255..568dc76ff5 100644 --- a/unsloth_cli/tests/test_claude_subagent_mcp.py +++ b/unsloth_cli/tests/test_claude_subagent_mcp.py @@ -43,6 +43,41 @@ def test_protocol_lists_and_calls_local_agent(): } +def test_protocol_exposes_read_only_agent_for_claude_plan_mode(): + listed = bridge._response( + {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, + run_read_only_agent = lambda task: task, + read_only_tool_name = "unsloth_plan_agent", + ) + tools = {tool["name"]: tool for tool in listed["result"]["tools"]} + assert tools["unsloth_agent"]["annotations"]["readOnlyHint"] is False + assert tools["unsloth_plan_agent"]["annotations"] == { + "readOnlyHint": True, + "destructiveHint": False, + "idempotentHint": True, + "openWorldHint": True, + } + + called = bridge._response( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "unsloth_plan_agent", + "arguments": {"task": " inspect this "}, + }, + }, + run_agent = lambda task: f"write: {task}", + run_read_only_agent = lambda task: f"plan: {task}", + read_only_tool_name = "unsloth_plan_agent", + ) + assert called["result"] == { + "content": [{"type": "text", "text": "plan: inspect this"}], + "isError": False, + } + + def test_protocol_returns_tool_errors_to_parent(): response = bridge._response( { @@ -212,6 +247,38 @@ def test_local_child_uses_unsloth_without_overwriting_parent_auth( assert "CLAUDE_CODE_OAUTH_TOKEN" not in child_env +def test_read_only_local_child_uses_plan_mode(monkeypatch, tmp_path): + captured = {} + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", "1") + monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path)) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(bridge, "_claude_flags", lambda model: []) + + class Process: + pid = 1234 + returncode = 0 + + def communicate(self, timeout): + return json.dumps({"is_error": False, "result": "PLAN_OK"}), "" + + def poll(self): + return self.returncode + + def popen(command, **kwargs): + captured["command"] = command + return Process() + + monkeypatch.setattr(bridge.subprocess, "Popen", popen) + assert bridge.run_local_agent("plan this", read_only = True) == "PLAN_OK" + command = captured["command"] + assert command[command.index("--permission-mode") + 1] == "plan" + prompt = command[command.index("--append-system-prompt") + 1] + assert "read-only local coding subagent" in prompt + + def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path): monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888") monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test") diff --git a/unsloth_cli/tests/test_codex_subagent_mcp.py b/unsloth_cli/tests/test_codex_subagent_mcp.py new file mode 100644 index 0000000000..c0c97ca123 --- /dev/null +++ b/unsloth_cli/tests/test_codex_subagent_mcp.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import io +import json +import os +import subprocess + +import pytest + +import unsloth_cli.codex_subagent_mcp as bridge + + +def _write_config(tmp_path, *, bypass_permissions = False): + path = tmp_path / "subagent.json" + path.write_text( + json.dumps( + { + "api_key": "sk-unsloth-test", + "codex_home": str(tmp_path / "child"), + "bypass_permissions": bypass_permissions, + } + ) + ) + return path + + +def test_protocol_uses_codex_specific_tool_name(): + requests = "\n".join( + [ + json.dumps({"jsonrpc": "2.0", "id": 0, "method": "initialize"}), + json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}), + json.dumps( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": bridge._CODEX_SUBAGENT_MCP_TOOL, + "arguments": {"task": " inspect this "}, + }, + } + ), + ] + ) + output = io.StringIO() + bridge.serve( + io.StringIO(requests), + output, + run_agent = lambda task, cancel_event: f"completed: {task}", + tool_name = bridge._CODEX_SUBAGENT_MCP_TOOL, + tool_description = bridge._CODEX_SUBAGENT_TOOL_DESCRIPTION, + instructions = bridge._SERVER_INSTRUCTIONS, + ) + responses = { + response["id"]: response for response in map(json.loads, output.getvalue().splitlines()) + } + assert responses[0]["result"]["instructions"] == bridge._SERVER_INSTRUCTIONS + assert len(bridge._SERVER_INSTRUCTIONS) <= 512 + assert responses[1]["result"]["tools"][0]["name"] == "spawn_local_agent" + assert ( + "Use this tool instead of the built-in spawn_agent tool" + in responses[1]["result"]["tools"][0]["description"] + ) + assert responses[1]["result"]["tools"][0]["annotations"]["destructiveHint"] is True + assert responses[2]["result"] == { + "content": [{"type": "text", "text": "completed: inspect this"}], + "isError": False, + } + + +@pytest.mark.parametrize("bypass_permissions", [False, True]) +@pytest.mark.parametrize("wsl_bridge", [False, True]) +def test_local_child_uses_explicit_unsloth_profile( + monkeypatch, tmp_path, bypass_permissions, wsl_bridge +): + config = _write_config(tmp_path, bypass_permissions = bypass_permissions) + monkeypatch.setenv(bridge._CODEX_SUBAGENT_CONFIG_ENV, str(config)) + credential_names = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN") + for name in credential_names: + monkeypatch.setenv(name, "cloud-key") + monkeypatch.setenv("CODEX_SQLITE_HOME", str(tmp_path / "parent-sqlite")) + if wsl_bridge: + monkeypatch.setattr( + bridge, + "_wsl_shim_env", + lambda command, env, unset: ( + env, + ( + bridge._CODEX_ENV_KEY, + "CODEX_HOME/p", + "CODEX_SQLITE_HOME/p", + *unset, + "PWD/p", + ), + ), + ) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = {} + + class Process: + pid = 1234 + returncode = 0 + + def communicate(self, timeout): + captured["timeout"] = timeout + return ( + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "LOCAL_OK"}, + } + ), + "", + ) + + def poll(self): + return self.returncode + + def popen(command, **kwargs): + captured["command"] = command + captured.update(kwargs) + return Process() + + monkeypatch.setattr(bridge.subprocess, "Popen", popen) + assert bridge.run_local_agent("reply exactly LOCAL_OK") == "LOCAL_OK" + command = captured["command"] + assert command[:4] == ["/usr/local/bin/codex", "--oss", "--profile", "unsloth_api"] + if bypass_permissions: + assert "--dangerously-bypass-approvals-and-sandbox" in command + else: + assert command[4:8] == ["--sandbox", "workspace-write", "--ask-for-approval", "never"] + assert command[command.index("exec") + 1 : command.index("exec") + 4] == [ + "--ephemeral", + "--json", + "--skip-git-repo-check", + ] + assert command[-1].endswith("Task: reply exactly LOCAL_OK") + assert captured["cwd"] == os.getcwd() + assert captured["stdin"] is subprocess.DEVNULL + assert captured["stdout"] is subprocess.PIPE + assert captured["stderr"] is subprocess.PIPE + if os.name == "nt": + assert captured["creationflags"] == subprocess.CREATE_NEW_PROCESS_GROUP + else: + assert captured["start_new_session"] is True + assert captured["env"]["CODEX_HOME"] == str(tmp_path / "child") + assert captured["env"]["CODEX_SQLITE_HOME"] == str(tmp_path / "child") + assert captured["env"][bridge._CODEX_ENV_KEY] == "sk-unsloth-test" + if wsl_bridge: + assert all(captured["env"][name] == "" for name in credential_names) + wslenv = captured["env"]["WSLENV"].split(":") + assert all( + name in {entry.split("/", 1)[0] for entry in wslenv} for name in bridge._CODEX_ENV_UNSET + ) + assert "CODEX_SQLITE_HOME/p" in wslenv + assert "PWD/p" in wslenv + else: + assert all(name not in captured["env"] for name in credential_names) + + +def test_local_child_returns_last_agent_message(): + output = "\n".join( + [ + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "intermediate"}, + } + ), + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "final"}, + } + ), + ] + ) + assert bridge._result_text(output) == "final" + + +def test_local_child_prioritizes_failed_turn_over_progress(): + output = "\n".join( + [ + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "still working"}, + } + ), + json.dumps({"type": "turn.failed", "error": {"message": "local failure"}}), + ] + ) + with pytest.raises(RuntimeError, match = "local failure"): + bridge._result_text(output) + + +def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path): + config = _write_config(tmp_path) + monkeypatch.setenv(bridge._CODEX_SUBAGENT_CONFIG_ENV, str(config)) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/codex") + cancel_event = bridge.threading.Event() + stopped = [] + + class Process: + pid = 1234 + returncode = None + + def communicate(self, timeout): + cancel_event.set() + raise subprocess.TimeoutExpired("codex", timeout) + + def poll(self): + return self.returncode + + process = Process() + monkeypatch.setattr(bridge.subprocess, "Popen", lambda *args, **kwargs: process) + + def stop(child): + stopped.append(child) + child.returncode = -15 + + monkeypatch.setattr(bridge, "_stop_child", stop) + with pytest.raises(RuntimeError, match = "cancelled"): + bridge.run_local_agent("wait", cancel_event) + assert stopped == [process] diff --git a/unsloth_cli/tests/test_pi_subagent.py b/unsloth_cli/tests/test_pi_subagent.py index beac6770df..0276366656 100644 --- a/unsloth_cli/tests/test_pi_subagent.py +++ b/unsloth_cli/tests/test_pi_subagent.py @@ -64,7 +64,12 @@ import {{ existsSync }} from "node:fs"; import {{ pathToFileURL }} from "node:url"; mock.module("typebox", () => ({{ - Type: {{ Object: (value) => value, String: (value) => value }}, + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, }})); test("cancellation stops the Pi child process group", async () => {{ @@ -138,10 +143,25 @@ def test_pi_child_error_events_fail_the_tool_call(tmp_path): driver = tmp_path / "pi-driver.js" driver.write_text( """ -const event = { - type: "message_end", - message: { role: "assistant", stopReason: "error", errorMessage: "backend unreachable", content: [] }, -}; +const task = process.argv.at(-1).replace(/^Task: /, ""); +const event = task === "pass" + ? { + type: "message_end", + message: { + role: "assistant", + stopReason: "stop", + content: [{ type: "text", text: "PASS_OK" }], + }, + } + : { + type: "message_end", + message: { + role: "assistant", + stopReason: "error", + errorMessage: "backend unreachable", + content: [], + }, + }; console.log(JSON.stringify(event)); """, encoding = "utf-8", @@ -154,7 +174,12 @@ import {{ expect, mock, test }} from "bun:test"; import {{ pathToFileURL }} from "node:url"; mock.module("typebox", () => ({{ - Type: {{ Object: (value) => value, String: (value) => value }}, + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, }})); test("child error events fail the tool call", async () => {{ @@ -168,14 +193,30 @@ test("child error events fail the tool call", async () => {{ registerTool(value) {{ tool = value; }}, }}); - const execution = tool.execute( + const singleExecution = tool.execute( "call", {{ task: "fail" }}, undefined, undefined, {{ cwd: {str(tmp_path)!r} }}, ); - await expect(execution).rejects.toThrow("backend unreachable"); + await expect(singleExecution).rejects.toThrow("backend unreachable"); + + const parallelExecution = tool.execute( + "call", + {{ tasks: ["pass", "fail"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const parallelError = await parallelExecution.then( + () => "", + (error) => String(error), + ); + expect(parallelError).toContain("Parallel: 1/2 local agents succeeded"); + expect(parallelError).toContain("PASS_OK"); + expect(parallelError).toContain("Agent 2 failed"); + expect(parallelError).toContain("backend unreachable"); }}, 10_000); """, encoding = "utf-8", @@ -189,3 +230,247 @@ test("child error events fail the tool call", async () => {{ ) assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script") +def test_pi_parallel_agents_run_together_and_preserve_transcripts(tmp_path): + bun = shutil.which("bun") + if bun is None: + pytest.skip("Bun is required to execute the bundled Pi extension") + + config = tmp_path / "subagent.json" + config.write_text( + json.dumps( + { + "baseUrl": "http://127.0.0.1:8000/v1", + "apiKey": "private-token", + "model": "local-model", + "contextWindow": 32768, + "maxTokens": 8192, + } + ), + encoding = "utf-8", + ) + starts = tmp_path / "starts" + driver = tmp_path / "pi-driver.js" + driver.write_text( + f""" +import * as fs from "node:fs"; + +const task = process.argv.at(-1).replace(/^Task: /, ""); +fs.appendFileSync({str(starts)!r}, `${{task}}\\n`); +for (let attempt = 0; attempt < 100; attempt++) {{ + const count = fs.readFileSync({str(starts)!r}, "utf8").trim().split("\\n").filter(Boolean).length; + if (count >= 2) break; + await Bun.sleep(20); +}} +const event = {{ + type: "message_end", + message: {{ + role: "assistant", + stopReason: "stop", + content: [{{ type: "text", text: `DONE_${{task}}` }}], + }}, +}}; +console.log(JSON.stringify(event)); +console.log(JSON.stringify({{ + type: "tool_execution_end", + toolCallId: `tool_${{task}}`, + toolName: "read", + result: {{ content: [{{ type: "text", text: `TOOL_${{task}}` }}] }}, + isError: false, +}})); +const toolResult = {{ + role: "toolResult", + toolCallId: `tool_${{task}}`, + toolName: "read", + content: [{{ type: "text", text: `TOOL_${{task}}` }}], + isError: false, +}}; +// Current Pi emits a completed tool result both as message_end and in the +// following turn_end. Preserve it once in the transcript. +console.log(JSON.stringify({{ + type: "message_end", + message: toolResult, +}})); +console.log(JSON.stringify({{ + type: "turn_end", + message: event.message, + toolResults: [toolResult], +}})); +""", + encoding = "utf-8", + ) + extension = Path(__file__).parents[1] / "pi_subagent.ts" + test_file = tmp_path / "pi-parallel.test.ts" + test_file.write_text( + f""" +import {{ expect, mock, test }} from "bun:test"; +import {{ pathToFileURL }} from "node:url"; + +mock.module("typebox", () => ({{ + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, +}})); + +test("parallel tasks launch one child each and retain their transcripts", async () => {{ + process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r}; + process.argv[1] = {str(driver)!r}; + + const loaded = await import(pathToFileURL({str(extension)!r}).href); + let tool; + loaded.default({{ + registerProvider() {{}}, + registerTool(value) {{ tool = value; }}, + }}); + + expect(tool.executionMode).toBe("parallel"); + const result = await tool.execute( + "call", + {{ tasks: ["ALPHA", "BETA"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + expect(result.content[0].text).toContain("Parallel: 2/2 local agents succeeded"); + expect(result.content[0].text).toContain("DONE_ALPHA"); + expect(result.content[0].text).toContain("DONE_BETA"); + expect(result.details.mode).toBe("parallel"); + expect(result.details.results).toHaveLength(2); + expect(result.details.results[0].transcript).toHaveLength(2); + expect(result.details.results[1].transcript).toHaveLength(2); + expect(result.details.results[0].transcript[0].content[0].text).toBe("DONE_ALPHA"); + expect(result.details.results[0].transcript[1].content[0].text).toBe("TOOL_ALPHA"); + expect(result.details.results[1].transcript[0].content[0].text).toBe("DONE_BETA"); + expect(result.details.results[1].transcript[1].content[0].text).toBe("TOOL_BETA"); +}}, 10_000); +""", + encoding = "utf-8", + ) + + completed = subprocess.run( + [bun, "test", str(test_file)], + capture_output = True, + text = True, + timeout = 15, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script") +def test_pi_parallel_agent_cap_spans_concurrent_tool_calls(tmp_path): + bun = shutil.which("bun") + if bun is None: + pytest.skip("Bun is required to execute the bundled Pi extension") + + config = tmp_path / "subagent.json" + config.write_text( + json.dumps( + { + "baseUrl": "http://127.0.0.1:8000/v1", + "apiKey": "private-token", + "model": "local-model", + "contextWindow": 32768, + "maxTokens": 8192, + } + ), + encoding = "utf-8", + ) + markers = tmp_path / "active" + markers.mkdir() + peaks = tmp_path / "peaks" + driver = tmp_path / "pi-driver.js" + driver.write_text( + f""" +import * as fs from "node:fs"; + +const task = process.argv.at(-1).replace(/^Task: /, ""); +const marker = `{str(markers)!s}/${{process.pid}}`; +fs.writeFileSync(marker, task); +await Bun.sleep(150); +fs.appendFileSync({str(peaks)!r}, `${{fs.readdirSync({str(markers)!r}).length}}\\n`); +await Bun.sleep(150); +fs.unlinkSync(marker); +console.log(JSON.stringify({{ + type: "message_end", + message: {{ + role: "assistant", + stopReason: "stop", + content: [{{ type: "text", text: `DONE_${{task}}` }}], + }}, +}})); +""", + encoding = "utf-8", + ) + extension = Path(__file__).parents[1] / "pi_subagent.ts" + test_file = tmp_path / "pi-global-cap.test.ts" + test_file.write_text( + f""" +import {{ expect, mock, test }} from "bun:test"; +import {{ pathToFileURL }} from "node:url"; + +mock.module("typebox", () => ({{ + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, +}})); + +test("concurrent tool calls share the four-agent cap", async () => {{ + process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r}; + process.argv[1] = {str(driver)!r}; + + const loaded = await import(pathToFileURL({str(extension)!r}).href); + let tool; + loaded.default({{ + registerProvider() {{}}, + registerTool(value) {{ tool = value; }}, + }}); + + const first = tool.execute( + "call-1", + {{ tasks: ["A1", "A2", "A3", "A4"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const second = tool.execute( + "call-2", + {{ tasks: ["B1", "B2", "B3", "B4"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const results = await Promise.all([first, second]); + expect(results[0].content[0].text).toContain("4/4 local agents succeeded"); + expect(results[1].content[0].text).toContain("4/4 local agents succeeded"); + const afterQueue = await tool.execute( + "call-3", + {{ task: "C" }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + expect(afterQueue.content[0].text).toContain("DONE_C"); +}}, 10_000); +""", + encoding = "utf-8", + ) + + completed = subprocess.run( + [bun, "test", str(test_file)], + capture_output = True, + text = True, + timeout = 15, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + observed = [int(value) for value in peaks.read_text().splitlines()] + assert max(observed) == 4 diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 8ce0bcc0d4..42745ef9c5 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -621,51 +621,223 @@ def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch): assert not (tmp_path / "model-catalog.json").exists() -def test_write_codex_subagent_config_keeps_parent_model_out(tmp_path, monkeypatch): +def test_write_codex_subagent_bridge_keeps_parent_credentials_out(tmp_path, monkeypatch): monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"} - path = start.write_codex_subagent_config(BASE, "private-token", local, tmp_path) - agent = _parse_toml(path.read_text()) - assert agent["name"] == "unsloth" - assert "local agent" in agent["description"].lower() - assert agent["model_provider"] == start._CODEX_PROFILE - assert agent["model"] == local["id"] - assert agent["model_context_window"] == MODEL["context_length"] - assert agent["model_providers"][start._CODEX_PROFILE] == { - "name": "Unsloth Studio", - "base_url": f"{BASE}/v1", - "wire_api": "responses", - "auth": { - "command": sys.executable, - "args": [ - "-c", - "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])", - str(tmp_path / "unsloth-auth.json"), - ], - "timeout_ms": 5000, - }, + path = start.write_codex_subagent_bridge( + BASE, + "private-token", + local, + tmp_path, + yolo = False, + ) + assert json.loads(path.read_text()) == { + "api_key": "private-token", + "codex_home": str(tmp_path / "child"), + "bypass_permissions": False, } - assert json.loads((tmp_path / "unsloth-auth.json").read_text()) == {"token": "private-token"} - catalog = json.loads((tmp_path / agent["model_catalog_json"]).read_text()) + assert path.stat().st_mode & 0o077 == 0 + profile = _parse_toml((tmp_path / "child" / "unsloth_api.config.toml").read_text()) + assert profile["model"] == local["id"] + assert profile["model_provider"] == start._CODEX_PROFILE + assert profile["model_context_window"] == MODEL["context_length"] + config = _parse_toml((tmp_path / "child" / "config.toml").read_text()) + assert config["model_providers"][start._CODEX_PROFILE]["base_url"] == f"{BASE}/v1" + catalog = json.loads((tmp_path / "child" / profile["model_catalog_json"]).read_text()) assert catalog["models"][0]["slug"] == local["id"] +def test_write_codex_parent_overlay_preserves_user_state_and_instructions(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "config.toml").write_text('model = "cloud-model"\n') + (source / "auth.json").write_text('{"auth": "cloud"}\n') + (source / "sessions").mkdir() + (source / "AGENTS.override.md").write_text("Keep my existing instructions.\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + + assert (overlay / "config.toml").read_text() == 'model = "cloud-model"\n' + assert (overlay / "auth.json").read_text() == '{"auth": "cloud"}\n' + assert (overlay / "sessions").is_dir() + instructions = (overlay / "AGENTS.override.md").read_text() + assert instructions.startswith("Keep my existing instructions.\n") + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in instructions + assert not (overlay / "AGENTS.md").exists() + assert (overlay / "AGENTS.override.md").stat().st_mode & 0o077 == 0 + assert (source / "AGENTS.override.md").read_text() == "Keep my existing instructions.\n" + + +def test_write_codex_parent_overlay_refreshes_reused_entries(tmp_path, monkeypatch): + first = tmp_path / "first-codex" + first.mkdir() + (first / "auth.json").write_text('{"auth": "old"}\n') + (first / "old-only.toml").write_text("old\n") + second = tmp_path / "second-codex" + second.mkdir() + (second / "auth.json").write_text('{"auth": "new"}\n') + overlay_path = tmp_path / "managed" / "parent" + + monkeypatch.setenv("CODEX_HOME", str(first)) + overlay = start.write_codex_parent_overlay(overlay_path) + assert (overlay / "auth.json").read_text() == '{"auth": "old"}\n' + assert (overlay / "old-only.toml").exists() + + monkeypatch.setenv("CODEX_HOME", str(second)) + overlay = start.write_codex_parent_overlay(overlay_path) + assert (overlay / "auth.json").read_text() == '{"auth": "new"}\n' + assert not (overlay / "old-only.toml").exists() + + +def test_write_codex_parent_overlay_does_not_use_itself_as_source(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "auth.json").write_text('{"auth": "cloud"}\n') + overlay_path = tmp_path / "managed" / "parent" + monkeypatch.setenv("CODEX_HOME", str(source)) + overlay = start.write_codex_parent_overlay(overlay_path) + + monkeypatch.setenv("CODEX_HOME", str(overlay)) + overlay = start.write_codex_parent_overlay(overlay_path) + + assert (overlay / "auth.json").read_text() == '{"auth": "cloud"}\n' + manifest = json.loads((overlay / start._CODEX_PARENT_OVERLAY_MANIFEST).read_text()) + assert manifest["source_home"] == str(source) + + +def test_write_codex_parent_overlay_refreshes_fallback_copies(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + config = source / "config.toml" + config.write_text('model = "first"\n') + sessions = source / "sessions" + sessions.mkdir() + (sessions / "existing.jsonl").write_text("existing session\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + + def deny_symlink(*args, **kwargs): + raise OSError("symlinks unavailable") + + monkeypatch.setattr(Path, "symlink_to", deny_symlink) + monkeypatch.setattr(start, "_create_directory_junction", lambda source, target: False) + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + (overlay / "history.jsonl").write_text("session state\n") + config.write_text('model = "second"\n') + + overlay = start.write_codex_parent_overlay(overlay) + + assert (overlay / "config.toml").read_text() == 'model = "second"\n' + assert (overlay / "sessions" / "existing.jsonl").read_text() == "existing session\n" + assert (overlay / "history.jsonl").read_text() == "session state\n" + + config.unlink() + overlay = start.write_codex_parent_overlay(overlay) + assert not (overlay / "config.toml").exists() + assert (overlay / "history.jsonl").read_text() == "session state\n" + + +def test_create_directory_junction_uses_windows_mklink(tmp_path, monkeypatch): + captured = {} + monkeypatch.setattr(start.os, "name", "nt") + + def run(command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + source = tmp_path / "source" + target = tmp_path / "target" + + assert start._create_directory_junction(source, target) is True + assert captured["command"] == [ + "cmd.exe", + "/d", + "/c", + "mklink", + "/J", + str(target), + str(source), + ] + assert captured["kwargs"] == { + "capture_output": True, + "text": True, + "timeout": 30, + "check": False, + } + + @pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") -def test_codex_subagent_auth_uses_wsl_for_windows_codex(monkeypatch, tmp_path): +def test_write_codex_parent_overlay_uses_windows_home_for_windows_codex(tmp_path, monkeypatch): + windows_profile = tmp_path / "windows-profile" + source = windows_profile / ".codex" + source.mkdir(parents = True) + (source / "auth.json").write_text('{"auth": "windows"}\n') + executable = "/mnt/c/Users/x/AppData/Roaming/npm/codex" + monkeypatch.delenv("CODEX_HOME", raising = False) + monkeypatch.delenv("USERPROFILE", raising = False) + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: executable) + + def check_output(command, **kwargs): + if command[0] == "cmd.exe": + assert kwargs["cwd"] == str(Path(executable).parent) + return r"C:\Users\x" + "\n" + assert command == ["wslpath", "-u", r"C:\Users\x"] + return str(windows_profile) + "\n" + + monkeypatch.setattr(start.subprocess, "check_output", check_output) + + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + + assert (overlay / "auth.json").read_text() == '{"auth": "windows"}\n' + + +def test_codex_parent_overlay_launch_uses_private_temp_root_and_cleans_up(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "auth.json").write_text("{}\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + agents_root = tmp_path / "agents" + monkeypatch.setattr(start, "_agents_config_root", lambda: agents_root) + + with start._codex_parent_overlay(tmp_path / "session", launch = True, persist = False) as overlay: + assert overlay.parent == agents_root / ".tmp" + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in (overlay / "AGENTS.md").read_text() + assert overlay.exists() + + assert not overlay.exists() + + +@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") +def test_codex_subagent_bridge_uses_wsl_for_windows_codex(monkeypatch, tmp_path): monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False) monkeypatch.setattr( start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex.exe", ) - - path = start.write_codex_subagent_config(BASE, "private-token", MODEL, tmp_path) - auth = _parse_toml(path.read_text())["model_providers"][start._CODEX_PROFILE]["auth"] - - assert auth["command"] == "wsl.exe" - assert auth["args"][:5] == ["-d", "Ubuntu", "--", sys.executable, "-c"] - assert auth["args"][-1] == str(tmp_path / "unsloth-auth.json") + flags = start._codex_subagent_flags(tmp_path / "subagent.json") + prefix = f"mcp_servers.{start._CODEX_SUBAGENT_MCP_SERVER}=" + override = next(value for value in flags if value.startswith(prefix)) + server = _parse_toml("server = " + override.removeprefix(prefix))["server"] + assert server["command"] == "wsl.exe" + assert server["args"] == [ + "-d", + "Ubuntu", + "--", + sys.executable, + "-c", + server["args"][5], + str(tmp_path / "subagent.json"), + ] + assert "sys.path.insert" in server["args"][5] + assert f"from {start._CODEX_SUBAGENT_MCP_MODULE} import main" in server["args"][5] + assert server["required"] is True + assert server["enabled_tools"] == [start._CODEX_SUBAGENT_MCP_TOOL] + assert server["default_tools_approval_mode"] == "approve" + assert not any(value.startswith("developer_instructions=") for value in flags) @pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") @@ -813,7 +985,7 @@ def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path "--plugin-dir", str(plugin), "--allowedTools", - start._CLAUDE_SUBAGENT_TOOL, + f"{start._CLAUDE_SUBAGENT_TOOL},{start._CLAUDE_SUBAGENT_PLAN_TOOL}", "hello", ] assert "--model" not in command @@ -841,6 +1013,7 @@ def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path } skill = (plugin / "skills" / "local-agent" / "SKILL.md").read_text() assert "spawn an Unsloth agent or local agent" in skill + assert "In plan mode" in skill assert "Ask Claude to spawn an Unsloth or local agent." in result.output @@ -1016,6 +1189,11 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch): monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) + source_home = tmp_path / "user-codex" + source_home.mkdir() + (source_home / "config.toml").write_text('model = "cloud-model"\n') + (source_home / "AGENTS.md").write_text("Keep the user's guidance.\n") + monkeypatch.setenv("CODEX_HOME", str(source_home)) result = CliRunner().invoke( start.start_app, [ @@ -1029,20 +1207,34 @@ def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, assert result.exit_code == 0, result.output command = _launch_command(result.output) assert command[0] == "codex" - assert command[1:3] == ["--enable", "multi_agent"] - assert "agents.max_depth=1" in command assert "--oss" not in command assert "--profile" not in command assert "--model" not in command - assert "CODEX_HOME" not in result.output + parent_home = tmp_path / "agents" / "codex-subagent" / "parent" + _assert_env_set(result.output, "CODEX_HOME", str(parent_home)) assert start._CODEX_ENV_KEY not in result.output assert "sk-unsloth-feedfacefeedface" not in result.output home = tmp_path / "agents" / "codex-subagent" - agent_path = home / "unsloth.toml" - agent = _parse_toml(agent_path.read_text()) - assert agent["model"] == MODEL["id"] + ":UD-Q4_K_XL" - assert "env_key" not in agent["model_providers"][start._CODEX_PROFILE] - assert f"agents.unsloth.config_file={json.dumps(str(agent_path))}" in command + bridge_path = home / "subagent.json" + bridge = json.loads(bridge_path.read_text()) + assert bridge["api_key"] == "sk-unsloth-feedfacefeedface" + assert bridge["codex_home"] == str(home / "child") + assert bridge["bypass_permissions"] is False + profile = _parse_toml((home / "child" / "unsloth_api.config.toml").read_text()) + assert profile["model"] == MODEL["id"] + ":UD-Q4_K_XL" + prefix = f"mcp_servers.{start._CODEX_SUBAGENT_MCP_SERVER}=" + override = next(value for value in command if value.startswith(prefix)) + assert override.startswith(prefix) + server = _parse_toml("server = " + override.removeprefix(prefix))["server"] + assert server["command"] == sys.executable + assert server["args"] == ["-c", server["args"][1], str(bridge_path)] + assert "sys.path.insert" in server["args"][1] + assert f"from {start._CODEX_SUBAGENT_MCP_MODULE} import main" in server["args"][1] + assert server["enabled_tools"] == [start._CODEX_SUBAGENT_MCP_TOOL] + assert not any(value.startswith("developer_instructions=") for value in command) + parent_instructions = (parent_home / "AGENTS.md").read_text() + assert parent_instructions.startswith("Keep the user's guidance.\n") + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in parent_instructions assert "Ask Codex to spawn an Unsloth or local agent." in result.output