feat: support embedding Agent with instance tools, skills, and storage

src/bub/builtin/__init__.py

@@ -0,0 +1,3 @@
+from .agent import Agent
+
+__all__ = ["Agent"]

src/bub/builtin/agent.py

@@ -25,15 +25,10 @@
 from bub.envelope import field_of
 from bub.framework import BubFramework
 from bub.skills import discover_skills, render_skills_prompt
-from bub.store import AsyncTapeStoreAdapter, InMemoryTapeStore, is_async_tape_store
+from bub.store import AsyncTapeStore, AsyncTapeStoreAdapter, InMemoryTapeStore, TapeStore, is_async_tape_store
 from bub.streaming import AsyncStreamEvents, StreamEvent, StreamState
 from bub.tape import Tape
-from bub.tools import (
-    REGISTRY,
-    Tool,
-    ToolContext,
-    model_tools,
-)
+from bub.tools import REGISTRY, Tool, ToolContext, model_tools
 from bub.turn import TurnState
 from bub.utils import workspace_from_state
 
@@ -44,18 +39,53 @@
 class Agent:
     """Agent that processes prompts using hooks, tools, tape, and any-llm-sdk."""
 
-    def __init__(self, framework: BubFramework) -> None:
+    def __init__(
+        self,
+        framework: BubFramework,
+        *,
+        tools: Collection[Tool] | None = None,
+        tape_store: TapeStore | AsyncTapeStore | None = None,
+        skill_dirs: Collection[Path] | None = None,
+    ) -> None:
+        """Create a builtin agent with instance-specific tools, skills, and storage.
+
+        Args:
+            framework: Configured hook runtime supplying prompts, tape context,
+                interception hooks, and optional shared resources.
+            tools: Tools available to this instance. None snapshots the global
+                registry; an empty collection disables tools.
+            tape_store: Explicit store, preferred over the framework's active
+                store. Without either, the agent uses an in-memory store.
+            skill_dirs: Skill roots in precedence order. None uses project, user,
+                and builtin discovery; an empty collection disables discovery.
+
+        Settings come from Bub's process-wide configuration. The caller owns the
+        lifecycle of an explicitly supplied store.
+        """
         self.settings = load_settings()
         self.framework = framework
+        self.tools = {tool.name: tool for tool in tools} if tools is not None else REGISTRY.copy()
+        self.tape_store = tape_store
+        self.skill_dirs = skill_dirs
         self.model_runner = ModelRunner(self.settings, hooks=framework.get_agent_hooks())
 
     @cached_property
     def tape(self) -> Tape:
+        """Return the lazily constructed, cached tape factory for this agent.
+
+        Select the explicit store, active framework store, or an in-memory fallback,
+        in that order. Adapt synchronous stores and use hook-provided context and
+        sidecars. Archive files use ``bub.home / 'tapes'`` independently of the store.
+        """
         import bub
 
