diff --git a/examples/member-authored-workflow/README.md b/examples/member-authored-workflow/README.md new file mode 100644 index 00000000..aef56bcf --- /dev/null +++ b/examples/member-authored-workflow/README.md @@ -0,0 +1,107 @@ +# Member-authored Workflows + +A Team member can write Python Workflow source with its normal editing tools, +then execute that file through `manage_workflow`. The source is imported by a +Python process inside that member's execution environment. The entry function +receives the usual `ctx` and `args`, so it can use arbitrary Python control flow +and the existing Workflow operations. + +Run the offline example from an installed checkout. It creates a local HTTP +model, has the parent member write `generated.py`, starts that file, and returns +the child result to the parent. Every model request stays on the local machine. + +```bash +uv run python examples/member-authored-workflow/demo.py +``` + +The result includes `Local answer`, a computed total of `15`, and `90` tokens +reported by the scripted model across the parent and child. The printed +workspace retains the generated source, parent transcripts, child Workflow +manifest, usage, and result. + +Declare the tool on the role that needs it, together with `bash` and the tools +its Workflow agents will use. + +```yaml +entry: author +roles: + author: + prompt: Write and run Workflows to complete the task. + tools: [file_write, file_read, bash, manage_workflow] +``` + +`manage_workflow` shares the role's shell execution permission. Local host +execution requires the Team caller to enable `allow_unisolated_shell`. +The target environment must provide Python with OpenCollab installed. The +`python` argument selects that environment's executable and defaults to +`python3`. A child session's requested tool names must be available to its +parent member. +Its output caps, shell isolation requirement, and file-creation restriction +also carry into child sessions. + +The five actions use the same tool. `launch` returns immediately with a child +identifier belonging to the calling member. `status` observes progress, `wait` +returns a result when execution finishes, `result` reads the retained outcome, +and `stop` waits for child cleanup. An optional `timeout` on `wait` bounds that +wait and leaves the child running. + +```json +{"action": "launch", "source": "generated.py", "entry": "main", "args": {"count": 6}} +{"action": "status", "id": ""} +{"action": "wait", "id": "", "timeout": 30} +{"action": "result", "id": ""} +{"action": "stop", "id": ""} +``` + +Model credentials and supported generation settings inherit the member's +effective configuration through the child process's standard input. Child usage is charged once to the parent's token +total. Finite child budgets reserve the member's remaining budget across +simultaneously active children and the parent's subsequent model requests. +A child with an omitted budget receives half the unreserved remainder so the +parent can coordinate its runs. Unbounded parents can launch unbounded +children. An explicit child `budget` limits that child. The parent's step limit +supplies the default for child sessions. A Workflow may select a per-call step +limit with `ctx.agent_run(max_steps=...)`. + +A launch lasts until its Workflow finishes or its parent stops. Ordinary +parent completion stops any remaining children before saving the final parent +snapshot. Active child control files live in +`.opencollab/child-runs/` inside the environment. When the parent run has an +artifacts directory, it receives each child's native Workflow manifest, +Session transcripts, result, and command output. The finished child's +`artifacts` field names that persistent host directory. Parent shutdown removes +child control files before the Team captures the member's patch. +Source errors, execution errors, and output serialization errors appear in the +result. Workflow agent failures retain the existing `agent_failures` field, +including cases where the Workflow chooses to return after a failed agent. + +Cancellation first asks the child Workflow runtime to stop its sessions and +save its state. If it fails to settle, the owning Environment cancels its +command and performs its existing process cleanup. A `wait` call uses the +managed lifecycle, so an ordinary `bash` timeout does not determine how long +the Workflow can run. + +An authored Workflow inherits its member's isolation strategy. Each +`ctx.agent` or `ctx.agent_run` call can select `shared`, `worktree`, +`docker-group`, `docker-session`, or a host-registered strategy through its +`isolation` argument. See the [session isolation guide](../../docs/session-isolation.md) +for provider configuration. For example, these calls create two independent +Docker execution environments when the host configures the `docker-session` +provider. + +```python +from opencollab.tools import builtin_tools + +async def main(ctx, args): + tools = builtin_tools("bash") + first = await ctx.agent("inspect the source", tools=tools, isolation="docker-session") + second = await ctx.agent("verify the source", tools=tools, isolation="docker-session") + return first, second +``` + +The host owns the configured provider factories and each authored run's child +environments. The member's Python process sends environment operations and +model requests through its private run directory. The host calls the member's +resolved model backend and returns its native response and usage. This also +supports a member container with Docker's default `none` network. Parent +cancellation waits for active child operations and releases their environments. diff --git a/examples/member-authored-workflow/demo.py b/examples/member-authored-workflow/demo.py new file mode 100644 index 00000000..ec5c8301 --- /dev/null +++ b/examples/member-authored-workflow/demo.py @@ -0,0 +1,103 @@ +"""Run a generated Workflow against an offline HTTP model.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +import tempfile +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +from opencollab import OpenCollab + +SOURCE = '''from opencollab.workflows import workflow + +@workflow +async def main(ctx, args): + answer = await ctx.agent(args["question"], tools=[]) + return {"answer": answer, "total": sum(range(args["count"]))} +''' + + +async def run(workspace: Path) -> dict: + observed = {} + parent_step = 0 + + class Model(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + nonlocal parent_step + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + messages = body["messages"] + is_parent = messages[0].get("content") == "Offline workflow author" + tool_call = None + if is_parent: + parent_step += 1 + if parent_step == 1: + name, args = "file_write", {"path": "generated.py", "mode": "create", "content": SOURCE} + elif parent_step == 2: + name, args = "manage_workflow", {"action": "launch", "source": "generated.py", + "python": sys.executable, + "args": {"question": "Give a local answer", "count": 6}} + elif parent_step == 3: + launched = json.loads([m for m in messages if m["role"] == "tool"][-1]["content"]) + name, args = "manage_workflow", {"action": "wait", "id": launched["id"]} + else: + observed.update(json.loads([m for m in messages if m["role"] == "tool"][-1]["content"])) + if parent_step <= 3: + tool_call = {"id": f"call-{parent_step}", "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}} + message = {"role": "assistant", "content": "Collected child result" if is_parent else "Local answer"} + if tool_call is not None: + message.update(content=None, tool_calls=[tool_call]) + payload = json.dumps({ + "id": "offline", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": message, + "finish_reason": "tool_calls" if tool_call else "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + workspace.mkdir(parents=True, exist_ok=True) + team = workspace / "team.yaml" + team.write_text("entry: author\nroles:\n author:\n prompt: Offline workflow author\n" + " tools: [file_write, bash, manage_workflow]\n", encoding="utf-8") + server = ThreadingHTTPServer(("127.0.0.1", 0), Model) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + result = await OpenCollab( + workspace, model="gpt-4o", provider="openai", + api_key="offline-model", # pragma: allowlist secret + base_url=f"http://127.0.0.1:{server.server_port}/v1", + config={"thinking": False, "budget": 10_000}, + ).team( + "Write and execute your own Workflow", config=team, trace=False, + artifacts=workspace / "parent-run", use_worktrees=False, allow_unisolated_shell=True, + ) + finally: + server.shutdown() + server.server_close() + thread.join() + result.raise_for_status() + if observed.get("status") != "completed": + raise RuntimeError(f"child failed: {observed}") + return {"parent_output": result.output, "total_tokens": result.tokens, "child": observed, + "workspace": str(workspace)} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workspace", type=Path, default=None) + options = parser.parse_args() + directory = options.workspace or Path(tempfile.mkdtemp(prefix="opencollab-authored-")) + print(json.dumps(asyncio.run(run(directory.resolve())), indent=2)) diff --git a/opencollab/README.md b/opencollab/README.md index 0eb0d2d4..a53f043e 100644 --- a/opencollab/README.md +++ b/opencollab/README.md @@ -282,6 +282,11 @@ another async workflow with the same context or use `ctx.candidate_workflow` for an isolated candidate. Team spawning continues to create members within the same Team. Workflow sessions receive the tools selected by their Python workflow. +Members that write new Workflow source during a run can use `manage_workflow` +to launch the file and later inspect, wait for, or stop its child run. +The [member-authored Workflow example](../examples/member-authored-workflow/README.md) +documents its permissions, budget sharing, lifecycle, and offline demonstration. + ### Workflow authoring A workflow is a plain async Python function tagged with `@workflow`. Create diff --git a/opencollab/adapters/_env_base.py b/opencollab/adapters/_env_base.py index b92cda08..d36ab371 100644 --- a/opencollab/adapters/_env_base.py +++ b/opencollab/adapters/_env_base.py @@ -58,7 +58,7 @@ async def setup(self, mount_dir: str | None = None) -> str: raise ValueError("mount_dir is supported only by container environments") return self.workspace - async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> ExecResult: + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> ExecResult: raise NotImplementedError async def read_file(self, path: str) -> str: diff --git a/opencollab/adapters/_env_docker.py b/opencollab/adapters/_env_docker.py index ffa57a54..49d30d55 100644 --- a/opencollab/adapters/_env_docker.py +++ b/opencollab/adapters/_env_docker.py @@ -548,8 +548,8 @@ async def _exec( ) return decoded - async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> ExecResult: - return await self._exec(cmd, timeout=timeout) + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> ExecResult: + return await self._exec(cmd, timeout=timeout, input_bytes=input_bytes) async def read_file(self, path: str) -> str: result = await self._exec( diff --git a/opencollab/adapters/_env_local.py b/opencollab/adapters/_env_local.py index 0bc50789..dc227f02 100644 --- a/opencollab/adapters/_env_local.py +++ b/opencollab/adapters/_env_local.py @@ -94,7 +94,7 @@ def _close_workspace(self) -> None: os.close(self._workspace_fd) self._workspace_fd = None - async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> ExecResult: + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> ExecResult: self._ensure_active() try: result = await run_process( @@ -104,6 +104,7 @@ async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> ExecResult: timeout=timeout, registry=self._processes, output_limit=PROCESS_OUTPUT_CAPTURE_BYTES, + input_bytes=input_bytes, ) except asyncio.TimeoutError as exc: return timed_out_result(exc, -1, timeout) diff --git a/opencollab/adapters/session_environments.py b/opencollab/adapters/session_environments.py index d4389d42..d0b769ea 100644 --- a/opencollab/adapters/session_environments.py +++ b/opencollab/adapters/session_environments.py @@ -48,6 +48,20 @@ def __init__( self._source = base_environment or LocalEnvironment(workspace) self._owned_source = self._source if base_environment is None else None + @property + def default_isolation(self) -> str: + return self._default + + @property + def strategy_names(self) -> tuple[str, ...]: + return tuple(dict.fromkeys(("shared", "worktree", *self._providers))) + + def child_pool(self, environment: EnvironmentPort, *, role: str) -> SessionEnvironmentPool: + return SessionEnvironmentPool( + self._workspace, isolation=self._roles.get(role) or self._default, + base_environment=environment, providers=self._providers, + ) + def _validate(self, name: str) -> None: isolation_name(name) if name not in {"shared", "worktree"} and name not in self._providers: diff --git a/opencollab/adapters/tools/_paths.py b/opencollab/adapters/tools/_paths.py index 45eac8b1..b71560f3 100644 --- a/opencollab/adapters/tools/_paths.py +++ b/opencollab/adapters/tools/_paths.py @@ -9,6 +9,9 @@ from __future__ import annotations +import os +import posixpath + from opencollab.application.tool_execution import ToolRuntime @@ -20,4 +23,12 @@ def checked_path(runtime: ToolRuntime, path: str) -> str: workspace — callers that prefer a soft error catch it and return a string. """ policy = runtime.safety_policy - return policy.check_path(path) if policy else path + checked = policy.check_path(path) if policy else path + environment = runtime.environment + workspace = getattr(environment, "workspace", None) + if policy and workspace is not None and not getattr(environment, "local_filesystem", False): + # Keep a remote workspace in its own namespace after host policy validation. + host_root = os.path.realpath(workspace) + if os.path.isabs(checked) and os.path.commonpath((host_root, checked)) == host_root: + checked = posixpath.join(workspace, os.path.relpath(checked, host_root)) + return checked diff --git a/opencollab/application/_session_run_completion.py b/opencollab/application/_session_run_completion.py index ae0c6990..9313cd13 100644 --- a/opencollab/application/_session_run_completion.py +++ b/opencollab/application/_session_run_completion.py @@ -693,13 +693,8 @@ def _request_output_limit( self._run_control is not None and self.hard_budget_tokens is not None ), ) - ceilings = [self.max_budget_tokens] - if self._run_control is not None: - ceilings.append(self.hard_budget_tokens) - for ceiling in ceilings: - if ceiling is None: - continue - remaining_budget = int(ceiling) - int(self.state.used_tokens) + remaining_budget = self.remaining_budget_tokens + if remaining_budget is not None: output_budget = remaining_budget - reserved_input_tokens if output_budget < minimum_output_tokens: raise _TokenBudgetStop( @@ -709,6 +704,16 @@ def _request_output_limit( max_output_tokens = min(max_output_tokens, output_budget) return max_output_tokens + @property + def remaining_budget_tokens(self) -> int | None: + ceilings = [self.max_budget_tokens] + if self._run_control is not None: + ceilings.append(self.hard_budget_tokens) + finite = [int(ceiling) for ceiling in ceilings if ceiling is not None] + if not finite: + return None + return min(finite) - int(self.state.used_tokens) - sum(reserve() for reserve in self.budget_reservations) + def _completion_accepts_output_limit(self) -> bool: """Include the default limit when the injected completion accepts it.""" try: diff --git a/opencollab/application/ports.py b/opencollab/application/ports.py index fdf3cb70..13e8234c 100644 --- a/opencollab/application/ports.py +++ b/opencollab/application/ports.py @@ -27,7 +27,7 @@ async def setup(self, mount_dir: str | None = None) -> str: """Prepare the environment and return its ready workspace identity.""" ... - async def exec_cmd(self, cmd: str, timeout: float = 120.0) -> Any: + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> Any: ... async def read_file(self, path: str) -> str: diff --git a/opencollab/application/session.py b/opencollab/application/session.py index 7805b1f8..aa6066fc 100644 --- a/opencollab/application/session.py +++ b/opencollab/application/session.py @@ -7,7 +7,7 @@ import threading from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Callable +from typing import TYPE_CHECKING, Any, Awaitable, Callable from opencollab.application.async_timeout import await_owned_operation from opencollab.application.autosave import AutoSaveSubscriber @@ -110,6 +110,7 @@ def __init__( # cleanup, but callers can otherwise interleave a new message or a # second run with the same mutable FSM. self._turn_lock = asyncio.Lock() + self._owned_runtime_closers: list[Callable[[], Awaitable[None]]] = [] self.state = runtime.state self.event_bus = runtime.event_bus @@ -257,6 +258,7 @@ def persistence_errors(self) -> tuple[Exception, ...]: async def aclose(self) -> None: """Close the model transport when this session created and owns it.""" + await self._close_owned_runtimes() if self._llm_close_error is not None: raise self._llm_close_error if not self._owns_llm or self._llm_closed: @@ -277,6 +279,18 @@ async def aclose(self) -> None: raise self.runner.finish_run_control(lambda: self.pending_cleanup_tasks) + def own_runtime( + self, close: Callable[[], Awaitable[None]], *, reserve_tokens: Callable[[], int] | None = None, + ) -> None: + """Stop attached work before terminal persistence and transport closure.""" + self._owned_runtime_closers.append(close) + if reserve_tokens is not None: + self.runner.budget_reservations.append(reserve_tokens) + + async def _close_owned_runtimes(self) -> None: + for close in self._owned_runtime_closers: + await await_owned_operation(close(), propagate_cancellation=False) + @property def messages(self) -> list[dict]: return self.state.messages @@ -362,6 +376,13 @@ async def run_loop(self, cancel_event: asyncio.Event | None = None) -> str: primary_error = exc raise finally: + if self.state.phase.is_terminal() or primary_error is not None: + try: + await self._close_owned_runtimes() + except BaseException as cleanup_error: + if primary_error is None: + raise + add_exception_note(primary_error, f"owned runtime cleanup also failed: {cleanup_error}") self.runner._record_run_terminal() try: await self._checkpoint_terminal_snapshot() diff --git a/opencollab/application/session_run.py b/opencollab/application/session_run.py index 634cf083..f9184545 100644 --- a/opencollab/application/session_run.py +++ b/opencollab/application/session_run.py @@ -118,6 +118,7 @@ def __init__( self.tool_execution = tool_execution self.tracer = tracer self.max_budget_tokens = max_budget_tokens + self.budget_reservations: list[Callable[[], int]] = [] self._initialize_run_control(run_control, run_id) self.max_steps = max_steps self.deferrable_tool_names = deferrable_tool_names diff --git a/opencollab/bootstrap/authored_workflow.py b/opencollab/bootstrap/authored_workflow.py new file mode 100644 index 00000000..53901b17 --- /dev/null +++ b/opencollab/bootstrap/authored_workflow.py @@ -0,0 +1,338 @@ +"""Own member-authored Workflow processes through an existing Environment.""" + +from __future__ import annotations + +import asyncio +import json +import math +import posixpath +import re +import shlex +import tempfile +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from opencollab.adapters.session_environments import SessionEnvironmentPool +from opencollab.adapters.tools._paths import checked_path +from opencollab.adapters.tools.base import Tool +from opencollab.application.async_timeout import await_owned_operation +from opencollab.application.tool_execution import ToolRuntime +from opencollab.bootstrap.authored_workflow_bridge import AuthoredWorkflowBridge +from opencollab.bootstrap.tool_registry import TOOL_LIMIT_FIELDS + +if TYPE_CHECKING: + from opencollab.application.ports import EnvironmentPort + from opencollab.application.session import Session + from opencollab.bootstrap.context_builder import SpawnConfig + +_RUNNER_STOP_GRACE_SECONDS = 5.0 + + +@dataclass +class _Child: + id: str + environment: EnvironmentPort + artifacts: str + budget: int | None + task: asyncio.Task[None] | None = None + execution: asyncio.Task[Any] | None = None + stop_requested: bool = False + status: str = "running" + tokens: int = 0 + result: dict[str, Any] | None = None + cleanup_error: Exception | None = None + bridge: AuthoredWorkflowBridge | None = None + saved_artifacts: str | None = None + artifacts_removed: bool = False + + +class AuthoredWorkflowTool(Tool): + name = "manage_workflow" + description = ( + "Execute a Workflow you wrote in a Python source file in your execution environment. " + "launch returns a child id immediately. Use status, wait, stop or result with that id. " + "entry defaults to main and accepts (ctx, args). Tokens share your Team budget. " + "The environment's Python must have OpenCollab installed." + ) + disable_outer_timeout = True + parameters = { + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["launch", "status", "wait", "stop", "result"]}, + "id": {"type": "string"}, + "source": {"type": "string", "description": "Python file inside your workspace."}, + "entry": {"type": "string", "default": "main"}, + "args": {"type": "object"}, + "python": {"type": "string", "default": "python3"}, + "budget": {"type": "integer", "minimum": 1}, + "concurrency": {"type": "integer", "minimum": 1, "default": 4}, + "timeout": {"type": "number", "minimum": 0, + "description": "Optional wait duration. Expiry leaves the child running."}, + }, + "required": ["action"], + "additionalProperties": False, + } + + def __init__( + self, get_session: Callable[[], Session], config: SpawnConfig, + save_dir: str | None, *, allow_unisolated_shell: bool = False, + environment_pool_factory: Callable[..., Any] | None = None, + ) -> None: + self._get_session = get_session + self._config = config + self._save_dir = save_dir + self._allow_unisolated_shell = allow_unisolated_shell + self._environment_pool_factory = environment_pool_factory + self._children: dict[str, _Child] = {} + self._registered = False + self._launch_lock = asyncio.Lock() + + async def execute_with_runtime(self, params: dict[str, Any], runtime: ToolRuntime) -> str: + parent = self._get_session() + if not self._registered: + parent.own_runtime(self.aclose, reserve_tokens=self._reserved_tokens) + self._registered = True + action = params["action"] + if action == "launch": + async with self._launch_lock: + child = await self._launch(params, runtime) + else: + child = self._children.get(params.get("id", "")) + if child is None: + raise ValueError("unknown child run for this Team member") + if action == "stop": + await self._stop(child) + elif action == "wait": + assert child.task is not None + timeout = params.get("timeout") + if timeout is not None and (not math.isfinite(timeout) or timeout <= 0): + raise ValueError("timeout must be finite and positive") + try: + await asyncio.wait_for(asyncio.shield(child.task), timeout=timeout) + except asyncio.TimeoutError: + pass + elif action not in {"status", "result"}: + raise ValueError("unknown Workflow action") + return json.dumps(self._snapshot(child, include_result=action in {"wait", "stop", "result"}), + ensure_ascii=False) + + def _snapshot(self, child: _Child, *, include_result: bool) -> dict[str, Any]: + record = {"id": child.id, "status": child.status, "tokens": child.tokens, + "artifacts": child.saved_artifacts or child.artifacts} + if include_result: + record["result"] = child.result + return record + + def _reserved_tokens(self) -> int: + return sum(max(0, (item.budget or 0) - item.tokens) + for item in self._children.values() if item.status == "running") + + async def _launch(self, params: dict[str, Any], runtime: ToolRuntime) -> _Child: + parent = self._get_session() + environment = runtime.environment + if environment is None or environment.revoked: + raise RuntimeError("no active execution environment available") + shell = parent.agent.find_tool("bash") + if shell is None: + raise PermissionError("manage_workflow source execution requires the member's bash permission") + isolated = bool(getattr(environment, "process_isolated", False)) + if (not self._allow_unisolated_shell or getattr(shell, "require_process_isolation", False)) and not isolated: + raise PermissionError("manage_workflow requires an OS process sandbox or allow_unisolated_shell") + if parent.runner._run_cancellation_requested() or parent.state.phase.is_terminal(): + raise RuntimeError("parent session is stopping") + source = params.get("source") + if not isinstance(source, str) or not source or "\x00" in source: + raise ValueError("source must name a Python file in the workspace") + source = checked_path(runtime, source) + await environment.read_file(source) + budget = params.get("budget") + if budget is not None and (isinstance(budget, bool) or not isinstance(budget, int) or budget < 1): + raise ValueError("budget must be a positive integer") + remaining = parent.runner.remaining_budget_tokens + if remaining is not None: + if remaining <= 0: + raise ValueError("Team session token budget is exhausted or reserved by child runs") + budget = max(1, remaining // 2) if budget is None else min(budget, remaining) + child_id = f"child-{uuid.uuid4().hex}" + artifacts = posixpath.join(environment.workspace, ".opencollab", "child-runs", child_id) + child = _Child(child_id, environment, artifacts, budget) + config = {key: value for key, value in {**vars(self._config), **vars(parent.agent)}.items() + if key in {"model", "provider", "wire_protocol", "api_key", "auth_token", "base_url", + "temperature", "top_p", "thinking", "thinking_params", "reasoning_effort", + "context_window", "llm_timeout", "llm_max_retries", "llm_connect_timeout", + "llm_first_event_timeout", "llm_stream_idle_timeout", "llm_stream_chat", + "provider_error_time_budget", "thinking_params_explicit", "model_routes", + "request_pool", "request_capacity", "logical_model"}} + config["max_output_tokens"] = parent.agent.max_tokens_per_step + tool_limits = { + tool.name: {key: getattr(tool, key) + for key in (*TOOL_LIMIT_FIELDS.get(tool.name, ()), "allow_create", "require_process_isolation") + if hasattr(tool, key)} + for tool in parent.agent.tools + } + request = {"id": child_id, "source": source, "entry": params.get("entry", "main"), + "args": params.get("args", {}), "config": config, "budget": budget, + "max_steps": parent.max_steps, "concurrency": params.get("concurrency", 4), + "allowed_tools": [tool.name for tool in parent.agent.tools], + "tool_limits": tool_limits, + "process_isolated": isolated, + "artifacts": artifacts} + command = " ".join([shlex.quote(params.get("python", "python3")), "-m", + "opencollab.bootstrap.authored_workflow_runner"]) + if runtime.safety_policy is not None: + await runtime.safety_policy.check_cmd_interactive(command, runtime.confirm_fn()) + pool = (self._environment_pool_factory(environment, role=parent.agent.name) + if self._environment_pool_factory is not None else SessionEnvironmentPool( + environment.workspace, isolation="shared", base_environment=environment, + )) + child.bridge = AuthoredWorkflowBridge( + environment, artifacts, pool, config, runtime, + allow_unisolated_shell=self._allow_unisolated_shell + and not getattr(shell, "require_process_isolation", False), + ) + request["bridge"] = child.bridge.options() + request_bytes = json.dumps(request, allow_nan=False).encode("utf-8") + self._children[child_id] = child + child.task = asyncio.create_task(self._execute(child, command, request_bytes)) + return child + + async def _refresh_usage(self, child: _Child) -> None: + try: + usage = json.loads(await child.environment.read_file(posixpath.join(child.artifacts, "usage.json"))) + except FileNotFoundError: + return + self._account_usage(child, usage["tokens"]) + + def _account_usage(self, child: _Child, tokens: Any) -> None: + if isinstance(tokens, bool) or not isinstance(tokens, int) or tokens < child.tokens: + raise ValueError("invalid child usage total") + delta = tokens - child.tokens + if delta: + parent = self._get_session() + parent.state.add_used_tokens(delta) + parent.enqueue_auto_save() + child.tokens = tokens + + async def _execute(self, child: _Child, command: str, request_bytes: bytes) -> None: + execution = asyncio.create_task(child.environment.exec_cmd( + command, timeout=math.inf, input_bytes=request_bytes, + )) + child.execution = execution + if child.stop_requested: + execution.cancel() + logs: dict[str, Any] = {} + status = "running" + try: + while not execution.done(): + await asyncio.wait({execution}, timeout=0.1) + assert child.bridge is not None + await child.bridge.poll() + await self._refresh_usage(child) + if self._get_session().runner._run_cancellation_requested(): + await child.environment.write_file(posixpath.join(child.artifacts, "stop"), "stop") + outcome = execution.result() + logs = {"returncode": outcome.returncode, "stdout": outcome.stdout, "stderr": outcome.stderr} + if outcome.returncode: + raise RuntimeError(f"Workflow runner exited {outcome.returncode}: {outcome.stderr[-2000:]}") + child.result = json.loads(await child.environment.read_file(posixpath.join(child.artifacts, "result.json"))) + self._account_usage(child, child.result["tokens"]) + status = child.result["status"] + except asyncio.CancelledError: + status = "stopped" + child.result = {"status": "stopped", "output": None, "reason": "cancelled"} + except Exception as exc: + status = "failed" + child.result = {"status": "failed", "output": None, "reason": f"{type(exc).__name__}: {exc}"} + finally: + if not execution.done(): + execution.cancel() + settled = await asyncio.gather(execution, return_exceptions=True) + if isinstance(settled[0], Exception): + status = "failed" + child.cleanup_error = settled[0] + child.result = {"status": "failed", "output": None, + "reason": f"child process cleanup failed: {settled[0]}"} + try: + assert child.bridge is not None + await await_owned_operation(child.bridge.aclose(), propagate_cancellation=False) + await self._refresh_usage(child) + self._account_usage(child, max(child.tokens, child.bridge.tokens)) + except Exception as exc: + status = "failed" + child.cleanup_error = exc + child.result = {"status": "failed", "output": None, + "reason": f"{(child.result or {}).get('reason', '')}; child accounting failed: {exc}"} + if child.result is not None: + child.result["tokens"] = child.tokens + child.status = status + if self._save_dir is not None: + try: + directory = Path(tempfile.mkdtemp(prefix=f"{child.id}-", dir=self._save_dir)) + native_files = (child.result or {}).get("artifact_files") + if native_files is None: + try: + native_files = json.loads(await child.environment.read_file( + posixpath.join(child.artifacts, "artifact_files.json"), + )) + except FileNotFoundError: + native_files = [] + for name in native_files: + if not isinstance(name, str) or not re.fullmatch( + r"(?:workflow\.json|\d+(?:_[^/\n]+)?\.json)(?:\.journal(?:\.lock)?)?", name, + ): + raise ValueError("invalid child artifact name") + try: + content = await child.environment.read_file(posixpath.join(child.artifacts, name)) + except FileNotFoundError: + continue + (directory / name).write_text(content, encoding="utf-8") + child.saved_artifacts = str(directory) + (directory / "result.json").write_text(json.dumps(self._snapshot(child, include_result=True)), + encoding="utf-8") + (directory / "process.json").write_text(json.dumps(logs), encoding="utf-8") + except Exception as exc: + child.status = "failed" + child.cleanup_error = exc + child.result = {"status": "failed", "output": None, + "reason": f"child result persistence failed: {exc}", "tokens": child.tokens} + + async def _stop(self, child: _Child) -> None: + if child.task is None or child.task.done(): + return + child.stop_requested = True + try: + await child.environment.write_file(posixpath.join(child.artifacts, "stop"), "stop") + await asyncio.wait_for(asyncio.shield(child.task), timeout=_RUNNER_STOP_GRACE_SECONDS) + except (OSError, RuntimeError, asyncio.TimeoutError): + if child.execution is not None and not child.execution.done(): + child.execution.cancel() + await await_owned_operation( + asyncio.gather(child.task, return_exceptions=True), propagate_cancellation=False, + ) + if child.task.cancelled(): + child.status = "stopped" + child.result = {"status": "stopped", "output": None, "reason": "cancelled", "tokens": child.tokens} + + async def aclose(self) -> None: + failures = [] + async with self._launch_lock: + for child in self._children.values(): + try: + await self._stop(child) + except Exception as exc: + failures.append(exc) + if child.cleanup_error is not None: + failures.append(child.cleanup_error) + elif not child.artifacts_removed: + try: + result = await child.environment.exec_cmd(f"rm -rf -- {shlex.quote(child.artifacts)}") + if result.returncode: + raise OSError("child runtime artifact cleanup failed") + child.artifacts_removed = True + except Exception as exc: + failures.append(exc) + if failures: + raise RuntimeError(f"member-authored Workflow cleanup failed: {failures[0]}") from failures[0] diff --git a/opencollab/bootstrap/authored_workflow_bridge.py b/opencollab/bootstrap/authored_workflow_bridge.py new file mode 100644 index 00000000..5fc2f6e7 --- /dev/null +++ b/opencollab/bootstrap/authored_workflow_bridge.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import asyncio +import base64 +import builtins +import json +import math +import posixpath +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from contextvars import ContextVar +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from opencollab.adapters._env_base import Environment, ExecResult, TextFileRange +from opencollab.adapters.llm import errors as llm_errors +from opencollab.adapters.llm.client import LLMClient +from opencollab.adapters.llm.types import LLMResponse, Usage, to_plain_data +from opencollab.adapters.tools.base import host_write_lock +from opencollab.application._session_run_usage import _normalize_completion_usage +from opencollab.application.async_timeout import await_owned_operation, consume_task_result +from opencollab.application.exception_notes import add_exception_note + +_METADATA = {"context_window", "estimate_request_tokens", "estimate_context_tokens", "minimum_output_tokens"} +_COMPLETION = {"messages", "tools", "temperature", "thinking", "thinking_params", "reasoning_effort", + "tool_choice", "top_p", "max_output_tokens", "response_session_id"} +_ENVIRONMENT = {"exec_cmd", "read_file", "read_text_range", "write_file", "write_temp_file", "remove_file"} + + +def _write_json(path: Path, value: Any) -> None: + temporary = path.with_suffix(".tmp") + temporary.write_text(json.dumps(value, ensure_ascii=False, allow_nan=False), encoding="utf-8") + temporary.replace(path) + + +def _restore_error(value: dict[str, Any]) -> Exception: + name = value.pop("type") + cls = getattr(llm_errors, name, getattr(builtins, name, RuntimeError)) + if not isinstance(cls, type) or not issubclass(cls, Exception): + cls = RuntimeError + error = cls(value.pop("message")) + for name, item in value.items(): + setattr(error, name, Usage(**item) if name == "usage" else item) + return error + + +class _FileEdit: + """Run native operations in the task holding their environment's edit lock.""" + + def __init__(self, environment: Any, path: str) -> None: + self.environment = environment + self.ready = asyncio.get_running_loop().create_future() + self.ready.add_done_callback(consume_task_result) + self.calls: asyncio.Queue[tuple[str, dict[str, Any], asyncio.Future[Any]] | None] = asyncio.Queue() + self.task = asyncio.create_task(self._run(path)) + + async def _run(self, path: str) -> None: + try: + async with host_write_lock(path, self.environment): + self.ready.set_result(None) + while (call := await self.calls.get()) is not None: + action, arguments, reply = call + if reply.cancelled(): + continue + try: + value = await getattr(self.environment, action)(**arguments) + except asyncio.CancelledError: + reply.cancel() + raise + except Exception as exc: + if not reply.done(): + reply.set_exception(exc) + else: + if not reply.done(): + reply.set_result(value) + except BaseException as exc: + if not self.ready.done(): + self.ready.set_exception(exc) + raise + finally: + while not self.calls.empty(): + call = self.calls.get_nowait() + if call is not None and not call[2].done(): + call[2].set_exception(RuntimeError("file edit lock closed")) + + async def call(self, action: str, arguments: dict[str, Any]) -> Any: + if self.task.done(): + self.task.result() + raise RuntimeError("file edit lock closed") + reply = asyncio.get_running_loop().create_future() + self.calls.put_nowait((action, arguments, reply)) + return await reply + + async def release(self) -> None: + self.calls.put_nowait(None) + await self.task + + +class AuthoredWorkflowBridge: + """Own only this authored run's strategy pools and admitted operations.""" + + def __init__(self, environment: Any, artifacts: str, pool: Any, llm: Any, runtime: Any, + *, allow_unisolated_shell: bool) -> None: + relative = posixpath.relpath(artifacts, environment.workspace) + if relative == ".." or relative.startswith("../") or posixpath.isabs(relative): + raise ValueError("child mailbox must be inside the member workspace") + self._environment = environment + self._directory = posixpath.join(artifacts, "bridge") + self._pool, self._config, self._runtime = pool, llm, runtime + self._llm = None + self._allow_unisolated_shell = allow_unisolated_shell + self._environments = {0: environment} + self._next_environment = 0 + self.tokens = 0 + self._tasks: dict[str, asyncio.Task] = {} + self._handles: dict[str, int] = {} + self._edits: dict[str, _FileEdit] = {} + self._cancelled_tasks: set[asyncio.Task[Any]] = set() + + def _cancel(self, task: asyncio.Task[Any]) -> None: + if not task.done() and task not in self._cancelled_tasks: + self._cancelled_tasks.add(task) + task.cancel() + + def _model(self) -> Any: + if self._llm is None: + aliases = {"llm_timeout": "request_timeout", "llm_max_retries": "max_retries", + "llm_connect_timeout": "connect_timeout", "llm_first_event_timeout": "first_event_timeout", + "llm_stream_idle_timeout": "stream_idle_timeout", "llm_stream_chat": "stream_chat"} + allowed = {"model", "provider", "api_key", "auth_token", "base_url", "wire_protocol", "context_window", + "provider_error_time_budget", "request_pool", "request_capacity", *aliases} + self._llm = LLMClient(**{aliases.get(key, key): value for key, value in self._config.items() + if key in allowed and (value is not None or key == "context_window")}) + return self._llm + + def options(self) -> dict[str, Any]: + return {"strategies": self._pool.strategy_names, "isolation": self._pool.default_isolation, + "metadata": tuple(_METADATA), + "requires_thinking_prefix": None, + "supports_request_admission": self._config.get("request_pool") is not None, + "supports_response_session_identity": self._config.get("wire_protocol") == "responses"} + + async def poll(self) -> None: + try: + requests = json.loads(await self._environment.read_file(posixpath.join(self._directory, "requests.json"))) + except FileNotFoundError: + return + if not isinstance(requests, dict): + raise ValueError("invalid child bridge requests") + for identity, request in requests.items(): + if not identity.isdecimal() or len(identity) > 12 or not isinstance(request, dict): + raise ValueError("invalid child bridge request identity") + handle = request.get("environment", 0) + if isinstance(handle, bool) or not isinstance(handle, int): + raise ValueError("invalid child bridge environment handle") + if identity not in self._tasks: + self._tasks[identity] = asyncio.create_task(self._respond(identity, request)) + self._handles[identity] = handle + elif request.get("cancel"): + self._cancel(self._tasks[identity]) + if identity in self._edits: + self._cancel(self._edits[identity].task) + for identity, task in tuple(self._tasks.items()): + if identity not in requests and task.done(): + task.result() + self._cancelled_tasks.discard(task) + del self._tasks[identity] + del self._handles[identity] + + async def _respond(self, identity: str, request: dict[str, Any]) -> None: + try: + if request.get("cancel"): + raise asyncio.CancelledError + payload = json.loads(await await_owned_operation(self._environment.read_file( + posixpath.join(self._directory, f"{identity}.request.json"), + ), propagate_cancellation=True)) + if not isinstance(payload, dict): + raise ValueError("bridge request payload must be an object") + reply = {"value": await self._dispatch(identity, {**payload, **request})} + except BaseException as exc: + if not isinstance(exc, (Exception, asyncio.CancelledError)): + raise + error = {"type": type(exc).__name__, "message": str(exc)} + for name in ("usage", "status_code", "code", "param", "body"): + value = getattr(exc, name, None) + if value is not None: + error[name] = to_plain_data(value) + reply = {"error": error} + await await_owned_operation(self._environment.write_file( + posixpath.join(self._directory, f"{identity}.json"), json.dumps(reply, ensure_ascii=False), + )) + + async def _dispatch(self, identity: str, request: dict[str, Any]) -> Any: + action, arguments = request["action"], request.get("arguments", {}) + if not isinstance(arguments, dict): + raise ValueError("bridge arguments must be an object") + handle = request.get("environment", 0) + if isinstance(handle, bool) or not isinstance(handle, int) or handle not in self._environments: + raise ValueError("unknown environment for this authored run") + environment = self._environments[handle] + if action == "file_write_lock": + if set(arguments) != {"path"}: + raise ValueError("file edit lock needs a path") + edit = _FileEdit(environment, arguments["path"]) + self._edits[identity] = edit + try: + await asyncio.shield(edit.ready) + except BaseException: + self._cancel(edit.task) + await await_owned_operation(asyncio.gather(edit.task, return_exceptions=True)) + del self._edits[identity] + raise + return identity + if action == "file_write_unlock": + if set(arguments) != {"id"}: + raise ValueError("file edit unlock needs its owned lock") + edit = self._edits.get(arguments["id"]) + if edit is None or edit.environment is not environment: + raise ValueError("unknown file edit lock for this environment") + await edit.release() + del self._edits[arguments["id"]] + return None + if action == "metadata": + name = arguments["name"] + if name == "requires_thinking_prefix": + return self._model().requires_thinking_prefix + if name not in _METADATA: + raise ValueError("unknown model metadata operation") + return getattr(self._model(), name)(*arguments.get("args", []), **arguments.get("kwargs", {})) + if action == "complete": + if set(arguments) - _COMPLETION: + raise ValueError("unknown completion controls") + notifications = [] + if getattr(self._model(), "supports_request_admission", False): + def admitted() -> None: + notifications.append(asyncio.create_task(self._environment.write_file( + posixpath.join(self._directory, f"{identity}.admitted"), "ready", + ))) + arguments["on_admitted"] = admitted + try: + try: + response = await self._model().complete(**arguments) + except BaseException as exc: + if getattr(exc, "usage", None) is not None: + self.tokens += _normalize_completion_usage(exc.usage)[1] + raise + self.tokens += _normalize_completion_usage(response.usage)[1] + return to_plain_data(response) + finally: + if notifications: + await await_owned_operation(asyncio.gather(*notifications), propagate_cancellation=True) + if action == "acquire": + if set(arguments) != {"strategy", "label"}: + raise ValueError("acquire accepts a strategy and label") + acquired = await self._pool.acquire(arguments["label"], source=environment, + isolation=arguments["strategy"]) + self._next_environment += 1 + self._environments[self._next_environment] = acquired + return {"handle": self._next_environment, "workspace": acquired.workspace, + "process_isolated": bool(getattr(acquired, "process_isolated", False))} + if action == "release": + if arguments or handle == 0: + raise ValueError("release needs an owned child environment") + pending = [task for key, task in self._tasks.items() if key != identity and not task.done() + and self._environments.get(self._handles[key]) is environment] + pending.extend(edit.task for edit in self._edits.values() + if edit.environment is environment and not edit.task.done()) + for task in pending: + self._cancel(task) + await asyncio.gather(*pending, return_exceptions=True) + await self._pool.release_env(environment) + del self._environments[handle] + return None + if action not in _ENVIRONMENT: + raise ValueError("unknown child environment operation") + if action == "exec_cmd": + if not self._allow_unisolated_shell and not getattr(environment, "process_isolated", False): + raise PermissionError("child command requires process isolation") + from opencollab.bootstrap.container import build_workspace_safety_policy + policy = build_workspace_safety_policy(environment) + if policy is not None: + await policy.check_cmd_interactive(arguments["cmd"], self._runtime.confirm_fn()) + if arguments.get("timeout", 120) is None: + arguments["timeout"] = math.inf + if "input_bytes" in arguments: + arguments["input_bytes"] = base64.b64decode(arguments["input_bytes"], validate=True) + edit_id = request.get("edit") + if edit_id is None: + result = await getattr(environment, action)(**arguments) + else: + edit = self._edits.get(edit_id) + if edit is None or edit.environment is not environment: + raise ValueError("unknown file edit lock for this environment") + result = await edit.call(action, arguments) + if isinstance(result, (ExecResult, TextFileRange)): + result = asdict(result) + if result.get("stdout_bytes") is not None: + result["stdout_bytes"] = base64.b64encode(result["stdout_bytes"]).decode() + return result + + async def aclose(self) -> None: + tasks = [*self._tasks.values(), *(edit.task for edit in self._edits.values())] + for task in tasks: + self._cancel(task) + results = await asyncio.gather(*tasks, return_exceptions=True) + if self._llm is not None: + await self._llm.close() + await self._pool.release() + failures = [item for item in results if isinstance(item, Exception)] + if failures: + raise OSError(f"child bridge response persistence failed: {failures[0]}") from failures[0] + + +class BridgeClient: + def __init__(self, artifacts: Path, options: dict[str, Any]) -> None: + self.directory = artifacts / "bridge" + self.directory.resolve().relative_to(Path.cwd().resolve()) + self.directory.mkdir(parents=True, exist_ok=True) + self.options = options + self._next = 0 + self._pending: dict[str, Any] = {} + + def _publish(self) -> None: + _write_json(self.directory / "requests.json", self._pending) + + def _begin(self, action: str, arguments: dict[str, Any], environment: int, edit: str | None = None) -> str: + self._next += 1 + identity = str(self._next) + payload: dict[str, Any] = {"action": action, "arguments": arguments} + if edit is not None: + payload["edit"] = edit + _write_json(self.directory / f"{identity}.request.json", payload) + self._pending[identity] = {"environment": environment} + self._publish() + return identity + + def _result(self, identity: str) -> Any: + path = self.directory / f"{identity}.json" + reply = json.loads(path.read_text(encoding="utf-8")) + del self._pending[identity] + self._publish() + (self.directory / f"{identity}.request.json").unlink(missing_ok=True) + if "error" in reply: + raise _restore_error(reply["error"]) + return reply["value"] + + def metadata(self, name: str, *args: Any, **kwargs: Any) -> Any: + identity = self._begin("metadata", {"name": name, "args": args, "kwargs": kwargs}, 0) + while not (self.directory / f"{identity}.json").exists(): + time.sleep(0.01) + return self._result(identity) + + async def call(self, action: str, arguments: dict[str, Any], environment: int = 0, + *, on_admitted: Any = None, edit: str | None = None) -> Any: + identity = self._begin(action, arguments, environment, edit) + try: + while not (self.directory / f"{identity}.json").exists(): + if on_admitted is not None and (self.directory / f"{identity}.admitted").exists(): + on_admitted() + on_admitted = None + await asyncio.sleep(0.01) + return self._result(identity) + except asyncio.CancelledError as cancellation: + self._pending[identity]["cancel"] = True + self._publish() + + async def acknowledged() -> Any: + while not (self.directory / f"{identity}.json").exists(): + await asyncio.sleep(0.01) + return self._result(identity) + + try: + value = await await_owned_operation(acknowledged(), propagate_cancellation=False) + except Exception as exc: + usage = getattr(exc, "usage", None) + else: + usage = Usage(**value["usage"]) if action == "complete" and isinstance(value, dict) else None + if action == "file_write_lock": + try: + await await_owned_operation(self.call("file_write_unlock", {"id": value}, environment)) + except Exception as exc: + add_exception_note(cancellation, f"file edit lock cleanup also failed: {exc}") + if usage is not None: + cancellation.usage = usage + raise + + def provider(self, strategy: str, **_kwargs: Any) -> Any: + client = self + + class Pool: + async def acquire(self, role: str, *, source: Any = None) -> Any: + handle = source.handle if isinstance(source, BridgeEnvironment) else 0 + result = await client.call("acquire", {"strategy": strategy, "label": role}, handle) + return BridgeEnvironment(client, result) + + async def release_env(self, environment: Any) -> None: + await environment.cleanup() + + async def release(self) -> None: + pass + + return Pool() + + +class BridgeEnvironment(Environment): + def __init__(self, client: BridgeClient, options: dict[str, Any]) -> None: + super().__init__() + self.client, self.handle = client, options["handle"] + self.workspace, self.process_isolated = options["workspace"], options["process_isolated"] + self._released = False + self._edit: ContextVar[str | None] = ContextVar("authored_file_edit", default=None) + + async def _call(self, action: str, arguments: dict[str, Any]) -> Any: + return await self.client.call(action, arguments, self.handle, edit=self._edit.get()) + + @asynccontextmanager + async def file_write_lock(self, path: str) -> AsyncIterator[None]: + identity = await self.client.call("file_write_lock", {"path": path}, self.handle) + token = self._edit.set(identity) + try: + yield None + finally: + self._edit.reset(token) + await await_owned_operation( + self.client.call("file_write_unlock", {"id": identity}, self.handle), + propagate_cancellation=True, + ) + + async def exec_cmd(self, cmd: str, timeout: float = 120.0, *, input_bytes: bytes | None = None) -> ExecResult: + arguments: dict[str, Any] = {"cmd": cmd, "timeout": timeout if math.isfinite(timeout) else None} + if input_bytes is not None: + arguments["input_bytes"] = base64.b64encode(input_bytes).decode() + result = await self._call("exec_cmd", arguments) + if result.get("stdout_bytes") is not None: + result["stdout_bytes"] = base64.b64decode(result["stdout_bytes"]) + return ExecResult(**result) + + async def read_file(self, path: str) -> str: + return await self._call("read_file", {"path": path}) + + async def read_text_range(self, path: str, **kwargs: Any) -> TextFileRange: + return TextFileRange(**await self._call("read_text_range", {"path": path, **kwargs})) + + async def write_file(self, path: str, content: str) -> None: + await self._call("write_file", {"path": path, "content": content}) + + async def write_temp_file(self, content: str, *, prefix: str, suffix: str = ".tmp") -> str: + return await self._call("write_temp_file", {"content": content, "prefix": prefix, "suffix": suffix}) + + async def remove_file(self, path: str) -> None: + await self._call("remove_file", {"path": path}) + + async def cleanup(self) -> None: + if not self._released: + async def release() -> None: + await self.client.call("release", {}, self.handle) + self._released = True + await await_owned_operation(release(), propagate_cancellation=True) + + +class BridgeLLM: + def __init__(self, client: BridgeClient, environment: Any) -> None: + self.client = client + self.handle = environment.handle if isinstance(environment, BridgeEnvironment) else 0 + self.supports_response_session_identity = client.options["supports_response_session_identity"] + self.supports_request_admission = client.options["supports_request_admission"] + + @property + def requires_thinking_prefix(self) -> bool: + return self.client.metadata("requires_thinking_prefix") + + def context_window(self) -> int | None: + return self.client.metadata("context_window") + + def estimate_request_tokens( + self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, + *, thinking: bool = False, thinking_params: dict[str, Any] | None = None, + ) -> int: + return self.client.metadata("estimate_request_tokens", messages, tools, + thinking=thinking, thinking_params=thinking_params) + + def estimate_context_tokens( + self, messages: list[dict[str, Any]], *, thinking: bool = False, + thinking_params: dict[str, Any] | None = None, + ) -> int: + return self.client.metadata("estimate_context_tokens", messages, + thinking=thinking, thinking_params=thinking_params) + + def minimum_output_tokens( + self, *, max_output_tokens: int, thinking: bool = False, + thinking_params: dict[str, Any] | None = None, + ) -> int: + return self.client.metadata("minimum_output_tokens", max_output_tokens=max_output_tokens, + thinking=thinking, thinking_params=thinking_params) + + async def complete(self, messages: list[dict[str, Any]], *, on_admitted: Any = None, **kwargs: Any) -> LLMResponse: + result = await self.client.call( + "complete", {"messages": messages, **kwargs}, self.handle, on_admitted=on_admitted, + ) + result["usage"] = Usage(**result["usage"]) + return LLMResponse(**result) + + async def close(self) -> None: + pass diff --git a/opencollab/bootstrap/authored_workflow_runner.py b/opencollab/bootstrap/authored_workflow_runner.py new file mode 100644 index 00000000..1a81c13c --- /dev/null +++ b/opencollab/bootstrap/authored_workflow_runner.py @@ -0,0 +1,160 @@ +"""Execute member-authored source in the member's Python environment.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import json +import re +import signal +import sys +from dataclasses import replace +from functools import partial +from pathlib import Path +from typing import Any + +from opencollab.adapters.env import LocalEnvironment +from opencollab.application.run_control import RunControl, RunEvent +from opencollab.application.structured_output import StructuredOutputTool +from opencollab.application.submit_findings import SubmitFindingsTool +from opencollab.bootstrap.authored_workflow_bridge import BridgeClient, BridgeLLM +from opencollab.bootstrap.session_factory import build_session +from opencollab.bootstrap.workflow_runtime import run_workflow + + +def _write_json(path: Path, value: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, ensure_ascii=False, allow_nan=False), encoding="utf-8") + temporary.replace(path) + + +async def _run(request: dict[str, Any]) -> dict[str, Any]: + artifacts = Path(request["artifacts"]) + artifacts.mkdir(parents=True, exist_ok=True) + tokens = 0 + artifact_files = ["workflow.json"] + _write_json(artifacts / "artifact_files.json", artifact_files) + bridge = BridgeClient(artifacts, request["bridge"]) if "bridge" in request else None + + def build_child(**kwargs: Any) -> Any: + forbidden = { + tool.name for tool in kwargs["agent"].tools + if not isinstance(tool, (StructuredOutputTool, SubmitFindingsTool)) + } - set(request["allowed_tools"]) + if forbidden: + raise PermissionError(f"workflow requests tools unavailable to its Team member: {sorted(forbidden)}") + for tool in kwargs["agent"].tools: + for name, ceiling in request["tool_limits"].get(tool.name, {}).items(): + if hasattr(tool, name): + if name == "require_process_isolation": + value = bool(getattr(tool, name) or ceiling) + elif name == "allow_create": + value = bool(getattr(tool, name) and ceiling) + else: + value = min(getattr(tool, name), ceiling) + setattr(tool, name, value) + control = kwargs.get("run_control") or RunControl() + + def on_event(event: RunEvent) -> None: + nonlocal tokens + if event.type == "usage": + tokens += event.data["total_tokens"] + _write_json(artifacts / "usage.json", {"tokens": tokens}) + if control.on_event is not None: + control.on_event(event) + + kwargs["run_control"] = replace(control, on_event=on_event) + if bridge is not None: + kwargs["llm"] = BridgeLLM(bridge, kwargs["env"]) + session = build_session(**kwargs) + if session.auto_save_path is not None: + name = Path(session.auto_save_path).name + artifact_files.extend((name, name + ".journal", name + ".journal.lock")) + _write_json(artifacts / "artifact_files.json", artifact_files) + return session + + async def execute() -> Any: + workspace = Path.cwd().resolve() + source = (workspace / request["source"]).resolve(strict=True) + source.relative_to(workspace) + if source.suffix != ".py": + raise ValueError("workflow source must be a Python file") + sys.dont_write_bytecode = True + spec = importlib.util.spec_from_file_location(f"_oc_child_{request['id']}", source) + if spec is None or spec.loader is None: + raise ValueError("workflow source cannot be loaded") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + sys.path.insert(0, str(source.parent)) + spec.loader.exec_module(module) + flow = getattr(module, request["entry"]) + if not callable(flow) and not callable(getattr(flow, "fn", None)): + raise TypeError("workflow entry must be callable") + environment = LocalEnvironment(str(workspace)) + environment.process_isolated = request["process_isolated"] + try: + return await run_workflow( + flow, request["args"], cfg=request["config"], workspace=str(workspace), + env=environment, cleanup_environment=True, budget=request["budget"], + max_steps=request["max_steps"], max_concurrency=request["concurrency"], + save_dir=str(artifacts), trace=False, run_id=request["id"], + session_builder=build_child, return_details=True, limit_mode="explicit", + **({"isolation": request["bridge"]["isolation"], + "environment_providers": {name: partial(bridge.provider, name) + for name in request["bridge"]["strategies"]}} + if bridge is not None else {}), + ) + finally: + await environment.cleanup() + + owner = asyncio.create_task(execute()) + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, owner.cancel) + + async def watch_stop() -> None: + while not owner.done(): + if (artifacts / "stop").exists(): + owner.cancel() + return + await asyncio.sleep(0.05) + + watcher = asyncio.create_task(watch_stop()) + result: dict[str, Any] = {"tokens": tokens} + try: + details = await owner + result.update( + status="stopped" if details.stop_reason else "completed", + output=details.output, reason=details.stop_reason, + agent_failures=details.agent_failures, + ) + except asyncio.CancelledError as exc: + notes = tuple(getattr(exc, "__notes__", ())) + result.update(status="failed" if notes else "stopped", output=None, + reason="cancelled", cleanup_errors=notes) + except Exception as exc: + result.update(status="failed", output=None, reason=f"{type(exc).__name__}: {exc}") + finally: + watcher.cancel() + await asyncio.gather(watcher, return_exceptions=True) + for sig in (signal.SIGINT, signal.SIGTERM): + loop.remove_signal_handler(sig) + result["tokens"] = tokens + result["artifact_files"] = [item.name for item in artifacts.iterdir() if item.is_file() and re.fullmatch( + r"(?:workflow\.json|\d+(?:_[^/\n]+)?\.json)(?:\.journal(?:\.lock)?)?", item.name, + )] + try: + _write_json(artifacts / "result.json", result) + except (TypeError, ValueError) as exc: + result.update(status="failed", output=None, reason=f"workflow output is not JSON serializable: {exc}") + _write_json(artifacts / "result.json", result) + return result + + +def main() -> None: + request = json.loads(sys.stdin.buffer.read()) + asyncio.run(_run(request)) + + +if __name__ == "__main__": + main() diff --git a/opencollab/bootstrap/context_builder.py b/opencollab/bootstrap/context_builder.py index 4b8db27f..39fa6dc3 100644 --- a/opencollab/bootstrap/context_builder.py +++ b/opencollab/bootstrap/context_builder.py @@ -208,6 +208,7 @@ def build_agent( *, scheduler: SchedulerPort | None = None, workflow_tool: Tool | None = None, + authored_workflow_tool: Tool | None = None, ask_user_available: bool = False, allow_unisolated_shell: bool = False, plan: ContextPlan | None = None, @@ -223,6 +224,7 @@ def build_agent( role.tools, scheduler=scheduler, workflow_tool=workflow_tool, + authored_workflow_tool=authored_workflow_tool, skill_store=self._skill_store, ask_user_available=ask_user_available, allow_unisolated_shell=allow_unisolated_shell, diff --git a/opencollab/bootstrap/scheduler_factory.py b/opencollab/bootstrap/scheduler_factory.py index 66d3e0d8..3a7e9d86 100644 --- a/opencollab/bootstrap/scheduler_factory.py +++ b/opencollab/bootstrap/scheduler_factory.py @@ -193,6 +193,7 @@ def build_scheduler( safety_policy_factory=build_workspace_safety_policy, ), team_cfg=team_cfg, + environment_pool_factory=getattr(worktree_pool, "child_pool", None), lead_workspace=ctx.workspace, lead_environment=environment, interactive=interactive, diff --git a/opencollab/bootstrap/session_factory.py b/opencollab/bootstrap/session_factory.py index 4f704310..3ae8c6c1 100644 --- a/opencollab/bootstrap/session_factory.py +++ b/opencollab/bootstrap/session_factory.py @@ -478,6 +478,7 @@ def __init__( cfg: SpawnConfig, *, team_cfg: TeamConfig | None = None, + environment_pool_factory: Callable[..., Any] | None = None, lead_workspace: str | None = None, lead_environment: Environment | None = None, interactive: bool = False, @@ -495,6 +496,7 @@ def __init__( else None ) self._team = team_cfg or default_team_config() + self._environment_pool_factory = environment_pool_factory self._workflows = None if any("run_workflow" in role.tools for role in self._team.roles.values()): from opencollab.bootstrap.workflow_runtime import discover_workflows @@ -643,6 +645,9 @@ def build_spawn_session( role, scheduler=scheduler, workflow_tool=self._workflow_tool(lambda: session), + authored_workflow_tool=self._authored_workflow_tool( + lambda: session, seated_at_start=self._prebuilt_roster, + ), # Spawned agents have no independent human to answer their questions. ask_user_available=False, allow_unisolated_shell=self._unisolated_shell_allowed( @@ -685,6 +690,17 @@ def _workflow_tool(self, get_session: Callable[[], Session]) -> Tool | None: return TeamWorkflowTool(get_session, self._workflows, self._cfg, self._save_dir) + def _authored_workflow_tool( + self, get_session: Callable[[], Session], *, seated_at_start: bool, + ) -> Tool: + from opencollab.bootstrap.authored_workflow import AuthoredWorkflowTool + + return AuthoredWorkflowTool( + get_session, self._cfg, self._save_dir, + environment_pool_factory=self._environment_pool_factory, + allow_unisolated_shell=self._unisolated_shell_allowed(seated_at_start=seated_at_start), + ) + def _role_profile(self, role_name: str) -> SingleAgentProfile | None: """The agent profile this seat runs under, or ``None`` for OpenCollab's. @@ -731,6 +747,7 @@ def create_lead_session( self._team.entry, scheduler=scheduler, workflow_tool=self._workflow_tool(lambda: session), + authored_workflow_tool=self._authored_workflow_tool(lambda: session, seated_at_start=True), ask_user_available=self._interactive, allow_unisolated_shell=self._unisolated_shell_allowed( seated_at_start=True diff --git a/opencollab/bootstrap/team_workflow.py b/opencollab/bootstrap/team_workflow.py index d85da530..82ed4e09 100644 --- a/opencollab/bootstrap/team_workflow.py +++ b/opencollab/bootstrap/team_workflow.py @@ -62,8 +62,8 @@ async def execute_with_runtime(self, params: dict[str, Any], runtime: ToolRuntim parent = self._get_session() cancel_event = parent.runner.default_cancel_event or parent.runner._run_cancel_event budget = params.get("budget") - if parent.max_budget_tokens is not None: - remaining = parent.max_budget_tokens - parent.used_tokens + remaining = parent.runner.remaining_budget_tokens + if remaining is not None: if remaining <= 0: return "Error: Team session token budget is exhausted." budget = remaining if budget is None else min(budget, remaining) diff --git a/opencollab/bootstrap/tool_registry.py b/opencollab/bootstrap/tool_registry.py index f07677f1..a51e3786 100644 --- a/opencollab/bootstrap/tool_registry.py +++ b/opencollab/bootstrap/tool_registry.py @@ -51,7 +51,7 @@ SKILL_TOOL_FACTORIES: dict[str, Callable[[SkillStorePort], Tool]] = { "use_skill": UseSkillTool, } -WORKFLOW_TOOL_NAMES = frozenset({"run_workflow"}) +WORKFLOW_TOOL_NAMES = frozenset({"run_workflow", "manage_workflow"}) KNOWN_TOOL_NAMES: frozenset[str] = ( frozenset(STATELESS_TOOL_FACTORIES) | frozenset(SCHEDULER_TOOL_FACTORIES) @@ -119,6 +119,7 @@ def build_tools_for_role( *, scheduler: SchedulerPort | None = None, workflow_tool: Tool | None = None, + authored_workflow_tool: Tool | None = None, skill_store: SkillStorePort | None = None, ask_user_available: bool = False, allow_unisolated_shell: bool = False, @@ -143,9 +144,10 @@ def build_tools_for_role( if name == "ask_user" and not ask_user_available: continue if name in WORKFLOW_TOOL_NAMES: - if workflow_tool is None: - raise ValueError("run_workflow requires a Team session") - tools.append(workflow_tool) + selected = workflow_tool if name == "run_workflow" else authored_workflow_tool + if selected is None: + raise ValueError(f"{name} requires a Team session") + tools.append(selected) elif name in STATELESS_TOOL_FACTORIES: tools.append( _instantiate( diff --git a/tests/interfaces/test_authored_workflow_isolation.py b/tests/interfaces/test_authored_workflow_isolation.py new file mode 100644 index 00000000..26ca374c --- /dev/null +++ b/tests/interfaces/test_authored_workflow_isolation.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import asyncio +import json +import os +import re +import subprocess +from pathlib import Path + +import pytest + +from opencollab import OpenCollab +from opencollab.adapters.env import DockerEnvironment +from opencollab.environments import DockerMount, DockerOptions, docker_isolation_provider +from tests.interfaces.test_member_authored_workflow import _call, _serve_model +from tests.support.paths import REPO_ROOT + + +def _containers(): + return set(subprocess.check_output([ + "docker", "ps", "--no-trunc", "-aq", "--filter", "label=opencollab.owner", + ], text=True).split()) + + +@pytest.mark.parametrize("cancel", [False, True]) +async def test_member_container_authors_workflow_with_host_owned_child_containers(tmp_path, monkeypatch, cancel): + image = os.environ.get("OPENCOLLAB_DOCKER_TEST_IMAGE") + wheels = os.environ.get("OPENCOLLAB_DOCKER_WHEEL_DIR") + if not image or not wheels: + pytest.skip("local Docker image and Linux wheel directory enable the authored container probe") + monkeypatch.setenv("OPENCOLLAB_API_USAGE_LOG", os.devnull) + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "seed.txt").write_text("source seed\n") + for args in (("init", "-q"), ("add", "."), ( + "-c", "user.name=Test", "-c", "user.email=test@example.invalid", "commit", "-qm", "initial", + )): + subprocess.run(["git", "-C", str(workspace), *args], check=True, capture_output=True) + team = tmp_path / "team.yaml" + team.write_text("entry: writer\nroles:\n writer:\n isolation: docker-group\n" + " prompt: Writer\n tools: [file_write, bash, manage_workflow]\n") + source = '''import asyncio +from opencollab.tools import builtin_tools + +async def main(ctx, args): + tools = builtin_tools("bash") + return await asyncio.gather( + ctx.agent("first worker", tools=tools), + ctx.agent("second worker", tools=tools, isolation="docker-session"), + ) +''' + observed = {} + workers = [] + owner_ids = set() + setup = DockerEnvironment.setup + + async def remember_owner(environment, *args, **kwargs): + identity = await setup(environment, *args, **kwargs) + if not environment._attached: + owner_ids.add(identity) + return identity + + monkeypatch.setattr(DockerEnvironment, "setup", remember_owner) + provider = docker_isolation_provider(DockerOptions(image=image, mounts=( + DockerMount(str(REPO_ROOT), "/oc-source"), DockerMount(str(Path(wheels).resolve()), "/oc-wheels"), + )), scope="group") + + def manager_provider(**kwargs): + pool = provider(**kwargs) + acquire = pool.acquire + + async def prepare(role, *, source=None): + environment = await acquire(role, source=source) + outcome = await environment.exec_cmd( + "python3 -m pip install --no-index --no-compile --find-links=/oc-wheels " + "openai anthropic typer rich prompt-toolkit filelock pydantic pyyaml " + ">/tmp/oc-install.log 2>&1 && " + "printf '/oc-source\\n' > /usr/local/lib/python3.12/site-packages/oc-source.pth" + ) + assert outcome.returncode == 0, (await environment.exec_cmd("cat /tmp/oc-install.log")).stdout + return environment + + pool.acquire = prepare + return pool + + child_provider = docker_isolation_provider(DockerOptions(image=image), scope="session") + + def isolated_provider(**kwargs): + pool = child_provider(**kwargs) + acquire = pool.acquire + + async def remember(role, *, source=None): + environment = await acquire(role, source=source) + workers.append(environment) + return environment + + pool.acquire = remember + return pool + + def respond(body): + replies = [item for item in body["messages"] if item["role"] == "tool"] + manager = body["messages"][0]["content"].startswith("Writer") + if manager: + if not replies: + call = _call("file_write", {"path": "authored.py", "content": source, "mode": "create"}) + elif len(replies) == 1: + call = _call("manage_workflow", {"action": "launch", "source": "authored.py"}) + elif len(replies) == 2: + call = _call("manage_workflow", {"action": "wait", "id": json.loads(replies[-1]["content"])["id"]}) + else: + observed.update(json.loads(replies[-1]["content"])) + return {"role": "assistant", "content": "collected"}, {"prompt_tokens": 3, "completion_tokens": 2} + elif not replies: + call = _call("bash", {"command": "cat seed.txt && cat /etc/hostname && pwd && " + "test ! -S /var/run/docker.sock && printf ready > entered" + + (" && sleep 300" if cancel else "")}) + else: + return ({"role": "assistant", "content": replies[-1]["content"]}, + {"prompt_tokens": 3, "completion_tokens": 2}) + return ({"role": "assistant", "content": None, "tool_calls": [call]}, + {"prompt_tokens": 3, "completion_tokens": 2}) + + with _serve_model(respond) as (endpoint, requests): + task = asyncio.create_task(OpenCollab( + workspace, model="gpt-4o", api_key="offline-test", base_url=endpoint, # pragma: allowlist secret + environment_providers={"docker-group": manager_provider, "docker-session": isolated_provider}, + config={"llm_max_retries": 0, "thinking": False}, + ).team("write and run your Workflow", config=team, isolation="docker-session", + artifacts=tmp_path / "artifacts", trace=False)) + if cancel: + async def entered(): + while len(workers) < 1: + await asyncio.sleep(0.05) + while True: + try: + if await workers[0].read_file("entered") == "ready": + return + except FileNotFoundError: + pass + await asyncio.sleep(0.05) + await asyncio.wait_for(entered(), 60) + task.cancel() + try: + cancelled = await task + except asyncio.CancelledError: + pass + else: + assert cancelled.status == "stopped", cancelled + children = list((tmp_path / "artifacts").glob("child-*/result.json")) + assert len(children) == 1 + stopped = json.loads(children[0].read_text()) + assert stopped["status"] == "stopped", stopped + assert not stopped["result"].get("cleanup_errors") + assert json.loads((children[0].parent / "workflow.json").read_text())["status"] == "stopped" + else: + result = await asyncio.wait_for(task, 90) + assert result.ok and result.output == "collected" + assert observed["status"] == "completed", observed + answers = observed["result"]["output"] + assert len(answers) == 2 and all("source seed" in answer for answer in answers) + hostnames = [re.search(r"(?m)^[0-9a-f]{12}$", answer).group() for answer in answers] + assert hostnames[0] != hostnames[1] + saved = Path(observed["artifacts"]) + assert json.loads((saved / "workflow.json").read_text())["sessions"] == 2 + assert len(list(saved.glob("[0-9]*.json"))) == 2 + patch = Path(result.metrics["entry_patch"]).read_text() + assert "child-runs" not in patch and "__pycache__" not in patch + assert all("api_key" not in path.read_text() for path in saved.glob("*.json")) + assert len(requests) >= 4 + assert owner_ids and owner_ids.isdisjoint(_containers()) + + +async def test_authored_worker_keeps_resolved_model_when_physical_name_is_an_alias(tmp_path, monkeypatch): + import sys + + from opencollab.adapters.env import LocalEnvironment + from opencollab.application.tool_execution import ToolRuntime + from tests.interfaces.test_member_authored_workflow import _manager + + monkeypatch.setenv("PYTHONPATH", str(REPO_ROOT)) + monkeypatch.setenv("OPENCOLLAB_API_USAGE_LOG", os.devnull) + (tmp_path / "authored.py").write_text( + "async def main(ctx, args):\n return await ctx.agent('child', tools=[])\n" + ) + + def respond(body): + return {"role": "assistant", "content": body["model"]}, {"prompt_tokens": 3, "completion_tokens": 2} + + with _serve_model(respond) as (endpoint, requests): + environment = LocalEnvironment(str(tmp_path)) + parent, tool, _runtime = _manager(environment) + parent.agent.model = "physical-model" + parent.agent.logical_model = "writer" + parent.agent.api_key = "offline-test" # pragma: allowlist secret + parent.agent.base_url = endpoint + tool._config.model_routes = { + "writer": {"model": "physical-model", "provider": "openai", "base_url": endpoint, + "api_key": "local"}, # pragma: allowlist secret + "physical-model": {"model": "wrong-model", "provider": "openai", "base_url": endpoint, + "api_key": "wrong"}, # pragma: allowlist secret + } + runtime = ToolRuntime(environment, None, None) + try: + launched = json.loads(await tool.execute_with_runtime({ + "action": "launch", "source": "authored.py", "python": sys.executable, + }, runtime)) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": launched["id"]}, runtime)) + assert result["status"] == "completed", result + assert result["result"]["output"] == "physical-model" + assert [request["model"] for request in requests] == ["physical-model"] + finally: + await parent.aclose() + await environment.cleanup() diff --git a/tests/interfaces/test_member_authored_workflow.py b/tests/interfaces/test_member_authored_workflow.py new file mode 100644 index 00000000..330ce25c --- /dev/null +++ b/tests/interfaces/test_member_authored_workflow.py @@ -0,0 +1,587 @@ +from __future__ import annotations + +import asyncio +import json +import os +import sys +import threading +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +from opencollab.adapters.llm.types import LLMResponse, Usage +from opencollab.adapters.tools.bash import BashTool +from opencollab.application.event_bus import EventBus +from opencollab.application.tool_execution import ToolRuntime +from opencollab.application.workflow_registry import Registry, WorkflowSpec +from opencollab.bootstrap import container +from opencollab.bootstrap.authored_workflow import AuthoredWorkflowTool +from opencollab.bootstrap.context_builder import SpawnConfig +from opencollab.bootstrap.session_factory import build_session +from opencollab.bootstrap.team_workflow import TeamWorkflowTool +from opencollab.bootstrap.workflow_runtime import run_workflow +from opencollab.domain.agent import Agent +from opencollab.sdk import OpenCollab +from tests.support.environment_test_support import local_environment_factory as local_environment_factory +from tests.support.paths import REPO_ROOT + + +@pytest.fixture(autouse=True) +def _child_import_path(monkeypatch): + monkeypatch.setenv("PYTHONPATH", str(REPO_ROOT)) + monkeypatch.setenv("OPENCOLLAB_API_USAGE_LOG", os.devnull) + + +def _call(name, args): + return {"id": "call", "type": "function", "function": {"name": name, "arguments": json.dumps(args)}} + + +@contextmanager +def _serve_model(reply): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + requests.append(body) + message, usage = reply(body) + payload = json.dumps({ + "id": "offline", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": message, "finish_reason": ( + "tool_calls" if message.get("tool_calls") else "stop" + )}], + "usage": usage, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/v1", requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +async def test_member_writes_workflow_source_and_collects_loopback_model_result( + tmp_path, monkeypatch, +): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + requests.append(body) + message = {"role": "assistant", "content": "loopback answer"} + finish = "stop" + if not any(item["role"] == "tool" for item in body["messages"]): + message.update(content=None, tool_calls=[_call("file_read", {"path": "authored.py"})]) + finish = "tool_calls" + payload = json.dumps({ + "id": "offline", "object": "chat.completion", "created": 0, "model": "gpt-4o", + "choices": [{"index": 0, "message": message, "finish_reason": finish}], + "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + observed = {} + source = '''from opencollab.workflows import workflow +from opencollab.tools import builtin_tools + +@workflow +async def main(ctx, args): + tools = builtin_tools("file_read", limits={"file_read": {"max_read_chars": 5000}}) + answer = await ctx.agent(args["question"], tools=tools) + return {"answer": answer, "number": sum(range(args["count"]))} +''' + + class ScriptedParent: + def __init__(self, **kwargs): + pass + + def context_window(self): + return 200_000 + + async def close(self): + pass + + async def complete(self, messages, **kwargs): + replies = [message for message in messages if message["role"] == "tool"] + if len(replies) == 0: + call = _call("file_write", {"path": "authored.py", "content": source, "mode": "create"}) + elif len(replies) == 1: + call = _call("manage_workflow", {"action": "launch", "source": "authored.py", + "args": {"question": "answer locally", "count": 6}, "python": sys.executable}) + elif len(replies) == 2: + launched = json.loads(replies[-1]["content"]) + assert launched["status"] == "running" + call = _call("manage_workflow", {"action": "wait", "id": launched["id"]}) + else: + observed.update(json.loads(replies[-1]["content"])) + return LLMResponse(content="collected child result", usage=Usage(input_tokens=3, output_tokens=2), + finish_reason="stop") + return LLMResponse(tool_calls=[call], usage=Usage(input_tokens=3, output_tokens=2), + finish_reason="tool_calls") + + monkeypatch.setattr(container, "LLMClient", ScriptedParent) + team = tmp_path / "team.yaml" + team.write_text("entry: writer\nroles:\n writer:\n prompt: Writer\n" + " tools: [file_write, file_read, bash, manage_workflow]\n" + "tool_limits:\n file_read:\n max_read_chars: 64\n") + try: + result = await OpenCollab(tmp_path, model="gpt-4o", api_key="offline-test", # pragma: allowlist secret + base_url=f"http://127.0.0.1:{server.server_port}/v1", + config={"budget": 10_000, "thinking": False}).team( + "write and execute your own Workflow", config=team, artifacts=tmp_path / "parent-artifacts", + trace=False, use_worktrees=False, allow_unisolated_shell=True, + ) + finally: + server.shutdown() + server.server_close() + thread.join() + assert result.output == "collected child result" + assert result.tokens == 56 + assert observed["status"] == "completed" + assert observed["result"]["output"] == {"answer": "loopback answer", "number": 15} + assert observed["tokens"] == 36 + assert len(requests) == 2 + assert requests[0]["model"] == "gpt-4o" + child_tool_reply = [item for item in requests[1]["messages"] if item["role"] == "tool"][-1]["content"] + assert "more lines below" in child_tool_reply + assert "return {" not in child_tool_reply + assert json.loads((Path(observed["artifacts"]) / "workflow.json").read_text())["status"] == "completed" + assert not list(tmp_path.glob(".oc-child-request-*")) + assert all("offline-test" not in path.read_text(errors="ignore") + for path in tmp_path.rglob("*") if path.is_file()) + + +def _manager(environment, *, budget=1000, shell=True, llm=None): + class UnusedModel: + def context_window(self): + return 200_000 + + agent = Agent(name="parent", system_prompt="parent", tools=[BashTool()] if shell else []) + config = SpawnConfig(model="gpt-4o", provider="openai", api_key=None, base_url=None, + llm_timeout=60, tracer=None, event_bus=EventBus(), permission_policy=None) + tool = AuthoredWorkflowTool(lambda: parent, config, None, allow_unisolated_shell=True) + agent.tools.append(tool) + parent = build_session(agent=agent, env=environment, max_budget_tokens=budget, llm=llm or UnusedModel()) + return parent, tool, ToolRuntime(environment, None, None) + + +async def _launch(tool, runtime, **kwargs): + return json.loads(await tool.execute_with_runtime( + {"action": "launch", "source": "authored.py", "python": sys.executable, **kwargs}, runtime, + )) + + +async def _wait_file(path): + async def wait(): + while not path.exists(): + await asyncio.sleep(0.02) + await asyncio.wait_for(wait(), 10) + + +@pytest.mark.parametrize("stop", ["explicit", "parent-close", "cancel-wait"]) +async def test_stop_cleans_child_before_returning(local_environment_factory, tmp_path, stop): + (tmp_path / "authored.py").write_text('''import asyncio +from pathlib import Path + +async def main(ctx, args): + Path("entered").write_text("ready") + try: + await asyncio.sleep(300) + finally: + Path("cleaned").write_text("yes") +''') + env = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(env) + launched = await _launch(tool, runtime) + await _wait_file(tmp_path / "entered") + waiting = asyncio.create_task(tool.execute_with_runtime({"action": "wait", "id": launched["id"]}, runtime)) + if stop == "explicit": + stopped = json.loads(await tool.execute_with_runtime({"action": "stop", "id": launched["id"]}, runtime)) + assert stopped["status"] == "stopped" + elif stop == "parent-close": + await parent.aclose() + else: + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + await parent.aclose() + if not waiting.cancelled(): + assert json.loads(await waiting)["status"] == "stopped" + assert (tmp_path / "cleaned").read_text() == "yes" + assert not env._processes._processes + assert not env.revoked + + +async def test_wait_timeout_keeps_child_and_result_stable(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text('''import asyncio +async def main(ctx, args): + await asyncio.sleep(0.5) + return {"value": 7} +''') + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path))) + child = await _launch(tool, runtime) + pending = json.loads(await tool.execute_with_runtime( + {"action": "wait", "id": child["id"], "timeout": 0.01}, runtime, + )) + assert pending["status"] == "running" and pending["result"] is None + completed = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert completed["result"]["output"] == {"value": 7} + assert json.loads(await tool.execute_with_runtime({"action": "result", "id": child["id"]}, runtime)) == completed + assert parent.used_tokens == 0 + await parent.aclose() + + +async def test_stop_grace_leaves_finished_runner_result_persistence_owned( + local_environment_factory, tmp_path, monkeypatch, +): + from opencollab.bootstrap import authored_workflow + + monkeypatch.setattr(authored_workflow, "_RUNNER_STOP_GRACE_SECONDS", 0.05) + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + environment = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(environment) + entered, finish = asyncio.Event(), asyncio.Event() + native_read = environment.read_file + read_cancelled = False + + async def delayed_result(path): + nonlocal read_cancelled + value = await native_read(path) + if path.endswith("/result.json"): + entered.set() + try: + await finish.wait() + except asyncio.CancelledError: + read_cancelled = True + raise + return value + + monkeypatch.setattr(environment, "read_file", delayed_result) + child = await _launch(tool, runtime) + stopping = None + try: + await asyncio.wait_for(entered.wait(), 10) + assert tool._children[child["id"]].execution.done() + stopping = asyncio.create_task(tool.execute_with_runtime({"action": "stop", "id": child["id"]}, runtime)) + await asyncio.sleep(0.1) + assert not stopping.done() and not read_cancelled + finish.set() + result = json.loads(await asyncio.wait_for(stopping, 2)) + assert result["status"] == "completed" and result["result"]["output"] == 1 + finally: + finish.set() + if stopping is not None: + await asyncio.gather(stopping, return_exceptions=True) + await parent.aclose() + + +async def test_child_budget_reservations_and_parent_ownership(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("import asyncio\nasync def main(ctx, args):\n await asyncio.sleep(300)\n") + env = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(env, budget=100) + first = await _launch(tool, runtime, budget=40) + second = await _launch(tool, runtime, budget=90) + assert tool._children[first["id"]].budget == 40 + assert tool._children[second["id"]].budget == 60 + with pytest.raises(ValueError, match="reserved"): + await _launch(tool, runtime) + other_parent, other_tool, other_runtime = _manager(env) + with pytest.raises(ValueError, match="unknown child"): + await other_tool.execute_with_runtime({"action": "stop", "id": first["id"]}, other_runtime) + await parent.aclose() + await other_parent.aclose() + + +@pytest.mark.parametrize("terminal", ["done", "failed", "cancelled"]) +async def test_parent_terminal_turn_stops_owned_children(local_environment_factory, tmp_path, terminal): + (tmp_path / "authored.py").write_text('''import asyncio +from pathlib import Path +async def main(ctx, args): + Path("entered").write_text("ready") + try: + await asyncio.sleep(300) + finally: + Path("cleaned").write_text("yes") +''') + + class ParentModel: + def context_window(self): + return 200_000 + + async def complete(self, messages, **kwargs): + if not any(m["role"] == "tool" for m in messages): + return LLMResponse(tool_calls=[_call("manage_workflow", { + "action": "launch", "source": "authored.py", "python": sys.executable, + })], usage=Usage(input_tokens=3, output_tokens=2), finish_reason="tool_calls") + await _wait_file(tmp_path / "entered") + if terminal == "failed": + raise ValueError("parent failed") + if terminal == "cancelled": + raise asyncio.CancelledError + return LLMResponse(content="parent finished", usage=Usage(input_tokens=3, output_tokens=2), + finish_reason="stop") + + env = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(env, budget=10_000, llm=ParentModel()) + await parent.add_user_message("launch then finish") + if terminal == "failed": + with pytest.raises(ValueError, match="parent failed"): + await parent.run_loop() + elif terminal == "cancelled": + with pytest.raises(asyncio.CancelledError): + await parent.run_loop() + else: + assert await parent.run_loop() == "parent finished" + assert (tmp_path / "cleaned").exists() + assert all(item.status == "stopped" for item in tool._children.values()) + assert not env._processes._processes + await parent.aclose() + + +async def test_unbounded_parent_can_launch_unbounded_children(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 7\n") + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path)), budget=None) + child = await _launch(tool, runtime) + assert tool._children[child["id"]].budget is None + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["result"]["output"] == 7 + await parent.aclose() + + +async def test_owned_child_cleanup_remains_idempotent_after_environment_closes(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 7\n") + environment = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(environment, budget=None) + child = await _launch(tool, runtime) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["status"] == "completed" and result["result"]["output"] == 7 + await parent.aclose() + assert not Path(child["artifacts"]).exists() + await environment.cleanup() + await parent.aclose() + + +async def test_child_reservations_limit_parent_requests_and_release_after_stop(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("import asyncio\nasync def main(ctx, args):\n await asyncio.sleep(300)\n") + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path))) + messages = [{"role": "user", "content": "brief"}] + + def limit(): + return parent.runner._request_output_limit(messages, None, thinking=False) + + before = limit() + launched = await _launch(tool, runtime, budget=400) + child = tool._children[launched["id"]] + assert limit() == before - 400 + tool._account_usage(child, 50) + assert limit() == before - 400 + tool._account_usage(child, 50) + assert parent.used_tokens == 50 + await tool.execute_with_runtime({"action": "stop", "id": child.id}, runtime) + assert limit() == before - 50 + await parent.aclose() + + +@pytest.mark.parametrize("source", [ + "async def main(ctx, args):\n raise ValueError('authored failure')\n", + "this is invalid python !", + "async def main(ctx, args):\n return object()\n", + "async def main(ctx, args):\n return float('nan')\n", +]) +async def test_authored_errors_are_explicit(local_environment_factory, tmp_path, source): + (tmp_path / "authored.py").write_text(source) + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path))) + child = await _launch(tool, runtime) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["status"] == "failed" + assert result["result"]["reason"] + await parent.aclose() + + +async def test_child_sessions_preserve_member_tool_permissions(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text( + "async def main(ctx, args):\n return await ctx.agent('read', tools='coding')\n", + ) + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path))) + child = await _launch(tool, runtime) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["result"]["output"] is None + assert result["result"]["agent_failures"][0]["exception_type"] == "PermissionError" + assert parent.used_tokens == 0 + await parent.aclose() + + +@pytest.mark.parametrize("capture", ["structured_output", "submit_findings"]) +async def test_managed_workflows_retain_internal_capture_tools(local_environment_factory, tmp_path, capture): + schema = { + "type": "object", "properties": {"value": {"type": "integer"}}, "required": ["value"], + } + options = {"schema": schema} if capture == "structured_output" else { + "enforcement_strength": "needs-enforcement", + } + payload = {"value": 7} if capture == "structured_output" else { + "findings": [], "summary": "No source evidence was requested.", "insufficient_evidence": True, + } + + def reply(body): + assert [tool["function"]["name"] for tool in body["tools"]] == [capture] + return {"role": "assistant", "content": None, "tool_calls": [_call(capture, payload)]}, { + "prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18, + } + + async def regular(ctx, args): + return await ctx.agent("Return a result", tools=[], **options) + + (tmp_path / "authored.py").write_text( + "async def main(ctx, args):\n" + " return await ctx.agent('Return a result', tools=[], **args)\n", + ) + with _serve_model(reply) as (url, requests): + config = {"model": "gpt-4o", "provider": "openai", "api_key": "offline-test", # pragma: allowlist secret + "base_url": url, "thinking": False} + control = await run_workflow( + regular, {}, cfg=config, workspace=str(tmp_path), trace=False, budget=None, + return_details=True, limit_mode="explicit", + ) + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path)), budget=None) + parent.agent.base_url = url + parent.agent.api_key = config["api_key"] + try: + child = await _launch(tool, runtime, args=options) + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["status"] == "completed" + assert result["result"]["output"] == control.output + assert result["result"]["agent_failures"] == [] + assert result["tokens"] == parent.used_tokens == control.tokens == 18 + assert len(requests) == 2 + if capture == "structured_output": + assert control.output == payload + else: + assert payload["summary"] in control.output + finally: + await parent.aclose() + + +async def test_read_only_member_cannot_execute_source(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path)), shell=False) + with pytest.raises(PermissionError, match="bash permission"): + await _launch(tool, runtime) + await parent.aclose() + + +async def test_stop_before_command_starts_releases_reservation(local_environment_factory, tmp_path, monkeypatch): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + env = local_environment_factory(str(tmp_path)) + parent, tool, runtime = _manager(env) + launched = await _launch(tool, runtime) + + async def unavailable_directory(*args): + raise FileNotFoundError("child directory is not ready") + + monkeypatch.setattr(env, "write_file", unavailable_directory) + result = json.loads(await tool.execute_with_runtime({"action": "stop", "id": launched["id"]}, runtime)) + assert result["status"] == "stopped" + assert tool._reserved_tokens() == 0 + assert not env._processes._processes + await parent.aclose() + + +async def test_member_shell_isolation_requirement_is_retained(local_environment_factory, tmp_path): + (tmp_path / "authored.py").write_text("async def main(ctx, args):\n return 1\n") + parent, tool, runtime = _manager(local_environment_factory(str(tmp_path))) + parent.agent.find_tool("bash").require_process_isolation = True + with pytest.raises(PermissionError, match="OS process sandbox"): + await _launch(tool, runtime) + await parent.aclose() + + +@pytest.mark.parametrize("budget", [1000, None]) +async def test_registered_workflow_honors_live_managed_reservations(local_environment_factory, tmp_path, budget): + (tmp_path / "authored.py").write_text('''import asyncio +from pathlib import Path +async def main(ctx, args): + Path("entered").write_text("ready") + while not Path("release").exists(): + await asyncio.sleep(0.01) + return await ctx.agent("Return a result", tools=[]) +''') + registered_limits = [] + + class RegisteredModel: + def context_window(self): + return 200_000 + + def estimate_request_tokens(self, messages, tools=None, **kwargs): + return 5 + + async def complete(self, messages, **kwargs): + registered_limits.append(kwargs["max_output_tokens"]) + return LLMResponse(content="registered answer", finish_reason="stop", usage=Usage( + input_tokens=5, output_tokens=min(695, kwargs["max_output_tokens"]), + )) + + def reply(body): + assert body["max_tokens"] >= 50 + return {"role": "assistant", "content": "managed answer"}, { + "prompt_tokens": 300, "completion_tokens": 50, "total_tokens": 350, + } + + async def regular(ctx, args): + return {"budget": ctx.budget.total, "answer": await ctx.agent("Return a result", tools=[])} + + with _serve_model(reply) as (url, requests): + parent, tool, runtime = _manager( + local_environment_factory(str(tmp_path)), budget=budget, llm=RegisteredModel(), + ) + parent.agent.base_url = url + parent.agent.api_key = "offline-test" # pragma: allowlist secret + registry = Registry() + registry.register(WorkflowSpec(name="regular", description="Return an answer", fn=regular)) + registered = TeamWorkflowTool(lambda: parent, registry, tool._config, None) + parent.agent.tools.append(registered) + try: + child = await _launch(tool, runtime, budget=400) + await _wait_file(tmp_path / "entered") + parent_limit = parent.runner._request_output_limit( + [{"role": "user", "content": "Return a result"}], None, thinking=False, + ) + value = json.loads(await registered.execute_with_runtime({"name": "regular", "args": {}}, runtime)) + assert value["output"]["budget"] == (600 if budget is not None else None) + assert value["output"]["answer"] == "registered answer" + assert registered_limits == [parent_limit] + assert tool._reserved_tokens() == 400 + (tmp_path / "release").write_text("go") + result = json.loads(await tool.execute_with_runtime({"action": "wait", "id": child["id"]}, runtime)) + assert result["result"]["output"] == "managed answer" + assert result["tokens"] == 350 + assert parent.used_tokens == value["tokens"] + result["tokens"] + assert parent.used_tokens == (950 if budget is not None else 1050) + assert len(requests) == 1 + finally: + await parent.aclose() diff --git a/tests/runtime/test_authored_workflow_bridge.py b/tests/runtime/test_authored_workflow_bridge.py new file mode 100644 index 00000000..f1ac07af --- /dev/null +++ b/tests/runtime/test_authored_workflow_bridge.py @@ -0,0 +1,453 @@ +from __future__ import annotations + +import asyncio +import os +from types import SimpleNamespace + +import pytest + +from opencollab.adapters.env import DockerEnvironment, LocalEnvironment +from opencollab.adapters.llm.client import LLMClient +from opencollab.adapters.llm.errors import TransientProviderError +from opencollab.adapters.llm.request_admission import request_slot +from opencollab.adapters.llm.types import LLMResponse, Usage +from opencollab.adapters.session_environments import SessionEnvironmentPool +from opencollab.adapters.tools.fs import FileWriteTool +from opencollab.application._session_run_shared import _TokenBudgetStop +from opencollab.application.ports import ( + ContextTokenEstimatorPort, + RequestOutputRequirementPort, + RequestTokenEstimatorPort, +) +from opencollab.application.tool_execution import ToolRuntime +from opencollab.bootstrap.authored_workflow_bridge import ( + AuthoredWorkflowBridge, + BridgeClient, + BridgeEnvironment, + BridgeLLM, +) +from opencollab.bootstrap.session_factory import build_session +from opencollab.domain.agent import Agent + + +@pytest.mark.parametrize("docker", [False, True]) +async def test_bridge_retains_native_whole_edit_lock(tmp_path, monkeypatch, docker): + image = os.environ.get("OPENCOLLAB_DOCKER_TEST_IMAGE") + if docker and not image: + pytest.skip("local Docker image enables the native edit-lock probe") + monkeypatch.chdir(tmp_path) + workspace = tmp_path / "workspace" + workspace.mkdir() + environment = DockerEnvironment(image=image) if docker else LocalEnvironment(str(workspace)) + mailbox = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + bridge = AuthoredWorkflowBridge( + mailbox, str(artifacts), SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=mailbox), + {}, ToolRuntime(mailbox, None, None), allow_unisolated_shell=True, + ) + bridge._environments[1] = environment + client = BridgeClient(artifacts, bridge.options()) + proxy = BridgeEnvironment(client, {"handle": 1, "workspace": environment.workspace, + "process_isolated": docker}) + native_read = environment.read_file + + async def snapshot_then_yield(path): + value = await native_read(path) + await asyncio.sleep(0.05) + return value + + monkeypatch.setattr(environment, "read_file", snapshot_then_yield) + stop = asyncio.Event() + + async def serve(): + while not stop.is_set(): + await bridge.poll() + await asyncio.sleep(0.005) + + server = asyncio.create_task(serve()) + try: + await environment.setup() + for target in (environment, proxy): + await environment.write_file("target.txt", "alpha\nbeta\n") + runtime = ToolRuntime(target, None, None) + replies = await asyncio.wait_for(asyncio.gather(*( + FileWriteTool().execute_with_runtime( + {"mode": "str_replace", "path": "target.txt", "old_str": word, "new_str": word.upper()}, + runtime, + ) for word in ("alpha", "beta") + )), 15) + assert all("Replaced" in reply for reply in replies) + assert await native_read("target.txt") == "ALPHA\nBETA\n" + with pytest.raises(FileNotFoundError): + async with proxy.file_write_lock("missing.txt"): + await proxy.read_file("missing.txt") + finally: + stop.set() + await server + await bridge.aclose() + await environment.cleanup() + await mailbox.cleanup() + + +async def test_cancelled_bridge_edit_settles_native_write_and_releases_waiters(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + environment = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), + SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment), + {}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + client = BridgeClient(artifacts, bridge.options()) + proxy = BridgeEnvironment(client, {"handle": 0, "workspace": environment.workspace, + "process_isolated": False}) + entered, finish, stop = asyncio.Event(), asyncio.Event(), asyncio.Event() + native_write = environment.write_file + + async def delayed_write(path, content): + if path == "target.txt" and content == "late write": + entered.set() + await finish.wait() + await native_write(path, content) + + monkeypatch.setattr(environment, "write_file", delayed_write) + + async def serve(): + while not stop.is_set(): + await bridge.poll() + await asyncio.sleep(0.005) + + async def write(): + async with proxy.file_write_lock("target.txt"): + await proxy.write_file("target.txt", "late write") + + async def queued(): + async with proxy.file_write_lock("target.txt"): + pytest.fail("cancelled waiter acquired the edit lock") + + server = asyncio.create_task(serve()) + writing = asyncio.create_task(write()) + waiting = None + try: + await asyncio.wait_for(entered.wait(), 2) + waiting = asyncio.create_task(queued()) + await asyncio.sleep(0.05) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(waiting, 2) + writing.cancel() + await asyncio.sleep(0.05) + assert not writing.done() + finish.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(writing, 2) + assert await environment.read_file("target.txt") == "late write" + async with proxy.file_write_lock("target.txt"): + await proxy.write_file("target.txt", "next edit") + assert await environment.read_file("target.txt") == "next edit" + finally: + finish.set() + for task in (writing, waiting): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(*(task for task in (writing, waiting) if task is not None), return_exceptions=True) + stop.set() + await server + await bridge.aclose() + await environment.cleanup() + + +async def test_cancelled_bridge_lock_consumes_and_releases_already_published_acquisition(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + environment = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), + SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment), + {}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + client = BridgeClient(artifacts, bridge.options()) + proxy = BridgeEnvironment(client, {"handle": 0, "workspace": environment.workspace, + "process_isolated": False}) + native_write = environment.write_file + stop = asyncio.Event() + + async def publish_then_cancel(path, content): + await native_write(path, content) + if path.endswith("/bridge/1.json"): + acquiring.cancel() + + monkeypatch.setattr(environment, "write_file", publish_then_cancel) + + async def acquire(): + async with proxy.file_write_lock("target.txt"): + pytest.fail("cancelled acquisition entered its edit") + + async def serve(): + while not stop.is_set(): + await bridge.poll() + await asyncio.sleep(0.005) + + acquiring = asyncio.create_task(acquire()) + server = asyncio.create_task(serve()) + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(acquiring, 2) + async with proxy.file_write_lock("target.txt"): + await proxy.write_file("target.txt", "next edit") + assert not bridge._edits + finally: + if not acquiring.done(): + acquiring.cancel() + await asyncio.gather(acquiring, return_exceptions=True) + stop.set() + await server + await bridge.aclose() + await environment.cleanup() + + +async def test_bridge_retains_native_budget_protocols_and_thinking_output_floor(tmp_path): + native = LLMClient(model="claude-sonnet-4-5", provider="anthropic", + api_key="offline-test") # pragma: allowlist secret + client = SimpleNamespace( + options={"supports_response_session_identity": False, "supports_request_admission": False}, + metadata=lambda name, *args, **kwargs: getattr(native, name)(*args, **kwargs), + ) + llm = BridgeLLM(client, None) + thinking_params = {"thinking": {"type": "enabled", "budget_tokens": 1024}} + environment = LocalEnvironment(str(tmp_path)) + session = build_session( + agent=Agent(name="child", system_prompt="child", tools=[], model=native.model, + provider="anthropic", max_tokens_per_step=4096, + thinking=True, thinking_params=thinking_params), + env=environment, llm=llm, max_budget_tokens=1024, + ) + messages = [ + {"role": "user", "content": "Return an answer."}, + {"role": "assistant", "content": "Prior answer", "reasoning_content": "discarded " * 1000}, + ] + try: + configured, minimum, reserved = session.runner._request_budget_requirements( + messages, None, thinking=True, + ) + assert configured == 4096 and minimum == 1025 + assert reserved == native.estimate_request_tokens(messages, thinking=True, thinking_params=thinking_params) + assert session.runner._estimate_context_message(messages[-1]) == native.estimate_context_tokens( + [messages[-1]], thinking=True, thinking_params=thinking_params, + ) + with pytest.raises(_TokenBudgetStop): + session.runner._request_output_limit(messages, None, thinking=True) + for protocol in (RequestTokenEstimatorPort, ContextTokenEstimatorPort, RequestOutputRequirementPort): + assert isinstance(native, protocol) and isinstance(llm, protocol) + finally: + await session.aclose() + await native.close() + await environment.cleanup() + + +async def test_bridge_acknowledges_request_cancelled_before_first_poll(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + environment = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), + SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment), + {}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + client = BridgeClient(artifacts, bridge.options()) + operation = asyncio.create_task(client.call("exec_cmd", {"cmd": "touch unwanted"})) + try: + await asyncio.sleep(0) + operation.cancel() + await asyncio.sleep(0) + assert client._pending["1"]["cancel"] + await bridge.poll() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(operation, 2) + assert not (tmp_path / "unwanted").exists() and not client._pending + finally: + if not operation.done(): + operation.cancel() + await bridge.aclose() + await asyncio.gather(operation, return_exceptions=True) + await environment.cleanup() + + +async def test_bridge_reads_concurrent_large_request_payloads_independently(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + environment = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), + SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment), + {}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + received = [] + + class Model: + async def complete(self, messages): + received.append(len(messages[0]["content"])) + return LLMResponse(content=str(received[-1]), usage=Usage(3, 2)) + + async def close(self): + pass + + bridge._llm = Model() + client = BridgeClient(artifacts, bridge.options()) + prompts = ["x" * (1536 * 1024), "x" * (1536 * 1024), "\u4e2d" * (512 * 1024)] + identities = [client._begin("complete", {"messages": [{"role": "user", "content": prompt}]}, 0) + for prompt in prompts] + payload_reads = [] + native_read = environment.read_file + + async def remember_read(path): + if path.endswith(".request.json"): + payload_reads.append(path) + return await native_read(path) + + monkeypatch.setattr(environment, "read_file", remember_read) + try: + payloads = list(client.directory.glob("*.request.json")) + assert len(payloads) == 3 + assert sum(path.stat().st_size for path in payloads) > 4 * 1024 * 1024 + assert (client.directory / f"{identities[-1]}.request.json").stat().st_size < 2 * 1024 * 1024 + assert (client.directory / "requests.json").stat().st_size < 1024 + await bridge.poll() + await asyncio.wait_for(asyncio.gather(*bridge._tasks.values()), 5) + await bridge.poll() + results = [client._result(identity) for identity in identities] + assert sorted(received) == sorted(map(len, prompts)) + assert [result["content"] for result in results] == [str(len(prompt)) for prompt in prompts] + assert len(payload_reads) == 3 + assert bridge.tokens == 15 and not client._pending + assert not list(client.directory.glob("*.request.json")) + finally: + await bridge.aclose() + await environment.cleanup() + + +@pytest.mark.parametrize("error_usage", [False, True]) +async def test_cancelled_bridge_call_waits_for_native_slot_and_retains_late_usage(tmp_path, monkeypatch, error_usage): + monkeypatch.chdir(tmp_path) + environment = LocalEnvironment(str(tmp_path)) + artifacts = tmp_path / "artifacts" + artifacts.mkdir() + pool_path = tmp_path / "request-pool" + pool_path.mkdir() + first_entered = asyncio.Event() + first_cancelled = asyncio.Event() + finish = asyncio.Event() + admissions = [] + entered = 0 + + class Model: + supports_request_admission = True + + async def complete(self, messages, on_admitted=None): + nonlocal entered + async with request_slot(str(pool_path), 1): + entered += 1 + index = entered + if on_admitted is not None: + on_admitted() + if index == 1: + first_entered.set() + try: + await finish.wait() + except asyncio.CancelledError: + first_cancelled.set() + await finish.wait() + if error_usage and index == 1: + error = TransientProviderError("reported provider failure") + error.usage = Usage(7, 3) + raise error + return LLMResponse(content="completed", usage=Usage(7, 3), + provider_items=[{"type": "reasoning", "encrypted_content": "opaque"}], + provider_state={"response_id": "native-response"}) + + async def close(self): + pass + + bridge = AuthoredWorkflowBridge( + environment, str(artifacts), SessionEnvironmentPool(str(tmp_path), isolation="shared", + base_environment=environment), + {"request_pool": str(pool_path)}, ToolRuntime(environment, None, None), allow_unisolated_shell=True, + ) + bridge._llm = Model() + client = BridgeClient(artifacts, bridge.options()) + llm = BridgeLLM(client, environment) + stop = asyncio.Event() + + async def serve(): + while not stop.is_set(): + await bridge.poll() + await asyncio.sleep(0.005) + + server = asyncio.create_task(serve()) + cancelled_usage = [] + + async def first_request(): + try: + return await llm.complete([], on_admitted=lambda: admissions.append("first")) + except asyncio.CancelledError as exc: + cancelled_usage.append(exc.usage.total_tokens) + raise + + first = asyncio.create_task(first_request()) + second = None + try: + await asyncio.wait_for(first_entered.wait(), 2) + first.cancel() + await asyncio.wait_for(first_cancelled.wait(), 2) + second = asyncio.create_task(llm.complete([], on_admitted=lambda: admissions.append("second"))) + await asyncio.sleep(0.1) + assert not first.done() and not second.done() + assert entered == 1 and "second" not in admissions + finish.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(first, 2) + assert cancelled_usage == [10] + response = await asyncio.wait_for(second, 2) + assert response.provider_items[0]["encrypted_content"] == "opaque" + assert response.provider_state["response_id"] == "native-response" + assert bridge.tokens == 20 and entered == 2 + finally: + finish.set() + for task in (first, second): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(*(task for task in (first, second) if task is not None), return_exceptions=True) + stop.set() + await server + await bridge.aclose() + await environment.cleanup() + + +async def test_bridge_rejects_foreign_handles_endpoint_controls_and_source_paths(tmp_path): + environment = LocalEnvironment(str(tmp_path)) + pool = SessionEnvironmentPool(str(tmp_path), isolation="shared", base_environment=environment) + bridge = AuthoredWorkflowBridge( + environment, str(tmp_path / "artifacts"), pool, {}, ToolRuntime(environment, None, None), + allow_unisolated_shell=True, + ) + try: + with pytest.raises(ValueError, match="inside the member workspace"): + AuthoredWorkflowBridge(environment, str(tmp_path.parent), pool, {}, None, allow_unisolated_shell=True) + with pytest.raises(ValueError, match="unknown environment"): + await bridge._dispatch("1", {"action": "exec_cmd", "environment": 99, + "arguments": {"cmd": "touch foreign"}}) + with pytest.raises(ValueError, match="strategy and label"): + await bridge._dispatch("2", {"action": "acquire", "arguments": { + "strategy": "worktree", "label": "child", "workspace": str(tmp_path.parent), + }}) + with pytest.raises(ValueError, match="unknown completion controls"): + await bridge._dispatch("3", {"action": "complete", "arguments": {"base_url": "http://localhost"}}) + assert bridge._llm is None and not (tmp_path / "foreign").exists() + finally: + await bridge.aclose() + await environment.cleanup() diff --git a/tests/runtime/test_docker_env.py b/tests/runtime/test_docker_env.py index dd00c3fa..d9565381 100644 --- a/tests/runtime/test_docker_env.py +++ b/tests/runtime/test_docker_env.py @@ -182,6 +182,27 @@ def respond(command, _kwargs): await env.cleanup() assert all(call[0][1] != "rm" for call in fake.calls) +async def test_command_stdin_is_forwarded_outside_argv(monkeypatch) -> None: + payload = b"private-command-input" + + def respond(command, kwargs): + if command[1] == "inspect": + return _result(stdout=f"{CONTAINER_ID}\t/swe-task\ttrue\n".encode()) + assert command[1] == "exec" + assert "-i" in command + assert payload.decode() not in " ".join(command) + assert kwargs["input_bytes"] == payload + return _result(stdout=b"received") + + fake = FakeDocker(respond) + _patch(monkeypatch, fake) + env = DockerEnvironment(container_id="swe-task", workspace="/repo") + try: + result = await env.exec_cmd("consume-input", input_bytes=payload) + assert result.stdout == "received" + finally: + await env.cleanup() + async def test_attached_binding_ignores_nonpositive_dropped_output_metadata(monkeypatch) -> None: fake = FakeDocker( lambda command, _kwargs: _result(