tau-coding-agent 0.1.0__py3-none-any.whl

tau/extensions/context.py

@@ -0,0 +1,462 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Awaitable, Callable, Literal
+
+if TYPE_CHECKING:
+    from tau.runtime.service import Runtime
+    from tau.session.manager import SessionManager
+    from tau.session.types import SessionEntry
+    from tau.settings.manager import SettingsManager
+    from tau.tui.ui_context import UIContext
+
+
+# ---------------------------------------------------------------------------
+# Session method option dataclasses
+# ---------------------------------------------------------------------------
+
+@dataclass
+class NewSessionOptions:
+    """Options for ExtensionContext.new_session()."""
+    parent_session: str | None = None
+    with_session: Callable[[ExtensionContext], Awaitable[None] | None] | None = None
+
+
+@dataclass
+class ForkOptions:
+    """Options for ExtensionContext.fork()."""
+    position: Literal["before", "at"] = "at"
+    with_session: Callable[[ExtensionContext], Awaitable[None] | None] | None = None
+
+
+@dataclass
+class NavigateTreeOptions:
+    """Options for ExtensionContext.navigate_tree()."""
+    summarize: bool = False
+    custom_instructions: str | None = None
+    replace_instructions: bool = False
+    label: str | None = None
+
+
+@dataclass
+class SwitchSessionOptions:
+    """Options for ExtensionContext.switch_session()."""
+    with_session: Callable[[ExtensionContext], Awaitable[None] | None] | None = None
+
+
+class ExtensionContext:
+    """
+    Runtime context passed to extension command handlers and event handlers.
+
+    Provides live session state rather than the static info available at
+    registration time (which lives on ExtensionAPI).
+
+    Handler signature:
+        async def my_cmd(ctx: ExtensionContext, args: list[str]) -> None:
+            print(ctx.cwd, ctx.model_id)
+    """
+
+    def __init__(
+        self,
+        cwd: Path,
+        settings: SettingsManager | None,
+        model_id: str,
+        provider_id: str,
+        session_manager: SessionManager | None = None,
+        layout: object | None = None,
+        runtime: Runtime | None = None,
+        model_thinking: bool = False,
+    ) -> None:
+        self._cwd = cwd
+        self._settings = settings
+        self._model_id = model_id
+        self._provider_id = provider_id
+        self._session_manager = session_manager
+        self._layout = layout
+        self._runtime = runtime
+        self._model_thinking = model_thinking
+
+    @classmethod
+    def from_runtime(cls, runtime: Runtime) -> ExtensionContext:
+        """Construct a context snapshot from the live runtime."""
+        sm = runtime.session_manager
+        agent = runtime.agent
+        llm = agent._engine.llm if agent is not None else None
+        return cls(
+            cwd=sm.cwd if sm is not None else Path("."),
+            settings=runtime.settings_manager,
+            model_id=llm.model.id if llm is not None else "",
+            provider_id=llm.provider_id if llm is not None else "",
+            session_manager=sm,
+            layout=getattr(runtime, "_layout", None),
+            runtime=runtime,
+            model_thinking=bool(llm.model.thinking) if llm is not None else False,
+        )
+
+    @property
+    def cwd(self) -> Path:
+        """Current working directory."""
+        return self._cwd
+
+    @property
+    def settings(self) -> SettingsManager | None:
+        """Settings manager, or None if not available."""
+        return self._settings
+
+    @property
+    def model_id(self) -> str:
+        """Active model identifier."""
+        return self._model_id
+
+    @property
+    def provider_id(self) -> str:
+        """Active provider identifier."""
+        return self._provider_id
+
+    @property
+    def model_thinking(self) -> bool:
+        """Whether the active model supports extended thinking."""
+        return self._model_thinking
+
+    @property
+    def mode(self) -> str:
+        """Execution mode: ``'tui'`` when running inside an interactive TUI session, ``'headless'`` otherwise."""
+        return "tui" if self._layout is not None else "headless"
+
+    @property
+    def session_entries(self) -> list[SessionEntry]:
+        """All entries in the session file across every branch (the full tree).
+
+        Use this when you need to inspect the complete history regardless of
+        which branch the user is currently on.  Most extensions should prefer
+        ``branch_entries`` instead.
+        """
+        if self._session_manager is None:
+            return []
+        return self._session_manager.get_entries()
+
+    @property
+    def branch_entries(self) -> list[SessionEntry]:
+        """Entries on the current branch only (root → current leaf, in order).
+
+        This is the correct view for restoring per-branch extension state — it
+        excludes entries that belong to abandoned branches, so you always see
+        the data that is actually in scope for the current conversation.
+
+        Example — restore saved state from the most recent custom entry::
+
+            for entry in reversed(ctx.branch_entries):
+                if isinstance(entry, CustomInfoEntry) and entry.custom_type == "my-ext":
+                    restore(entry.data)
+                    break
+        """
+        if self._session_manager is None:
+            return []
+        return self._session_manager.get_branch()
+
+    @property
+    def ui(self) -> "UIContext | None":
+        """TUI customization API. None when running outside a TUI session."""
+        if self._layout is None:
+            return None
+        from tau.tui.ui_context import UIContext
+        return UIContext(self._layout, settings=self._settings)  # type: ignore[arg-type]
+
+    @property
+    def has_ui(self) -> bool:
+        """True when dialog-capable UI is available (TUI mode)."""
+        return self._layout is not None
+
+    # ── Agent state ───────────────────────────────────────────────────────────
+
+    def is_idle(self) -> bool:
+        """Return True when the agent is not currently streaming a response."""
+        if self._runtime is None:
+            return True
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return True
+        return not getattr(agent, "_running", False)
+
+    def abort(self) -> None:
+        """Abort the current agent operation (no-op if idle)."""
+        if self._runtime is None:
+            return
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return
+        cancel_fn = getattr(agent, "cancel", None)
+        if callable(cancel_fn):
+            cancel_fn()
+
+    def shutdown(self) -> None:
+        """Gracefully shut down tau and exit."""
+        import sys
+        sys.exit(0)
+
+    def get_context_usage(self) -> dict | None:
+        """Return current context usage info or None if unavailable.
+
+        Keys: ``tokens`` (int | None), ``context_window`` (int), ``percent`` (float | None).
+        """
+        if self._runtime is None:
+            return None
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return None
+        usage = agent.get_context_usage()
+        if usage is None:
+            return None
+        tokens = getattr(usage, "tokens", None)
+        window = getattr(usage, "context_window", None) or 0
+        percent = getattr(usage, "percent", None)
+        return {"tokens": tokens, "context_window": window, "percent": percent}
+
+    def compact(self, custom_instructions: str | None = None) -> None:
+        """Trigger context compaction without waiting for completion."""
+        import asyncio
+        if self._runtime is None:
+            return
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return
+        compact_fn = getattr(agent, "compact", None)
+        if callable(compact_fn):
+            import inspect
+            result = compact_fn(custom_instructions=custom_instructions)
+            if inspect.isawaitable(result):
+                asyncio.ensure_future(result)  # type: ignore[arg-type]
+
+    def get_system_prompt(self) -> str:
+        """Return the current effective system prompt."""
+        if self._runtime is None:
+            return ""
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return ""
+        return getattr(agent, "_system_prompt", "") or ""
+
+    # ── Session control (command context) ─────────────────────────────────────
+
+    async def wait_for_idle(self) -> None:
+        """Suspend until the agent finishes its current turn."""
+        if self._runtime is None:
+            return
+        agent = self._runtime.agent
+        if agent is None:
+            return
+        await agent._engine.state.idle_event.wait()
+
+    async def new_session(self, options: NewSessionOptions | None = None) -> dict:
+        """Start a new session.  Returns ``{"cancelled": bool}``.
+
+        ``options.with_session(ctx)`` runs inside the new session before the UI
+        transitions — use it to inject initial messages via ``ctx.send_user_message``.
+        """
+        if self._runtime is None:
+            return {"cancelled": True}
+        opts = options or NewSessionOptions()
+        await self._runtime.new_session(with_session=opts.with_session)
+        return {"cancelled": False}
+
+    async def fork(self, entry_id: str, options: ForkOptions | None = None) -> dict:
+        """Fork from a specific entry.  Returns ``{"cancelled": bool}``."""
+        if self._runtime is None:
+            return {"cancelled": True}
+        opts = options or ForkOptions()
+        await self._runtime.fork_session(
+            entry_id, position=opts.position, with_session=opts.with_session
+        )
+        return {"cancelled": False}
+
+    async def navigate_tree(
+        self,
+        target_id: str,
+        *,
+        summarize: bool = False,
+        custom_instructions: str | None = None,
+        options: NavigateTreeOptions | None = None,
+    ) -> dict:
+        """Navigate the session tree.  Returns ``{"cancelled": bool}``."""
+        if self._runtime is None:
+            return {"cancelled": True}
+        opts = options or NavigateTreeOptions(
+            summarize=summarize, custom_instructions=custom_instructions
+        )
+        ok = await self._runtime.navigate_tree(
+            target_id,
+            summarize=opts.summarize,
+            custom_instructions=opts.custom_instructions,
+            replace_instructions=opts.replace_instructions,
+            label=opts.label,
+        )
+        return {"cancelled": not ok}
+
+    async def switch_session(
+        self,
+        session_path: str,
+        options: SwitchSessionOptions | None = None,
+    ) -> dict:
+        """Switch to a different session file.  Returns ``{"cancelled": bool}``."""
+        if self._runtime is None:
+            return {"cancelled": True}
+        opts = options or SwitchSessionOptions()
+        await self._runtime.resume_session(Path(session_path), with_session=opts.with_session)
+        return {"cancelled": False}
+
+    async def send_message(self, content: str) -> None:
+        """Append a plain text message to the current session as a user turn.
+
+        Useful inside ``with_session`` callbacks to seed a new session with
+        context before the agent processes its first real turn.
+        """
+        if self._runtime is None:
+            return
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return
+        engine = getattr(agent, "_engine", None)
+        if engine is None:
+            return
+        from tau.message.types import UserMessage, TextContent
+        msg = UserMessage(contents=[TextContent(content=content)])
+        await engine.steer(msg)
+
+    async def send_user_message(
+        self,
+        content: str,
+        deliver_as: Literal["steer", "follow_up"] = "steer",
+    ) -> None:
+        """Inject a user message into the active engine queue.
+
+        ``deliver_as='steer'`` inserts immediately (mid-turn interception).
+        ``deliver_as='follow_up'`` queues for after the current turn completes.
+        """
+        if self._runtime is None:
+            return
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return
+        engine = getattr(agent, "_engine", None)
+        if engine is None:
+            return
+        from tau.message.types import UserMessage, TextContent
+        msg = UserMessage(contents=[TextContent(content=content)])
+        if deliver_as == "follow_up":
+            await engine.follow_up(msg)
+        else:
+            await engine.steer(msg)
+
+    async def is_project_trusted(self) -> bool | None:
+        """Return the trust status of the current project directory.
+
+        Checks in order:
+        1. The ``project_trust`` hook (extensions can override).
+        2. The settings manager's cached trust state (set at startup or by the prompt).
+        3. Returns ``None`` if undecided.
+        """
+        if self._runtime is None:
+            return None
+        sm_session = getattr(self._runtime, "session_manager", None)
+        cwd = str(sm_session.cwd) if sm_session is not None else ""
+
+        from tau.hooks.types import ProjectTrustEvent, ProjectTrustResult
+        results = await self._runtime.hooks.emit(ProjectTrustEvent(project_dir=cwd))
+        for r in results:
+            if isinstance(r, ProjectTrustResult) and r.trusted is not None:
+                return r.trusted
+
+        # Fall back to the settings manager's runtime trust state
+        sm = getattr(self._runtime, "settings_manager", None)
+        if sm is not None:
+            return sm.is_project_trusted()
+        return None
+
+    def set_project_trusted(self, trusted: bool, *, remember: bool = False) -> None:
+        """Set the project's trust state for this session.
+
+        ``remember=True`` persists the decision to ``~/.tau/trust.json``
+        so it survives future sessions.
+        """
+        sm = getattr(self._runtime, "settings_manager", None) if self._runtime else None
+        if sm is None:
+            return
+        sm.set_project_trusted(trusted)
+        if remember:
+            from tau.trust.manager import trust_store
+            sm_session = getattr(self._runtime, "session_manager", None)
+            cwd = sm_session.cwd if sm_session is not None else None
+            if cwd is not None:
+                trust_store.set(str(cwd), trusted)
+
+    async def reload(self) -> None:
+        """Reload extensions, skills, prompts, and settings."""
+        if self._runtime is None:
+            return
+        await self._runtime.reload_extensions()
+
+    # ── Context inspection helpers ─────────────────────────────────────────────
+
+    def has_pending_messages(self) -> bool:
+        """Return True if any steering or follow-up messages are queued.
+
+        Useful to avoid injecting a duplicate message when one is already
+        waiting to be processed by the engine.
+        """
+        if self._runtime is None:
+            return False
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return False
+        has_fn = getattr(agent, "has_pending_messages", None)
+        if callable(has_fn):
+            return bool(has_fn())
+        return False
+
+    @property
+    def signal(self) -> "object | None":
+        """The current abort signal (``asyncio.Event``) while the agent is streaming.
+
+        The event is *set* when the current operation has been aborted.
+        Returns ``None`` when the agent is idle or unavailable.
+        """
+        if self._runtime is None:
+            return None
+        agent = getattr(self._runtime, "agent", None)
+        if agent is None:
+            return None
+        return getattr(agent, "_signal", None)
+
+    def get_system_prompt_options(self) -> dict:
+        """Return metadata about how the active system prompt was assembled.
+
+        Keys:
+        - ``skills`` — list of skill names currently loaded
+        - ``prompts`` — list of prompt template names currently loaded
+        - ``tools`` — list of tool names registered with the engine
+        - ``system_prompt_length`` — character length of the built prompt
+        """
+        from tau.skills.registry import skill_registry
+        from tau.prompts.registry import prompt_registry
+
+        skill_names = [s.name for s in skill_registry.list()]
+        prompt_names = [p.name for p in prompt_registry.list()]
+
+        tool_names: list[str] = []
+        if self._runtime is not None:
+            agent = getattr(self._runtime, "agent", None)
+            if agent is not None:
+                engine = getattr(agent, "_engine", None)
+                if engine is not None:
+                    tools = getattr(engine, "tools", None) or []
+                    tool_names = [t.name for t in tools]
+
+        prompt_len = len(self.get_system_prompt())
+
+        return {
+            "skills": skill_names,
+            "prompts": prompt_names,
+            "tools": tool_names,
+            "system_prompt_length": prompt_len,
+        }

