pythinker-code 0.8.0__py3-none-any.whl

pythinker_code/extensions.py

@@ -0,0 +1,151 @@
+"""In-process extension API.
+
+This module exposes Pythinker's extension hooks for the registries that
+already exist in the codebase:
+
+* tool renderers           — :func:`register_tool_renderer`
+* keybindings              — :func:`register_keybinding`
+* footer status badges     — :func:`register_footer_status`
+* event handlers           — :func:`register_event_handler`
+* lifecycle hooks          — :func:`register_extension`
+
+The footer status registry is a thin placeholder for the day the card-style
+:class:`FooterState` is wired into the live prompt; until then it just
+stores entries that callers can read out.
+
+the ``register_command`` overlaps with Pythinker's existing
+``SlashCommandRegistry`` — extensions should use that registry directly.
+
+Extensions are registered with::
+
+    from pythinker_code.extensions import register_extension
+
+    def setup(ctx):
+        ctx.register_tool_renderer(my_renderer)
+        ctx.register_keybinding("app.tools.expand", "alt+e")
+        ctx.bus.on("assistant.message", on_assistant_message)
+
+    register_extension(name="my-ext", setup=setup)
+
+The host calls :func:`run_pending_extensions` once at startup; setup
+errors are logged but do not crash the host.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from pythinker_code.events import EventBus, EventHandler, create_event_bus
+from pythinker_code.ui.shell.keymap import register_keybinding as _register_keybinding
+from pythinker_code.ui.shell.tool_renderers import (
+    ToolRenderDefinition,
+)
+from pythinker_code.ui.shell.tool_renderers import (
+    register_tool_renderer as _register_tool_renderer,
+)
+from pythinker_code.utils.logging import logger
+
+__all__ = [
+    "ExtensionContext",
+    "ExtensionSetup",
+    "footer_statuses",
+    "register_extension",
+    "registered_extensions",
+    "run_pending_extensions",
+    "shared_event_bus",
+]
+
+
+_GLOBAL_BUS: EventBus = create_event_bus()
+"""Process-wide bus extensions write to. Independent from any per-session
+event streams to keep the surface stable."""
+
+_FOOTER_STATUSES: dict[str, str] = {}
+"""Read-only-from-outside map of extension id → status text."""
+
+
+def shared_event_bus() -> EventBus:
+    """Return the global :class:`EventBus` extensions subscribe to."""
+    return _GLOBAL_BUS
+
+
+def footer_statuses() -> dict[str, str]:
+    """Return a copy of the current extension footer-status map."""
+    return dict(_FOOTER_STATUSES)
+
+
+@dataclass(slots=True)
+class ExtensionContext:
+    """Handle passed into an extension's setup function.
+
+    Each method delegates to the appropriate global registry. Method names
+    mirror the ``register*`` API surface so porting an extension between
+    runtimes is mostly mechanical.
+    """
+
+    name: str
+    bus: EventBus = field(default_factory=lambda: _GLOBAL_BUS)
+
+    def register_tool_renderer(self, definition: ToolRenderDefinition) -> None:
+        _register_tool_renderer(definition)
+
+    def register_keybinding(self, semantic_id: str, *keys: str) -> None:
+        _register_keybinding(semantic_id, *keys)
+
+    def register_footer_status(self, key: str, text: str) -> None:
+        """Add or replace an extension's footer status badge.
+
+        ``key`` is normally the extension name. Pass an empty *text* to
+        clear it. The footer reads :func:`footer_statuses` when rendering.
+        """
+        if not text:
+            _FOOTER_STATUSES.pop(key, None)
+            return
+        _FOOTER_STATUSES[key] = text
+
+    def register_event_handler(self, channel: str, handler: EventHandler) -> Callable[[], None]:
+        return self.bus.on(channel, handler)
+
+
+ExtensionSetup = Callable[[ExtensionContext], None]
+
+
+@dataclass(frozen=True, slots=True)
+class _Extension:
+    name: str
+    setup: ExtensionSetup
+    started: bool = False
+
+
+_PENDING: list[_Extension] = []
+_LOADED: list[str] = []
+
+
+def register_extension(*, name: str, setup: ExtensionSetup) -> None:
+    """Queue *setup* to run when :func:`run_pending_extensions` is next called."""
+    _PENDING.append(_Extension(name=name, setup=setup))
+
+
+def run_pending_extensions() -> list[str]:
+    """Execute every pending extension's setup. Returns the names that ran.
+
+    Errors are logged and the rest of the queue continues — the same behavior.
+    """
+    started: list[str] = []
+    while _PENDING:
+        ext = _PENDING.pop(0)
+        ctx = ExtensionContext(name=ext.name)
+        try:
+            ext.setup(ctx)
+        except Exception:  # noqa: BLE001 — extension error must not crash host
+            logger.exception("Extension '{name}' setup failed", name=ext.name)
+            continue
+        _LOADED.append(ext.name)
+        started.append(ext.name)
+    return started
+
+
+def registered_extensions() -> list[str]:
+    """Snapshot of extension names that have been successfully started."""
+    return list(_LOADED)

pythinker_code/hooks/__init__.py

@@ -0,0 +1,4 @@
+from pythinker_code.hooks.config import HOOK_EVENT_TYPES, HookDef, HookEventType
+from pythinker_code.hooks.engine import HookEngine
+
+__all__ = ["HookDef", "HookEventType", "HOOK_EVENT_TYPES", "HookEngine"]

pythinker_code/hooks/config.py

@@ -0,0 +1,34 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+HookEventType = Literal[
+    "PreToolUse",
+    "PostToolUse",
+    "PostToolUseFailure",
+    "UserPromptSubmit",
+    "Stop",
+    "StopFailure",
+    "SessionStart",
+    "SessionEnd",
+    "SubagentStart",
+    "SubagentStop",
+    "PreCompact",
+    "PostCompact",
+    "Notification",
+]
+
+HOOK_EVENT_TYPES: list[str] = list(HookEventType.__args__)  # type: ignore[attr-defined]
+
+
+class HookDef(BaseModel):
+    """A single hook definition in config.toml."""
+
+    event: HookEventType
+    """Which lifecycle event triggers this hook."""
+    command: str
+    """Shell command to execute. Receives JSON on stdin."""
+    matcher: str = ""
+    """Regex pattern to filter. Empty matches everything."""
+    timeout: int = Field(default=30, ge=1, le=600)
+    """Timeout in seconds. Fail-open on timeout."""

pythinker_code/hooks/engine.py

@@ -0,0 +1,383 @@
+from __future__ import annotations
+
+import asyncio
+import re
+import time
+import uuid
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from pythinker_code.hooks.config import HookDef, HookEventType
+from pythinker_code.hooks.runner import HookResult, run_hook
+from pythinker_code.utils.logging import logger
+
+# Callback signatures for wire integration
+type OnTriggered = Callable[[str, str, int], None]
+"""(event, target, hook_count) -> None"""
+
+type OnResolved = Callable[[str, str, str, str, int], None]
+"""(event, target, action, reason, duration_ms) -> None"""
+
+type OnWireHookRequest = Callable[[WireHookHandle], Awaitable[None]]
+"""Called when a wire hook needs client handling. The callback should send
+the request over the wire and resolve the handle when the client responds."""
+
+
+@dataclass
+class WireHookSubscription:
+    """A client-side hook subscription registered via wire initialize."""
+
+    id: str
+    event: str
+    matcher: str = ""
+    timeout: int = 30
+
+
+@dataclass
+class WireHookHandle:
+    """A pending wire hook request waiting for client response."""
+
+    id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
+    subscription_id: str = ""
+    event: str = ""
+    target: str = ""
+    input_data: dict[str, Any] = field(default_factory=lambda: {})
+    _future: asyncio.Future[HookResult] | None = field(default=None, repr=False)
+
+    def _get_future(self) -> asyncio.Future[HookResult]:
+        if self._future is None:
+            self._future = asyncio.get_running_loop().create_future()
+        return self._future
+
+    async def wait(self) -> HookResult:
+        """Wait for the client to respond."""
+        return await self._get_future()
+
+    def resolve(self, action: str = "allow", reason: str = "") -> None:
+        """Resolve with client's decision."""
+        result = HookResult(action=action, reason=reason)  # type: ignore[arg-type]
+        future = self._get_future()
+        if not future.done():
+            future.set_result(result)
+
+
+class HookEngine:
+    """Loads hook definitions and executes matching hooks in parallel.
+
+    Supports two hook sources:
+    - Server-side (config.toml): shell commands executed locally
+    - Client-side (wire subscriptions): forwarded to client via HookRequest
+    """
+
+    def __init__(
+        self,
+        hooks: list[HookDef] | None = None,
+        cwd: str | None = None,
+        *,
+        on_triggered: OnTriggered | None = None,
+        on_resolved: OnResolved | None = None,
+        on_wire_hook: OnWireHookRequest | None = None,
+    ):
+        self._hooks: list[HookDef] = list(hooks) if hooks else []
+        self._wire_subs: list[WireHookSubscription] = []
+        self._cwd = cwd
+        self._on_triggered = on_triggered
+        self._on_resolved = on_resolved
+        self._on_wire_hook = on_wire_hook
+        self._by_event: dict[str, list[HookDef]] = {}
+        self._wire_by_event: dict[str, list[WireHookSubscription]] = {}
+        self._pending_fire_and_forget: set[asyncio.Task[Any]] = set()
+        self._rebuild_index()
+
+    def fire_and_forget_trigger(
+        self,
+        event: HookEventType,
+        *,
+        matcher_value: str = "",
+        input_data: dict[str, Any],
+    ) -> asyncio.Task[list[HookResult]]:
+        """Trigger a hook in the background and keep a strong reference to the
+        task. asyncio holds tasks in a WeakSet, so naively writing
+        ``asyncio.create_task(engine.trigger(...))`` and discarding the local
+        variable lets Python's GC collect the still-pending task — which fires
+        ``loop.call_exception_handler`` with no exception field and surfaces as
+        ``Unhandled exception in event loop / Exception None`` in the
+        prompt_toolkit terminal. Use this helper any time the caller wants to
+        fire a hook without awaiting its completion.
+        """
+        task: asyncio.Task[list[HookResult]] = asyncio.create_task(
+            self.trigger(event, matcher_value=matcher_value, input_data=input_data)
+        )
+        self._pending_fire_and_forget.add(task)
+        task.add_done_callback(self._pending_fire_and_forget.discard)
+        task.add_done_callback(self._log_fire_and_forget_failure)
+        return task
+
+    @staticmethod
+    def _log_fire_and_forget_failure(task: asyncio.Task[Any]) -> None:
+        if task.cancelled():
+            return
+        exc = task.exception()
+        if exc is not None:
+            logger.opt(exception=exc).warning("Fire-and-forget hook task failed")
+
+    def _rebuild_index(self) -> None:
+        self._by_event.clear()
+        for h in self._hooks:
+            self._by_event.setdefault(h.event, []).append(h)
+        self._wire_by_event.clear()
+        for s in self._wire_subs:
+            self._wire_by_event.setdefault(s.event, []).append(s)
+
+    def add_hooks(self, hooks: list[HookDef]) -> None:
+        """Add server-side hooks at runtime. Rebuilds index."""
+        self._hooks.extend(hooks)
+        self._rebuild_index()
+
+    def add_wire_subscriptions(self, subs: list[WireHookSubscription]) -> None:
+        """Register client-side hook subscriptions from wire initialize."""
+        self._wire_subs.extend(subs)
+        self._rebuild_index()
+
+    def set_callbacks(
+        self,
+        on_triggered: OnTriggered | None = None,
+        on_resolved: OnResolved | None = None,
+        on_wire_hook: OnWireHookRequest | None = None,
+    ) -> None:
+        """Set wire event callbacks."""
+        self._on_triggered = on_triggered
+        self._on_resolved = on_resolved
+        self._on_wire_hook = on_wire_hook
+
+    @property
+    def has_hooks(self) -> bool:
+        return bool(self._hooks) or bool(self._wire_subs)
+
+    def has_hooks_for(self, event: HookEventType) -> bool:
+        return bool(self._by_event.get(event)) or bool(self._wire_by_event.get(event))
+
+    @property
+    def summary(self) -> dict[str, int]:
+        """Event -> total count of hooks (server + wire)."""
+        counts: dict[str, int] = {}
+        for event, hooks in self._by_event.items():
+            counts[event] = counts.get(event, 0) + len(hooks)
+        for event, subs in self._wire_by_event.items():
+            counts[event] = counts.get(event, 0) + len(subs)
+        return counts
+
+    def details(self) -> dict[str, list[dict[str, str]]]:
+        """Event -> list of {matcher, command/type} for display."""
+        result: dict[str, list[dict[str, str]]] = {}
+        for event, hooks in self._by_event.items():
+            entries = result.setdefault(event, [])
+            for h in hooks:
+                entries.append(
+                    {
+                        "matcher": h.matcher or "(all)",
+                        "source": "server",
+                        "command": h.command,
+                    }
+                )
+        for event, subs in self._wire_by_event.items():
+            entries = result.setdefault(event, [])
+            for s in subs:
+                entries.append(
+                    {
+                        "matcher": s.matcher or "(all)",
+                        "source": "wire",
+                        "command": "(client-side)",
+                    }
+                )
+        return result
+
+    def _match_regex(self, pattern: str, value: str) -> bool:
+        if not pattern:
+            return True
+        try:
+            return bool(re.search(pattern, value))
+        except re.error:
+            logger.warning("Invalid regex in hook matcher: {}", pattern)
+            return False
+
+    async def trigger(
+        self,
+        event: HookEventType,
+        *,
+        matcher_value: str = "",
+        input_data: dict[str, Any],
+    ) -> list[HookResult]:
+        """Run all matching hooks (server + wire) in parallel."""
+
+        # --- Match server-side hooks ---
+        seen_commands: set[str] = set()
+        server_matched: list[HookDef] = []
+        for h in self._by_event.get(event, []):
+            if not self._match_regex(h.matcher, matcher_value):
+                continue
+            if h.command in seen_commands:
+                continue
+            seen_commands.add(h.command)
+            server_matched.append(h)
+
+        # --- Match wire subscriptions ---
+        wire_matched: list[WireHookSubscription] = []
+        for s in self._wire_by_event.get(event, []):
+            if not self._match_regex(s.matcher, matcher_value):
+                continue
+            wire_matched.append(s)
+
+        total = len(server_matched) + len(wire_matched)
+        if total == 0:
+            return []
+
+        try:
+            results = await self._execute_hooks(
+                event, matcher_value, server_matched, wire_matched, input_data
+            )
+        except Exception as exc:
+            from pythinker_code.telemetry.errors import report_handled_error
+
+            report_handled_error(exc, site="hooks.engine.run")
+            logger.warning("Hook engine error for {}, failing open", event)
+            return []
+
+        # Telemetry runs outside the fail-open try: a telemetry failure
+        # must NEVER discard hook results. For security-critical hooks
+        # (PreToolUse block), treating a sink failure as fail-open would
+        # silently bypass the block.
+        try:
+            from pythinker_code.telemetry import track
+
+            has_block = any(r.action == "block" for r in results)
+            track("hook_triggered", event_type=event, action="block" if has_block else "allow")
+        except Exception:
+            logger.debug("Telemetry for hook_triggered failed")
+
+        return results
+
+    async def _execute_hooks(
+        self,
+        event: str,
+        matcher_value: str,
+        server_matched: list[HookDef],
+        wire_matched: list[WireHookSubscription],
+        input_data: dict[str, Any],
+    ) -> list[HookResult]:
+        """Run matched hooks and emit wire events. Separated for fail-open wrapping."""
+        total = len(server_matched) + len(wire_matched)
+        logger.debug(
+            "Triggering {} hooks for {} ({} server, {} wire)",
+            total,
+            event,
+            len(server_matched),
+            len(wire_matched),
+        )
+
+        # --- HookTriggered ---
+        if self._on_triggered:
+            try:
+                self._on_triggered(event, matcher_value, total)
+            except Exception as e:
+                from pythinker_code.telemetry.errors import report_handled_error
+
+                report_handled_error(e, site="hooks.engine.triggered_cb")
+                logger.warning(
+                    "HookTriggered callback failed for {event}: {error}, continuing",
+                    event=event,
+                    error=e,
+                )
+
+        t0 = time.monotonic()
+        tasks: list[asyncio.Task[HookResult]] = []
+
+        # Server-side: run shell commands
+        for h in server_matched:
+            tasks.append(
+                asyncio.create_task(
+                    run_hook(h.command, input_data, timeout=h.timeout, cwd=self._cwd)
+                )
+            )
+
+        # Wire-side: send request to client, wait for response
+        for s in wire_matched:
+            tasks.append(
+                asyncio.create_task(
+                    self._dispatch_wire_hook(
+                        s.id, event, matcher_value, input_data, timeout=s.timeout
+                    )
+                )
+            )
+
+        results = list(await asyncio.gather(*tasks))
+        duration_ms = int((time.monotonic() - t0) * 1000)
+
+        # Aggregate: block if any hook blocked
+        action = "allow"
+        reason = ""
+        for r in results:
+            if r.action == "block":
+                action = "block"
+                reason = r.reason
+                logger.warning(
+                    "Hook blocked {event} (matcher={matcher}): {reason}",
+                    event=event,
+                    matcher=matcher_value,
+                    reason=reason,
+                )
+                break
+
+        # --- HookResolved ---
+        if self._on_resolved:
+            try:
+                self._on_resolved(event, matcher_value, action, reason, duration_ms)
+            except Exception as e:
+                from pythinker_code.telemetry.errors import report_handled_error
+
+                report_handled_error(e, site="hooks.engine.resolved_cb")
+                logger.warning(
+                    "HookResolved callback failed for {event}: {error}, continuing",
+                    event=event,
+                    error=e,
+                )
+
+        return results
+
+    async def _dispatch_wire_hook(
+        self,
+        subscription_id: str,
+        event: str,
+        target: str,
+        input_data: dict[str, Any],
+        *,
+        timeout: int = 30,
+    ) -> HookResult:
+        """Send a hook request to the wire client and wait for response."""
+        if not self._on_wire_hook:
+            return HookResult(action="allow")
+
+        handle = WireHookHandle(
+            subscription_id=subscription_id,
+            event=event,
+            target=target,
+            input_data=input_data,
+        )
+        # Run the callback in background so timeout applies to the
+        # full client round-trip, not just handle.wait().
+        hook_task: asyncio.Task[None] = asyncio.ensure_future(self._on_wire_hook(handle))
+        hook_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
+        try:
+            return await asyncio.wait_for(handle.wait(), timeout=timeout)
+        except TimeoutError:
+            hook_task.cancel()
+            logger.warning("Wire hook timed out: {} {}", event, target)
+            return HookResult(action="allow", timed_out=True)
+        except Exception as e:
+            from pythinker_code.telemetry.errors import report_handled_error
+
+            report_handled_error(e, site="hooks.engine.wire")
+            hook_task.cancel()
+            logger.warning("Wire hook failed: {} {}: {}", event, target, e)
+            return HookResult(action="allow")