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

tau/extensions/loader.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import importlib.util
+import inspect
+import json
+import logging
+import sys
+import traceback
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from tau.extensions.api import (
+    Extension, ExtensionAPI, ExtensionError,
+    LoadExtensionsResult, _RuntimeRef,
+)
+from tau.extensions.events import EventBus
+
+if TYPE_CHECKING:
+    from tau.inference.api.text.service import TextLLM
+    from tau.settings.manager import SettingsManager
+    from tau.settings.types import ExtensionEntry
+
+logger = logging.getLogger(__name__)
+
+_ENTRY_POINT = "register"
+
+
+class ExtensionLoader:
+    """
+    Discovers and loads Python extension files from configured directories.
+
+    Discovery order (project first, then global, then explicit paths):
+      - <cwd>/.tau/extensions/
+      - ~/.tau/extensions/
+      - Any entries in settings.extensions.list that are enabled
+
+    Within each directory, files and subdirectories are sorted alphabetically.
+    Files whose names start with ``_`` are skipped.  Files whose stem appears
+    in ``disabled_stems`` are also skipped.  Duplicate resolved paths are
+    skipped silently.  A failing extension is logged and recorded in the
+    returned ``LoadExtensionsResult`` — it never crashes startup.
+
+    Subdirectory support:
+      Given ``extensions/my_pkg/``:
+        1. ``my_pkg/manifest.json`` with ``{"tau": {"extensions": ["./main.py"]}}``
+           → load the declared paths.
+        2. ``my_pkg/__init__.py``
+           → load as a Python package entry point.
+      If neither is present the subdirectory is ignored.
+
+    Dependency installation:
+      ``manifest.json`` may also declare ``{"tau": {"dependencies": ["pkg>=1.0"]}}``.
+      Before the entry file is executed, these specs are installed (once — cached by
+      a hash of the dependency list) via ``uv pip install``, falling back to the
+      venv's own ``pip`` when ``uv`` is not on PATH. No dedicated venv is created
+      under ``.tau/``: project extensions install into the project's own ``.venv``
+      when one already exists, otherwise everything falls back to the global
+      packages venv at ``~/.tau/venv``. The resolved venv's site-packages directory
+      is then added to ``sys.path`` so the extension's own imports resolve.
+    """
+
+    def __init__(
+        self,
+        project_dir: Path | None = None,
+        global_dir: Path | None = None,
+        extra_entries: list[ExtensionEntry] | None = None,
+        extra_sources: dict[str, str] | None = None,
+        disabled_stems: set[str] | None = None,
+        entry_configs: dict[str, dict] | None = None,
+        builtins_dir: Path | None = None,
+        # Dependencies injected for per-extension ExtensionAPI creation
+        llm: TextLLM | None = None,
+        settings: SettingsManager | None = None,
+        cwd: Path | None = None,
+        runtime_ref: _RuntimeRef | None = None,
+        events: EventBus | None = None,
+    ) -> None:
+        self._builtins_dir: Path | None = builtins_dir if (builtins_dir is not None and builtins_dir.is_dir()) else None
+        self._dirs: list[Path] = []
+        self._dir_sources: dict[Path, str] = {}
+        self._extra_paths: list[Path] = []
+        self._extra_sources: dict[str, str] = extra_sources or {}
+        self._disabled_stems: set[str] = disabled_stems or set()
+        self._entry_configs: dict[str, dict] = entry_configs or {}
+        self._subdir_deps: dict[Path, list[str]] = {}
+        self._subdir_settings: dict[Path, dict] = {}
+        self._llm = llm
+        self._settings = settings
+        self._cwd = cwd or Path(".")
+        self._runtime_ref = runtime_ref
+        self._events = events or EventBus()
+
+        # Project first, then global
+        if project_dir is not None and project_dir.is_dir():
+            self._dirs.append(project_dir)
+            self._dir_sources[project_dir] = "project"
+        if global_dir is not None and global_dir.is_dir():
+            self._dirs.append(global_dir)
+            self._dir_sources[global_dir] = "global"
+
+        for entry in (extra_entries or []):
+            resolved = Path(entry.path).expanduser().resolve()
+            if resolved.exists():
+                self._extra_paths.append(resolved)
+                if entry.settings:
+                    self._entry_configs.setdefault(resolved.stem, entry.settings)
+
+    # ── Discovery ──────────────────────────────────────────────────────────────
+
+    def _subdir_entries(self, subdir: Path) -> list[Path]:
+        """Return extension entry files for a subdirectory, or empty list."""
+        manifest = subdir / "manifest.json"
+        if manifest.is_file():
+            try:
+                data = json.loads(manifest.read_text(encoding="utf-8"))
+                from tau.settings.paths import get_app_name
+                app_data = data.get(get_app_name().lower(), {})
+
+                deps = app_data.get("dependencies", [])
+                if deps:
+                    self._subdir_deps[subdir.resolve()] = list(deps)
+
+                settings_schema = app_data.get("settings")
+                if isinstance(settings_schema, dict):
+                    self._subdir_settings[subdir.resolve()] = settings_schema
+
+                declared = app_data.get("extensions", [])
+                if declared:
+                    entries = [(subdir / p).resolve() for p in declared]
+                    valid = [e for e in entries if e.is_file()]
+                    if valid:
+                        return valid
+            except (json.JSONDecodeError, OSError):
+                pass
+
+        init = subdir / "__init__.py"
+        if init.is_file():
+            return [init]
+
+        return []
+
+    def _discover_in_dir(self, d: Path, seen: set[Path]) -> list[Path]:
+        """Collect extension files from one directory (one level deep)."""
+        found: list[Path] = []
+        for entry in sorted(d.iterdir(), key=lambda e: e.name):
+            if entry.name.startswith("_"):
+                continue
+            if entry.is_file() and entry.suffix == ".py":
+                if entry.stem in self._disabled_stems:
+                    continue
+                r = entry.resolve()
+                if r not in seen:
+                    seen.add(r)
+                    found.append(entry)
+            elif entry.is_dir():
+                if entry.name in self._disabled_stems:
+                    continue
+                for sub_entry in self._subdir_entries(entry):
+                    r = sub_entry.resolve()
+                    if r not in seen:
+                        seen.add(r)
+                        found.append(sub_entry)
+        return found
+
+    def _discover(self) -> list[tuple[Path, str]]:
+        """Return deduplicated list of (extension file, source label) pairs to load."""
+        found: list[tuple[Path, str]] = []
+        seen: set[Path] = set()
+
+        # Builtins always first, before any user dirs
+        if self._builtins_dir is not None:
+            for path in self._discover_in_dir(self._builtins_dir, seen):
+                found.append((path, "builtin"))
+
+        for d in self._dirs:
+            source = self._dir_sources.get(d, "unknown")
+            for path in self._discover_in_dir(d, seen):
+                found.append((path, source))
+
+        for p in self._extra_paths:
+            source = self._extra_sources.get(str(p.resolve()), "explicit")
+            if p.is_file() and p.suffix == ".py":
+                r = p.resolve()
+                if r not in seen:
+                    seen.add(r)
+                    found.append((p, source))
+            elif p.is_dir():
+                entries = self._subdir_entries(p)
+                if entries:
+                    for e in entries:
+                        r = e.resolve()
+                        if r not in seen:
+                            seen.add(r)
+                            found.append((e, source))
+                else:
+                    for path in self._discover_in_dir(p, seen):
+                        found.append((path, source))
+
+        return found
+
+    # ── Dependency installation ──────────────────────────────────────────────────
+
+    def _resolve_venv_dir(self, source: str) -> Path:
+        """Pick a venv to install extension dependencies into.
+
+        Project extensions always use the project's own ``.venv`` — reusing it
+        if it already exists, creating it at ``<cwd>/.venv`` if not. No dedicated
+        venv is ever created under ``.tau/``. Everything else (global, explicit,
+        unknown sources) uses the global packages venv at ``~/.tau/venv``.
+        """
+        from tau.settings.paths import get_packages_venv
+
+        if source == "project":
+            return self._cwd / ".venv"
+
+        return get_packages_venv(None)
+
+    def _ensure_dependencies(self, subdir: Path, deps: list[str], source: str) -> None:
+        """Install an extension's declared dependencies, once per dependency set.
+
+        Runs synchronously (called via ``asyncio.to_thread``) so the blocking
+        subprocess install doesn't stall the event loop.
+        """
+        from tau.packages.manager import PackageManager
+
+        venv_dir = self._resolve_venv_dir(source)
+        pkg_mgr = PackageManager(venv_dir)
+
+        digest = hashlib.sha256("\n".join(sorted(deps)).encode("utf-8")).hexdigest()
+        cache_file = venv_dir / ".tau_ext_deps.json"
+        cache: dict[str, str] = {}
+        if cache_file.is_file():
+            try:
+                cache = json.loads(cache_file.read_text(encoding="utf-8"))
+            except (json.JSONDecodeError, OSError):
+                cache = {}
+
+        key = str(subdir.resolve())
+        if cache.get(key) != digest:
+            logger.info("Installing dependencies for extension %r: %s", subdir.name, ", ".join(deps))
+            pkg_mgr.install_requirements(deps)
+            cache[key] = digest
+            cache_file.parent.mkdir(parents=True, exist_ok=True)
+            cache_file.write_text(json.dumps(cache, indent=2), encoding="utf-8")
+
+        site = pkg_mgr.site_packages()
+        if site and str(site) not in sys.path:
+            sys.path.insert(0, str(site))
+
+    # ── Loading ────────────────────────────────────────────────────────────────
+
+    async def load(self) -> LoadExtensionsResult:
+        """Discover and load all extensions in parallel.
+
+        Returns a ``LoadExtensionsResult`` with loaded ``Extension`` objects and
+        any ``ExtensionError`` entries for files that failed to load.
+        """
+        discovered = self._discover()
+
+        async def _load(path: Path, source: str) -> tuple[Extension | None, list[ExtensionError]]:
+            config_key = path.parent.name if path.name == "__init__.py" else path.stem
+            config = self._entry_configs.get(config_key, {})
+            ext, errs = await self._load_one(path, config, source=source)
+            if ext is not None:
+                logger.debug("Loaded extension: %s (%s)", path, source)
+            return ext, errs
+
+        results = await asyncio.gather(*[_load(p, s) for p, s in discovered])
+
+        extensions: list[Extension] = []
+        errors: list[ExtensionError] = []
+        for ext, errs in results:
+            if ext is not None:
+                extensions.append(ext)
+            errors.extend(errs)
+
+        return LoadExtensionsResult(extensions=extensions, errors=errors)
+
+    async def _load_one(
+        self, path: Path, config: dict, *, source: str = "unknown"
+    ) -> tuple[Extension | None, list[ExtensionError]]:
+        """Load a single extension file and call its ``register(tau)`` factory."""
+        str_path = str(path)
+        errors: list[ExtensionError] = []
+
+        try:
+            deps = self._subdir_deps.get(path.parent.resolve())
+            if deps and source != "builtin":
+                await asyncio.to_thread(self._ensure_dependencies, path.parent, deps, source)
+
+            module_name = f"_tau_ext_{hashlib.sha1(str(path.resolve()).encode()).hexdigest()[:16]}"
+            spec = importlib.util.spec_from_file_location(module_name, path)
+            if spec is None or spec.loader is None:
+                errors.append(ExtensionError(
+                    extension_path=str_path, event="load",
+                    error=f"Cannot create module spec for {path}",
+                ))
+                return None, errors
+
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[module_name] = module
+            spec.loader.exec_module(module)  # type: ignore[union-attr]
+
+            register_fn = getattr(module, _ENTRY_POINT, None)
+            if register_fn is None or not callable(register_fn):
+                errors.append(ExtensionError(
+                    extension_path=str_path, event="load",
+                    error=f"No '{_ENTRY_POINT}(tau)' function in {path.name}",
+                ))
+                return None, errors
+
+            ext = Extension(path=str_path, config=config, source=source)
+            api = ExtensionAPI(
+                extension=ext,
+                llm=self._llm,  # type: ignore[arg-type]
+                settings=self._settings,  # type: ignore[arg-type]
+                cwd=self._cwd,
+                runtime_ref=self._runtime_ref,
+                events=self._events,
+            )
+
+            result = register_fn(api)
+            if inspect.isawaitable(result):
+                await result
+
+            self._attach_manifest_panel(ext, path)
+
+            return ext, errors
+
+        except Exception:
+            tb = traceback.format_exc()
+            errors.append(ExtensionError(
+                extension_path=str_path, event="load",
+                error=tb.strip().splitlines()[-1],
+                stack=tb,
+            ))
+            return None, errors
+
+    def _attach_manifest_panel(self, ext: Extension, path: Path) -> None:
+        """Auto-build a /settings sub-panel from the manifest ``settings`` schema.
+
+        Skipped when the extension already registered a panel itself (manual
+        ``register_settings`` wins) or when no schema is declared.
+        """
+        if ext.settings_registrations:
+            return
+        if self._settings is None:
+            return  # no settings manager → nothing to persist changes to
+        settings = self._settings
+        schema = self._subdir_settings.get(path.parent.resolve())
+        if not schema:
+            return
+
+        from tau.tui.components.settings_schema import build_manifest_panel
+
+        ext_dir = path.parent
+        try:
+            settings_path = str(ext_dir.relative_to(self._cwd))
+        except ValueError:
+            settings_path = str(ext_dir)
+        module_path = str(path)  # absolute path of the loaded entry file
+
+        def _apply(key: str, value: Any) -> None:
+            # Persist to settings.json, then reload just this extension so the
+            # change takes effect live without re-running other extensions.
+            settings.set_extension_config_key(settings_path, key, value)
+            runtime = self._runtime_ref.runtime if self._runtime_ref is not None else None
+            if runtime is None:
+                return
+            reload_one = getattr(runtime, "reload_extension", None)
+            if reload_one is not None:
+                asyncio.ensure_future(reload_one(module_path))
+            else:  # fall back to full reload if single-extension reload is unavailable
+                reload_all = getattr(runtime, "reload_extensions", None)
+                if reload_all is not None:
+                    asyncio.ensure_future(reload_all())
+
+        reg = build_manifest_panel(
+            schema, ext.config,
+            default_title=ext_dir.name,
+            apply=_apply,
+        )
+        if reg is not None:
+            ext.settings_registrations.append(reg)

