diff --git a/raven/agent/loop/main.py b/raven/agent/loop/main.py index 1ddadc64..ccaf0b66 100644 --- a/raven/agent/loop/main.py +++ b/raven/agent/loop/main.py @@ -553,6 +553,21 @@ def __init__( self._consolidation_tasks: set[asyncio.Task] = set() + # A switch that arrives mid-turn is parked here until no turn is + # running, so a turn in flight finishes on the provider it started + # with. A depth counter, not a flag: OriginPools gates USER and + # system origins on independent semaphores with no global cap + # (spine/scheduler.py), so a user turn and a cron turn overlap on + # this loop under the TUI defaults. + self._pending_provider: tuple[LLMProvider, str] | None = None + self._turns_in_flight = 0 + + # ``self.subagents``, ``self.context_engine`` and + # ``self.memory_consolidator`` were each handed ``provider`` earlier in + # this constructor and hold their own reference; ``_adopt_provider`` is + # what keeps them from outliving a live model switch. Add the call + # there when adding another holder. + # Phase B-3: the L4 facade (``DefaultMemoryEngine`` / # ``MemoryEngine`` ABC) has been retired. AgentLoop now holds # the underlying subsystems directly: @@ -607,6 +622,62 @@ def _apply_disabled_tools(self) -> None: if self.tools.has(name): self.tools.unregister(name) + def set_provider(self, provider: LLMProvider, model: str) -> None: + """Point the loop and everything it built at a new provider/model. + + ``config.set model`` builds a provider from the prospective config + and hands it here. Assigning ``self.provider`` alone is not enough: + the subagent manager, the context engine's LLM-backed segments and + the consolidator each captured the provider handed to them in + ``__init__``. Left behind, they keep calling the old endpoint for + the rest of the process -- which is how switching away from a dead + credential fixed the main loop while subagents and the skill + rewriter/gate went on failing to authenticate. + + A switch that lands while any turn is running is parked rather than + applied: the loop reads ``self.provider`` at call time (eight sites + in this module, plus the context engine and consolidator + underneath), so adopting mid-turn would relay one conversation + across two vendors. How that surfaces depends on the path: the + ``chat_with_retry`` sites turn a rejected request into + ``finish_reason="error"`` content, so the turn reports a failure + with no sign that its endpoint moved, while ``_llm_call_stream`` + (which a TUI turn takes) catches only ``TimeoutError`` and lets the + rejection propagate. Neither is a diagnosis the user can act on. + + The park is the second line of defence, not the first: the RPC + rejects a switch outright when the caller's own session has a turn + in flight (``is_turn_active`` in ``tui_rpc.methods.config``). This + covers what that guard cannot see -- a caller that passes no + ``session_id``, and the proactive turns that run in their own lanes. + Note the RPC still answers ``applied: True`` and the config file is + already written, so a parked switch is applied on disk while the + loop reports the old model until the last turn drains. + + Detached subagents are not covered by that park -- they outlive the + turn that spawned them -- so ``SubagentManager`` snapshots instead. + """ + if self._turns_in_flight: + self._pending_provider = (provider, model) + return + self._adopt_provider(provider, model) + + def _adopt_provider(self, provider: LLMProvider, model: str) -> None: + """Hand a provider to the loop and every subsystem holding the old one.""" + self.provider = provider + self.model = model + self.subagents.set_provider(provider, model) + self.context_engine.set_provider(provider, model) + self.memory_consolidator.set_provider(provider, model) + + def _adopt_pending_provider(self) -> None: + """Apply a parked switch. Callers must check that no turn is running.""" + pending = self._pending_provider + if pending is None: + return + self._pending_provider = None + self._adopt_provider(*pending) + def configure_personalization(self, enable: bool) -> None: """Global switch for the 4-step personalization flow (PAHF-inspired). @@ -2558,14 +2629,60 @@ async def run_turn( inline_tool_stream: bool = False, usage_sink: dict[str, Any] | None = None, text_sink: dict[str, Any] | None = None, + ) -> TurnOutcome: + """Turn boundary for a live ``/model`` switch; see ``_run_turn``. + + A parked switch is adopted here, before the turn reads + ``self.provider`` for the first time, and the count kept for the + duration is what parks the next one. Wrapping rather than + snapshotting because the provider is read from ``self`` at eight + sites in this module and by the context engine and consolidator + underneath them -- one boundary covers all of them, a snapshot + would have to be threaded through each. + + Both ends gate on zero, because turns overlap: a user turn and a + proactive turn hold slots in separate pools. Adopting on the way in + would otherwise land a switch parked for a turn that is still + running, and clearing a flag on the way out would unpark it just as + wrongly. Adopting again on the last exit is what keeps a park from + outliving the turns it was waiting on. + """ + if self._turns_in_flight == 0: + self._adopt_pending_provider() + self._turns_in_flight += 1 + try: + return await self._run_turn( + req, + emit, + drain, + stream=stream, + inline_tool_stream=inline_tool_stream, + usage_sink=usage_sink, + text_sink=text_sink, + ) + finally: + self._turns_in_flight -= 1 + if self._turns_in_flight == 0: + self._adopt_pending_provider() + + async def _run_turn( + self, + req: TurnRequest, + emit: Emit, + drain: Drain, + *, + stream: bool = True, + inline_tool_stream: bool = False, + usage_sink: dict[str, Any] | None = None, + text_sink: dict[str, Any] | None = None, ) -> TurnOutcome: """Spine-native turn entry: consume a TurnRequest, fan the agent's output onto the single ``emit``, return a TurnOutcome. Collapses the legacy output paths (a str return + the five callbacks) onto one boundary. Named ``run_turn`` rather than ``run``: ``run`` is the runtime keep-alive - (executor / debug server / MCP up, then idle). A spine runner wraps this - method to satisfy the TurnRunner protocol. + (executor / debug server / MCP up, then idle). A spine runner calls the + public ``run_turn`` to satisfy the TurnRunner protocol. ``stream`` is the canon Q2-D assembly switch: a streaming outlet (TUI) wires it True so the reply goes out as StreamDelta and dissolves (b2 — no diff --git a/raven/agent/subagent/manager.py b/raven/agent/subagent/manager.py index 1d17e90e..3e6ef382 100644 --- a/raven/agent/subagent/manager.py +++ b/raven/agent/subagent/manager.py @@ -75,6 +75,19 @@ def __init__( # pruned to the rolling window on access, so it self-bounds. self._session_spawn_times: dict[str, deque[float]] = {} + def set_provider(self, provider: LLMProvider, model: str) -> None: + """Adopt the provider a live ``/model`` switch just built. + + Subagents run on the parent's provider, so a switch that is not + propagated here leaves every spawn calling the credential the loop + has already abandoned. Only spawns requested after this call are + affected: a subagent is a detached task that outlives the turn that + spawned it, so the loop's park cannot cover it and ``spawn`` + snapshots the pair it was asked for. + """ + self.provider = provider + self.model = model + async def spawn( self, task: str, @@ -107,7 +120,13 @@ async def spawn( display_label = label or task[:30] + ("..." if len(task) > 30 else "") origin = {"channel": origin_channel, "chat_id": origin_chat_id, "session_key": quota_key} - bg_task = asyncio.create_task(self._run_subagent(task_id, task, display_label, origin)) + # Snapshot here rather than where the task starts running: it queues + # behind the concurrency gate and a sandbox boot first, and a switch + # landing in that window would hand this task an endpoint the user + # chose after asking for it. + bg_task = asyncio.create_task( + self._run_subagent(task_id, task, display_label, origin, self.provider, self.model) + ) self._running_tasks[task_id] = bg_task if session_key: self._session_tasks.setdefault(session_key, set()).add(task_id) @@ -131,6 +150,8 @@ async def _run_subagent( task: str, label: str, origin: dict[str, str], + provider: LLMProvider, + model: str, ) -> None: """Execute the subagent task and announce the result.""" logger.info("Subagent [{}] starting task: {}", task_id, label) @@ -141,7 +162,7 @@ async def _run_subagent( async with self._gate: executor = build_executor(self._sandbox_config, self.workspace, self._owned_ids) async with executor: - await self._run_subagent_inner(task_id, task, label, origin, executor) + await self._run_subagent_inner(task_id, task, label, origin, executor, provider, model) except Exception as e: error_msg = f"Error: {str(e)}" logger.error("Subagent [{}] failed: {}", task_id, e) @@ -154,6 +175,8 @@ async def _run_subagent_inner( label: str, origin: dict[str, str], executor: Any, + provider: LLMProvider, + model: str, ) -> None: try: # Build subagent tools (no message tool, no spawn tool) @@ -191,10 +214,10 @@ async def _run_subagent_inner( while iteration < max_iterations: iteration += 1 - response = await self.provider.chat_with_retry( + response = await provider.chat_with_retry( messages=messages, tools=tools.get_definitions(), - model=self.model, + model=model, ) if response.has_tool_calls: diff --git a/raven/context_engine/assembler.py b/raven/context_engine/assembler.py index a2c990dd..debf165d 100644 --- a/raven/context_engine/assembler.py +++ b/raven/context_engine/assembler.py @@ -37,6 +37,7 @@ if TYPE_CHECKING: from raven.context_engine.curator import TurnContext + from raven.providers.base import LLMProvider class ContextAssembler(ContextEngine): @@ -64,6 +65,15 @@ def owns_compaction(self) -> bool: # the full append-only log and skips the host MemoryConsolidator. return True + def set_provider(self, provider: "LLMProvider", model: str) -> None: + # Duck-typed on purpose: only the builders that actually call an LLM + # implement it, and putting it on the SegmentBuilder protocol would + # force an empty override onto every purely textual builder. + for builder in self._builders: + setter = getattr(builder, "set_provider", None) + if callable(setter): + setter(provider, model) + async def assemble( self, session_key: str, diff --git a/raven/context_engine/base.py b/raven/context_engine/base.py index 398387ff..3b8e8722 100644 --- a/raven/context_engine/base.py +++ b/raven/context_engine/base.py @@ -35,6 +35,7 @@ # for ``ContextEngine``, so referencing ``TurnContext`` only in type # hints keeps the loop unbroken. from raven.context_engine.curator import TurnContext + from raven.providers.base import LLMProvider # --------------------------------------------------------------------------- @@ -142,6 +143,17 @@ def owns_compaction(self) -> bool: lets the engine manage history compaction itself (Curator archives messages out-of-band).""" + def set_provider(self, provider: "LLMProvider", model: str) -> None: + """Adopt the provider a live ``/model`` switch just built. + + Segments that call an LLM hold the provider handed to them at + construction; without this they keep calling the old one for the + rest of the process. Concrete rather than abstract so a future + implementation with no LLM-backed segment is not forced to write an + empty override; ``ContextAssembler`` is the only one today and does + override it. + """ + @abstractmethod async def assemble( self, diff --git a/raven/context_engine/curator.py b/raven/context_engine/curator.py index 9864b213..732ee683 100644 --- a/raven/context_engine/curator.py +++ b/raven/context_engine/curator.py @@ -352,6 +352,12 @@ def __init__( # CuratorSegmentBuilder before any build/validate call. self.prefix: "AssembledPrefix | None" = None + def set_provider(self, provider: LLMProvider, model: str) -> None: + """Adopt the provider a live ``/model`` switch just built.""" + self.provider = provider + self.model = model + self.trimmer.set_provider(provider, model) + @staticmethod def working_state_segment(working_state: str | None) -> str: """Render segment 6 text (``# Curator Working State``) or ``""``.""" diff --git a/raven/context_engine/history_trimmer.py b/raven/context_engine/history_trimmer.py index d40d08ba..b4a62cac 100644 --- a/raven/context_engine/history_trimmer.py +++ b/raven/context_engine/history_trimmer.py @@ -80,6 +80,12 @@ def __init__( self.get_tool_definitions = get_tool_definitions self.context_window_tokens = context_window_tokens + def set_provider(self, provider: LLMProvider, model: str) -> None: + """Adopt the provider a live ``/model`` switch just built, so token + estimates keep matching the model actually being called.""" + self.provider = provider + self.model = model + # ------------------------------------------------------------------ # Pure history-shaping helpers (no token estimation / no I/O) # ------------------------------------------------------------------ diff --git a/raven/context_engine/segments/curator.py b/raven/context_engine/segments/curator.py index c443021b..1c2055d7 100644 --- a/raven/context_engine/segments/curator.py +++ b/raven/context_engine/segments/curator.py @@ -82,6 +82,20 @@ def __init__( ) self._turn_ids: dict[str, str] = {} + def set_provider(self, provider: LLMProvider, model: str) -> None: + """Adopt the provider a live ``/model`` switch just built. + + ``curator_model`` is re-derived with the constructor's own + expression, so a switch cannot make it mean something it did not + mean at build time. The default is non-empty, so in practice it is + a pin; an explicitly empty ``context.curator_model`` is the one + config that follows the agent model, and it follows it here too. + """ + self.provider = provider + self.model = model + self.curator_model = self.config.curator_model or model + self.assembler.set_provider(provider, model) + async def build(self, ctx: AssemblyContext) -> Segment | None: if ctx.prefix is None: raise RuntimeError("CuratorSegmentBuilder requires ctx.prefix (phase B)") diff --git a/raven/context_engine/segments/skills.py b/raven/context_engine/segments/skills.py index 948a8465..dfd53b7a 100644 --- a/raven/context_engine/segments/skills.py +++ b/raven/context_engine/segments/skills.py @@ -41,6 +41,7 @@ from raven.memory_engine.skill_forge.gate import LLMGateFilter from raven.memory_engine.skill_forge.rewriter import QueryRewriter from raven.memory_engine.skill_forge.types import RouterHit + from raven.providers.base import LLMProvider from raven.skill_hub import SkillHubClient log = logging.getLogger(__name__) @@ -73,6 +74,14 @@ def __init__( self._hub_client = hub_client self._get_tool_definitions = get_tool_definitions + def set_provider(self, provider: "LLMProvider", model: str) -> None: + """Hand a live ``/model`` switch down to the two LLM users in this + segment. The router itself holds no provider.""" + if self._rewriter is not None: + self._rewriter.set_provider(provider, model) + if self._gate is not None: + self._gate.set_provider(provider, model) + @trace.instrument("skill.inject", kind="skill", detached=True, extract=semconv.skill_inject_skills) async def build(self, ctx: AssemblyContext) -> Segment | None: if self._router is None: diff --git a/raven/memory_engine/consolidate/consolidator.py b/raven/memory_engine/consolidate/consolidator.py index 853ce093..1a9daa32 100644 --- a/raven/memory_engine/consolidate/consolidator.py +++ b/raven/memory_engine/consolidate/consolidator.py @@ -1723,6 +1723,11 @@ def __init__( self.enable_foresight = enable_foresight self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + def set_provider(self, provider: "LLMProvider", model: str) -> None: + """Adopt the provider a live ``/model`` switch just built.""" + self.provider = provider + self.model = model + def get_lock(self, session_key: str) -> asyncio.Lock: """Return the shared consolidation lock for one session.""" return self._locks.setdefault(session_key, asyncio.Lock()) diff --git a/raven/memory_engine/skill_forge/gate.py b/raven/memory_engine/skill_forge/gate.py index d9fbe33d..10184dc9 100644 --- a/raven/memory_engine/skill_forge/gate.py +++ b/raven/memory_engine/skill_forge/gate.py @@ -61,6 +61,21 @@ def __init__( self._temperature = temperature self._max_tokens = max_tokens + def set_provider(self, provider: "LLMProvider", model: str) -> None: + """Adopt the provider a live ``/model`` switch just built. + + Unset, ``_model`` follows whatever the new provider defaults to. + + Set, it is a pin, and this leaves it pinned -- which is what a + restart on the new model would produce, since the gate is built + with the agent's provider and the pin regardless of which vendor + the pin names. Note that a pin is only a model id: the credential + comes from the provider, so a pin naming a vendor the provider does + not serve was already broken at boot, and stays broken here. + """ + del model + self._provider = provider + @trace.instrument("skill.gate", kind="skill", extract=semconv.skill_gate) async def filter( self, diff --git a/raven/memory_engine/skill_forge/rewriter.py b/raven/memory_engine/skill_forge/rewriter.py index 7009e03b..5c3dd3b6 100644 --- a/raven/memory_engine/skill_forge/rewriter.py +++ b/raven/memory_engine/skill_forge/rewriter.py @@ -74,6 +74,17 @@ def __init__( self._max_tokens = max_tokens self._temperature = temperature + def set_provider(self, provider: "LLMProvider", model: str) -> None: + """Adopt the provider a live ``/model`` switch just built. + + Held rather than looked up per call, so without this the rewriter + keeps calling the provider captured at construction after the loop + has moved on -- a switch away from an unusable credential fixes the + main path and leaves this one failing. + """ + del model # the rewriter runs on the provider's default model + self._provider = provider + @trace.instrument("skill.rewrite", kind="skill", extract=semconv.skill_rewrite) async def analyze(self, query: str) -> RewriteResult: truncated = (query or "").strip()[:_QUERY_MAX_LENGTH] diff --git a/raven/tui_rpc/methods/config.py b/raven/tui_rpc/methods/config.py index 91239a26..15567c17 100644 --- a/raven/tui_rpc/methods/config.py +++ b/raven/tui_rpc/methods/config.py @@ -389,8 +389,10 @@ def _set_model( _save_config(payload) if loop is not None: - loop.provider = built_provider - loop.model = raw_value + # Not a two-attribute assignment: the loop hands its provider to the + # subagent manager, the context engine and the consolidator at build + # time, and each keeps it. set_provider is what reaches them. + loop.set_provider(built_provider, raw_value) return {"applied": True, "previous": previous, "value": raw_value} diff --git a/tests/test_agent_loop_model_switch.py b/tests/test_agent_loop_model_switch.py new file mode 100644 index 00000000..365ddcd5 --- /dev/null +++ b/tests/test_agent_loop_model_switch.py @@ -0,0 +1,483 @@ +"""A live model switch must reach everything holding the old provider. + +``config.set key="model"`` builds a fresh provider and hands it to +``AgentLoop.set_provider``. The loop is not the only holder: the subagent +manager, the context engine's LLM-backed segments and the consolidator +each captured the provider they were built with. When the switch stopped +at ``loop.provider``, those three kept calling the abandoned endpoint -- +subagent spawns and the skill rewriter/gate failed to authenticate while +the main loop worked fine. + +Work already in flight is the other half: every call site reads the +provider from ``self`` at call time, so an unconditional swap would relay +one conversation across two vendors. The loop parks a switch until no +turn is running -- a count, not a flag, because a user turn and a +proactive turn hold slots in separate pools and overlap. A detached +subagent outlives that window entirely, so ``spawn`` snapshots the pair +it was asked for. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from raven.agent.loop.main import AgentLoop +from raven.agent.subagent.manager import SubagentManager +from raven.config.raven import ContextConfig, SkillForgeConfig +from raven.context_engine.assembler import ContextAssembler +from raven.context_engine.segments.curator import CuratorSegmentBuilder +from raven.context_engine.segments.skills import SkillsSegmentBuilder +from raven.memory_engine.skill_forge.gate import LLMGateFilter +from raven.memory_engine.skill_forge.rewriter import QueryRewriter +from raven.providers.base import LLMResponse, ToolCallRequest + +NEW_MODEL = "anthropic/claude-opus-4-8" + + +class _Provider: + """Minimal provider stand-in; identity is what the assertions track.""" + + def __init__(self, name: str = "old") -> None: + self.name = name + + def get_default_model(self) -> str: + return "fake/model" + + async def chat_with_retry(self, **kwargs) -> LLMResponse: + return LLMResponse(content="ok", finish_reason="stop") + + +class _Recorder: + """Holder that remembers the provider/model it was last pointed at.""" + + def __init__(self) -> None: + self.provider: object = "old-provider" + self.model = "old-model" + + def set_provider(self, provider: object, model: str) -> None: + self.provider = provider + self.model = model + + +class _StubExecutor: + """``_run_subagent_inner`` only passes this to ExecTool; no command runs.""" + + @property + def is_sandboxed(self) -> bool: + return False + + async def exec(self, command: str, **kwargs): # pragma: no cover - unused + raise NotImplementedError + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +def _noop_submit(*args, **kwargs) -> None: + """``_announce_result`` calls the spine submit without awaiting it.""" + return None + + +class _TextOnlyBuilder: + """A segment that never calls an LLM, so it has no set_provider.""" + + name = "identity" + order = 1 + needs_prefix = False + + async def build(self, ctx): # pragma: no cover - never invoked here + return None + + +def _loop(tmp_path) -> AgentLoop: + return AgentLoop( + provider=_Provider(), + workspace=tmp_path, + model="fake/model", + context_config=ContextConfig(), + skill_forge_config=SkillForgeConfig(), + ) + + +# --------------------------------------------------------------------------- +# Fan-out +# --------------------------------------------------------------------------- + + +def test_set_provider_reaches_every_holder() -> None: + loop = object.__new__(AgentLoop) + loop.provider = "old-provider" + loop.model = "old-model" + loop.subagents = _Recorder() + loop.context_engine = _Recorder() + loop.memory_consolidator = _Recorder() + loop._turns_in_flight = 0 + loop._pending_provider = None + + new_provider = SimpleNamespace(name="new-provider") + loop.set_provider(new_provider, NEW_MODEL) + + assert loop.provider is new_provider + assert loop.model == NEW_MODEL + for holder in (loop.subagents, loop.context_engine, loop.memory_consolidator): + assert holder.provider is new_provider + assert holder.model == NEW_MODEL + + +def test_switch_reaches_the_real_holders_a_loop_builds(tmp_path) -> None: + """The stubbed fan-out above proves the dispatcher; this proves the + receivers. Every holder here is the class a real run uses, reached by + walking the engine the factory actually assembled -- so a setter that is + renamed, dropped, or quietly wrong fails here instead of in production. + """ + loop = _loop(tmp_path) + engine = loop.context_engine + assert isinstance(engine, ContextAssembler) + + skills = next(b for b in engine._builders if isinstance(b, SkillsSegmentBuilder)) + curator = next(b for b in engine._builders if isinstance(b, CuratorSegmentBuilder)) + assert skills._gate is not None, "llm_gate_enabled defaults True; the gate is a holder" + assert skills._rewriter is not None + + new_provider = _Provider("new") + loop.set_provider(new_provider, NEW_MODEL) + + assert loop.provider is new_provider + assert loop.subagents.provider is new_provider + assert loop.subagents.model == NEW_MODEL + assert loop.memory_consolidator.provider is new_provider + assert skills._gate._provider is new_provider + assert skills._rewriter._provider is new_provider + assert curator.provider is new_provider + assert curator.assembler.provider is new_provider + assert curator.assembler.trimmer.provider is new_provider + + +def test_fan_out_targets_still_exist_on_a_real_loop(tmp_path) -> None: + """Guard against silent drift. ``ContextAssembler.set_provider`` walks its + builders duck-typed and skips anything without the method, so a renamed + holder or setter would leave the fan-out green while it quietly stops + covering that subsystem. + """ + loop = _loop(tmp_path) + for attr in ("subagents", "context_engine", "memory_consolidator"): + holder = getattr(loop, attr, None) + assert holder is not None, f"AgentLoop.{attr} is gone; _adopt_provider still calls it" + assert callable(getattr(holder, "set_provider", None)), f"AgentLoop.{attr} lost set_provider" + + +def test_assembler_forwards_to_llm_backed_builders_only() -> None: + llm_builder = _Recorder() + llm_builder.name = "skills" + llm_builder.order = 5 + llm_builder.needs_prefix = False + text_builder = _TextOnlyBuilder() + + assembler = ContextAssembler([llm_builder, text_builder], lambda: []) + new_provider = SimpleNamespace(name="new-provider") + + # The text-only builder has no set_provider; walking must skip it rather + # than blow up, which is why the fan-out is duck-typed. + assembler.set_provider(new_provider, NEW_MODEL) + + assert llm_builder.provider is new_provider + assert llm_builder.model == NEW_MODEL + + +# --------------------------------------------------------------------------- +# Pins +# --------------------------------------------------------------------------- + + +def test_rewriter_and_gate_adopt_the_new_provider() -> None: + new_provider = SimpleNamespace(name="new-provider") + + rewriter = QueryRewriter("old-provider") + rewriter.set_provider(new_provider, NEW_MODEL) + assert rewriter._provider is new_provider + + gate = LLMGateFilter("old-provider") + gate.set_provider(new_provider, NEW_MODEL) + assert gate._provider is new_provider + # Unset means the gate follows the provider's default model, so the + # switch must not pin it to the agent's model behind the user's back. + assert gate._model is None + + +def test_pinned_gate_model_survives_the_switch() -> None: + gate = LLMGateFilter("old-provider", model="openai/gpt-5-mini") + gate.set_provider(SimpleNamespace(name="new-provider"), NEW_MODEL) + assert gate._model == "openai/gpt-5-mini" + + +def test_the_default_curator_model_is_a_pin(tmp_path) -> None: + """``ContextConfig.curator_model`` has a non-empty default, so in practice + it does not follow the agent model -- at construction or across a switch. + """ + loop = _loop(tmp_path) + curator = next(b for b in loop.context_engine._builders if isinstance(b, CuratorSegmentBuilder)) + pinned = curator.curator_model + assert pinned == ContextConfig().curator_model + + loop.set_provider(_Provider("new"), NEW_MODEL) + assert curator.curator_model == pinned + + +def test_an_empty_curator_model_follows_the_agent_model_both_times(tmp_path) -> None: + """The field has no ``min_length``, so ``curator_model: ""`` validates and + the constructor's ``or model`` follows the agent model. A switch has to + follow it too, or the same config means one thing at build time and + another after. + """ + loop = AgentLoop( + provider=_Provider(), + workspace=tmp_path, + model="fake/model", + context_config=ContextConfig(curator_model=""), + skill_forge_config=SkillForgeConfig(), + ) + curator = next(b for b in loop.context_engine._builders if isinstance(b, CuratorSegmentBuilder)) + assert curator.curator_model == "fake/model" + + loop.set_provider(_Provider("new"), NEW_MODEL) + assert curator.curator_model == NEW_MODEL + + +# --------------------------------------------------------------------------- +# In-flight work +# --------------------------------------------------------------------------- + + +def test_a_switch_during_a_turn_is_parked(tmp_path) -> None: + """Every call site reads ``self.provider`` at call time, so adopting + mid-turn would send the rest of one turn to a different vendor. + """ + loop = _loop(tmp_path) + started = _Provider("started-with") + loop.set_provider(started, "started/model") + + loop._turns_in_flight = 1 + switched = _Provider("switched-to") + loop.set_provider(switched, NEW_MODEL) + + assert loop.provider is started, "the turn in flight must keep what it started with" + assert loop.subagents.provider is started + assert loop._pending_provider == (switched, NEW_MODEL) + + loop._turns_in_flight = 0 + loop._adopt_pending_provider() + + assert loop.provider is switched + assert loop.subagents.provider is switched + assert loop._pending_provider is None + + +def test_a_switch_between_turns_applies_immediately(tmp_path) -> None: + loop = _loop(tmp_path) + switched = _Provider("switched-to") + loop.set_provider(switched, NEW_MODEL) + + assert loop.provider is switched + assert loop._pending_provider is None + + +@pytest.mark.asyncio +async def test_a_running_subagent_keeps_the_provider_it_started_with(tmp_path) -> None: + """A subagent is a detached task that outlives the turn that spawned it, + so the loop's park cannot cover it -- ``spawn`` snapshots instead. Without + that, iteration k+1 calls the new vendor carrying k iterations of the old + vendor's message shapes. + """ + seen: list[tuple[str, str]] = [] + release = asyncio.Event() + + class _TwoStepProvider(_Provider): + async def chat_with_retry(self, **kwargs) -> LLMResponse: + seen.append((self.name, kwargs.get("model"))) + if len(seen) == 1: + # Hold the task open across the switch, then ask for one more + # iteration so a re-read of self.provider would show up. + await release.wait() + return LLMResponse( + content="", + tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})], + finish_reason="tool_calls", + ) + return LLMResponse(content="done", finish_reason="stop") + + manager = SubagentManager( + provider=_TwoStepProvider("started-with"), + workspace=tmp_path, + model="started/model", + ) + manager._submit = _noop_submit + + task = asyncio.create_task( + manager._run_subagent_inner( + "t1", + "do the thing", + "thing", + {"channel": "cli", "chat_id": "direct", "session_key": "s"}, + _StubExecutor(), + manager.provider, + manager.model, + ) + ) + await asyncio.sleep(0) + manager.set_provider(_TwoStepProvider("switched-to"), NEW_MODEL) + release.set() + await task + + assert [name for name, _ in seen] == ["started-with", "started-with"] + assert [model for _, model in seen] == ["started/model", "started/model"] + # The next spawn does get the new one -- the snapshot is per task, not a freeze. + assert manager.provider.name == "switched-to" + assert manager.model == NEW_MODEL + + +@pytest.mark.asyncio +async def test_run_turn_adopts_on_entry_and_releases_on_exit(tmp_path) -> None: + """The wrapper is the whole mechanism, so drive the real one. Asserting on + ``_turns_in_flight`` alone would pass with the wrapper deleted. + """ + loop = _loop(tmp_path) + started = _Provider("started-with") + loop.set_provider(started, "started/model") + + switched = _Provider("switched-to") + loop._pending_provider = (switched, NEW_MODEL) + + seen: dict[str, object] = {} + + async def _fake_run_turn(*args, **kwargs): + seen["provider"] = loop.provider + seen["depth"] = loop._turns_in_flight + return "outcome" + + loop._run_turn = _fake_run_turn + assert await loop.run_turn(None, None, None) == "outcome" + + assert seen["provider"] is switched, "a parked switch must land before the turn reads it" + assert seen["depth"] == 1 + assert loop._turns_in_flight == 0 + + +@pytest.mark.asyncio +async def test_run_turn_releases_its_slot_when_the_turn_raises(tmp_path) -> None: + """A turn that fails must not leave the loop looking busy forever -- every + later switch would be parked and never adopted. + """ + loop = _loop(tmp_path) + + async def _boom(*args, **kwargs): + raise RuntimeError("turn failed") + + loop._run_turn = _boom + with pytest.raises(RuntimeError): + await loop.run_turn(None, None, None) + + assert loop._turns_in_flight == 0 + + switched = _Provider("switched-to") + loop.set_provider(switched, NEW_MODEL) + assert loop.provider is switched + + +@pytest.mark.asyncio +async def test_a_second_turn_does_not_unpark_a_switch_under_the_first(tmp_path) -> None: + """OriginPools gates USER and system origins on independent semaphores with + no global cap, so a user turn and a cron turn overlap on one loop. A count + is what makes the park survive the shorter of the two. + """ + loop = _loop(tmp_path) + started = _Provider("started-with") + loop.set_provider(started, "started/model") + + long_turn_running = asyncio.Event() + release_long = asyncio.Event() + switched = _Provider("switched-to") + during_short: dict[str, object] = {} + + async def _long(*args, **kwargs): + long_turn_running.set() + await release_long.wait() + during_short["provider_at_end_of_long"] = loop.provider + return "long" + + async def _short(*args, **kwargs): + during_short["provider_during_short"] = loop.provider + return "short" + + loop._run_turn = _long + long_task = asyncio.create_task(loop.run_turn(None, None, None)) + await long_turn_running.wait() + + # The switch lands while only the long turn is running. + loop.set_provider(switched, NEW_MODEL) + assert loop.provider is started + assert loop._pending_provider is not None + + # A second, shorter turn starts and finishes underneath it. + loop._run_turn = _short + await loop.run_turn(None, None, None) + + assert during_short["provider_during_short"] is started, "the short turn must not adopt the park" + assert loop.provider is started, "the long turn is still running" + assert loop._turns_in_flight == 1 + + release_long.set() + await long_task + + assert loop.provider is switched, "the last turn out adopts it" + assert loop._pending_provider is None + + +@pytest.mark.asyncio +async def test_spawn_snapshots_before_the_task_queues(tmp_path) -> None: + """The snapshot must be taken in ``spawn``, not where the task starts + running: a spawn waits on the concurrency gate and a sandbox boot first, + and a switch landing in that window would hand the task an endpoint the + user chose after asking for it. + + Driven through the real ``_run_subagent`` so the window is genuinely + open -- stubbing it would prove only that ``spawn`` passes *a* pair, which + a snapshot taken later would also satisfy. + """ + import raven.agent.subagent.manager as manager_mod + + served: list[str] = [] + + class _RecordingProvider(_Provider): + async def chat_with_retry(self, **kwargs) -> LLMResponse: + served.append(self.name) + return LLMResponse(content="done", finish_reason="stop") + + manager = SubagentManager( + provider=_RecordingProvider("started-with"), + workspace=tmp_path, + model="started/model", + ) + manager._submit = _noop_submit + # Hold every spawn in exactly the window the snapshot exists for. + manager._gate = asyncio.Semaphore(0) + + original_build = manager_mod.build_executor + manager_mod.build_executor = lambda cfg, workspace, owned_ids=None: _StubExecutor() + try: + await manager.spawn("do the thing", label="thing", session_key="s") + manager.set_provider(_RecordingProvider("switched-to"), NEW_MODEL) + manager._gate.release() + for _ in range(50): + await asyncio.sleep(0) + if served: + break + finally: + manager_mod.build_executor = original_build + + assert served == ["started-with"], "the spawn ran on the provider it was asked for" diff --git a/tests/test_sandbox_unit.py b/tests/test_sandbox_unit.py index a66561b0..93945e63 100644 --- a/tests/test_sandbox_unit.py +++ b/tests/test_sandbox_unit.py @@ -1168,12 +1168,17 @@ def _patched_build(cfg, workspace, owned_ids=None): subagent_mod.build_executor = _patched_build try: # Patch the inner method so the agent loop completes quickly - async def _fast_inner(task_id, task, label, origin, executor): + async def _fast_inner(task_id, task, label, origin, executor, provider, model): await manager._announce_result(task_id, label, task, "done", origin, "ok") manager._run_subagent_inner = _fast_inner await manager._run_subagent( - "t1", "test task", "test", {"channel": "cli", "chat_id": "direct", "session_key": "cli:direct"} + "t1", + "test task", + "test", + {"channel": "cli", "chat_id": "direct", "session_key": "cli:direct"}, + manager.provider, + manager.model, ) finally: subagent_mod.build_executor = original diff --git a/tests/test_subagent_manager.py b/tests/test_subagent_manager.py index 1f65a772..15298c3f 100644 --- a/tests/test_subagent_manager.py +++ b/tests/test_subagent_manager.py @@ -94,7 +94,7 @@ async def _drive(monkeypatch, *, max_concurrent: int, spawn_n: int) -> int: state = {"current": 0, "peak": 0} release = asyncio.Event() - async def _stub_inner(task_id, task, label, origin, executor) -> None: + async def _stub_inner(task_id, task, label, origin, executor, provider, model) -> None: state["current"] += 1 state["peak"] = max(state["peak"], state["current"]) await release.wait() @@ -146,6 +146,8 @@ async def _capture_announcement(task_id, label, task, result, origin, status) -> "delete", {"channel": "tui", "chat_id": "default", "session_key": "tui:session-a"}, executor, + manager.provider, + manager.model, ) assert executor.commands == [] @@ -247,7 +249,7 @@ async def test_cancel_by_session_cancels_live_task(monkeypatch): entered = asyncio.Event() release = asyncio.Event() - async def _blocking_inner(task_id, task, label, origin, executor) -> None: + async def _blocking_inner(task_id, task, label, origin, executor, provider, model) -> None: entered.set() await release.wait() # never set — keeps the task live until cancelled diff --git a/tests/test_tui_rpc_config.py b/tests/test_tui_rpc_config.py index caa18cd3..66e031a7 100644 --- a/tests/test_tui_rpc_config.py +++ b/tests/test_tui_rpc_config.py @@ -137,10 +137,30 @@ async def test_config_set_missing_key_param_raises_validation(fake_home: Path) - # ---------------------------------------------------------------------------- +class _FakeLoop: + """Stand-in for AgentLoop's half of the switch contract. + + Records the ``set_provider`` calls: assigning ``provider``/``model`` + directly would leave the subagent manager, the context engine and the + consolidator on the old provider, so the handler must go through the + method, not the attributes. + """ + + def __init__(self, provider: object, model: str) -> None: + self.provider = provider + self.model = model + self.switches: list[tuple[object, str]] = [] + + def set_provider(self, provider: object, model: str) -> None: + self.provider = provider + self.model = model + self.switches.append((provider, model)) + + async def test_config_set_model_reassigns_loop_and_persists(fake_home: Path, monkeypatch) -> None: import raven.tui_rpc.methods.config as config_mod - loop = SimpleNamespace(provider="old-prov", model="old-model") + loop = _FakeLoop("old-prov", "old-model") new_provider = SimpleNamespace(name="new-prov") monkeypatch.setattr(config_mod, "is_turn_active", lambda _key: False) @@ -165,6 +185,9 @@ async def test_config_set_model_reassigns_loop_and_persists(fake_home: Path, mon assert result["value"] == "anthropic/claude-opus-4-8" assert loop.model == "anthropic/claude-opus-4-8" assert loop.provider is new_provider + # Routed through set_provider, so everything holding the old provider + # (subagents, context-engine segments, consolidator) gets told too. + assert loop.switches == [(new_provider, "anthropic/claude-opus-4-8")] cfg = json.loads((fake_home / ".raven" / "config.json").read_text()) assert cfg["agents"]["defaults"]["model"] == "anthropic/claude-opus-4-8" @@ -219,7 +242,7 @@ def _boom(_cfg): lambda *a, **k: SimpleNamespace(agents=SimpleNamespace(defaults=SimpleNamespace(model="", provider="auto"))), ) - loop = SimpleNamespace(provider="keep-prov", model="anthropic/claude-sonnet-4-5") + loop = _FakeLoop("keep-prov", "anthropic/claude-sonnet-4-5") with pytest.raises(ModelNotAvailableError): await config_set( { @@ -233,6 +256,7 @@ def _boom(_cfg): # Loop untouched and on-disk model preserved. assert loop.model == "anthropic/claude-sonnet-4-5" assert loop.provider == "keep-prov" + assert loop.switches == [] cfg = json.loads((fake_home / ".raven" / "config.json").read_text()) assert cfg["agents"]["defaults"]["model"] == "anthropic/claude-sonnet-4-5"