tau/extensions/events.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections import defaultdict
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+class EventBus:
+    """
+    Lightweight pub/sub bus for cross-extension communication.
+
+    Separate from the main hook bus — intended for extensions to send
+    signals to each other rather than to the core runtime.
+
+    Usage:
+        # Publisher
+        await tau.events.emit("my-ext:done", {"count": 42})
+
+        # Subscriber (in another extension or same one)
+        @tau.events.on("my-ext:done")
+        async def handler(data):
+            print(data["count"])
+    """
+
+    def __init__(self) -> None:
+        self._handlers: dict[str, list[Callable]] = defaultdict(list)
+
+    def on(self, event: str, handler: Callable | None = None) -> Any:
+        """Subscribe to a named event.
+
+        Direct call:
+            tau.events.on("my-ext:done", handler)
+
+        Decorator:
+            @tau.events.on("my-ext:done")
+            async def handler(data): ...
+        """
+        if handler is not None:
+            self._handlers[event].append(handler)
+            return lambda: self._handlers[event].remove(handler)
+
+        def decorator(fn: Callable) -> Callable:
+            """Register the function as an event handler."""
+            self._handlers[event].append(fn)
+            return fn
+
+        return decorator
+
+    async def emit(self, event: str, data: Any = None) -> list[Any]:
+        """Fire all handlers subscribed to *event*. Returns collected results."""
+        results: list[Any] = []
+        for handler in list(self._handlers.get(event, [])):
+            try:
+                result = handler(data)
+                if asyncio.iscoroutine(result):
+                    result = await result
+                results.append(result)
+            except Exception:
+                logger.exception("EventBus handler %r raised on event %r", handler, event)
+        return results
+
+    def clear(self, event: str | None = None) -> None:
+        """Remove all handlers for one event, or all events if None."""
+        if event is None:
+            self._handlers.clear()
+        else:
+            self._handlers.pop(event, None)