feat(builtin): ,model <id> command for per-session model switching (#244)

* feat(builtin): add `,model <id>` command for per-session model switching

Allow switching the model for the current session from chat via the
`,model <provider:model>` command, persisted across restarts.

- BuiltinImpl.run_model_stream now forwards state["model"] as the per-call
  model override to Agent.run_stream (one-line pass-through).
- New builtin tool `model` (set_model) writes state["model"]; as a registered
  tool it is directly invocable as the `,model <id>` chat command, and also
  remains callable by the agent.
- A minimal SessionStateStore persists state["model"] per session under
  bub.home/sessions/<id>.json via load_state/save_state (atomic write). Empty
  values clear the override; nothing is written to the process environment.
- state["model"] is read once at turn start, so a mid-turn switch applies from
  the next turn (documented in the run_model_stream hookspec).

Tests (LLM mocked) cover the pass-through, the tool/command, the end-to-end
per-call override, persistence round-trip, and concurrent-session isolation.

* Persist per-session model on the tape instead of a parallel store

Drops the SessionStateStore (a second JSON persistence module under
bub.home/sessions/) and reuses the existing per-session tape to carry
state["model"]:

- set_model now records each switch as a `model_switch` event on the
  session tape (context.tape), which merges back to the persisted store
  at end of turn — no new store.
- load_state recovers the latest `model_switch` event from the session
  tape (reading the persisted store before the per-turn fork exists) and
  injects state["model"]; nothing recorded -> no injection (fresh session
  never inherits another's model).
- save_state reverts to lifespan-only; the tool owns persistence now.

Removes src/bub/builtin/session_state.py + tests/test_session_state.py;
rewrites the persistence tests against the tape. Reuses two existing
mechanisms (the in-memory state dict + the tape) rather than adding one.

`uv run pytest` (205 passed) and `uv run ruff check .` are green.

---------

Co-authored-by: yiling <yiling@ebay.com>

Author: Lynskylate

src/bub/builtin/hook_impl.py

@@ -70,6 +70,23 @@ def _get_agent(self) -> Agent:
             self._agent = Agent(self.framework)
         return self._agent
 
+    async def _recover_session_model(self, session_id: str) -> str | None:
+        """Recover the latest per-session model override recorded on the session tape.
+
+        The ``model`` tool records each switch as a ``model_switch`` event on the
+        session's tape. Scanning that tape here (before the per-turn fork exists)
+        reads the persisted store, so a choice from a prior turn or restart is
+        restored. Returns ``None`` when nothing was recorded, so a fresh session
+        never inherits another session's model.
+        """
+        session = self._get_agent().tape.session_tape(session_id, self.framework.workspace)
+        entries = list(await session.store.fetch_all(session.query().kinds("event")))
+        for entry in reversed(entries):
+            if entry.kind == "event" and entry.payload.get("name") == "model_switch":
+                model = (entry.payload.get("data") or {}).get("model")
+                return str(model) if model else None
+        return None
+
     @staticmethod
     async def _discard_message(_: ChannelMessage) -> None:
         return
@@ -118,6 +135,11 @@ async def load_state(self, message: ChannelMessage, session_id: str) -> State:
         state = {"session_id": session_id, "_runtime_agent": self._get_agent()}
         if context := field_of(message, "context_str"):
             state["context"] = context
+        # Carry over a previously recorded per-session model override from the
+        # session tape. Only set when a prior turn actually recorded one, so a
+        # fresh/unknown session never inherits another session's model.
+        if model := await self._recover_session_model(session_id):
+            state["model"] = model
         return state
 
     @hookimpl
@@ -126,6 +148,9 @@ async def save_state(self, session_id: str, state: State, message: ChannelMessag
         lifespan = field_of(message, "lifespan")
         if lifespan is not None:
             await lifespan.__aexit__(tp, value, traceback)
+        # The per-session model override is persisted on the session tape by the
+        # ``model`` tool itself (a ``model_switch`` event, merged back at end of
+        # turn), so nothing to write here — this hook only closes the lifespan.
 
     @hookimpl
     async def build_prompt(self, message: ChannelMessage, session_id: str, state: State) -> str | list[dict]:
@@ -158,7 +183,12 @@ async def build_prompt(self, message: ChannelMessage, session_id: str, state: St
 
     @hookimpl
     async def run_model_stream(self, prompt: str | list[dict], session_id: str, state: State) -> AsyncStreamEvents:
-        return await self._get_agent().run_stream(session_id=session_id, prompt=prompt, state=state)
+        return await self._get_agent().run_stream(
+            session_id=session_id,
+            prompt=prompt,
+            state=state,
+            model=state.get("model"),
+        )
 
     @hookimpl
     def register_cli_commands(self, app: typer.Typer) -> None:

src/bub/builtin/tools.py

@@ -403,6 +403,22 @@ async def quit_tool(*, context: ToolContext) -> str:
     return "Session tasks stopped."
 
 
+@tool(name="model", context=True)
+async def set_model(model_id: str, *, context: ToolContext) -> str:
+    """Switch the model for THIS session. Invoke as the `,model <model_id>` command.
+
+    Takes effect on the NEXT turn and persists across restarts. Pass any
+    ``provider:model`` string (for example ``openai:gpt-4o`` or
+    ``openrouter:openrouter/free``). An invalid model surfaces as an error on the
+    next turn — run `,model <valid_id>` again to recover.
+    """
+    context.state["model"] = model_id
+    # Persist on the session tape (merged back at end of turn); load_state
+    # recovers the latest `model_switch` event next turn / after restart.
+    await context.tape.append_event("model_switch", {"model": model_id})
+    return f"Session model set to {model_id} (applies from the next turn)."
+
+
 def _resolve_path(context: ToolContext, raw_path: str) -> Path:
     workspace = context.state.get("_runtime_workspace")
     path = Path(raw_path).expanduser()

src/bub/hookspecs.py

@@ -43,7 +43,13 @@ def run_model(self, prompt: str | list[dict], session_id: str, state: State) ->
 
     @hookspec(firstresult=True)
     def run_model_stream(self, prompt: str | list[dict], session_id: str, state: State) -> AsyncStreamEvents:
-        """Run model for one turn and return a stream of events. Should not be implemented if `run_model` is implemented."""
+        """Run model for one turn and return a stream of events. Should not be implemented if `run_model` is implemented.
+
+        Implementations may honor a runtime model override by reading
+        ``state["model"]`` (any ``provider:model`` string). The value takes
+        effect on the turn in which it is read, so a model switched mid-turn via
+        the `,model <id>` command applies from the *next* turn.
+        """
         raise NotImplementedError
 
     @hookspec

tests/test_builtin_agent.py

@@ -8,6 +8,7 @@
 import pytest
 from any_llm.types.completion import ChatCompletionChunk
 
+import bub.builtin.tools  # noqa: F401  — registers builtin tools (incl. `model`)
 from bub.builtin.agent import Agent
 from bub.builtin.model_runner import ModelRunner
 from bub.builtin.settings import AgentSettings
@@ -243,6 +244,35 @@ async def test_agent_run_model_defaults_to_none() -> None:
     assert completion_kwargs["model"] == "test:model"
 
 
+@pytest.mark.asyncio
+async def test_agent_run_model_override_does_not_mutate_default() -> None:
+    """A per-call model override must not leak into the agent's configured model.
+
+    The override is resolved per turn (``model or self.settings.model``) and
+    forwarded to any-llm; it must never be written back to ``settings.model``.
+    This is the agent-layer half of the guarantee that a session-scoped model
+    switch (state['model'] -> run_stream(model=...)) cannot bleed across
+    sessions the way a process-global env var would.
+    """
+    agent = _make_agent()
+    fork_capture = _ForkCapture()
+    agent.tape = _FakeTapeFactory(fork_capture)  # type: ignore[assignment]
+    default_model = agent.settings.model
+
+    result = await agent.run_stream(
+        session_id="user/s1",
+        prompt="hello",
+        state={"_runtime_workspace": "/tmp"},  # noqa: S108
+        model="openai:gpt-4o",
+    )
+    [event async for event in result]
+
+    completion_kwargs = _model_runner(agent).completion_kwargs
+    assert completion_kwargs is not None
+    assert completion_kwargs["model"] == "openai:gpt-4o"
+    assert agent.settings.model == default_model
+
+
 @pytest.mark.asyncio
 async def test_agent_run_resolves_allowed_tool_aliases_and_limits_prompt() -> None:
     allowed_name = "tests.allowed_agent_tool"
@@ -295,3 +325,27 @@ async def test_agent_run_rejects_unknown_allowed_tools() -> None:
 
     with pytest.raises(ValueError, match="tests_missing_agent_tool"):
         [event async for event in stream]
+
+
+@pytest.mark.asyncio
+async def test_run_command_model_switches_session_model_directly() -> None:
+    """,model <model_id> runs the `model` builtin as a chat command with no LLM call.
+
+    The command writes state['model'] on the same state object the framework
+    hands to run_stream, so the override is picked up by run_model_stream on the
+    next turn. Persistence is a `model_switch` event the tool records on the
+    session tape (merged back at end of turn), not a side effect of save_state.
+    """
+    agent = _make_agent()
+    fork_capture = _ForkCapture()
+    agent.tape = _FakeTapeFactory(fork_capture)  # type: ignore[assignment]
+    state: dict[str, Any] = {"_runtime_workspace": "/tmp"}  # noqa: S108
+    assert "model" not in REGISTRY or REGISTRY["model"].context is True
+
+    stream = await agent.run_stream(session_id="user/s1", prompt=",model openai:gpt-4o", state=state)
+    events = [event async for event in stream]
+
+    # state["model"] is mutated in place (tool context shares the framework state).
+    assert state["model"] == "openai:gpt-4o"
+    deltas = [event.data.get("delta", "") for event in events if event.kind == "text"]
+    assert any("Session model set to openai:gpt-4o" in delta for delta in deltas)

tests/test_builtin_hook_impl.py

@@ -9,9 +9,11 @@
 
 from bub.builtin.hook_impl import AGENTS_FILE_NAME, DEFAULT_SYSTEM_PROMPT, BuiltinImpl
 from bub.builtin.store import FileTapeStore
+from bub.builtin.tape import Tape
 from bub.channels.message import ChannelMessage
 from bub.framework import BubFramework
 from bub.runtime import AsyncStreamEvents, StreamEvent
+from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, TapeContext
 
 
 class RecordingLifespan:
@@ -26,18 +28,36 @@ async def __aexit__(self, exc_type, exc, traceback) -> None:
         self.exit_args = (exc_type, exc, traceback)
 
 
+def _fake_tape(home: Path) -> Tape:
+    return Tape(
+        archive_path=home / "tapes",
+        store=AsyncTapeStoreAdapter(InMemoryTapeStore()),
+        context=TapeContext(),
+    )
+
+
 class FakeAgent:
-    def __init__(self, home: Path) -> None:
+    def __init__(self, home: Path, *, tape: Tape | None = None) -> None:
         self.settings = SimpleNamespace(home=home)
+        # A real in-memory async tape so load_state's recovery path runs against
+        # the same store the tests write `model_switch` events to.
+        self.tape = tape if tape is not None else _fake_tape(home)
         self.run_calls: list[tuple[str, str, dict[str, object]]] = []
-        self.run_stream_calls: list[tuple[str, str, dict[str, object]]] = []
+        self.run_stream_calls: list[tuple[str, str, dict[str, object], str | None]] = []
 
     async def run(self, *, session_id: str, prompt: str, state: dict[str, object]) -> str:
         self.run_calls.append((session_id, prompt, state))
         return "agent-output"
 
-    async def run_stream(self, *, session_id: str, prompt: str, state: dict[str, object]) -> AsyncStreamEvents:
-        self.run_stream_calls.append((session_id, prompt, state))
+    async def run_stream(
+        self,
+        *,
+        session_id: str,
+        prompt: str,
+        state: dict[str, object],
+        model: str | None = None,
+    ) -> AsyncStreamEvents:
+        self.run_stream_calls.append((session_id, prompt, state, model))
 
         async def iterator():
             yield StreamEvent("text", {"delta": "agent-output"})
@@ -108,6 +128,43 @@ async def test_load_state_and_save_state_manage_lifespan_and_context(tmp_path: P
     assert isinstance(lifespan.exit_args[1], ValueError)
 
 
+@pytest.mark.asyncio
+async def test_load_state_injects_model_recorded_on_session_tape(tmp_path: Path) -> None:
+    """A model_switch event recorded on the session tape is restored into state on load."""
+    _, impl, agent = _build_impl(tmp_path)
+    session = agent.tape.session_tape("resolved-session", impl.framework.workspace)
+    await session.append_event("model_switch", {"model": "openai:gpt-4o"})
+
+    message = ChannelMessage(session_id="session", channel="cli", chat_id="room", content="hello")
+
+    state = await impl.load_state(message=message, session_id="resolved-session")
+
+    assert state["model"] == "openai:gpt-4o"
+
+
+@pytest.mark.asyncio
+async def test_load_state_does_not_inject_model_for_unknown_session(tmp_path: Path) -> None:
+    """A session with nothing recorded on its tape must not inherit any model (no leakage)."""
+    _, impl, _ = _build_impl(tmp_path)
+
+    message = ChannelMessage(session_id="session", channel="cli", chat_id="room", content="hello")
+
+    state = await impl.load_state(message=message, session_id="fresh-session")
+
+    assert "model" not in state
+
+
+@pytest.mark.asyncio
+async def test_recover_session_model_returns_latest_recorded(tmp_path: Path) -> None:
+    """When several switches were recorded, the most recent one wins."""
+    _, impl, agent = _build_impl(tmp_path)
+    session = agent.tape.session_tape("resolved-session", impl.framework.workspace)
+    await session.append_event("model_switch", {"model": "openai:gpt-4o"})
+    await session.append_event("model_switch", {"model": "anthropic:claude-3"})
+
+    assert await impl._recover_session_model("resolved-session") == "anthropic:claude-3"
+
+
 @pytest.mark.asyncio
 async def test_build_prompt_marks_commands_and_prefixes_context(tmp_path: Path) -> None:
     _, impl, _ = _build_impl(tmp_path)
@@ -160,10 +217,32 @@ async def test_run_model_stream_delegates_to_agent(tmp_path: Path) -> None:
     events = [event async for event in stream]
 
     assert [(event.kind, event.data) for event in events] == [("text", {"delta": "agent-output"})]
-    assert agent.run_stream_calls == [("session", "prompt", state)]
+    assert agent.run_stream_calls == [("session", "prompt", state, None)]
     assert agent.run_calls == []
 
 
+@pytest.mark.asyncio
+async def test_run_model_stream_forwards_state_model_override(tmp_path: Path) -> None:
+    """state['model'] must be forwarded as the per-call model override."""
+    _, impl, agent = _build_impl(tmp_path)
+    state = {"model": "openai:gpt-4o"}
+
+    await impl.run_model_stream(prompt="prompt", session_id="session", state=state)
+
+    assert agent.run_stream_calls == [("session", "prompt", state, "openai:gpt-4o")]
+
+
+@pytest.mark.asyncio
+async def test_run_model_stream_passes_none_when_state_has_no_model(tmp_path: Path) -> None:
+    """Without state['model'] the agent must fall back to its configured model."""
+    _, impl, agent = _build_impl(tmp_path)
+    state = {"context": "ctx"}
+
+    await impl.run_model_stream(prompt="prompt", session_id="session", state=state)
+
+    assert agent.run_stream_calls[-1][3] is None
+
+
 def test_system_prompt_appends_workspace_agents_file(tmp_path: Path) -> None:
     _, impl, _ = _build_impl(tmp_path)
     (tmp_path / AGENTS_FILE_NAME).write_text("local rules", encoding="utf-8")

tests/test_builtin_tools.py

@@ -20,6 +20,7 @@
     quit_tool,
     render_tools_prompt,
     resolve_tool_names,
+    set_model,
 )
 from bub.runtime import ErrorKind
 from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, TapeContext
@@ -151,6 +152,44 @@ def test_resolve_tool_names_rejects_unknown_names() -> None:
         resolve_tool_names(None, exclude={" tests_missing_tool "})
 
 
+def test_set_model_is_registered_with_context() -> None:
+    assert "model" in REGISTRY
+    tool_obj = REGISTRY["model"]
+    assert tool_obj.context is True
+    assert tool_obj.parameters == {
+        "type": "object",
+        "properties": {"model_id": {"type": "string"}},
+        "required": ["model_id"],
+    }
+
+
+@pytest.mark.asyncio
+async def test_set_model_writes_model_into_state_and_records_on_tape(tmp_path) -> None:
+    context = _tool_context(tmp_path)
+    assert "model" not in context.state
+
+    result = await set_model.run(model_id="openai:gpt-4o", context=context)
+
+    assert context.state["model"] == "openai:gpt-4o"
+    assert "openai:gpt-4o" in result
+    assert "next turn" in result.lower()
+    # The switch is also persisted as a `model_switch` event on the session
+    # tape, which load_state recovers on the next turn / after restart.
+    entries = list(await context.tape.store.fetch_all(context.tape.query().kinds("event")))
+    switches = [entry for entry in entries if entry.kind == "event" and entry.payload.get("name") == "model_switch"]
+    assert len(switches) == 1
+    assert switches[0].payload.get("data") == {"model": "openai:gpt-4o"}
+
+
+@pytest.mark.asyncio
+async def test_set_model_overwrites_previous_model(tmp_path) -> None:
+    context = _tool_context(tmp_path, model="openai:gpt-4o")
+
+    await set_model.run(model_id="anthropic:claude-3", context=context)
+
+    assert context.state["model"] == "anthropic:claude-3"
+
+
 @pytest.mark.asyncio
 async def test_bash_returns_stdout_for_foreground_command(tmp_path) -> None:
     result = await bash.run(cmd=_python_shell("print('hello')"), context=_tool_context(tmp_path))