-        tape_store = self.framework.get_tape_store()
-        if tape_store is None:
-            tape_store = InMemoryTapeStore()
+        tape_store: TapeStore | AsyncTapeStore | None
+        if self.tape_store is not None:
+            tape_store = self.tape_store
+        else:
+            tape_store = self.framework.get_tape_store()
+            if tape_store is None:
+                tape_store = InMemoryTapeStore()
         if not is_async_tape_store(tape_store):
             tape_store = AsyncTapeStoreAdapter(tape_store)
         return Tape(
@@ -91,17 +121,49 @@ async def run_stream(
         *,
         session_id: str,
         prompt: str | list[dict],
-        state: TurnState,
+        state: TurnState | None = None,
         model: str | None = None,
         allowed_skills: Collection[str] | None = None,
         allowed_tools: Collection[str] | None = None,
+        reasoning_effort: str | None = None,
     ) -> AsyncStreamEvents:
+        """Prepare a turn and return its stream; await this method before iterating.
+
+        Args:
+            session_id: Session identity within the workspace. A ``temp/`` prefix
+                prevents the turn's fork from merging back into its parent tape.
+            prompt: Text or multimodal content parts. Text beginning with a comma
+                after stripping whitespace invokes a builtin command.
+            state: Mutable turn state. None loads state through framework hooks
+                using this agent's store; supplied state skips that loading.
+                The current agent is always bound into the state.
+            model: Per-turn override, ahead of the state and configured model.
+            allowed_skills: Case-insensitive skill names available to this turn;
+                None leaves discovery unrestricted.
+            allowed_tools: Instance tool names or model aliases for the agent loop;
+                None allows all instance tools and an empty collection allows none.
+                Command execution uses the instance's tools directly.
+            reasoning_effort: Per-turn override of the value in state.
+
+        Consume the stream to completion to finish execution and tape merging.
+        A ``final`` event ends a model step, not necessarily the whole turn.
+        The returned object exposes ``error`` and ``usage``; execution can also
+        raise exceptions. This method does not render or dispatch outbound messages,
+        call save-state hooks, or serialize concurrent turns in the same session.
+        """
         if not prompt:
             return self._events_from_iterable([
                 StreamEvent("text", {"delta": "error: empty prompt"}),
                 StreamEvent("final", {"text": "error: empty prompt", "ok": False}),
             ])
 
+        if state is None:
+            state = await self.framework.build_state({"_runtime_agent": self}, session_id)
+        state["_runtime_agent"] = self  # Override the agent to the current instance.
+        if model is None:
+            model = state.get("model")
+        if reasoning_effort is not None:
+            state["reasoning_effort"] = reasoning_effort
         state.setdefault("session_id", session_id)
         tape = self.tape.session_tape(
             session_id, workspace_from_state(state), context=replace(self.tape.context, state=state)
@@ -139,13 +201,15 @@ async def _run_command(self, tape: Tape, *, line: str) -> str:
         output = ""
         status = "ok"
         try:
-            if name not in REGISTRY:
-                output = await REGISTRY["bash"].run(context=context, cmd=line)
+            if name not in self.tools:
+                if "bash" not in self.tools:
+                    raise ValueError("bash tool is not available")  # noqa: TRY301
+                output = await self.tools["bash"].run(context=context, cmd=line)
             else:
                 args = _parse_args(arg_tokens)
-                if REGISTRY[name].context:
+                if self.tools[name].context:
                     args.kwargs["context"] = context
-                output = REGISTRY[name].run(*args.positional, **args.kwargs)
+                output = self.tools[name].run(*args.positional, **args.kwargs)
                 if inspect.isawaitable(output):
                     output = await output
         except Exception as exc:
@@ -311,7 +375,7 @@ async def _stream_events_with_auto_handoff(
     def _load_skills_prompt(self, prompt: str, workspace: Path, allowed_skills: set[str] | None = None) -> str:
         skill_index = {
             skill.name.casefold(): skill
-            for skill in discover_skills(workspace)
+            for skill in discover_skills(workspace, skill_dirs=self.skill_dirs)
             if allowed_skills is None or skill.name.casefold() in allowed_skills
         }
         expanded_skills = set(HINT_RE.findall(prompt)) & set(skill_index.keys())
@@ -330,14 +394,14 @@ async def _run_once(
         if allowed_tools is not None:
             from bub.builtin.tools import resolve_tool_names
 
-            allowed_tools = resolve_tool_names(allowed_tools)
+            allowed_tools = resolve_tool_names(allowed_tools, all_names=self.tools)
         if allowed_skills is not None:
             allowed_skills = {name.casefold() for name in allowed_skills}
             tape.context.state["allowed_skills"] = list(allowed_skills)
         if allowed_tools is not None:
-            tools = [tool for tool in REGISTRY.values() if tool.name in allowed_tools]
+            tools = [tool for tool in self.tools.values() if tool.name in allowed_tools]
         else:
-            tools = list(REGISTRY.values())
+            tools = list(self.tools.values())
         return await self._run_once_stream(
             tape=tape,
             prompt=prompt,
@@ -394,7 +458,7 @@ def _system_prompt(
         blocks: list[str] = []
         if result := self.framework.get_system_prompt(prompt=prompt, state=state):
             blocks.append(result)
-        tools_prompt = render_tools_prompt(tools if tools is not None else REGISTRY.values())
+        tools_prompt = render_tools_prompt(tools if tools is not None else self.tools.values())
         if tools_prompt:
             blocks.append(tools_prompt)
         workspace = workspace_from_state(state)

src/bub/builtin/hook_impl.py

@@ -83,12 +83,14 @@ def __init__(self, framework: BubFramework) -> None:
         self.framework = framework
         self._agent: Agent | None = None
 
-    def _get_agent(self) -> Agent:
+    def _get_agent(self, state: TurnState | None = None) -> Agent:
+        if state and "_runtime_agent" in state:
+            return cast("Agent", state["_runtime_agent"])
         if self._agent is None:
             self._agent = Agent(self.framework)
         return self._agent
 
-    async def _recover_session_model(self, session_id: str) -> str | None:
+    async def _recover_session_model(self, session_id: str, *, agent: Agent) -> 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
@@ -97,17 +99,17 @@ async def _recover_session_model(self, session_id: str) -> str | None:
         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)
+        session = 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
 
-    async def _recover_session_reasoning_effort(self, session_id: str) -> str | None:
+    async def _recover_session_reasoning_effort(self, session_id: str, *, agent: Agent) -> str | None:
         """Recover the latest per-session reasoning effort override."""
-        session = self._get_agent().tape.session_tape(session_id, self.framework.workspace)
+        session = 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") == "reasoning_effort_switch":
@@ -166,15 +168,19 @@ async def load_state(self, message: ChannelMessage, session_id: str) -> TurnStat
         lifespan = field_of(message, "lifespan")
         if lifespan is not None:
             await lifespan.__aenter__()
-        state = {"session_id": session_id, "_runtime_agent": self._get_agent()}
+        # SDK calls supply their agent before recovery so state comes from its store.
+        agent = field_of(message, "_runtime_agent")
+        if agent is None:
+            agent = self._get_agent()
+        state = {"session_id": session_id, "_runtime_agent": 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):
+        if model := await self._recover_session_model(session_id, agent=agent):
             state["model"] = model
-        if reasoning_effort := await self._recover_session_reasoning_effort(session_id):
+        if reasoning_effort := await self._recover_session_reasoning_effort(session_id, agent=agent):
             state["reasoning_effort"] = reasoning_effort
         if model := field_of(message, "context", {}).get("model"):
             state["model"] = model
@@ -227,7 +233,7 @@ async def build_prompt(self, message: ChannelMessage, session_id: str, state: Tu
 
     @hookimpl
     async def run_model_stream(self, prompt: str | list[dict], session_id: str, state: TurnState) -> AsyncStreamEvents:
-        return await self._get_agent().run_stream(
+        return await self._get_agent(state).run_stream(
             session_id=session_id,
             prompt=prompt,
             state=state,
@@ -425,9 +431,11 @@ async def before_tool_call(
         replace it with a guidance ``tool_result`` so the model can re-issue a
         valid call on the next step.
         """
-        from bub.tools import REGISTRY, model_tools
+        from bub.tools import model_tools
 
-        available_tools = tuple(tool_item.name for tool_item in model_tools(REGISTRY.values()))
+        agent = self._get_agent(state)
+
+        available_tools = tuple(tool_item.name for tool_item in model_tools(agent.tools.values()))
         if call.tool in available_tools:
             return None
 

src/bub/builtin/tools.py

@@ -25,9 +25,10 @@ def _to_model_name(name: str) -> str:
     return name.replace(".", "_")
 
 
-def _tool_name_index() -> dict[str, str]:
-    real_names = {tool_name.casefold(): tool_name for tool_name in REGISTRY}
-    alias_names = {_to_model_name(tool_name).casefold(): tool_name for tool_name in REGISTRY}
+def _tool_name_index(all_names: Iterable[str]) -> dict[str, str]:
+    names = tuple(all_names)
+    real_names = {tool_name.casefold(): tool_name for tool_name in names}
+    alias_names = {_to_model_name(tool_name).casefold(): tool_name for tool_name in names}
     return {**alias_names, **real_names}
 
 
@@ -36,15 +37,15 @@ def resolve_tool_name(name: str) -> str | None:
     key = name.strip().casefold()
     if not key:
         return None
-    return _tool_name_index().get(key)
+    return _tool_name_index(REGISTRY).get(key)
 
 
-def _resolve_explicit_tool_names(names: Iterable[str]) -> tuple[set[str], set[str]]:
+def _resolve_explicit_tool_names(names: Iterable[str], index: dict[str, str]) -> tuple[set[str], set[str]]:
     resolved: set[str] = set()
     unknown: set[str] = set()
     for name in names:
         normalized_name = name.strip()
-        if resolved_name := resolve_tool_name(normalized_name):
+        if resolved_name := index.get(normalized_name.casefold()):
             resolved.add(resolved_name)
         else:
             unknown.add(normalized_name)
@@ -56,15 +57,19 @@ def _raise_unknown_tool_names(names: set[str]) -> None:
     raise ValueError(f"unknown tool name(s): {formatted}")
 
 
-def resolve_tool_names(names: Iterable[str] | None = None, *, exclude: Iterable[str] = ()) -> set[str]:
+def resolve_tool_names(
+    names: Iterable[str] | None = None, *, exclude: Iterable[str] = (), all_names: Iterable[str] | None = None
+) -> set[str]:
     """Resolve tool names from either runtime names or model-facing aliases."""
-    excluded, unknown_excluded = _resolve_explicit_tool_names(exclude)
+    available = tuple(REGISTRY if all_names is None else all_names)
+    index = _tool_name_index(available)
+    excluded, unknown_excluded = _resolve_explicit_tool_names(exclude, index)
     if unknown_excluded:
         _raise_unknown_tool_names(unknown_excluded)
     if names is None:
-        return set(REGISTRY) - excluded
+        return set(available) - excluded
 
-    resolved, unknown = _resolve_explicit_tool_names(names)
+    resolved, unknown = _resolve_explicit_tool_names(names, index)
     if unknown:
         _raise_unknown_tool_names(unknown)
     return resolved - excluded
@@ -239,12 +244,13 @@ def skill_describe(name: str | None = None, *, context: ToolContext) -> str:
     """
     from bub.utils import workspace_from_state
 
+    agent = _get_agent(context)
     allowed_skills = context.state.get("allowed_skills")
     if allowed_skills is not None and name and name.casefold() not in allowed_skills:
         return f"(skill '{name}' is not allowed in this context)"
 
     workspace = workspace_from_state(context.state)
-    skill_index = {skill.name: skill for skill in discover_skills(workspace)}
+    skill_index = {skill.name: skill for skill in discover_skills(workspace, skill_dirs=agent.skill_dirs)}
     if name is None:
         return "Available skills:\n" + "\n".join(f"- {skill.name}" for skill in skill_index.values())
     if name.casefold() not in skill_index:
@@ -338,7 +344,7 @@ async def run_subagent(param: SubAgentInput, *, context: ToolContext) -> str:
     else:
         subagent_session = param.session
     state = {**context.state, "session_id": subagent_session}
-    allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"})
+    allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.tools)
     output = ""
     async for event in await agent.run_stream(
         session_id=subagent_session,