bub 0.3.8__tar.gz → 0.3.9__tar.gz

{bub-0.3.8 → bub-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bub
-Version: 0.3.8
+Version: 0.3.9
 Summary: A common shape for agents that live alongside people.
 Project-URL: Homepage, https://bub.build
 Project-URL: Repository, https://github.com/bubbuild/bub

{bub-0.3.8 → bub-0.3.9}/src/bub/__init__.py

@@ -13,8 +13,20 @@ from bub.configure import Settings, config, ensure_config
 from bub.framework import DEFAULT_HOME, BubFramework
 from bub.hookspecs import hookimpl
 from bub.tools import tool
-
-__all__ = ["BubFramework", "Settings", "config", "ensure_config", "home", "hookimpl", "tool"]
+from bub.turn_admission import AdmitDecision, SteeringBuffer, TurnSnapshot
+
+__all__ = [
+    "AdmitDecision",
+    "BubFramework",
+    "Settings",
+    "SteeringBuffer",
+    "TurnSnapshot",
+    "config",
+    "ensure_config",
+    "home",
+    "hookimpl",
+    "tool",
+]
 
 try:
     __version__ = import_module("bub._version").version

{bub-0.3.8 → bub-0.3.9}/src/bub/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.3.8'
-__version_tuple__ = version_tuple = (0, 3, 8)
+__version__ = version = '0.3.9'
+__version_tuple__ = version_tuple = (0, 3, 9)
 
 __commit_id__ = commit_id = None

{bub-0.3.8 → bub-0.3.9}/src/bub/builtin/agent.py

@@ -24,6 +24,7 @@ from republic import (
     StreamEvent,
     StreamState,
     TapeContext,
+    Tool,
     ToolAutoResult,
     ToolContext,
 )
@@ -34,11 +35,11 @@ from bub.builtin.store import ForkTapeStore
 from bub.builtin.tape import TapeService
 from bub.framework import BubFramework
 from bub.skills import discover_skills, render_skills_prompt
-from bub.tools import REGISTRY, model_tools, render_tools_prompt
+from bub.tools import REGISTRY, model_tools, render_tools_prompt, resolve_tool_names
 from bub.types import State
 from bub.utils import workspace_from_state
 
-CONTINUE_PROMPT = "Continue the task."
+CONTINUE_PROMPT = "Continue the task until all targets are completed."
 HINT_RE = re.compile(r"\$([A-Za-z0-9_.-]+)")
 _CONTEXT_LENGTH_PATTERNS = re.compile(
     r"context.{0,20}(?:length|window)|maximum.{0,20}context|token.{0,10}limit|prompt.{0,10}too long|tokens? > \d+ maximum",
@@ -530,12 +531,12 @@ class Agent:
     ) -> AsyncStreamEvents | ToolAutoResult:
         prompt_text = prompt if isinstance(prompt, str) else _extract_text_from_parts(prompt)
         if allowed_tools is not None:
-            allowed_tools = {name.casefold() for name in allowed_tools}
+            allowed_tools = resolve_tool_names(allowed_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.casefold() in allowed_tools]
+            tools = [tool for tool in REGISTRY.values() if tool.name in allowed_tools]
         else:
             tools = list(REGISTRY.values())
         async with asyncio.timeout(self.settings.model_timeout_seconds):
@@ -543,7 +544,7 @@ class Agent:
                 return await tape.stream_events_async(
                     prompt=prompt,
                     system_prompt=self._system_prompt(
-                        prompt_text, state=tape.context.state, allowed_skills=allowed_skills
+                        prompt_text, state=tape.context.state, allowed_skills=allowed_skills, tools=tools
                     ),
                     max_tokens=self.settings.max_tokens,
                     tools=model_tools(tools),
@@ -553,18 +554,20 @@ class Agent:
                 return await tape.run_tools_async(
                     prompt=prompt,
                     system_prompt=self._system_prompt(
-                        prompt_text, state=tape.context.state, allowed_skills=allowed_skills
+                        prompt_text, state=tape.context.state, allowed_skills=allowed_skills, tools=tools
                     ),
                     max_tokens=self.settings.max_tokens,
                     tools=model_tools(tools),
                     model=model,
                 )
 
-    def _system_prompt(self, prompt: str, state: State, allowed_skills: set[str] | None = None) -> str:
+    def _system_prompt(
+        self, prompt: str, state: State, allowed_skills: set[str] | None = None, tools: Iterable[Tool] | None = None
+    ) -> str:
         blocks: list[str] = []
         if result := self.framework.get_system_prompt(prompt=prompt, state=state):
             blocks.append(result)
-        tools_prompt = render_tools_prompt(REGISTRY.values())
+        tools_prompt = render_tools_prompt(tools if tools is not None else REGISTRY.values())
         if tools_prompt:
             blocks.append(tools_prompt)
         workspace = workspace_from_state(state)

{bub-0.3.8 → bub-0.3.9}/src/bub/builtin/hook_impl.py

@@ -1,5 +1,5 @@
 import sys
-from datetime import UTC, datetime
+from datetime import datetime
 from pathlib import Path
 from typing import cast
 
@@ -134,7 +134,7 @@ class BuiltinImpl:
             message.kind = "command"
             return content
         context = field_of(message, "context_str")
-        now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+        now = datetime.now().astimezone().isoformat(timespec="seconds")
         context_prefix = f"{context}\n---Date: {now}---\n" if context else ""
         text = f"{context_prefix}{content}"
 

{bub-0.3.8 → bub-0.3.9}/src/bub/builtin/store.py

@@ -101,16 +101,22 @@ class ForkTapeStore:
     @contextlib.asynccontextmanager
     async def fork(self, tape: str, merge_back: bool = True) -> AsyncGenerator[None, None]:
         store = InMemoryTapeStore()
-        token = current_store.set(store)
-        tape_token = current_fork_tape.set(tape)
-        reset_token = current_tape_was_reset.set(False)
+        # Save/restore instead of ContextVar.reset(token) to avoid
+        # "Token was created in a different Context" when cleanup
+        # runs in a different asyncio Task (e.g. cancellation, TaskGroup).
+        prev_store = current_store.get(_empty_store)
+        prev_fork_tape = current_fork_tape.get()
+        prev_was_reset = current_tape_was_reset.get()
+        current_store.set(store)
+        current_fork_tape.set(tape)
+        current_tape_was_reset.set(False)
         try:
             yield
         finally:
             was_reset = current_tape_was_reset.get()
-            current_store.reset(token)
-            current_fork_tape.reset(tape_token)
-            current_tape_was_reset.reset(reset_token)
+            current_store.set(prev_store)
+            current_fork_tape.set(prev_fork_tape)
+            current_tape_was_reset.set(prev_was_reset)
             if merge_back:
                 if was_reset:
                     await self._parent.reset(tape)

bub-0.3.9/src/bub/channels/__init__.py

@@ -0,0 +1,5 @@
+from .base import Channel, Interface, Lifecycle
+from .manager import ChannelManager
+from .message import ChannelMessage
+
+__all__ = ["Channel", "ChannelManager", "ChannelMessage", "Interface", "Lifecycle"]

{bub-0.3.8 → bub-0.3.9}/src/bub/channels/base.py

@@ -39,3 +39,11 @@ class Channel(ABC):
     def stream_events(self, message: ChannelMessage, stream: AsyncIterable[StreamEvent]) -> AsyncIterable[StreamEvent]:
         """Optionally wrap the output stream for this channel."""
         return stream
+
+
+class Interface(Channel):
+    """User-facing inbound/outbound surface managed by the channel runtime."""
+
+
+class Lifecycle(Channel):
+    """Background runtime managed alongside channels."""

{bub-0.3.8 → bub-0.3.9}/src/bub/channels/cli/__init__.py

@@ -14,12 +14,11 @@ from prompt_toolkit.key_binding import KeyBindings
 from prompt_toolkit.patch_stdout import patch_stdout
 from republic import StreamEvent
 from rich import get_console
-from rich.live import Live
 
 import bub
 from bub.builtin.agent import Agent
 from bub.builtin.tape import TapeInfo
-from bub.channels.base import Channel
+from bub.channels.base import Interface
 from bub.channels.cli.renderer import CliRenderer
 from bub.channels.message import ChannelMessage
 from bub.envelope import field_of
@@ -27,7 +26,7 @@ from bub.tools import REGISTRY
 from bub.types import MessageHandler
 
 
-class CliChannel(Channel):
+class CliChannel(Interface):
     """A simple CLI channel for testing and debugging."""
 
     name = "cli"
@@ -141,23 +140,21 @@ class CliChannel(Channel):
     async def stream_events(
         self, message: ChannelMessage, stream: AsyncIterable[StreamEvent]
     ) -> AsyncIterable[StreamEvent]:
-        live: Live | None = None
-        text = ""
-        try:
-            async for event in stream:
-                if event.kind == "text":
-                    content = str(event.data.get("delta", ""))
-                    if not content.strip() and not text:
-                        continue  # skip leading whitespace-only events
-                    text += content
-                    if live is None:
-                        live = self._renderer.start_stream(message.kind, text)
-                    else:
-                        self._renderer.update_stream(live, kind=message.kind, text=text)
-                yield event
-        finally:
-            if live is not None:
-                self._renderer.finish_stream(live, kind=message.kind, text=text)
+        head_printed = False
+        console = get_console()
+        async for event in stream:
+            if event.kind == "text":
+                content = str(event.data.get("delta", ""))
+                if not content.strip() and not head_printed:
+                    continue  # skip leading whitespace-only events
+                if not head_printed:
+                    self._renderer.print_head(message.kind)
+                    head_printed = True
+
+                console.print(content, end="", highlight=False)
+            elif event.kind == "final":
+                console.print("\n")  # ensure we end with a newline after the final event
+            yield event
 
     def _build_prompt(self, workspace: Path) -> PromptSession[str]:
         kb = KeyBindings()

{bub-0.3.8 → bub-0.3.9}/src/bub/channels/cli/renderer.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 from rich.console import Console
-from rich.live import Live
 from rich.panel import Panel
 from rich.text import Text
 
@@ -33,54 +32,30 @@ class CliRenderer:
             return
         self.console.print(Text(text, style="bright_black"))
 
-    def panel(self, kind: MessageKind, text: str) -> Panel:
-        title, border_style = self._panel_style(kind)
-        return Panel(text, title=title, border_style=border_style)
-
     def command_output(self, text: str) -> None:
         if not text.strip():
             return
-        self.console.print(self.panel("command", text))
+        self.console.print(f"[cyan bold]Command >[/]\n{text}")
 
     def assistant_output(self, text: str) -> None:
         if not text.strip():
             return
-        self.console.print(self.panel("normal", text))
+        self.console.print(f"[blue bold]Assistant >[/]\n{text}")
 
     def error(self, text: str) -> None:
         if not text.strip():
             return
-        self.console.print(self.panel("error", text))
+        self.console.print(f"[red bold]Error >[/]\n{text}")
+
+    def print_head(self, kind: MessageKind) -> None:
+        if kind == "command":
+            self.console.print("[cyan bold]Command >[/]")
+        elif kind == "error":
+            self.console.print("[red bold]Error >[/]")
+        else:
+            self.console.print("[blue bold]Assistant >[/]")
 
     def log(self, message: object) -> None:
         text = str(message).rstrip()
         if text:
             self.console.print(text)
-
-    def start_stream(self, kind: MessageKind, text: str) -> Live:
-        live = Live(
-            self.panel(kind, text),
-            console=self.console,
-            auto_refresh=False,
-            transient=False,
-            vertical_overflow="visible",
-        )
-        live.start(refresh=True)
-        return live
-
-    def update_stream(self, live: Live, *, kind: MessageKind, text: str) -> None:
-        live.update(self.panel(kind, text), refresh=True)
-
-    def finish_stream(self, live: Live, *, kind: MessageKind, text: str) -> None:
-        live.update(self.panel(kind, text), refresh=True)
-        live.stop()
-
-    @staticmethod
-    def _panel_style(kind: MessageKind) -> tuple[str, str]:
-        match kind:
-            case "error":
-                return "Error", "red"
-            case "command":
-                return "Command", "green"
-            case _:
-                return "Assistant", "blue"

bub-0.3.9/src/bub/channels/manager.py

@@ -0,0 +1,343 @@
+import asyncio
+import contextlib
+import functools
+from collections.abc import AsyncIterable, Collection
+
+from loguru import logger
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
+from republic import StreamEvent
+
+from bub import config
+from bub.channels.base import Channel, Interface, Lifecycle
+from bub.channels.handler import BufferedMessageHandler
+from bub.channels.message import ChannelMessage
+from bub.configure import Settings, ensure_config
+from bub.envelope import content_of, field_of
+from bub.framework import BubFramework
+from bub.turn_admission import AdmitDecision, SessionTurnController
+from bub.types import Envelope, MessageHandler
+from bub.utils import wait_until_stopped
+
+
+@config()
+class ChannelSettings(Settings):
+    model_config = SettingsConfigDict(env_prefix="BUB_", extra="ignore", env_file=".env")
+
+    enabled_channels: str = Field(
+        default="all", description="Comma-separated list of enabled channels, or 'all' for all channels."
+    )
+    debounce_seconds: float = Field(
+        default=1.0,
+        description="Minimum seconds between processing two messages from the same channel to prevent overload.",
+    )
+    max_wait_seconds: float = Field(
+        default=10.0,
+        description="Maximum seconds to wait for processing before new messages reach the channel.",
+    )
+    active_time_window: float = Field(
+        default=60.0,
+        description="Time window in seconds to consider a channel active for processing messages.",
+    )
+    stream_output: bool = Field(default=False, description="Whether to stream model output to channels in real-time.")
+
+
+class ChannelManager:
+    def __init__(
+        self,
+        framework: BubFramework,
+        enabled_channels: Collection[str] | None = None,
+        stream_output: bool | None = None,
+    ) -> None:
+        self.framework = framework
+        self._channels: dict[str, Channel] = self.framework.get_channels(self.on_receive)
+        self._settings = ensure_config(ChannelSettings)
+        self._stream_output = stream_output if stream_output is not None else self._settings.stream_output
+        if enabled_channels is not None:
+            self._enabled_channels = list(enabled_channels)
+        else:
+            self._enabled_channels = self._settings.enabled_channels.split(",")
+        self._messages = asyncio.Queue[ChannelMessage]()
+        self._session_controllers: dict[str, SessionTurnController] = {}
+        self._session_handlers: dict[str, MessageHandler] = {}
+
+    async def on_receive(self, message: ChannelMessage) -> None:
+        channel = message.channel
+        session_id = message.session_id
+        if channel not in self._channels:
+            logger.warning(f"Received message from unknown channel '{channel}', ignoring.")
+            return
+        if session_id not in self._session_handlers:
+            handler: MessageHandler
+            if self._channels[channel].needs_debounce:
+                handler = BufferedMessageHandler(
+                    self._messages.put,
+                    active_time_window=self._settings.active_time_window,
+                    max_wait_seconds=self._settings.max_wait_seconds,
+                    debounce_seconds=self._settings.debounce_seconds,
+                )
+            else:
+                handler = self._messages.put
+            self._session_handlers[session_id] = handler
+        await self._session_handlers[session_id](message)
+
+    def get_channel(self, name: str) -> Channel | None:
+        return self._channels.get(name)
+
+    async def dispatch_output(self, message: Envelope) -> bool:
+        channel_name = field_of(message, "output_channel", field_of(message, "channel"))
+        if channel_name is None:
+            return False
+
+        channel_key = str(channel_name)
+        channel = self.get_channel(channel_key)
+        if channel is None:
+            return False
+
+        outbound = ChannelMessage(
+            session_id=str(field_of(message, "session_id", f"{channel_key}:default")),
+            channel=channel_key,
+            chat_id=str(field_of(message, "chat_id", "default")),
+            content=content_of(message),
+            context=field_of(message, "context", {}),
+            kind=field_of(message, "kind", "normal"),
+        )
+        await channel.send(outbound)
+        return True
+
+    def wrap_stream(self, message: Envelope, stream: AsyncIterable[StreamEvent]) -> AsyncIterable[StreamEvent]:
+        channel_name = field_of(message, "output_channel", field_of(message, "channel"))
+        if channel_name is None:
+            return stream
+
+        channel_key = str(channel_name)
+        channel = self.get_channel(channel_key)
+        if channel is None:
+            return stream
+
+        return channel.stream_events(message, stream)
+
+    async def quit(self, session_id: str) -> None:
+        controller = self._session_controllers.get(session_id)
+        if controller is None:
+            self.framework.clear_steering(session_id)
+            logger.info(f"channel.manager quit session_id={session_id}, cancelled 0 tasks")
+            return
+        controller.clear_pending()
+        controller.steering.drain_nowait()
+        tasks = set(controller.active_tasks)
+        current_task = asyncio.current_task()
+        cancelled_count = 0
+        for task in tasks:
+            if task is current_task:
+                continue
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+            cancelled_count += 1
+        controller.active_tasks.difference_update(task for task in tasks if task is not current_task)
+        self._drop_empty_controller(session_id)
+        logger.info(f"channel.manager quit session_id={session_id}, cancelled {cancelled_count} tasks")
+
+    def enabled_channels(self) -> list[Channel]:
+        """Channels are enabled by the following rules:
+        - Interfaces are enabled only if explicitly *included*.
+        - Lifecycles are always enabled unless explicitly *excluded*.
+        - Regular channels depend on the values of include and exclude.
+        """
+        included_channels = [
+            name.strip() for name in self._enabled_channels if name.strip() and not name.strip().startswith("!")
+        ]
+        excluded_channels = {name.strip()[1:] for name in self._enabled_channels if name.strip().startswith("!")}
+
+        if "all" in included_channels:
+            return [
+                channel
+                for channel in self._channels.values()
+                if channel.name not in excluded_channels and channel.enabled and not isinstance(channel, Interface)
+            ]
+        channels = [
+            channel
+            for name, channel in self._channels.items()
+            if name in included_channels and name not in excluded_channels and channel.enabled
+        ]
+        if not any(not isinstance(channel, Lifecycle) for channel in channels):
+            return channels
+        enabled_names = {channel.name for channel in channels}
+        channels.extend(
+            channel
+            for channel in self._channels.values()
+            if channel.name not in enabled_names
+            and channel.name not in excluded_channels
+            and channel.enabled
+            and isinstance(channel, Lifecycle)
+        )
+        return channels
+
+    def _controller(self, session_id: str) -> SessionTurnController:
+        controller = self._session_controllers.get(session_id)
+        if controller is None:
+            controller = SessionTurnController(session_id=session_id, steering=self.framework.steering(session_id))
+            self._session_controllers[session_id] = controller
+        return controller
+
+    def _drop_empty_controller(self, session_id: str) -> None:
+        controller = self._session_controllers.get(session_id)
+        if controller is None:
+            return
+        if controller.active() or controller.pending_queue or controller.steering.count > 0:
+            return
+        self._session_controllers.pop(session_id, None)
+        self.framework.clear_steering(session_id)
+
+    def _on_task_done(self, session_id: str, task: asyncio.Task) -> None:
+        if task.cancelled():
+            logger.info("channel.manager task cancelled session_id={}", session_id)
+        else:
+            task.exception()  # to log any exception
+        controller = self._session_controllers.get(session_id)
+        if controller is None:
+            return
+        controller.active_tasks.discard(task)
+        if not controller.active():
+            controller.promote_steering_to_pending()
+        self._schedule_pending(session_id)
+        self._drop_empty_controller(session_id)
+
+    async def _admit_message(self, message: ChannelMessage) -> bool:
+        try:
+            session_id = await self._resolve_message_session(message)
+        except Exception as exc:
+            logger.exception("channel.manager resolve_session failed")
+            await self.framework._hook_runtime.notify_error(stage="resolve_session", error=exc, message=message)
+            return False
+        controller = self._controller(session_id)
+        try:
+            decision = await self.framework.admit_message(
+                session_id=session_id,
+                message=message,
+                turn=controller.snapshot(),
+            )
+        except Exception as exc:
+            logger.exception("channel.manager admission hook failed")
+            await self.framework._hook_runtime.notify_error(stage="admit_message", error=exc, message=message)
+            return True
+        if decision is None:
+            self._drop_empty_controller(session_id)
+            return True
+        admitted = await self._apply_admission_decision(controller, message, decision)
+        if not admitted or not controller.active():
+            self._drop_empty_controller(session_id)
+        return admitted
+
+    async def _apply_admission_decision(
+        self,
+        controller: SessionTurnController,
+        message: ChannelMessage,
+        decision: AdmitDecision,
+    ) -> bool:
+        action = decision.action
+        if action == "process":
+            return True
+        if action == "drop":
+            logger.info(
+                "channel.manager admission drop session_id={} reason={}",
+                message.session_id,
+                decision.reason,
+            )
+            return False
+        if action == "follow_up":
+            return self._queue_pending(controller, message, decision.reason)
+        if action == "steer":
+            if controller.active():
+                controller.steering.put_nowait(message)
+                logger.info(
+                    "channel.manager admission steer session_id={} reason={}",
+                    message.session_id,
+                    decision.reason,
+                )
+                return False
+            return self._queue_pending(controller, message, decision.reason)
+        logger.warning("channel.manager admission unknown action={} session_id={}", decision.action, message.session_id)
+        return True
+
+    def _queue_pending(
+        self,
+        controller: SessionTurnController,
+        message: ChannelMessage,
+        reason: str | None,
+    ) -> bool:
+        if not controller.active():
+            return True
+        controller.add_pending(message)
+        logger.info(
+            "channel.manager admission follow_up session_id={} pending_count={} reason={}",
+            message.session_id,
+            len(controller.pending_queue),
+            reason,
+        )
+        return False
+
+    def _schedule_message(self, message: ChannelMessage) -> asyncio.Task:
+        controller = self._controller(message.session_id)
+        task = asyncio.create_task(self._run_message(message))
+        task.add_done_callback(functools.partial(self._on_task_done, message.session_id))
+        controller.active_tasks.add(task)
+        return task
+
+    def _schedule_pending(self, session_id: str) -> None:
+        controller = self._session_controllers.get(session_id)
+        if controller is None or controller.active():
+            return
+        message = controller.pop_pending()
+        if message is not None:
+            self._schedule_message(message)
+
+    async def _resolve_message_session(self, message: ChannelMessage) -> str:
+        session_id = await self.framework.resolve_session(message)
+        message.session_id = session_id
+        return session_id
+
+    async def _run_message(self, message: ChannelMessage) -> None:
+        await self.framework.process_inbound(message, self._stream_output)
+
+    async def listen_and_run(self) -> None:
+        stop_event = asyncio.Event()
+        self.framework.bind_outbound_router(self)
+        async with self.framework.running():
+            for channel in self.enabled_channels():
+                await channel.start(stop_event)
+            logger.info("channel.manager started listening")
+            try:
+                while True:
+                    message = await wait_until_stopped(self._messages.get(), stop_event)
+                    if not await self._admit_message(message):
+                        continue
+                    self._schedule_message(message)
+            except asyncio.CancelledError:
+                logger.info("channel.manager received shutdown signal")
+            except Exception:
+                logger.exception("channel.manager error")
+                raise
+            finally:
+                self.framework.bind_outbound_router(None)
+                await self.shutdown()
+                logger.info("channel.manager stopped")
+
+    async def shutdown(self) -> None:
+        count = 0
+        session_ids = list(self._session_controllers)
+        for controller in list(self._session_controllers.values()):
+            controller.clear_pending()
+            controller.steering.drain_nowait()
+            for task in set(controller.active_tasks):
+                task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+                count += 1
+        self._session_controllers.clear()
+        for session_id in session_ids:
+            self.framework.clear_steering(session_id)
+        logger.info(f"channel.manager cancelled {count} in-flight tasks")
+        for channel in self.enabled_channels():
+            await channel.stop()