tau/extensions/runtime.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import inspect
+import traceback
+from typing import TYPE_CHECKING, Any, Callable
+
+from tau.extensions.api import (
+    Extension, ExtensionError, LoadExtensionsResult,
+    ShortcutRegistration, _RuntimeRef,
+)
+
+if TYPE_CHECKING:
+    from tau.commands.types import CommandInfo
+    from tau.hooks.service import Hooks
+    from tau.tool.types import Tool
+
+
+# Events where handler return values matter for interception.
+# These are registered directly on the hooks bus (not via the catch-all subscriber)
+# so that Hooks.emit() collects their results.
+_INTERCEPTABLE_EVENTS: frozenset[str] = frozenset({
+    "before_compaction",
+    "user_terminal",
+    "resources_discover",
+    "project_trust",
+    "input",
+    "tool_result",
+    "context",
+})
+
+
+class ExtensionRuntime:
+    """
+    Owns all loaded extensions and dispatches lifecycle events to their handlers.
+
+    Subscribed as a catch-all listener on the hooks bus so that every event
+    emitted by the agent or runtime automatically reaches extension handlers —
+    no changes to the agent or engine are required.
+
+    Each handler is always called as ``handler(event, ctx)`` where ``ctx`` is a
+    fresh ``ExtensionContext`` snapshot built from the live runtime at dispatch
+    time.  Handler exceptions are caught per-handler and appended to ``errors``.
+
+    Interceptable events (``_INTERCEPTABLE_EVENTS``) are registered directly on
+    the hooks bus rather than going through the catch-all subscriber, so their
+    return values are collected by ``Hooks.emit()`` and available for inspection
+    by the caller (e.g. ``before_compaction`` handlers that return
+    ``BeforeCompactionResult``).
+    """
+
+    def __init__(
+        self,
+        load_result: LoadExtensionsResult,
+        hooks: Hooks,
+        runtime_ref: _RuntimeRef,
+    ) -> None:
+        self._extensions: list[Extension] = load_result.extensions
+        self._errors: list[ExtensionError] = list(load_result.errors)
+        self.runtime_ref: _RuntimeRef = runtime_ref
+        self._unsub = hooks.subscribe(self._dispatch)
+
+        # Register interceptable handlers directly so their results flow back
+        # through Hooks.emit() rather than being discarded by the subscriber path.
+        self._interceptable_unsubs: list[Callable[[], None]] = []
+        for ext in self._extensions:
+            for event_type in _INTERCEPTABLE_EVENTS:
+                for handler in ext.handlers.get(event_type, []):
+                    wrapped = self._make_interceptable_handler(ext, handler)
+                    self._interceptable_unsubs.append(hooks.register(event_type, wrapped))
+
+    # ── Dispatch ──────────────────────────────────────────────────────────────
+
+    def _make_interceptable_handler(self, ext: Extension, handler: Callable) -> Callable:
+        """Return a hooks-compatible wrapper that injects ctx and propagates the return value."""
+        async def wrapped(event: Any) -> Any:
+            """Invoke handler with extension context."""
+            from tau.extensions.context import ExtensionContext
+            runtime = self.runtime_ref.runtime
+            ctx = ExtensionContext.from_runtime(runtime) if runtime is not None else None
+            try:
+                result = handler(event, ctx)
+                if inspect.isawaitable(result):
+                    result = await result
+                return result
+            except Exception:
+                tb = traceback.format_exc()
+                self._errors.append(ExtensionError(
+                    extension_path=ext.path,
+                    event=getattr(event, "type", "unknown"),
+                    error=tb.strip().splitlines()[-1],
+                    stack=tb,
+                ))
+                return None
+
+        return wrapped
+
+    async def _dispatch(self, event: Any) -> None:
+        """Catch-all hooks subscriber — re-dispatches every event to extension handlers.
+
+        Interceptable events are skipped here; they are already handled by
+        directly-registered hooks so their return values reach Hooks.emit().
+        """
+        event_type: str | None = getattr(event, "type", None)
+        if not event_type:
+            return
+        if event_type in _INTERCEPTABLE_EVENTS:
+            return
+
+        from tau.extensions.context import ExtensionContext
+        runtime = self.runtime_ref.runtime
+        ctx = ExtensionContext.from_runtime(runtime) if runtime is not None else None
+
+        for ext in self._extensions:
+            for handler in ext.handlers.get(event_type, []):
+                try:
+                    result = handler(event, ctx)
+                    if inspect.isawaitable(result):
+                        await result
+                except Exception:
+                    tb = traceback.format_exc()
+                    self._errors.append(ExtensionError(
+                        extension_path=ext.path,
+                        event=event_type,
+                        error=tb.strip().splitlines()[-1],
+                        stack=tb,
+                    ))
+
+    def unsubscribe(self) -> None:
+        """Detach from the hooks bus (called before hot-reload replaces this runtime)."""
+        self._unsub()
+        for unsub in self._interceptable_unsubs:
+            unsub()
+        self._interceptable_unsubs.clear()
+
+    # ── Errors ────────────────────────────────────────────────────────────────
+
+    @property
+    def errors(self) -> list[ExtensionError]:
+        """All accumulated load and dispatch errors."""
+        return self._errors
+
+    # ── Accessors ─────────────────────────────────────────────────────────────
+
+    def get_tools(self) -> list[Tool]:
+        """Return all tools registered by extensions (last-writer-wins on name)."""
+        tools: dict[str, Any] = {}
+        for ext in self._extensions:
+            tools.update(ext.tools)
+        return list(tools.values())
+
+    def get_commands(self) -> list[CommandInfo]:
+        """Return all slash commands registered by extensions (last-writer-wins on name)."""
+        commands: dict[str, Any] = {}
+        for ext in self._extensions:
+            commands.update(ext.commands)
+        return list(commands.values())
+
+    def get_shortcuts(self) -> list[ShortcutRegistration]:
+        """Return all keyboard shortcuts registered by extensions."""
+        result: list[ShortcutRegistration] = []
+        for ext in self._extensions:
+            result.extend(ext.shortcuts)
+        return result
+
+    def get_prompt_appends(self) -> list[str]:
+        """Return all system-prompt additions registered by extensions."""
+        result: list[str] = []
+        for ext in self._extensions:
+            result.extend(ext.prompt_appends)
+        return result
+
+    def get_message_renderers(self) -> dict[str, Any]:
+        """Return merged message renderer registry (last-registered wins per type)."""
+        result: dict[str, Any] = {}
+        for ext in self._extensions:
+            result.update(ext.message_renderers)
+        return result
+
+    def get_autocomplete_providers(self) -> list[Any]:
+        """Return all autocomplete providers registered by extensions."""
+        result: list[Any] = []
+        for ext in self._extensions:
+            result.extend(ext.autocomplete_providers)
+        return result