From 0410b40f57d9cdc81021fc62b5f1b48f48c78d20 Mon Sep 17 00:00:00 2001 From: jfilipiuk Date: Fri, 11 Sep 2026 18:58:50 +0200 Subject: [PATCH] feat: inherit the caller's model for async sub-agent launch and update (#446) * feat: inherit the caller's model for async sub-agent launch and update * docs: tighten middleware related docstrings --- EvoScientist/llm/patches.py | 54 +++++- .../middleware/expert_async_subagent.py | 93 ++++++++-- tests/test_expert_async_subagent.py | 174 ++++++++++++++++++ tests/test_model_passthrough_patch.py | 68 +++++++ 4 files changed, 374 insertions(+), 15 deletions(-) diff --git a/EvoScientist/llm/patches.py b/EvoScientist/llm/patches.py index c239517..bcc38fe 100644 --- a/EvoScientist/llm/patches.py +++ b/EvoScientist/llm/patches.py @@ -36,6 +36,7 @@ from collections.abc import ( Mapping, Sequence, ) +from contextvars import ContextVar from typing import Any from langchain_core.messages import AIMessage, BaseMessage @@ -1086,6 +1087,41 @@ def _patch_anthropic_structured_output() -> None: # --------------------------------------------------------------------------- _model_passthrough_patched = False +# The launching run's per-run (model, provider), set by the async-task tools +# from ``runtime.config`` for the duration of a single ``runs.create`` call. +# ``_merge_runs_config_kwargs`` reads it as the highest-precedence source. +# +# Why a ContextVar rather than passing ``config=`` through each tool: the fix +# has to hold at every ``runs.create`` site (``start_async_task`` and +# ``update_async_task``), but only the tool functions can see the caller's +# per-run model (via ``runtime.config``, the config langgraph's ToolNode +# injects into tool calls). Threading it through a ContextVar +# lets the single merge point below inject it, so ``update_async_task`` (whose +# body we delegate to upstream unchanged) is covered without reimplementing it. +_caller_configurable: ContextVar[dict[str, str] | None] = ContextVar( + "_evo_caller_configurable", default=None +) + + +def _extract_caller_configurable(config: Any) -> dict[str, str]: + """Pull ``(model, model_provider)`` out of a launching run's ``config``. + + ``config`` is the tool's ``runtime.config`` (a ``RunnableConfig``-shaped + dict). Returns a dict suitable for ``configurable``; empty when the caller + carries no model override, so the config-default fallback still applies. + ``model_provider`` is only forwarded alongside a model — a bare provider + without a model is meaningless to the deployed graph's resolver. + """ + configurable = (config or {}).get("configurable") or {} + model = configurable.get("model") + provider = configurable.get("model_provider") + out: dict[str, str] = {} + if isinstance(model, str) and model: + out["model"] = model + if isinstance(provider, str) and provider: + out["model_provider"] = provider + return out + def _read_cfg_configurable() -> dict[str, str]: """Read live ``(model, provider)`` from EvoScientist config. @@ -1114,11 +1150,21 @@ def _read_cfg_configurable() -> dict[str, str]: def _merge_runs_config_kwargs(kwargs: dict) -> dict: """Merge the live model override into ``kwargs`` for ``runs.create``. - Preserves any caller-supplied ``config.configurable`` keys. EvoScientist's - keys take precedence on conflict (callers shouldn't be passing model - overrides — the CLI is the source of truth). + Model source, in increasing precedence: + + 1. ``_read_cfg_configurable()`` — the effective EvoScientist config. On the + in-process backend this is the CLI's live per-run model; on the + ``langgraph_server`` backend this proxy runs inside the dev-server + process, where it reports the *server's* config-default model, not the + CLI's per-run choice. + 2. ``_caller_configurable`` — the launching run's actual model, forwarded + by the async-task tools from ``runtime.config``. This wins so that a + sub-agent launched (or continued) while the caller runs on, say, a free + model does not silently fall back to a billed config-default. + + Any unrelated caller-supplied ``config.configurable`` keys are preserved. """ - overrides = _read_cfg_configurable() + overrides = {**_read_cfg_configurable(), **(_caller_configurable.get() or {})} if not overrides: return kwargs existing = kwargs.get("config") diff --git a/EvoScientist/middleware/expert_async_subagent.py b/EvoScientist/middleware/expert_async_subagent.py index f6a3e1b..673d09e 100644 --- a/EvoScientist/middleware/expert_async_subagent.py +++ b/EvoScientist/middleware/expert_async_subagent.py @@ -59,6 +59,7 @@ then called without ``runtime`` and raises ``TypeError``. import asyncio import logging import threading +from contextlib import contextmanager from datetime import UTC, datetime from typing import Any, NotRequired @@ -83,6 +84,74 @@ from langgraph.types import Command _logger = logging.getLogger(__name__) +@contextmanager +def _caller_model_scope(runtime: ToolRuntime): + """Forward the launching run's model to any ``runs.create`` in the block. + + An async sub-agent is launched via a bare ``runs.create`` inside the + caller's run. Without help it falls back to the server's config-default + model rather than the model the caller is running on — so a run started on + a free model silently bills the config-default. Read the caller's per-run + model from ``runtime.config`` — the config langgraph's ToolNode injects + into every tool call, the same ``configurable`` channel that carries + ``thread_id`` (``runtime`` is already injected into these tool + signatures, so it is the channel already in hand) — and publish it, for + the duration of the block, + to the contextvar the ``runs.create`` proxy reads. A sync ``with`` around an + ``await`` is fine: the value is set before the await and reset after, and + contextvars propagate across awaits within the same task. Empty when the + caller has no override, preserving the default-model behaviour. + """ + from ..llm.patches import _caller_configurable, _extract_caller_configurable + + token = _caller_configurable.set( + _extract_caller_configurable(getattr(runtime, "config", None)) + ) + try: + yield + finally: + _caller_configurable.reset(token) + + +def _build_expert_update_tool( + agent_map: dict[str, AsyncSubAgent], + clients: Any, +) -> StructuredTool: + """``update_async_task`` wrapped to inherit the caller's model. + + Delegates to upstream's tool body verbatim — preserving its + ``multitask_strategy`` and task-envelope semantics — inside + ``_caller_model_scope`` so the follow-up ``runs.create`` reaches the + sub-agent on the caller's model, not the config-default. The explicit + ``runtime: ToolRuntime`` signature is required: langchain decides runtime + injection from ``inspect.signature``, so a ``*args`` wrapper would strip it. + """ + base = _build_update_tool(agent_map, clients) + orig_func = base.func + orig_coro = base.coroutine + + def update_async_task( + task_id: str, message: str, runtime: ToolRuntime + ) -> str | Command: + with _caller_model_scope(runtime): + return orig_func(task_id=task_id, message=message, runtime=runtime) + + async def aupdate_async_task( + task_id: str, message: str, runtime: ToolRuntime + ) -> str | Command: + with _caller_model_scope(runtime): + return await orig_coro(task_id=task_id, message=message, runtime=runtime) + + return StructuredTool.from_function( + name=base.name, + func=update_async_task, + coroutine=aupdate_async_task, + description=base.description, + infer_schema=False, + args_schema=base.args_schema, + ) + + class ExpertAsyncSubAgent(AsyncSubAgent): """AsyncSubAgent spec extended with the expert-dispatch marker. @@ -307,11 +376,12 @@ def _build_expert_start_tool( try: client = clients.get_sync(subagent_type) thread = client.threads.create() - run = client.runs.create( - thread_id=thread["thread_id"], - assistant_id=spec["graph_id"], - input=input_dict, - ) + with _caller_model_scope(runtime): + run = client.runs.create( + thread_id=thread["thread_id"], + assistant_id=spec["graph_id"], + input=input_dict, + ) except Exception as e: _logger.warning( "Failed to launch async subagent '%s': %s", subagent_type, e @@ -349,11 +419,12 @@ def _build_expert_start_tool( try: client = clients.get_async(subagent_type) thread = await client.threads.create() - run = await client.runs.create( - thread_id=thread["thread_id"], - assistant_id=spec["graph_id"], - input=input_dict, - ) + with _caller_model_scope(runtime): + run = await client.runs.create( + thread_id=thread["thread_id"], + assistant_id=spec["graph_id"], + input=input_dict, + ) except Exception as e: _logger.warning( "Failed to launch async subagent '%s': %s", subagent_type, e @@ -454,7 +525,7 @@ class EvoAsyncSubAgentMiddleware(AsyncSubAgentMiddleware): self._resolve_lock, ), _build_check_tool(clients), - _build_update_tool(agent_map, clients), + _build_expert_update_tool(agent_map, clients), _build_cancel_tool(clients), _build_list_tasks_tool(clients), ] diff --git a/tests/test_expert_async_subagent.py b/tests/test_expert_async_subagent.py index 0142ea4..07beeb4 100644 --- a/tests/test_expert_async_subagent.py +++ b/tests/test_expert_async_subagent.py @@ -807,3 +807,177 @@ class TestResolveOnMissLocking: assert b_done.is_set() assert not t1.is_alive() assert not t2.is_alive() + + +class TestCallerModelInheritance: + """start / update forward the *caller's* per-run model into ``runs.create``, + beating the config-default. + + This is the bill-the-config-default bug on the ``langgraph_server`` backend: + the model-passthrough proxy runs inside the dev-server process, where + ``_ensure_config()`` reports the server's config-default (e.g. a billed + ``gemini-3-flash-preview``) rather than the CLI's per-run choice. The + launching run's real model reaches the tool as + ``runtime.config.configurable.model``, so it must win — otherwise a + sub-agent launched (or continued) while the caller is on a free model + silently bills the config-default. + """ + + def _runtime(self, *, model="free", provider="openrouter", state=None): + ns = SimpleNamespace( + tool_call_id="tc1", + config={"configurable": {"model": model, "model_provider": provider}}, + ) + if state is not None: + ns.state = state + return ns + + def _cfg_default(self): + from EvoScientist.config.settings import EvoScientistConfig + + return EvoScientistConfig(model="gemini-3-flash-preview", provider="openrouter") + + def _tracked_task(self, agent_name="writing-agent"): + return { + "task_id": "task-abc", + "agent_name": agent_name, + "thread_id": "task-abc", + "run_id": "old-run", + "status": "running", + "created_at": "2026-05-07T00:00:00Z", + "last_checked_at": "2026-05-07T00:00:00Z", + "last_updated_at": "2026-05-07T00:00:00Z", + } + + def test_start_forwards_caller_model_over_cfg(self): + mw = EvoAsyncSubAgentMiddleware(async_subagents=[_expert_spec()]) + start = next(t for t in mw.tools if t.name == "start_async_task") + + client = _fake_sync_client() + with ( + patch( + "EvoScientist.middleware.expert_async_subagent._ClientCache.get_sync", + return_value=client, + ), + patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=self._cfg_default(), + ), + ): + start.func( + description="w", + subagent_type="literature-review", + runtime=self._runtime(), + ) + + configurable = client.runs.create.call_args.kwargs["config"]["configurable"] + assert configurable["model"] == "free" + assert configurable["model_provider"] == "openrouter" + + @pytest.mark.asyncio + async def test_astart_forwards_caller_model_over_cfg(self): + mw = EvoAsyncSubAgentMiddleware(async_subagents=[_expert_spec()]) + start = next(t for t in mw.tools if t.name == "start_async_task") + + client = _fake_async_client() + with ( + patch( + "EvoScientist.middleware.expert_async_subagent._ClientCache.get_async", + return_value=client, + ), + patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=self._cfg_default(), + ), + ): + await start.coroutine( + description="w", + subagent_type="literature-review", + runtime=self._runtime(), + ) + + configurable = client.runs.create.await_args.kwargs["config"]["configurable"] + assert configurable["model"] == "free" + assert configurable["model_provider"] == "openrouter" + + def test_update_forwards_caller_model_over_cfg(self): + mw = EvoAsyncSubAgentMiddleware(async_subagents=[_standard_spec()]) + update = next(t for t in mw.tools if t.name == "update_async_task") + + client = _fake_sync_client() + state = {"async_tasks": {"task-abc": self._tracked_task()}} + with ( + patch( + "EvoScientist.middleware.expert_async_subagent._ClientCache.get_sync", + return_value=client, + ), + patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=self._cfg_default(), + ), + ): + update.func( + task_id="task-abc", + message="keep going", + runtime=self._runtime(state=state), + ) + + kwargs = client.runs.create.call_args.kwargs + assert kwargs["config"]["configurable"]["model"] == "free" + # Upstream update semantics preserved by delegation. + assert kwargs["multitask_strategy"] == "interrupt" + + @pytest.mark.asyncio + async def test_aupdate_forwards_caller_model_over_cfg(self): + mw = EvoAsyncSubAgentMiddleware(async_subagents=[_standard_spec()]) + update = next(t for t in mw.tools if t.name == "update_async_task") + + client = _fake_async_client() + state = {"async_tasks": {"task-abc": self._tracked_task()}} + with ( + patch( + "EvoScientist.middleware.expert_async_subagent._ClientCache.get_async", + return_value=client, + ), + patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=self._cfg_default(), + ), + ): + await update.coroutine( + task_id="task-abc", + message="keep going", + runtime=self._runtime(state=state), + ) + + kwargs = client.runs.create.await_args.kwargs + assert kwargs["config"]["configurable"]["model"] == "free" + assert kwargs["multitask_strategy"] == "interrupt" + + def test_caller_scope_reset_after_start(self): + """The contextvar must not leak past the tool call — a later launch + with no override falls back to the config-default, not the prior + caller's model.""" + from EvoScientist.llm import patches as patches_mod + + mw = EvoAsyncSubAgentMiddleware(async_subagents=[_expert_spec()]) + start = next(t for t in mw.tools if t.name == "start_async_task") + + client = _fake_sync_client() + with ( + patch( + "EvoScientist.middleware.expert_async_subagent._ClientCache.get_sync", + return_value=client, + ), + patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=self._cfg_default(), + ), + ): + start.func( + description="w", + subagent_type="literature-review", + runtime=self._runtime(), + ) + # Reset restores the default (None) — nothing leaks to the next launch. + assert not patches_mod._caller_configurable.get() diff --git a/tests/test_model_passthrough_patch.py b/tests/test_model_passthrough_patch.py index cafc36c..dd3e0d0 100644 --- a/tests/test_model_passthrough_patch.py +++ b/tests/test_model_passthrough_patch.py @@ -507,3 +507,71 @@ class TestPreserveExistingConfig: assert merged["config"] == { "configurable": {"model": "gpt-5", "model_provider": "openai"} } + + +# ============================================================================= +# 8. Caller's per-run model (contextvar) beats the config-default +# ============================================================================= + + +class TestExtractCallerConfigurable: + """``_extract_caller_configurable`` pulls (model, provider) from a config.""" + + def test_none_config_is_empty(self): + assert patches_mod._extract_caller_configurable(None) == {} + + def test_no_configurable_key_is_empty(self): + assert patches_mod._extract_caller_configurable({"tags": ["x"]}) == {} + + def test_model_only(self): + cfg = {"configurable": {"model": "free"}} + assert patches_mod._extract_caller_configurable(cfg) == {"model": "free"} + + def test_model_and_provider(self): + cfg = {"configurable": {"model": "free", "model_provider": "openrouter"}} + assert patches_mod._extract_caller_configurable(cfg) == { + "model": "free", + "model_provider": "openrouter", + } + + def test_bare_provider_without_model_dropped(self): + # A provider with no model is meaningless to the deployed resolver. + cfg = {"configurable": {"model_provider": "openrouter"}} + assert patches_mod._extract_caller_configurable(cfg) == {} + + +class TestCallerConfigurableWinsOverCfg: + """The launching run's model (contextvar) beats ``_ensure_config()``. + + This is the whole point on the ``langgraph_server`` backend: the proxy + runs in the dev-server process where ``_ensure_config()`` reports the + server's config-default, so the caller's per-run model must win. + """ + + def test_caller_model_overrides_config_default(self): + token = patches_mod._caller_configurable.set( + {"model": "free", "model_provider": "openrouter"} + ) + try: + with patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=_stub_cfg( + model="gemini-3-flash-preview", provider="openrouter" + ), + ): + merged = patches_mod._merge_runs_config_kwargs({"thread_id": "t1"}) + finally: + patches_mod._caller_configurable.reset(token) + assert merged["config"]["configurable"]["model"] == "free" + assert merged["config"]["configurable"]["model_provider"] == "openrouter" + + def test_no_caller_model_falls_back_to_config_default(self): + # Contextvar unset (default {}) → config-default applies as before. + with patch( + "EvoScientist.EvoScientist._ensure_config", + return_value=_stub_cfg( + model="gemini-3-flash-preview", provider="openrouter" + ), + ): + merged = patches_mod._merge_runs_config_kwargs({"thread_id": "t1"}) + assert merged["config"]["configurable"]["model"] == "gemini-3-flash-preview"