diff --git a/configs/README.md b/configs/README.md
index e1cdde36..4ae33415 100644
--- a/configs/README.md
+++ b/configs/README.md
@@ -62,6 +62,11 @@ environment credentials. Native bearer tokens use the `Authorization` header.
Base URLs also accept `OPENAI_BASE_URL` or `ANTHROPIC_BASE_URL` for the
selected provider after `OPENCOLLAB_BASE_URL`.
+Logical model names and request capacity shared by independent Team or Workflow
+processes are configured through [shared model routing](../docs/model-routing.md).
+The guide covers local backend URLs, API credentials, same-host request pools,
+and common gateways for multiple hosts.
+
## Wire protocol
`OPENCOLLAB_WIRE_PROTOCOL=chat_completions` is the default for OpenAI-compatible
diff --git a/docs/model-routing.md b/docs/model-routing.md
new file mode 100644
index 00000000..e70156e6
--- /dev/null
+++ b/docs/model-routing.md
@@ -0,0 +1,94 @@
+# Shared model routing and request capacity
+
+[Chinese](model-routing.zh-CN.md)
+
+A logical model name selects an explicit backend through `model_routes`. Each
+route supplies the physical model identifier, provider, endpoint and credentials.
+The runtime resolves the route before selecting model capabilities, reasoning
+controls, token estimators and context limits. Requests and provider usage records
+keep the physical model identifier. `OpenCollab.configuration` exposes the logical
+to physical mapping alongside provider, protocol, context and capacity metadata.
+
+```python
+import os
+from pathlib import Path
+
+from opencollab import OpenCollab
+
+pool = Path(os.environ["MODEL_REQUEST_POOL"]).resolve()
+pool.mkdir(parents=True, exist_ok=True)
+
+client = OpenCollab(
+ ".",
+ model="worker",
+ model_routes={
+ "worker": {
+ "model": "served-model-id",
+ "provider": "openai",
+ "wire_protocol": "chat_completions",
+ "base_url": os.environ["MODEL_BASE_URL"],
+ "api_key_env": "MODEL_API_KEY", # pragma: allowlist secret
+ "context_window": 131072,
+ "request_pool": str(pool),
+ "request_capacity": 2,
+ },
+ },
+)
+```
+
+The same mapping is accepted through `config={"model_routes": ...}`. CLI and SDK
+configuration also accept `OPENCOLLAB_MODEL_ROUTES` as a JSON object in the process
+environment or a [configuration env file](../configs/README.md). SDK constructor
+arguments take precedence. A Team role's `model` selects either a logical route or
+an ordinary physical model. Names absent from `model_routes` retain the existing
+default provider configuration. A route resolves once to its declared physical
+model, including when that identifier is also another logical name.
+
+`provider` defaults to `openai` and `wire_protocol` to `chat_completions`. Native
+Anthropic routes use `provider="anthropic"`. Route credentials follow the existing
+provider and endpoint credential selection rules. `api_key_env` selects a specific
+secret from the environment or selected env file and raises an error if unset.
+An explicit route `api_key` takes precedence. `auth_token` supports native
+Anthropic bearer authentication. Each route owns its endpoint credentials.
+`OpenCollab(...)` resolves env-file credentials at construction time. Explicit
+route `context_window` values apply to the physical model. Sampling, thinking,
+retry, timeout and output settings continue through the ordinary provider path.
+Responses requests retain their response items, provider identity and optional
+`response_session_id` handling.
+
+`request_pool` and `request_capacity` are supplied together. The pool is an
+existing absolute local directory shared by all participating processes. The
+runtime verifies the directory before acquiring a request slot. Its slot files
+hold OS locks until native completion or cancellation cleanup finishes. Queued
+cancellation releases only the waiter's resources. A terminated process releases
+its OS locks. Admission covers provider retries and releases before tools execute,
+so a parent waiting for a child leaves request capacity available to that child.
+The pool may be configured at the top level with `OPENCOLLAB_REQUEST_POOL` and
+`OPENCOLLAB_REQUEST_CAPACITY`, or the corresponding SDK configuration fields.
+A route-specific pool pair overrides that shared default. Multiple routes can
+share one pool when their requests consume the same backend capacity.
+
+The pool's `capacity` file records its configured size. Conflicting capacities
+raise an error because independent processes using different slot counts can
+oversubscribe the same backend. OS locks serialize that configuration decision.
+Keep the directory and its lock files in place while any participating process
+is running. To change capacity, stop those processes, empty the pool directory,
+then restart them with the new value. A local filesystem with working OS locks is
+required. Network filesystem locking and machines with separate filesystems
+require admission at the backend or gateway.
+
+A container must mount the same host directory at the configured absolute path.
+A missing directory produces an error explaining the mount requirement. Containers
+with independently created directories have independent capacity, even when the
+path strings match. Pool paths belong to the process running the Python model
+client, which may differ from its tools' execution environment.
+
+For multiple hosts, configure every route's `base_url` to the same externally
+managed gateway and configure its backend request capacity there. The gateway
+receives the route's physical model identifier through the existing provider wire
+protocol and owns admission across hosts. Gateway protocol support must cover
+the selected Chat Completions, Responses or Anthropic Messages path and preserve
+model identity, reasoning controls and continuation data. Backend work that
+continues after an HTTP disconnect is governed by the backend or gateway's
+cancellation policy. Local request pools coordinate client request lifetimes on
+one host.
diff --git a/docs/model-routing.zh-CN.md b/docs/model-routing.zh-CN.md
new file mode 100644
index 00000000..d7b66424
--- /dev/null
+++ b/docs/model-routing.zh-CN.md
@@ -0,0 +1,46 @@
+# 共享模型路由与请求容量
+
+[English](model-routing.md)
+
+`model_routes` 将逻辑模型名映射到明确配置的后端。每条路由包含实际模型标识、服务商、端点和凭据。运行时先解析路由，再选择模型能力、推理参数、token 估算方式和上下文限制。请求及服务商用量记录保留实际模型标识。`OpenCollab.configuration` 提供逻辑名与实际模型的对应关系，以及服务商、协议、上下文和容量信息。
+
+```python
+import os
+from pathlib import Path
+
+from opencollab import OpenCollab
+
+pool = Path(os.environ["MODEL_REQUEST_POOL"]).resolve()
+pool.mkdir(parents=True, exist_ok=True)
+
+client = OpenCollab(
+ ".",
+ model="worker",
+ model_routes={
+ "worker": {
+ "model": "served-model-id",
+ "provider": "openai",
+ "wire_protocol": "chat_completions",
+ "base_url": os.environ["MODEL_BASE_URL"],
+ "api_key_env": "MODEL_API_KEY", # pragma: allowlist secret
+ "context_window": 131072,
+ "request_pool": str(pool),
+ "request_capacity": 2,
+ },
+ },
+)
+```
+
+同一映射也可以通过 `config={"model_routes": ...}` 传入。CLI 和 SDK 均支持将 JSON 对象写入进程环境或[配置文件](../configs/README.md)中的 `OPENCOLLAB_MODEL_ROUTES`。SDK 构造参数优先。Team 角色的 `model` 可以选择逻辑路由，也可以填写普通模型标识。未出现在 `model_routes` 中的名称沿用已有的默认服务商配置。路由只解析一次，并使用其声明的实际模型，即使该标识同时也是另一个逻辑名称。
+
+`provider` 默认为 `openai`，`wire_protocol` 默认为 `chat_completions`。Anthropic 原生路由使用 `provider="anthropic"`。路由凭据遵循已有的服务商和端点选择规则。`api_key_env` 指定从环境或所选配置文件中读取的密钥变量，变量缺失时会报错。路由中明确提供的 `api_key` 优先。`auth_token` 支持 Anthropic 原生 bearer 认证。每条路由持有自己的端点凭据。`OpenCollab(...)` 在构造时完成环境文件凭据解析。路由中的 `context_window` 设置作用于实际模型。采样、思考、重试、超时和输出设置继续经过现有服务商实现。Responses 请求保留响应条目、服务商身份信息和可选的 `response_session_id` 行为。
+
+`request_pool` 和 `request_capacity` 需要成对提供。请求池使用已经存在的本地绝对路径目录，所有参与进程共享这个目录。运行时在取得请求名额前验证目录。每个名额通过文件上的操作系统锁持有，直到原生请求完成或取消清理结束。取消排队中的请求只释放该等待者自己的资源。进程结束时，操作系统会释放其文件锁。请求名额覆盖服务商重试，并在执行工具前释放，因此父任务等待子任务期间，子任务可以使用模型请求容量。
+
+顶层 `OPENCOLLAB_REQUEST_POOL` 和 `OPENCOLLAB_REQUEST_CAPACITY`，以及对应的 SDK 配置字段，可以设置共用请求池。路由中成对提供的池配置优先。当多个路由消耗同一个后端的容量时，它们可以共享一个池。
+
+池目录的 `capacity` 文件记录配置的容量。独立进程如果使用不同的名额数，就可能超过后端容量，因此冲突配置会报错。操作系统锁保证这些配置检查按顺序进行。参与进程运行期间，需要保留池目录及其锁文件。调整容量时，先停止这些进程，清空池目录，再使用新值启动。请求池依赖支持操作系统锁的本地文件系统。网络文件系统的锁语义，以及文件系统彼此独立的多台机器，需要由后端或网关管理请求容量。
+
+容器需要将同一个主机目录挂载到配置中的绝对路径。目录缺失时，错误信息会提示挂载要求。各自创建的独立目录拥有独立容量，即使它们的路径字符串相同。池路径属于运行 Python 模型客户端的进程，其所在环境可能与工具执行环境不同。
+
+多台主机共享容量时，将各条路由的 `base_url` 配置为同一个外部网关，并在网关中设置后端请求容量。网关通过已有服务商协议接收实际模型标识，负责跨主机分配名额。网关需要支持所选的 Chat Completions、Responses 或 Anthropic Messages 路径，并保留模型身份、推理参数和续接数据。HTTP 断开后仍在执行的后端工作由后端或网关的取消策略管理。本地请求池协调一台主机上的客户端请求生命周期。
diff --git a/opencollab/adapters/cli/config_resolve.py b/opencollab/adapters/cli/config_resolve.py
index 477b8a3f..1cdf124b 100644
--- a/opencollab/adapters/cli/config_resolve.py
+++ b/opencollab/adapters/cli/config_resolve.py
@@ -9,6 +9,7 @@
from rich.console import Console
from opencollab.bootstrap.config import accepted_api_key_envs, missing_api_key
+from opencollab.bootstrap.config import resolve_model_route as resolve_model_route
def missing_api_key_for(
diff --git a/opencollab/adapters/cli/main.py b/opencollab/adapters/cli/main.py
index 7c9fd261..0e4c1484 100644
--- a/opencollab/adapters/cli/main.py
+++ b/opencollab/adapters/cli/main.py
@@ -24,6 +24,7 @@
missing_api_key_for,
print_missing_key_hint,
resolve_config,
+ resolve_model_route,
)
from opencollab.adapters.cli.live_prompt import LivePrompt
from opencollab.adapters.cli.team import build_team_app
@@ -114,8 +115,11 @@ def main_callback(
raise typer.BadParameter("--hold requires --prompt or --prompt-file.")
cfg = resolve_config(workspace, model, provider, api_key, base_url, budget)
- if missing_api_key_for(cfg["provider"], cfg["api_key"], cfg["base_url"], cfg.get("auth_token")):
- print_missing_key_hint(console, cfg["provider"], cfg["base_url"])
+ model_cfg = _initial_model_config(cfg, team_config)
+ if missing_api_key_for(
+ model_cfg["provider"], model_cfg["api_key"], model_cfg["base_url"], model_cfg.get("auth_token"),
+ ):
+ print_missing_key_hint(console, model_cfg["provider"], model_cfg["base_url"])
raise typer.Exit(code=1)
run_with_bounded_shutdown(
@@ -133,6 +137,13 @@ def main_callback(
)
)
+def _initial_model_config(cfg: dict[str, Any], team_config_path: str | None) -> dict[str, Any]:
+ role_model = None
+ if cfg.get("model_routes"):
+ team = load_team_config(path=team_config_path)
+ role_model = team.role_for(team.entry).model
+ return resolve_model_route(cfg, role_model)
+
def _resolve_one_shot_prompt(prompt: str | None, prompt_file: str | None) -> str | None:
return resolve_text_input(
prompt,
@@ -338,9 +349,10 @@ async def _run(
# and lines are read straight from stdin.
has_screen = sys.stdin.isatty() and sys.stdout.isatty()
prompt = LivePrompt(tui, console=console, interactive=has_screen)
+ model_cfg = _initial_model_config(cfg, team_config_path)
tui.print_welcome(
- model=cfg["model"],
- provider=cfg["provider"],
+ model=model_cfg["model"],
+ provider=model_cfg["provider"],
workspace=workspace,
budget=cfg.get("budget"),
interactive=one_shot_prompt is None or (hold_after_run and has_screen),
diff --git a/opencollab/adapters/cli/workflow.py b/opencollab/adapters/cli/workflow.py
index 360ab12d..2183b3fe 100644
--- a/opencollab/adapters/cli/workflow.py
+++ b/opencollab/adapters/cli/workflow.py
@@ -20,6 +20,7 @@
missing_api_key_for,
print_missing_key_hint,
resolve_config,
+ resolve_model_route,
)
from opencollab.adapters.cli.text_input import resolve_text_input
from opencollab.application.async_timeout import run_with_bounded_shutdown
@@ -145,8 +146,11 @@ def run_cmd(
parsed_args["goal"] = task_text
cfg = resolve_config(workspace, model, provider, api_key, base_url, budget)
- if missing_api_key_for(cfg["provider"], cfg["api_key"], cfg["base_url"], cfg.get("auth_token")):
- print_missing_key_hint(console, cfg["provider"], cfg["base_url"])
+ model_cfg = resolve_model_route(cfg)
+ if missing_api_key_for(
+ model_cfg["provider"], model_cfg["api_key"], model_cfg["base_url"], model_cfg.get("auth_token"),
+ ):
+ print_missing_key_hint(console, model_cfg["provider"], model_cfg["base_url"])
raise typer.Exit(code=1)
# Fan-out workflows raise the fallback budget. Leave an explicit budget
diff --git a/opencollab/adapters/llm/client.py b/opencollab/adapters/llm/client.py
index 18e1da4d..67ece79f 100644
--- a/opencollab/adapters/llm/client.py
+++ b/opencollab/adapters/llm/client.py
@@ -9,7 +9,7 @@
import threading
import time
import uuid
-from typing import Any
+from typing import Any, Callable
from urllib.parse import urlsplit
import openai
@@ -32,6 +32,7 @@
normalize_wire_protocol,
warn_provider_near_miss,
)
+from opencollab.adapters.llm.request_admission import request_slot
from opencollab.adapters.llm.responses_messages import messages_to_input
from opencollab.adapters.llm.responses_provider import complete_responses
from opencollab.adapters.llm.retry import RetryTimeBudget, _scoped_native_retry_budget
@@ -93,6 +94,8 @@ def __init__(
user_agent: str | None = None,
stream_chat: bool = False,
auth_token: str | None = None,
+ request_pool: str | None = None,
+ request_capacity: int | None = None,
):
if context_window is not None and (
isinstance(context_window, bool)
@@ -108,6 +111,17 @@ def __init__(
raise ValueError(
"provider_error_time_budget must be a finite non-negative number"
)
+ if (request_pool is None) != (request_capacity is None):
+ raise ValueError("request_pool and request_capacity must be supplied together")
+ if request_capacity is not None and (
+ isinstance(request_capacity, bool) or not isinstance(request_capacity, int) or request_capacity <= 0
+ ):
+ raise ValueError("request_capacity must be a positive integer")
+ if request_pool is not None and not os.path.isabs(request_pool):
+ raise ValueError("request_pool must be an absolute local directory path")
+ self.request_pool = request_pool
+ self.request_capacity = request_capacity
+ self.supports_request_admission = request_pool is not None
self.model = model
self.provider = provider
self.max_retries = max(0, max_retries)
@@ -275,6 +289,7 @@ async def complete(
top_p: float | None = None,
max_output_tokens: int | None = None,
response_session_id: str | None = None,
+ on_admitted: Callable[[], None] | None = None,
) -> LLMResponse:
"""Single-shot completion. Returns full response.
@@ -300,66 +315,69 @@ async def complete(
# Measure here so every provider path records the call timing.
with recording() as first_token:
try:
- if self._anthropic:
- response = await complete_anthropic(
- self._anthropic,
- self.model,
- messages,
- tools,
- temperature,
- self.max_retries,
- thinking=thinking,
- thinking_params=thinking_params,
- tool_choice=tool_choice,
- top_p=top_p,
- max_output_tokens=max_output_tokens,
- reasoning_effort=reasoning_effort,
- provider_error_time_budget=retry_budget,
- )
- elif self.wire_protocol == RESPONSES:
- response = await complete_responses(
- self._openai,
- self.model,
- messages,
- tools,
- temperature,
- self.max_retries,
- thinking=thinking,
- thinking_params=thinking_params,
- tool_choice=tool_choice,
- top_p=top_p,
- max_output_tokens=max_output_tokens,
- reasoning_effort=reasoning_effort,
- prompt_cache_namespace=self._responses_prompt_cache_namespace,
- response_session_id=response_session_id or uuid.uuid4().hex,
- native_openai=(
- normalize_provider(self.provider) == "openai"
- and urlsplit(str(self._openai.base_url)).hostname in _NATIVE_OPENAI_HOSTS
- ),
- first_event_timeout=self.first_event_timeout,
- stream_idle_timeout=self.stream_idle_timeout,
- round_timeout=self.request_timeout,
- provider_error_time_budget=retry_budget,
- )
- else:
- response = await complete_openai(
- self._openai,
- self.model,
- messages,
- tools,
- temperature,
- self.max_retries,
- thinking=thinking,
- thinking_params=thinking_params,
- tool_choice=tool_choice,
- top_p=top_p,
- max_output_tokens=max_output_tokens,
- reasoning_effort=reasoning_effort,
- provider_error_time_budget=retry_budget,
- stream=self.stream_chat,
- first_event_timeout=self.first_event_timeout,
- stream_idle_timeout=self.stream_idle_timeout,
- )
+ async with request_slot(self.request_pool, self.request_capacity):
+ if on_admitted is not None:
+ on_admitted()
+ if self._anthropic:
+ response = await complete_anthropic(
+ self._anthropic,
+ self.model,
+ messages,
+ tools,
+ temperature,
+ self.max_retries,
+ thinking=thinking,
+ thinking_params=thinking_params,
+ tool_choice=tool_choice,
+ top_p=top_p,
+ max_output_tokens=max_output_tokens,
+ reasoning_effort=reasoning_effort,
+ provider_error_time_budget=retry_budget,
+ )
+ elif self.wire_protocol == RESPONSES:
+ response = await complete_responses(
+ self._openai,
+ self.model,
+ messages,
+ tools,
+ temperature,
+ self.max_retries,
+ thinking=thinking,
+ thinking_params=thinking_params,
+ tool_choice=tool_choice,
+ top_p=top_p,
+ max_output_tokens=max_output_tokens,
+ reasoning_effort=reasoning_effort,
+ prompt_cache_namespace=self._responses_prompt_cache_namespace,
+ response_session_id=response_session_id or uuid.uuid4().hex,
+ native_openai=(
+ normalize_provider(self.provider) == "openai"
+ and urlsplit(str(self._openai.base_url)).hostname in _NATIVE_OPENAI_HOSTS
+ ),
+ first_event_timeout=self.first_event_timeout,
+ stream_idle_timeout=self.stream_idle_timeout,
+ round_timeout=self.request_timeout,
+ provider_error_time_budget=retry_budget,
+ )
+ else:
+ response = await complete_openai(
+ self._openai,
+ self.model,
+ messages,
+ tools,
+ temperature,
+ self.max_retries,
+ thinking=thinking,
+ thinking_params=thinking_params,
+ tool_choice=tool_choice,
+ top_p=top_p,
+ max_output_tokens=max_output_tokens,
+ reasoning_effort=reasoning_effort,
+ provider_error_time_budget=retry_budget,
+ stream=self.stream_chat,
+ first_event_timeout=self.first_event_timeout,
+ stream_idle_timeout=self.stream_idle_timeout,
+ )
transport_timing = first_token.snapshot()
try:
response.transport_timing = transport_timing
diff --git a/opencollab/adapters/llm/request_admission.py b/opencollab/adapters/llm/request_admission.py
new file mode 100644
index 00000000..2c167093
--- /dev/null
+++ b/opencollab/adapters/llm/request_admission.py
@@ -0,0 +1,66 @@
+"""Request slots shared by independent processes on one local filesystem."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import AsyncIterator
+
+from filelock import FileLock, Timeout
+
+from opencollab.adapters.safe_files import (
+ create_regular_bytes_atomic,
+ ensure_directory_no_symlinks,
+ read_regular_text,
+)
+
+
+@asynccontextmanager
+async def _locked(lock: FileLock) -> AsyncIterator[None]:
+ while True:
+ try:
+ lock.acquire(timeout=0)
+ break
+ except Timeout:
+ await asyncio.sleep(0.05)
+ try:
+ yield
+ finally:
+ lock.release()
+
+
+@asynccontextmanager
+async def request_slot(pool: str | None, capacity: int | None) -> AsyncIterator[None]:
+ if pool is None:
+ yield
+ return
+ directory = Path(pool)
+ if not directory.is_dir():
+ raise ValueError(
+ f"request_pool must be an existing shared local directory: {pool}. "
+ "Mount the same directory into every participating process or container."
+ )
+ ensure_directory_no_symlinks(directory)
+ async with _locked(FileLock(str(directory / "capacity.lock"))):
+ record = directory / "capacity"
+ try:
+ configured_capacity = int(read_regular_text(record, max_bytes=64))
+ except FileNotFoundError:
+ create_regular_bytes_atomic(record, str(capacity).encode())
+ else:
+ if configured_capacity != capacity:
+ raise ValueError(f"request pool {pool!r} already has capacity {configured_capacity}, got {capacity}")
+ while True:
+ for index in range(capacity):
+ lock = FileLock(str(directory / f"slot-{index}.lock"))
+ try:
+ lock.acquire(timeout=0)
+ except Timeout:
+ continue
+ try:
+ yield
+ finally:
+ lock.release()
+ return
+ await asyncio.sleep(0.05)
diff --git a/opencollab/application/_session_run_completion.py b/opencollab/application/_session_run_completion.py
index ce83187c..ae0c6990 100644
--- a/opencollab/application/_session_run_completion.py
+++ b/opencollab/application/_session_run_completion.py
@@ -746,6 +746,9 @@ async def _invoke_summary(
kwargs["max_output_tokens"] = max_output_tokens
start = time.monotonic()
protected_call = self.state.wind_down_done
+ admitted = asyncio.Event() if getattr(self.llm, "supports_request_admission", False) else None
+ if admitted is not None:
+ kwargs["on_admitted"] = admitted.set
abandoned = False
if self.tracer:
self.tracer.log_step(step_type="llm_call_started", payload={
@@ -789,6 +792,7 @@ async def complete_owned() -> CompletionResponse:
owner, self._per_call_timeout,
late_task_tracker=self._mark_provider_task_draining,
late_result_handler=self._provider_task_done,
+ start_event=admitted,
)
except CallerTimeoutError as exc:
abandoned = True
@@ -813,6 +817,9 @@ async def _invoke_llm(self, **kwargs: Any) -> CompletionResponse:
if self._run_control is not None:
self._raise_if_run_cancelled()
protected_call = self.state.wind_down_done
+ admitted = asyncio.Event() if getattr(self.llm, "supports_request_admission", False) else None
+ if admitted is not None:
+ kwargs["on_admitted"] = admitted.set
operation = _complete_with_error_usage(self, self.llm.complete, protected_call=protected_call, **kwargs)
try:
response = await operation if self._per_call_timeout is None else await abandon_on_timeout(
@@ -824,6 +831,7 @@ async def _invoke_llm(self, **kwargs: Any) -> CompletionResponse:
task,
protected_call=protected_call,
),
+ start_event=admitted,
)
except CallerTimeoutError as exc:
raise GenerationTimeoutError(
diff --git a/opencollab/application/async_timeout.py b/opencollab/application/async_timeout.py
index 80afd17c..9dacdf93 100644
--- a/opencollab/application/async_timeout.py
+++ b/opencollab/application/async_timeout.py
@@ -115,6 +115,7 @@ async def abandon_on_timeout(
task_tracker: Callable[[asyncio.Task[T]], None] | None = None,
late_task_tracker: Callable[[asyncio.Task[T]], None] | None = None,
late_result_handler: Callable[[asyncio.Future[T]], None] | None = None,
+ start_event: asyncio.Event | None = None,
) -> T:
"""Await ``awaitable`` with a hard caller-side timeout.
@@ -124,6 +125,8 @@ async def abandon_on_timeout(
optional handler account for a delayed result in the background, and
immediately raises TimeoutError to the caller. Caller cancellation retains
the same background ownership and delayed-result handling.
+
+ ``start_event`` defers the deadline until request admission finishes.
"""
if timeout is None:
return await awaitable
@@ -139,7 +142,12 @@ async def abandon_on_timeout(
task = asyncio.ensure_future(awaitable)
if task_tracker is not None and isinstance(task, asyncio.Task):
task_tracker(task)
+ admission = asyncio.create_task(start_event.wait()) if start_event is not None else None
try:
+ if admission is not None:
+ await asyncio.wait({task, admission}, return_when=asyncio.FIRST_COMPLETED)
+ if task.done():
+ return task.result()
done, _ = await asyncio.wait({task}, timeout=normalized_timeout)
except asyncio.CancelledError:
if late_task_tracker is not None and isinstance(task, asyncio.Task):
@@ -147,6 +155,9 @@ async def abandon_on_timeout(
task.cancel()
task.add_done_callback(late_result_handler or consume_task_result)
raise
+ finally:
+ if admission is not None:
+ admission.cancel()
if task in done:
return task.result()
diff --git a/opencollab/bootstrap/_workflow_runtime_session.py b/opencollab/bootstrap/_workflow_runtime_session.py
index 98fd70b8..5e0ceab1 100644
--- a/opencollab/bootstrap/_workflow_runtime_session.py
+++ b/opencollab/bootstrap/_workflow_runtime_session.py
@@ -30,6 +30,7 @@
DEFAULT_TEMPERATURE,
DEFAULT_THINKING,
DEFAULT_TOP_P,
+ resolve_model_route,
resolve_thinking_params,
)
from opencollab.bootstrap.session_factory import (
@@ -114,6 +115,10 @@ def __init__(
run_id: str | None = None,
isolation: str | None = None,
environment_providers: Any | None = None,
+ model_routes: dict[str, Any] | None = None,
+ logical_model: str | None = None,
+ request_pool: str | None = None,
+ request_capacity: int | None = None,
) -> None:
self._llm = llm
self._session_builder = session_builder
@@ -122,6 +127,18 @@ def __init__(
self._environment_providers = environment_providers
self._strategy_pool = None
self._strategy_envs = []
+ settings = resolve_model_route({
+ "model": model, "provider": provider, "wire_protocol": wire_protocol,
+ "api_key": api_key, "auth_token": auth_token, "base_url": base_url,
+ "context_window": context_window, "model_routes": model_routes or {}, "logical_model": logical_model,
+ "request_pool": request_pool, "request_capacity": request_capacity,
+ })
+ model, provider, wire_protocol = (settings[key] for key in ("model", "provider", "wire_protocol"))
+ api_key, auth_token, base_url = (settings[key] for key in ("api_key", "auth_token", "base_url"))
+ context_window = settings["context_window"]
+ self._logical_model = settings.get("logical_model")
+ self._request_pool = settings.get("request_pool")
+ self._request_capacity = settings.get("request_capacity")
self._model = model
self._provider = provider
self._wire_protocol = wire_protocol
@@ -363,6 +380,9 @@ def build_workflow_session(
api_key=self._api_key,
auth_token=self._auth_token,
base_url=self._base_url,
+ logical_model=self._logical_model,
+ request_pool=self._request_pool,
+ request_capacity=self._request_capacity,
temperature=self._temperature,
top_p=self._top_p,
max_tokens_per_step=self._max_output_tokens,
@@ -516,6 +536,10 @@ def build_workflow_context(
isolation=isolation, environment_providers=environment_providers,
model=cfg["model"],
provider=cfg["provider"],
+ model_routes=cfg.get("model_routes"),
+ logical_model=cfg.get("logical_model"),
+ request_pool=cfg.get("request_pool"),
+ request_capacity=cfg.get("request_capacity"),
wire_protocol=cfg.get("wire_protocol", "chat_completions"),
api_key=cfg.get("api_key"),
auth_token=cfg.get("auth_token"),
diff --git a/opencollab/bootstrap/config.py b/opencollab/bootstrap/config.py
index e8da5931..a57fd31a 100644
--- a/opencollab/bootstrap/config.py
+++ b/opencollab/bootstrap/config.py
@@ -33,6 +33,7 @@
import json
import os
import stat
+from collections.abc import Mapping
from pathlib import Path
from typing import Any
@@ -120,17 +121,73 @@ def _optional_string(value: Any) -> Any:
return value.strip() or None
return value
+class ModelRoute(BaseModel):
+ model_config = ConfigDict(extra="forbid", hide_input_in_errors=True)
+
+ model: str = Field(min_length=1)
+ provider: str = Field(default="openai", min_length=1)
+ wire_protocol: str = "chat_completions"
+ base_url: str | None = None
+ api_key: str | None = Field(default=None, repr=False)
+ auth_token: str | None = Field(default=None, repr=False)
+ api_key_env: str | None = None
+ context_window: int | None = Field(default=None, ge=1)
+ request_pool: str | None = None
+ request_capacity: int | None = Field(default=None, ge=1)
+
+ @field_validator("model", "provider", "wire_protocol", mode="before")
+ @classmethod
+ def _strip_strings(cls, value: Any) -> Any:
+ return value.strip() if isinstance(value, str) else value
+
+ @field_validator("provider")
+ @classmethod
+ def _normalize_provider(cls, value: str) -> str:
+ return value.lower()
+
+ @field_validator("wire_protocol")
+ @classmethod
+ def _normalize_protocol(cls, value: str) -> str:
+ return normalize_wire_protocol(value)
+
+ @field_validator("base_url", "api_key", "auth_token", "api_key_env", "request_pool", mode="before")
+ @classmethod
+ def _optional_strings(cls, value: Any) -> Any:
+ return _optional_string(value)
+
+ @field_validator("context_window", "request_capacity", mode="before")
+ @classmethod
+ def _reject_boolean(cls, value: Any) -> Any:
+ if isinstance(value, bool):
+ raise ValueError("numeric configuration values must not be booleans")
+ return value
+
+ @model_validator(mode="after")
+ def _validate_pool(self) -> ModelRoute:
+ if (self.request_pool is None) != (self.request_capacity is None):
+ raise ValueError("request_pool and request_capacity must be supplied together")
+ if self.request_pool is not None and not os.path.isabs(self.request_pool):
+ raise ValueError("request_pool must be an absolute local directory path")
+ if is_anthropic(self.provider) and self.wire_protocol == "responses":
+ raise ValueError("Responses wire protocol requires an OpenAI-compatible provider")
+ return self
+
+
class OpenCollabConfig(BaseModel):
"""Runtime configuration for OpenCollab."""
- model_config = ConfigDict(extra="forbid")
+ model_config = ConfigDict(extra="forbid", hide_input_in_errors=True)
model: str = Field(default="gpt-4o", min_length=1)
+ logical_model: str | None = None
provider: str = Field(default="openai", min_length=1)
wire_protocol: str = Field(default="chat_completions", min_length=1)
api_key: str | None = Field(default=None, repr=False)
auth_token: str | None = Field(default=None, repr=False)
base_url: str | None = None
+ model_routes: dict[str, ModelRoute] = Field(default_factory=dict, repr=False)
+ request_pool: str | None = None
+ request_capacity: int | None = Field(default=None, ge=1)
budget: int = Field(default=1_000_000, ge=1)
temperature: float = Field(default=DEFAULT_TEMPERATURE, ge=0.0, le=2.0)
top_p: float | None = Field(default=DEFAULT_TOP_P, ge=0.0, le=1.0)
@@ -168,6 +225,7 @@ def _record_edited_thinking_params_source(self) -> OpenCollabConfig:
return self
@field_validator(
+ "request_capacity",
"budget",
"temperature",
"top_p",
@@ -187,6 +245,27 @@ def _reject_boolean_numbers(cls, value: Any) -> Any:
raise ValueError("numeric configuration values must not be booleans")
return value
+ @field_validator("model_routes", mode="before")
+ @classmethod
+ def _parse_routes(cls, value: Any) -> Any:
+ if value is None:
+ return {}
+ if isinstance(value, str):
+ value = json.loads(value)
+ if not isinstance(value, dict) or any(
+ not isinstance(key, str) or not key.strip() or key != key.strip() for key in value
+ ):
+ raise ValueError("model_routes must map non-empty model names to routes")
+ return value
+
+ @model_validator(mode="after")
+ def _validate_request_pool(self) -> OpenCollabConfig:
+ if (self.request_pool is None) != (self.request_capacity is None):
+ raise ValueError("request_pool and request_capacity must be supplied together")
+ if self.request_pool is not None and not os.path.isabs(self.request_pool):
+ raise ValueError("request_pool must be an absolute local directory path")
+ return self
+
@field_validator("top_p", mode="before")
@classmethod
def _coerce_top_p(cls, value: Any) -> Any:
@@ -244,7 +323,7 @@ def _normalize_reasoning_effort(cls, value: Any) -> Any:
)
return normalized
- @field_validator("api_key", "auth_token", "base_url", mode="before")
+ @field_validator("api_key", "auth_token", "base_url", "request_pool", mode="before")
@classmethod
def _empty_string_as_none(cls, value: Any) -> Any:
return _optional_string(value)
@@ -440,6 +519,9 @@ def resolve_ordered(*keys: str, default: str | None = None) -> str | None:
defaults = {
"model": "gpt-4o",
+ "model_routes": None,
+ "request_pool": None,
+ "request_capacity": None,
"wire_protocol": "chat_completions",
"budget": "1000000",
"temperature": str(DEFAULT_TEMPERATURE),
@@ -468,8 +550,44 @@ def resolve_ordered(*keys: str, default: str | None = None) -> str | None:
# Defaults share the same payload as an explicit DashScope protocol. Keep
# its source before validation replaces None with that default mapping.
values.setdefault("thinking_params_explicit", values["thinking_params"] is not None)
- return OpenCollabConfig.model_validate(values)
+ config = OpenCollabConfig.model_validate(values)
+ for route in config.model_routes.values():
+ if route.api_key_env and not route.api_key:
+ route.api_key = resolve_ordered(route.api_key_env)
+ if not route.api_key:
+ raise ValueError(f"model route API key environment variable is unset: {route.api_key_env}")
+ if not route.api_key and not route.auth_token:
+ names = api_key_env_precedence(route.provider, route.base_url)
+ if is_anthropic(route.provider):
+ names = tuple(name for name in names if name != "OPENCOLLAB_API_KEY")
+ route.api_key = resolve_ordered(*names)
+ if is_anthropic(route.provider) and not route.api_key:
+ route.auth_token = resolve_ordered("ANTHROPIC_AUTH_TOKEN")
+ if not route.auth_token:
+ route.api_key = resolve_ordered("OPENCOLLAB_API_KEY")
+ return config
def get_config(workspace: str | None = None) -> dict[str, Any]:
"""Get resolved configuration as a dict for existing callers."""
return build_config(workspace).model_dump()
+
+
+def resolve_model_route(config: Mapping[str, Any], model: str | None = None) -> dict[str, Any]:
+ resolved = dict(config)
+ if model is None and config.get("logical_model") is not None:
+ return resolved
+ if model is not None:
+ resolved["logical_model"] = None
+ selected = config["model"] if model is None else model
+ resolved["model"] = selected
+ route = config.get("model_routes", {}).get(selected)
+ if route is None:
+ return resolved
+ route = route if isinstance(route, ModelRoute) else ModelRoute.model_validate(route)
+ values = route.model_dump(exclude={"api_key_env", "request_pool", "request_capacity"})
+ resolved.update(values)
+ if route.request_pool is not None:
+ resolved.update(request_pool=route.request_pool, request_capacity=route.request_capacity)
+ resolved["base_url"] = resolve_provider_base_url(route.provider, route.base_url)
+ resolved["logical_model"] = selected
+ return resolved
diff --git a/opencollab/bootstrap/container.py b/opencollab/bootstrap/container.py
index 6abdd700..86f15ee8 100644
--- a/opencollab/bootstrap/container.py
+++ b/opencollab/bootstrap/container.py
@@ -132,6 +132,8 @@ def _resolve_llm(
context_window=getattr(agent, "context_window", None),
provider_error_time_budget=getattr(agent, "provider_error_time_budget", 0.0),
provider_retry_budget=provider_retry_budget,
+ **({"request_pool": agent.request_pool, "request_capacity": getattr(agent, "request_capacity", None)}
+ if getattr(agent, "request_pool", None) is not None else {}),
)
diff --git a/opencollab/bootstrap/context_builder.py b/opencollab/bootstrap/context_builder.py
index b92267cc..4b8db27f 100644
--- a/opencollab/bootstrap/context_builder.py
+++ b/opencollab/bootstrap/context_builder.py
@@ -3,6 +3,7 @@
from __future__ import annotations
from dataclasses import dataclass, field
+from typing import Any
from opencollab.adapters.llm.providers import RESPONSES
from opencollab.adapters.llm.types import model_capabilities
@@ -26,6 +27,7 @@
DEFAULT_THINKING,
DEFAULT_THINKING_PARAMS,
DEFAULT_TOP_P,
+ resolve_model_route,
)
from opencollab.bootstrap.team_config import RoleConfig, TeamConfig
from opencollab.bootstrap.tool_registry import COORDINATION_TOOL_NAMES, build_tools_for_role
@@ -72,6 +74,19 @@ class SpawnConfig:
llm_stream_chat: bool = False
thinking_params_explicit: bool | None = None
auth_token: str | None = field(default=None, repr=False)
+ model_routes: dict[str, Any] = field(default_factory=dict, repr=False)
+ request_pool: str | None = None
+ request_capacity: int | None = None
+ logical_model: str | None = None
+
+ def model_settings(self, model: str | None = None) -> dict[str, Any]:
+ return resolve_model_route(
+ {key: getattr(self, key) for key in (
+ "model", "provider", "api_key", "auth_token", "base_url", "wire_protocol",
+ "context_window", "model_routes", "request_pool", "request_capacity", "logical_model",
+ )},
+ model,
+ )
def __post_init__(self) -> None:
if self.thinking_params_explicit is None:
@@ -79,7 +94,7 @@ def __post_init__(self) -> None:
self.thinking_params = dict(DEFAULT_THINKING_PARAMS if self.thinking_params is None else self.thinking_params)
def reasoning_effort_for_role(self, role: RoleConfig) -> str | None:
- capabilities = model_capabilities(role.model or self.model)
+ capabilities = model_capabilities(self.model_settings(role.model)["model"])
if role.thinking is False and (
not self.thinking or capabilities.honors_workflow_thinking_override
):
@@ -214,9 +229,10 @@ def build_agent(
tool_limits=self._tool_limits_for(role),
)
cfg = self._cfg
- model = role.model or cfg.model
+ model_settings = cfg.model_settings(role.model)
+ model = model_settings["model"]
if (
- cfg.wire_protocol == RESPONSES
+ model_settings["wire_protocol"] == RESPONSES
and tools
and not model_capabilities(model).supports_responses_tools
):
@@ -240,15 +256,18 @@ def build_agent(
system_prompt=plan.system_prompt(),
tools=tools,
model=model,
- provider=cfg.provider,
- wire_protocol=cfg.wire_protocol,
- api_key=cfg.api_key,
- auth_token=cfg.auth_token,
- base_url=cfg.base_url,
+ provider=model_settings["provider"],
+ wire_protocol=model_settings["wire_protocol"],
+ api_key=model_settings["api_key"],
+ auth_token=model_settings["auth_token"],
+ base_url=model_settings["base_url"],
+ logical_model=model_settings.get("logical_model"),
+ request_pool=model_settings.get("request_pool"),
+ request_capacity=model_settings.get("request_capacity"),
temperature=temperature,
top_p=cfg.top_p,
max_tokens_per_step=cfg.max_output_tokens,
- context_window=cfg.context_window,
+ context_window=model_settings["context_window"],
thinking=thinking,
thinking_params=thinking_params,
thinking_params_explicit=(
diff --git a/opencollab/bootstrap/inspection.py b/opencollab/bootstrap/inspection.py
index 6e50adae..89ead0a1 100644
--- a/opencollab/bootstrap/inspection.py
+++ b/opencollab/bootstrap/inspection.py
@@ -8,18 +8,22 @@
from opencollab.adapters.llm.client import LLMClient
from opencollab.adapters.storage import SessionStore
+from opencollab.bootstrap.config import resolve_model_route
def configured_model_client(
config: Mapping[str, Any], *, client_type: type[Any] | None = None,
) -> Any:
"""Create an independently owned transport with the resolved configuration."""
+ config = resolve_model_route(config)
return (client_type or LLMClient)(
model=config["model"],
provider=config["provider"],
api_key=config.get("api_key"),
**({"auth_token": config["auth_token"]} if config.get("auth_token") else {}),
base_url=config.get("base_url"),
+ **({"request_pool": config["request_pool"], "request_capacity": config["request_capacity"]}
+ if config.get("request_pool") is not None else {}),
wire_protocol=config.get("wire_protocol", "chat_completions"),
context_window=config.get("context_window"),
max_retries=config.get("llm_max_retries", 3),
diff --git a/opencollab/bootstrap/programmatic.py b/opencollab/bootstrap/programmatic.py
index 209f82d7..056f723f 100644
--- a/opencollab/bootstrap/programmatic.py
+++ b/opencollab/bootstrap/programmatic.py
@@ -48,7 +48,7 @@
from opencollab.bootstrap.agent_runtime import (
run_agent as _run_agent,
)
-from opencollab.bootstrap.config import resolve_thinking_params
+from opencollab.bootstrap.config import resolve_model_route, resolve_thinking_params
from opencollab.bootstrap.scheduler_factory import build_scheduler # noqa: F401
from opencollab.bootstrap.session_factory import SESSION_MAX_STEPS
from opencollab.bootstrap.tool_registry import build_tools_for_role
@@ -345,6 +345,7 @@ async def run_agent(
run_id = run_id if run_id is not None else f"agent-{uuid.uuid4().hex}"
resolved_profile = agent_profile if agent_profile is not None else resolve_agent_profile(None)
resolved_tools = resolved_profile.resolve_tools(tools)
+ config = resolve_model_route(config)
agent = Agent(
name=name,
system_prompt=system_prompt if system_prompt is not None else resolved_profile.system_prompt,
@@ -355,6 +356,9 @@ async def run_agent(
api_key=config.get("api_key"),
auth_token=config.get("auth_token"),
base_url=config.get("base_url"),
+ logical_model=config.get("logical_model"),
+ request_pool=config.get("request_pool"),
+ request_capacity=config.get("request_capacity"),
context_window=config.get("context_window"),
max_tokens_per_step=config.get("max_output_tokens", 8_192),
temperature=config.get("temperature", 0.2),
diff --git a/opencollab/bootstrap/scheduler_factory.py b/opencollab/bootstrap/scheduler_factory.py
index dbd51a40..66d3e0d8 100644
--- a/opencollab/bootstrap/scheduler_factory.py
+++ b/opencollab/bootstrap/scheduler_factory.py
@@ -161,6 +161,10 @@ def build_scheduler(
SpawnConfig(
model=cfg["model"],
provider=cfg["provider"],
+ model_routes=cfg.get("model_routes", {}),
+ logical_model=cfg.get("logical_model"),
+ request_pool=cfg.get("request_pool"),
+ request_capacity=cfg.get("request_capacity"),
wire_protocol=cfg.get("wire_protocol", "chat_completions"),
api_key=cfg["api_key"],
auth_token=cfg.get("auth_token"),
diff --git a/opencollab/bootstrap/session_factory.py b/opencollab/bootstrap/session_factory.py
index e22d0b17..4f704310 100644
--- a/opencollab/bootstrap/session_factory.py
+++ b/opencollab/bootstrap/session_factory.py
@@ -521,9 +521,6 @@ def __init__(
def _validate_responses_tool_support(self) -> None:
"""Reject statically incompatible team roles before opening a workspace."""
- if self._cfg.wire_protocol != RESPONSES:
- return
-
# A restricted topology may name an undeclared destination. Such a
# destination is still spawnable and ``TeamConfig.role_for`` resolves it
# to the generic BASE_TOOL_NAMES role, so it must participate in the
@@ -538,7 +535,10 @@ def _validate_responses_tool_support(self) -> None:
)
for role_name, role in reachable_roles.items():
- model = role.model or self._cfg.model
+ settings = self._cfg.model_settings(role.model)
+ if settings["wire_protocol"] != RESPONSES:
+ continue
+ model = settings["model"]
tool_names = set(role.tools)
if role_name != self._team.entry or not self._interactive:
tool_names.discard("ask_user")
@@ -550,8 +550,9 @@ def _validate_responses_tool_support(self) -> None:
reasoning_effort=self._cfg.reasoning_effort_for_role(role),
)
fallback_tool_names = set(BASE_TOOL_NAMES) - {"ask_user"}
- if self._team.topology.allow_all:
- capabilities = model_capabilities(self._cfg.model)
+ default_settings = self._cfg.model_settings()
+ if self._team.topology.allow_all and default_settings["wire_protocol"] == RESPONSES:
+ capabilities = model_capabilities(default_settings["model"])
if fallback_tool_names and not capabilities.supports_responses_tools:
raise ValueError(
f"default model {self._cfg.model!r} does not support Responses "
diff --git a/opencollab/domain/agent.py b/opencollab/domain/agent.py
index 2420ca6d..3f617897 100644
--- a/opencollab/domain/agent.py
+++ b/opencollab/domain/agent.py
@@ -39,6 +39,9 @@ class Agent:
auth_token: str | None = field(default=None, repr=False, kw_only=True)
base_url: str | None = None
context_window: int | None = None
+ logical_model: str | None = field(default=None, kw_only=True)
+ request_pool: str | None = field(default=None, kw_only=True)
+ request_capacity: int | None = field(default=None, kw_only=True)
max_tokens_per_step: int = DEFAULT_MAX_TOKENS_PER_STEP
temperature: float = 0.0
top_p: float | None = None
diff --git a/opencollab/sdk/client.py b/opencollab/sdk/client.py
index 74c9cf05..160a9332 100644
--- a/opencollab/sdk/client.py
+++ b/opencollab/sdk/client.py
@@ -118,6 +118,7 @@ def __init__(
config: Mapping[str, Any] | None = None,
environment: Any | None = None,
environment_providers: Mapping[str, Any] | None = None,
+ model_routes: Mapping[str, Mapping[str, Any]] | None = None,
) -> None:
raw_workspace = os.fspath(workspace)
if not raw_workspace or "\x00" in raw_workspace:
@@ -128,6 +129,10 @@ def __init__(
if config is not None and not isinstance(config, Mapping):
raise TypeError("config must be a mapping")
overrides = dict(config or {})
+ if model is not None:
+ overrides.pop("logical_model", None)
+ if model_routes is not None:
+ overrides["model_routes"] = dict(model_routes)
overrides.update(
{
key: value
@@ -170,12 +175,20 @@ def configuration(self) -> Mapping[str, Any]:
"context_window",
"llm_max_retries",
"provider_error_time_budget",
+ "request_capacity",
)
public = {
key: copy.deepcopy(self._config[key])
for key in public_keys
if key in self._config
}
+ if self._config.get("model_routes"):
+ public["model_routes"] = {
+ name: {key: route[key] for key in (
+ "model", "provider", "wire_protocol", "context_window", "request_capacity"
+ )}
+ for name, route in self._config["model_routes"].items()
+ }
base_url = self._config.get("base_url")
public["base_url_sha256"] = (
hashlib.sha256(base_url.encode("utf-8")).hexdigest()
diff --git a/tests/interfaces/test_sdk_api.py b/tests/interfaces/test_sdk_api.py
index 452ad4a0..bf6bff54 100644
--- a/tests/interfaces/test_sdk_api.py
+++ b/tests/interfaces/test_sdk_api.py
@@ -50,7 +50,8 @@ def test_public_class_and_method_shapes_stay_lean() -> None:
positional: {"workspace": "."},
keyword: {
"model": None, "provider": None, "api_key": None,
- "base_url": None, "config": None, "environment": None, "environment_providers": None,
+ "base_url": None, "config": None, "environment": None,
+ "environment_providers": None, "model_routes": None,
},
},
sdk.OpenCollab.agent: {
diff --git a/tests/interfaces/test_shared_model_routing.py b/tests/interfaces/test_shared_model_routing.py
new file mode 100644
index 00000000..410cfff2
--- /dev/null
+++ b/tests/interfaces/test_shared_model_routing.py
@@ -0,0 +1,522 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import subprocess
+import sys
+import threading
+import time
+from contextlib import ExitStack, contextmanager
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+
+from opencollab import OpenCollab
+from opencollab.adapters.llm import client as transport
+from opencollab.adapters.llm.types import LLMResponse, Usage
+from opencollab.bootstrap.config import build_config
+from tests.support.paths import PACKAGE_ROOT
+from tests.support.provider_sdk_http import completion_http_response, install_sdk_transport
+
+
+@pytest.fixture(autouse=True)
+def clean_provider_environment(monkeypatch, tmp_path):
+ for name in tuple(os.environ):
+ if name.startswith("OPENCOLLAB_") or name in {
+ "OPENAI_API_KEY", "OPENAI_BASE_URL", "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_BASE_URL",
+ }:
+ monkeypatch.delenv(name)
+ monkeypatch.chdir(tmp_path)
+
+
+@pytest.mark.parametrize("entry", ["agent", "workflow", "team", "create_model_client"])
+async def test_sdk_routes_physical_model_provider_and_controls(tmp_path, monkeypatch, entry):
+ requests = []
+
+ def handler(request):
+ requests.append((str(request.url), dict(request.headers), json.loads(request.content)))
+ return completion_http_response(request, text="finished")
+
+ install_sdk_transport(monkeypatch, "anthropic", handler)
+ client = OpenCollab(tmp_path, model="writer",
+ api_key="default-endpoint-secret", model_routes={ # pragma: allowlist secret
+ "writer": {
+ "model": "claude-sonnet-4-6", "provider": "anthropic", "base_url": "https://writer.example",
+ "auth_token": "route-token", "context_window": 12345,
+ },
+ }, config={"max_output_tokens": 4096, "thinking_params": {}, "temperature": 0.4})
+ if entry == "create_model_client":
+ async with client.create_model_client() as llm:
+ assert llm.model == "claude-sonnet-4-6"
+ assert llm.provider == "anthropic"
+ assert llm.context_window() == 12345
+ assert llm.estimate_request_tokens([{"role": "user", "content": "finish"}]) > 0
+ assert llm.estimate_context_tokens([{"role": "user", "content": "finish"}]) > 0
+ response = await llm.complete([{"role": "user", "content": "finish"}], temperature=0.4)
+ assert response.content == "finished"
+ else:
+ if entry == "agent":
+ result = await client.agent("finish", tools=[], trace=False)
+ elif entry == "workflow":
+ async def flow(ctx, _inputs):
+ return await ctx.agent("finish", tools=[])
+
+ result = await client.workflow(flow, trace=False)
+ else:
+ path = tmp_path / "team.yaml"
+ path.write_text("entry: writer\nroles:\n writer:\n prompt: Finish.\n tools: []\n")
+ result = await client.team("finish", config=path, use_worktrees=False, trace=False)
+ assert result.ok
+ assert result.output == "finished"
+ assert len(requests) == 1
+ url, headers, body = requests[0]
+ assert url == "https://writer.example/v1/messages"
+ assert headers["authorization"] == "Bearer route-token"
+ assert "x-api-key" not in headers
+ assert body["model"] == "claude-sonnet-4-6"
+ assert body["temperature"] == 0.4
+ metadata = json.dumps(dict(client.configuration))
+ assert "claude-sonnet-4-6" in metadata
+ assert "route-token" not in metadata
+ assert "default-endpoint-secret" not in metadata
+ assert "writer.example" not in metadata
+
+
+async def test_routed_responses_preserve_identity_and_optional_session(tmp_path, monkeypatch):
+ response = LLMResponse(
+ content="answer", provider_model="gpt-5", usage=Usage(4, 2),
+ provider_items=[{"id": "item-1", "type": "message"}],
+ provider_state={"response_id": "resp-1"}, reasoning="reasoning",
+ )
+ captured = []
+
+ async def complete(*args, **kwargs):
+ captured.append((args, kwargs))
+ return response
+
+ monkeypatch.setattr(transport, "complete_responses", complete)
+ client = OpenCollab(tmp_path, model="reasoner", model_routes={
+ "reasoner": {
+ "model": "gpt-5", "api_key": "local", "wire_protocol": "responses", # pragma: allowlist secret
+ "base_url": "https://gateway.example/v1", "context_window": 10000,
+ },
+ })
+ async with client.create_model_client() as llm:
+ assert llm.supports_response_session_identity
+ assert llm.context_window() == 10000
+ actual = await llm.complete(
+ [{"role": "user", "content": "answer"}], reasoning_effort="high", response_session_id="session-1",
+ )
+ assert actual is response
+ assert response.provider_items == [{"id": "item-1", "type": "message"}]
+ assert response.provider_state == {"response_id": "resp-1"}
+ args, kwargs = captured[0]
+ assert args[1] == "gpt-5"
+ assert kwargs["reasoning_effort"] == "high"
+ assert kwargs["response_session_id"] == "session-1"
+
+
+def test_routes_resolve_env_file_secrets_once(tmp_path, monkeypatch):
+ config_dir = tmp_path / "configs"
+ config_dir.mkdir()
+ routes = {"worker": {"model": "local-model", "base_url": "http://localhost:9999/v1",
+ "api_key_env": "WORKER_KEY"}} # pragma: allowlist secret
+ (config_dir / ".env").write_text(f"WORKER_KEY=initial-key\nOPENCOLLAB_MODEL_ROUTES='{json.dumps(routes)}'\n")
+ client = OpenCollab(tmp_path, model="worker")
+ monkeypatch.setenv("WORKER_KEY", "later-key")
+ assert client._config["model_routes"]["worker"]["api_key"] == "initial-key" # pragma: allowlist secret
+ monkeypatch.setenv("OPENCOLLAB_MODEL_ROUTES", json.dumps({"worker": {"model": "other",
+ "api_key_env": "MISSING"}})) # pragma: allowlist secret
+ with pytest.raises(ValueError, match="MISSING"):
+ build_config(str(tmp_path))
+
+
+@pytest.mark.parametrize("bad", [
+ {"request_pool": "relative", "request_capacity": 2},
+ {"request_pool": "/tmp/pool"}, {"request_capacity": 2},
+ {"request_pool": "/tmp/pool", "request_capacity": True},
+ {"model_routes": {"worker": {"model": "physical", "request_capacity": 2}}},
+ {"model_routes": {"worker": {"model": "physical", "unexpected": True}}},
+])
+def test_invalid_pool_and_route_configuration_fails_explicitly(tmp_path, bad):
+ with pytest.raises(ValueError):
+ OpenCollab(tmp_path, config=bad)
+
+
+_WORKER = """
+import asyncio, json, sys, time
+from pathlib import Path
+from opencollab import OpenCollab
+root, mode, index = Path(sys.argv[1]), sys.argv[2], sys.argv[3]
+client = OpenCollab(root / index, model='workflow', config=json.loads((root / 'config.json').read_text()))
+(root / ('ready-' + index)).touch()
+while not (root / 'start').exists():
+ time.sleep(.01)
+async def main():
+ if mode == 'team':
+ result = await client.team('finish', config=root / 'team.yaml', use_worktrees=False, trace=False)
+ else:
+ async def flow(ctx, inputs):
+ return await ctx.agent('finish', tools=[])
+ result = await client.workflow(flow, trace=False)
+ print(json.dumps({'ok': result.ok, 'output': result.output}))
+asyncio.run(main())
+"""
+
+
+@contextmanager
+def _backend(state, expected_model):
+ class Handler(BaseHTTPRequestHandler):
+ def log_message(self, *_args):
+ pass
+
+ def do_POST(self):
+ body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+ with state["lock"]:
+ state["requests"].append((expected_model, body["model"]))
+ state["active"] += 1
+ state["peak"] = max(state["peak"], state["active"])
+ try:
+ time.sleep(0.2)
+ payload = json.dumps({
+ "id": "chat-test", "object": "chat.completion", "created": 1, "model": body["model"],
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": body["model"]},
+ "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4},
+ }).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)
+ finally:
+ with state["lock"]:
+ state["active"] -= 1
+
+ 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"
+ finally:
+ server.shutdown()
+ server.server_close()
+ thread.join(timeout=2)
+
+
+def test_independent_team_and_workflow_processes_share_actual_http_capacity(tmp_path):
+ state = {"lock": threading.Lock(), "requests": [], "active": 0, "peak": 0}
+ pool = tmp_path / "request-pool"
+ pool.mkdir()
+ processes = []
+ with ExitStack() as stack:
+ team_url = stack.enter_context(_backend(state, "team-physical"))
+ workflow_url = stack.enter_context(_backend(state, "workflow-physical"))
+ (tmp_path / "config.json").write_text(json.dumps({
+ "thinking_params": {}, "request_pool": str(pool), "request_capacity": 2,
+ "model_routes": {
+ "team": {"model": "team-physical", "base_url": team_url,
+ "api_key": "local"}, # pragma: allowlist secret
+ "workflow": {"model": "workflow-physical", "base_url": workflow_url,
+ "api_key": "local"}, # pragma: allowlist secret
+ },
+ }))
+ (tmp_path / "team.yaml").write_text(
+ "entry: reader\nroles:\n reader:\n prompt: Finish.\n model: team\n tools: []\n",
+ )
+ env = {**os.environ, "PYTHONPATH": str(PACKAGE_ROOT)}
+ try:
+ for index, mode in enumerate(["team", "workflow", "team", "workflow"]):
+ (tmp_path / str(index)).mkdir()
+ processes.append(subprocess.Popen(
+ [sys.executable, "-c", _WORKER, str(tmp_path), mode, str(index)], env=env,
+ stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
+ ))
+ deadline = time.monotonic() + 15
+ while len(list(tmp_path.glob("ready-*"))) < 4 and time.monotonic() < deadline:
+ for process in processes:
+ if process.poll() is not None:
+ stdout, stderr = process.communicate()
+ pytest.fail(f"worker exited before starting requests\n{stdout}\n{stderr}")
+ time.sleep(0.01)
+ assert len(list(tmp_path.glob("ready-*"))) == 4
+ (tmp_path / "start").touch()
+ results = []
+ for process in processes:
+ stdout, stderr = process.communicate(timeout=20)
+ assert process.returncode == 0, stderr
+ results.append(json.loads(stdout))
+ assert results == [
+ {"ok": True, "output": "team-physical"}, {"ok": True, "output": "workflow-physical"},
+ {"ok": True, "output": "team-physical"}, {"ok": True, "output": "workflow-physical"},
+ ]
+ finally:
+ for process in processes:
+ if process.poll() is None:
+ process.kill()
+ process.communicate(timeout=5)
+ assert sorted(state["requests"]) == [("team-physical", "team-physical")] * 2 + [
+ ("workflow-physical", "workflow-physical")
+ ] * 2
+ assert state["peak"] == 2
+ assert state["active"] == 0
+
+
+async def test_parent_tool_waiting_for_child_workflow_releases_request_slot(tmp_path, monkeypatch):
+ from opencollab.adapters.tools.base import Tool
+ from tests.support.provider_sdk_http import http_response
+
+ calls = []
+
+ def handler(request):
+ body = json.loads(request.content)
+ parent = bool(body.get("tools"))
+ calls.append("parent" if parent else "child")
+ if parent and not any(message["role"] == "tool" for message in body["messages"]):
+ return http_response(request, 200, json={
+ "id": "parent-call", "object": "chat.completion", "created": 1, "model": body["model"],
+ "choices": [{"index": 0, "finish_reason": "tool_calls", "message": {
+ "role": "assistant", "content": None,
+ "tool_calls": [{"id": "child-1", "type": "function", "function": {
+ "name": "run_child", "arguments": "{}",
+ }}],
+ }}],
+ "usage": {"prompt_tokens": 2, "completion_tokens": 2, "total_tokens": 4},
+ })
+ return completion_http_response(request, text="parent finished" if parent else "child finished")
+
+ install_sdk_transport(monkeypatch, "openai", handler)
+ client = OpenCollab(tmp_path, model="worker", model_routes={
+ "worker": {
+ "model": "physical-worker",
+ "api_key": "local", "request_pool": str(tmp_path), "request_capacity": 1, # pragma: allowlist secret
+ },
+ }, config={"thinking_params": {}})
+
+ class ChildTool(Tool):
+ name = "run_child"
+ description = "Run the child task."
+
+ async def execute_with_runtime(self, params, runtime):
+ async def flow(ctx, inputs):
+ return await ctx.agent("child", tools=[])
+
+ result = await client.workflow(flow, trace=False)
+ assert result.ok
+ return result.output
+
+ result = await asyncio.wait_for(client.agent("parent", tools=[ChildTool()], trace=False), 5)
+ assert result.ok
+ assert result.output == "parent finished"
+ assert calls == ["parent", "child", "parent"]
+
+
+@pytest.mark.parametrize("entry", ["workflow", "team"])
+async def test_route_capability_checks_use_physical_model(tmp_path, monkeypatch, entry):
+ client = OpenCollab(tmp_path, model="reasoner", model_routes={
+ "reasoner": {"model": "o3", "wire_protocol": "responses", "api_key": "local"}, # pragma: allowlist secret
+ }, config={"top_p": 0.5})
+ requests = []
+
+ def handler(request):
+ requests.append(request)
+ return completion_http_response(request, text="finished")
+
+ install_sdk_transport(monkeypatch, "openai", handler)
+ if entry == "workflow":
+ async def flow(ctx, _inputs):
+ return await ctx.agent("finish", tools=[])
+
+ result = await client.workflow(flow, trace=False)
+ assert result.output is None
+ assert result.tokens == 0
+ assert result.agent_failures[0]["exception_type"] == "ValueError"
+ else:
+ path = tmp_path / "team.yaml"
+ path.write_text("entry: reader\nroles:\n reader:\n prompt: Finish.\n tools: []\n")
+ with pytest.raises(ValueError, match="Responses sampling"):
+ await client.team("finish", config=path, use_worktrees=False, trace=False)
+ assert requests == []
+
+
+def test_cli_keeps_role_defaults_while_resolving_default_credentials(tmp_path, monkeypatch):
+ from opencollab.adapters.cli.config_resolve import missing_api_key_for, resolve_config, resolve_model_route
+
+ monkeypatch.setenv("OPENCOLLAB_MODEL", "worker")
+ monkeypatch.setenv("OPENCOLLAB_MODEL_ROUTES", json.dumps({
+ "worker": {"model": "local-physical",
+ "api_key": "local", "base_url": "http://localhost:9999/v1"}, # pragma: allowlist secret
+ }))
+ config = resolve_config(str(tmp_path), None, None, None, None, None)
+ assert config["model"] == "worker"
+ assert config["logical_model"] is None
+ selected = resolve_model_route(config)
+ assert selected["model"] == "local-physical"
+ assert selected["logical_model"] == "worker"
+ assert not missing_api_key_for(selected["provider"], selected["api_key"], selected["base_url"])
+
+
+def test_resolved_config_keeps_physical_identity_across_bootstrap_paths(tmp_path):
+ from opencollab.bootstrap._workflow_runtime_session import WorkflowSessionFactory
+ from opencollab.bootstrap.config import resolve_model_route
+ from opencollab.bootstrap.context_builder import ContextBuilder, SpawnConfig
+ from opencollab.bootstrap.team_config import RoleConfig, TeamConfig
+
+ client = OpenCollab(tmp_path, model="worker", model_routes={
+ "worker": {"model": "physical-worker",
+ "api_key": "local", "base_url": "https://worker.example/v1"}, # pragma: allowlist secret
+ "physical-worker": {"model": "unrelated-model",
+ "api_key": "other", "base_url": "https://other.example/v1"}, # pragma: allowlist secret
+ })
+ config = resolve_model_route(client._config)
+ assert resolve_model_route(config) == config
+ clone = OpenCollab(tmp_path, config=config)
+ assert clone._config["logical_model"] == "worker"
+ team = TeamConfig(entry="reader", roles={"reader": RoleConfig(prompt="Finish.", tools=[])})
+ spawn = SpawnConfig(
+ model=config["model"], provider=config["provider"], api_key=config["api_key"],
+ base_url=config["base_url"], llm_timeout=10, tracer=None, event_bus=None, permission_policy=None,
+ model_routes=config["model_routes"], logical_model=config["logical_model"],
+ )
+ agent = ContextBuilder(team, spawn).build_agent("reader")
+ assert agent.model == "physical-worker"
+ assert agent.base_url == "https://worker.example/v1"
+ captured = []
+ factory = WorkflowSessionFactory(
+ model=config["model"], provider=config["provider"], api_key=config["api_key"],
+ base_url=config["base_url"], model_routes=config["model_routes"], logical_model=config["logical_model"],
+ session_builder=lambda **kwargs: captured.append(kwargs),
+ )
+ factory.build_workflow_session(prompt="finish", budget=100, tools=[])
+ assert captured[0]["agent"].model == "physical-worker"
+ assert captured[0]["agent"].base_url == "https://worker.example/v1"
+ explicit = OpenCollab(tmp_path, config=config, model="physical-worker")
+ assert resolve_model_route(explicit._config)["model"] == "unrelated-model"
+
+
+@pytest.mark.parametrize("entry", ["sdk_team", "cli_team", "sdk_workflow", "cli_workflow"])
+@pytest.mark.parametrize("physical_model", [None, "gpt-4o"])
+def test_cli_and_sdk_keep_unrouted_models_on_the_default_backend(
+ tmp_path, monkeypatch, entry, physical_model,
+):
+ from typer.testing import CliRunner
+
+ from opencollab.adapters.cli.main import app
+
+ requests = []
+
+ def respond(request):
+ requests.append((str(request.url), dict(request.headers), json.loads(request.content)))
+ return completion_http_response(request, text="finished")
+
+ install_sdk_transport(monkeypatch, "openai", respond)
+ install_sdk_transport(monkeypatch, "anthropic", respond)
+ routes = {"writer": {
+ "model": "claude-sonnet-4-6", "provider": "anthropic", "base_url": "https://writer.example",
+ "auth_token": "route-token",
+ }}
+ model = "writer" if entry.endswith("team") else physical_model or "writer"
+ monkeypatch.setenv("OPENCOLLAB_MODEL", model)
+ monkeypatch.setenv("OPENCOLLAB_API_KEY", "default-key")
+ monkeypatch.setenv("OPENCOLLAB_BASE_URL", "https://default.example/v1")
+ monkeypatch.setenv("OPENCOLLAB_MODEL_ROUTES", json.dumps(routes))
+ monkeypatch.setenv("OPENCOLLAB_THINKING_PARAMS", "{}")
+ monkeypatch.setenv("OPENCOLLAB_MAX_OUTPUT_TOKENS", "4096")
+ team = tmp_path / "team.yaml"
+ team.write_text(
+ "entry: reader\nroles:\n reader:\n prompt: Finish.\n tools: []\n"
+ + (f" model: {physical_model}\n" if physical_model else ""),
+ )
+ if entry.startswith("sdk"):
+ client = OpenCollab(tmp_path)
+ if entry == "sdk_team":
+ result = asyncio.run(client.team("finish", config=team, use_worktrees=False, trace=False))
+ else:
+ async def flow(ctx, _inputs):
+ return await ctx.agent("finish", tools=[])
+
+ result = asyncio.run(client.workflow(flow, trace=False))
+ assert result.ok
+ assert result.output == "finished"
+ else:
+ if entry == "cli_team":
+ arguments = ["--team-config", str(team), "--no-worktrees", "--prompt", "finish"]
+ else:
+ workflows = tmp_path / "workflows"
+ workflows.mkdir()
+ (workflows / "finish.py").write_text(
+ "from opencollab.application.workflow_registry import workflow\n"
+ "@workflow(name='finish')\n"
+ "async def run(ctx, args):\n"
+ " return await ctx.agent('finish', tools=[])\n",
+ )
+ arguments = ["workflow", "run", "finish", "--no-save", "--no-trace"]
+ result = CliRunner().invoke(app, [*arguments, "--workspace", str(tmp_path)])
+ assert result.exit_code == 0, result.output
+ assert "finished" in result.output
+ assert len(requests) == 1
+ url, headers, body = requests[0]
+ if physical_model:
+ assert url == "https://default.example/v1/chat/completions"
+ assert headers["authorization"] == "Bearer default-key"
+ assert body["model"] == physical_model
+ else:
+ assert url == "https://writer.example/v1/messages"
+ assert headers["authorization"] == "Bearer route-token"
+ assert body["model"] == "claude-sonnet-4-6"
+
+
+@pytest.mark.parametrize("command", [[], ["workflow", "run", "duo"]])
+@pytest.mark.parametrize("secret_field", ["api_key", "auth_token"])
+def test_cli_route_validation_keeps_credentials_out_of_diagnostics(tmp_path, monkeypatch, command, secret_field):
+ from typer.testing import CliRunner
+
+ from opencollab.adapters.cli.main import app
+
+ secret = "DEMO" # pragma: allowlist secret
+ monkeypatch.setenv("OPENCOLLAB_MODEL_ROUTES", json.dumps({"worker": {
+ secret_field: secret, "model": "physical", "request_pool": str(tmp_path),
+ }}))
+ result = CliRunner().invoke(app, [*command, "--workspace", str(tmp_path)])
+ diagnostic = "".join(result.output.replace("\u2502", "").split())
+ assert result.exit_code == 2
+ assert "model_routes.worker" in diagnostic
+ assert "request_poolandrequest_capacitymustbesuppliedtogether" in diagnostic
+ assert secret not in result.output
+
+
+@pytest.mark.parametrize("entry", ["sdk", "cli"])
+def test_team_entry_route_supplies_its_own_cli_credentials(tmp_path, monkeypatch, entry):
+ from typer.testing import CliRunner
+
+ from opencollab.adapters.cli.main import app
+
+ requests = []
+
+ def respond(request):
+ requests.append((str(request.url), dict(request.headers), json.loads(request.content)))
+ return completion_http_response(request, text="finished")
+
+ install_sdk_transport(monkeypatch, "openai", respond)
+ monkeypatch.setenv("OPENCOLLAB_MODEL_ROUTES", json.dumps({"writer": {
+ "model": "physical-writer", "base_url": "https://writer.example/v1",
+ "api_key": "route-key", # pragma: allowlist secret
+ }}))
+ team = tmp_path / "team.yaml"
+ team.write_text("entry: reader\nroles:\n reader:\n prompt: Finish.\n model: writer\n tools: []\n")
+ if entry == "sdk":
+ result = asyncio.run(OpenCollab(tmp_path).team("finish", config=team, use_worktrees=False, trace=False))
+ assert result.ok
+ assert result.output == "finished"
+ else:
+ result = CliRunner().invoke(app, [
+ "--workspace", str(tmp_path), "--team-config", str(team), "--no-worktrees", "--prompt", "finish",
+ ])
+ assert result.exit_code == 0, result.output
+ assert "finished" in result.output
+ assert len(requests) == 1
+ url, headers, body = requests[0]
+ assert url == "https://writer.example/v1/chat/completions"
+ assert headers["authorization"] == "Bearer route-key"
+ assert body["model"] == "physical-writer"
diff --git a/tests/runtime/test_model_request_admission.py b/tests/runtime/test_model_request_admission.py
new file mode 100644
index 00000000..108b3306
--- /dev/null
+++ b/tests/runtime/test_model_request_admission.py
@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import subprocess
+import sys
+
+import pytest
+
+from opencollab.adapters.llm import client as transport
+from opencollab.adapters.llm.client import LLMClient
+from opencollab.adapters.llm.request_admission import request_slot
+from opencollab.adapters.llm.types import LLMResponse
+from tests.support.paths import PACKAGE_ROOT
+
+
+async def test_cancellation_retains_slot_until_native_cleanup_finishes(tmp_path, monkeypatch):
+ entered = asyncio.Event()
+ cleaning = asyncio.Event()
+ release = asyncio.Event()
+ calls = []
+
+ async def complete(*_args, **_kwargs):
+ calls.append(len(calls))
+ if len(calls) == 1:
+ entered.set()
+ try:
+ await asyncio.Event().wait()
+ except asyncio.CancelledError:
+ cleaning.set()
+ await release.wait()
+ raise
+ return LLMResponse(content="finished")
+
+ monkeypatch.setattr(transport, "complete_openai", complete)
+ kwargs = {"api_key": "local-test", "request_pool": str(tmp_path), "request_capacity": 1}
+ async with LLMClient(**kwargs) as first, LLMClient(**kwargs) as second:
+ active = asyncio.create_task(first.complete([]))
+ await asyncio.wait_for(entered.wait(), 2)
+ active.cancel()
+ await asyncio.wait_for(cleaning.wait(), 2)
+ queued = asyncio.create_task(second.complete([]))
+ await asyncio.sleep(0.12)
+ assert calls == [0]
+ assert not queued.done()
+ release.set()
+ with pytest.raises(asyncio.CancelledError):
+ await active
+ assert (await asyncio.wait_for(queued, 2)).content == "finished"
+ assert calls == [0, 1]
+
+
+async def test_cancelled_waiter_never_releases_someone_elses_slot(tmp_path):
+ async with request_slot(str(tmp_path), 1):
+ async def wait():
+ async with request_slot(str(tmp_path), 1):
+ pytest.fail("a cancelled waiter entered an occupied pool")
+
+ queued = asyncio.create_task(wait())
+ await asyncio.sleep(0.1)
+ queued.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await queued
+ following = asyncio.create_task(wait())
+ await asyncio.sleep(0.1)
+ assert not following.done()
+ following.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await following
+ async with request_slot(str(tmp_path), 1):
+ pass
+
+
+async def test_native_failure_keeps_slot_until_provider_returns(tmp_path, monkeypatch):
+ entered = asyncio.Event()
+ release = asyncio.Event()
+ calls = []
+
+ async def complete(*_args, **_kwargs):
+ calls.append(len(calls))
+ if len(calls) == 1:
+ entered.set()
+ await release.wait()
+ raise ValueError("terminal provider failure")
+ return LLMResponse(content="next")
+
+ monkeypatch.setattr(transport, "complete_openai", complete)
+ kwargs = {"api_key": "local-test", "request_pool": str(tmp_path), "request_capacity": 1}
+ async with LLMClient(**kwargs) as first, LLMClient(**kwargs) as second:
+ active = asyncio.create_task(first.complete([]))
+ await asyncio.wait_for(entered.wait(), 2)
+ queued = asyncio.create_task(second.complete([]))
+ await asyncio.sleep(0.1)
+ assert calls == [0]
+ release.set()
+ with pytest.raises(ValueError, match="terminal provider failure"):
+ await active
+ assert (await asyncio.wait_for(queued, 2)).content == "next"
+
+
+async def test_pool_rejects_capacity_disagreement_and_missing_mount(tmp_path):
+ async with request_slot(str(tmp_path), 2):
+ with pytest.raises(ValueError, match="already has capacity 2"):
+ async with request_slot(str(tmp_path), 3):
+ pass
+ with pytest.raises(ValueError, match="Mount the same directory"):
+ async with request_slot(str(tmp_path / "missing-mount"), 1):
+ pass
+ assert not (tmp_path / "missing-mount").exists()
+
+
+async def test_process_death_releases_os_slot(tmp_path):
+ script = """
+import asyncio, sys
+from opencollab.adapters.llm.request_admission import request_slot
+async def main():
+ async with request_slot(sys.argv[1], 1):
+ print('entered', flush=True)
+ await asyncio.Event().wait()
+asyncio.run(main())
+"""
+ env = {**os.environ, "PYTHONPATH": str(PACKAGE_ROOT)}
+ process = subprocess.Popen(
+ [sys.executable, "-c", script, str(tmp_path)], stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+ text=True, env=env,
+ )
+ try:
+ assert await asyncio.wait_for(asyncio.to_thread(process.stdout.readline), 10) == "entered\n"
+ process.kill()
+ await asyncio.wait_for(asyncio.to_thread(process.wait), 5)
+
+ async def acquire():
+ async with request_slot(str(tmp_path), 1):
+ return "recovered"
+
+ assert await asyncio.wait_for(acquire(), 2) == "recovered"
+ finally:
+ if process.poll() is None:
+ process.kill()
+ process.wait()
+ process.stdout.close()
+ process.stderr.close()
+
+
+async def test_native_retry_stays_inside_the_same_request_slot(tmp_path, monkeypatch):
+ from tests.support.provider_sdk_http import completion_http_response, http_response, install_sdk_transport
+
+ failed = asyncio.Event()
+ retrying = asyncio.Event()
+ release = asyncio.Event()
+ calls = []
+
+ async def handler(request):
+ import json
+
+ model = json.loads(request.content)["model"]
+ calls.append(model)
+ if len(calls) == 1:
+ failed.set()
+ return http_response(request, 503, json={"error": {"message": "temporarily unavailable"}})
+ if model == "first":
+ retrying.set()
+ await release.wait()
+ return completion_http_response(request)
+
+ install_sdk_transport(monkeypatch, "openai", handler)
+ kwargs = {"api_key": "local", "request_pool": str(tmp_path), "request_capacity": 1, "max_retries": 1}
+ async with LLMClient(model="first", **kwargs) as first, LLMClient(model="second", **kwargs) as second:
+ active = asyncio.create_task(first.complete([]))
+ await asyncio.wait_for(failed.wait(), 2)
+ queued = asyncio.create_task(second.complete([]))
+ await asyncio.wait_for(retrying.wait(), 3)
+ assert calls == ["first", "first"]
+ assert not queued.done()
+ release.set()
+ await asyncio.wait_for(asyncio.gather(active, queued), 3)
+ assert calls == ["first", "first", "second"]
diff --git a/tests/runtime/test_pooled_generation_timeout.py b/tests/runtime/test_pooled_generation_timeout.py
new file mode 100644
index 00000000..ceefde22
--- /dev/null
+++ b/tests/runtime/test_pooled_generation_timeout.py
@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from opencollab.adapters.llm import client as transport
+from opencollab.adapters.llm.client import LLMClient
+from opencollab.adapters.llm.request_admission import request_slot
+from opencollab.adapters.llm.types import LLMResponse, Usage
+from opencollab.application.session_run import GenerationTimeoutError
+from opencollab.bootstrap import build_session
+from opencollab.domain.agent import Agent
+from tests.runtime.test_async_compaction import history
+from tests.support.environment_test_support import local_environment_factory as local_environment_factory
+
+
+@pytest.mark.parametrize("summary", [False, True])
+async def test_queue_wait_leaves_the_full_provider_timeout(
+ tmp_path, monkeypatch, local_environment_factory, summary,
+):
+ calls = []
+
+ async def complete(*_args, **_kwargs):
+ calls.append("started")
+ await asyncio.sleep(0.06)
+ return LLMResponse(content="Earlier context." if summary and len(calls) == 1 else "answer")
+
+ monkeypatch.setattr(transport, "complete_openai", complete)
+ async with LLMClient(
+ api_key="local", request_pool=str(tmp_path), request_capacity=1, context_window=40_000,
+ ) as client:
+ session = build_session(
+ agent=Agent(name="queued", system_prompt="system"), llm=client,
+ env=local_environment_factory(str(tmp_path)), llm_timeout=0.1,
+ )
+ session.messages = history(summary)
+ async with request_slot(str(tmp_path), 1):
+ owner = asyncio.create_task(session.run_loop())
+ await asyncio.sleep(0.18)
+ assert calls == []
+ assert not owner.done()
+ assert await asyncio.wait_for(owner, 2) == "answer"
+ assert len(calls) == (2 if summary else 1)
+ await session.aclose()
+
+
+@pytest.mark.parametrize("summary", [False, True])
+async def test_admitted_timeout_retains_slot_and_late_usage_until_cleanup(
+ tmp_path, monkeypatch, local_environment_factory, summary,
+):
+ cancelled = asyncio.Event()
+ release = asyncio.Event()
+ calls = []
+
+ async def complete(*_args, **_kwargs):
+ calls.append("started")
+ if len(calls) == 1:
+ try:
+ await asyncio.Event().wait()
+ except asyncio.CancelledError:
+ cancelled.set()
+ await release.wait()
+ return LLMResponse(content="Earlier context.", usage=Usage(13, 24))
+
+ monkeypatch.setattr(transport, "complete_openai", complete)
+ async with LLMClient(
+ api_key="local", request_pool=str(tmp_path), request_capacity=1, context_window=40_000,
+ ) as client:
+ session = build_session(
+ agent=Agent(name="slow", system_prompt="system"), llm=client,
+ env=local_environment_factory(str(tmp_path)), llm_timeout=0.08,
+ )
+ session.messages = history(summary)
+ async with request_slot(str(tmp_path), 1):
+ owner = asyncio.create_task(session.run_loop())
+ await asyncio.sleep(0.15)
+ assert not owner.done()
+ assert calls == []
+ try:
+ with pytest.raises(GenerationTimeoutError):
+ await asyncio.wait_for(owner, 1)
+ await asyncio.wait_for(cancelled.wait(), 1)
+ assert session.pending_cleanup_tasks
+ following = asyncio.create_task(client.complete([]))
+ await asyncio.sleep(0.1)
+ assert calls == ["started"]
+ assert not following.done()
+ finally:
+ release.set()
+ await asyncio.wait_for(asyncio.gather(*session.pending_cleanup_tasks, return_exceptions=True), 1)
+ await asyncio.wait_for(following, 1)
+ assert session.used_tokens == 37
+ assert session.runner.late_provider_usage == (37,)
+ assert session.pending_cleanup_tasks == ()
+ await session.aclose()
+
+
+@pytest.mark.parametrize("summary", [False, True])
+@pytest.mark.parametrize("stop", ["cancel", "deadline"])
+async def test_outer_cancellation_removes_the_waiter_before_admission(
+ tmp_path, monkeypatch, local_environment_factory, summary, stop,
+):
+ calls = []
+
+ async def complete(*_args, **_kwargs):
+ calls.append("started")
+ return LLMResponse(content="answer")
+
+ monkeypatch.setattr(transport, "complete_openai", complete)
+ async with LLMClient(
+ api_key="local", request_pool=str(tmp_path), request_capacity=1, context_window=40_000,
+ ) as client:
+ session = build_session(
+ agent=Agent(name="cancelled", system_prompt="system"), llm=client,
+ env=local_environment_factory(str(tmp_path)), llm_timeout=0.08,
+ )
+ session.messages = history(summary)
+ async with request_slot(str(tmp_path), 1):
+ owner = asyncio.create_task(session.run_loop())
+ if stop == "cancel":
+ await asyncio.sleep(0.12)
+ assert not owner.done()
+ owner.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await asyncio.wait_for(owner, 1)
+ else:
+ with pytest.raises(asyncio.TimeoutError):
+ await asyncio.wait_for(owner, 0.12)
+ await asyncio.wait_for(asyncio.gather(*session.pending_cleanup_tasks, return_exceptions=True), 1)
+ await asyncio.sleep(0.1)
+ assert calls == []
+ assert session.used_tokens == 0
+ assert session.pending_cleanup_tasks == ()
+ assert (await asyncio.wait_for(client.complete([]), 1)).content == "answer"
+ await session.aclose()
diff --git a/tests/workflows/test_workflow_task_input.py b/tests/workflows/test_workflow_task_input.py
index 3510c1f6..698f8f34 100644
--- a/tests/workflows/test_workflow_task_input.py
+++ b/tests/workflows/test_workflow_task_input.py
@@ -89,6 +89,7 @@ async def run(spec, args, **kwargs):
workflow_cli,
"resolve_config",
lambda *args: {
+ "model": "fixture-model",
"provider": "openai",
"api_key": "fixture-key", # pragma: allowlist secret
"base_url": None,