Skip to content

feat: add Pi-aligned dynamic extension providers#417

Draft
alejandro-ao wants to merge 3 commits into
mainfrom
feat/llama-cpp-extension-api
Draft

feat: add Pi-aligned dynamic extension providers#417
alejandro-ao wants to merge 3 commits into
mainfrom
feat/llama-cpp-extension-api

Conversation

@alejandro-ao

@alejandro-ao alejandro-ao commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR adds a first-class dynamic-provider foundation for Tau extensions. It implements the architecture planned in #443 and provides the generic Tau core capabilities required by external integrations such as the llama.cpp extension proposed in #415.

The central change is that an extension can now register a provider object that owns its transport configuration, authentication strategy, model snapshot, asynchronous model discovery, and optional runtime factory. Tau then integrates that provider into startup resolution, /model, session switching, reload, resume, diagnostics, and resource cleanup.

No llama.cpp-specific behavior is added to Tau core. The same API is intended to support external extensions for llama.cpp, vLLM, Ollama, LM Studio, remote OpenAI-compatible gateways, and future non-OpenAI transports.

Why this matters

Before this work, an extension could implement provider-specific commands or discovery, but it could not robustly participate in Tau's normal provider lifecycle. The original prototype in this PR proved the basic idea, but modeled providers as mutable host configuration and did not safely handle overlays, asynchronous discovery, reload, runtime replacement, or restoration of built-in providers.

This implementation gives extensions a stable boundary:

  1. The extension owns provider-specific behavior.
  2. Tau owns provider lifecycle and session integration.
  3. tau_ai continues to own provider transports and streaming.
  4. tau_agent remains provider-agnostic.

That separation lets Tau add provider integrations without hardcoding each provider into core or weakening the reusable agent harness.

What was implemented

1. Provider objects instead of host-setting mutation

Extensions register a DynamicProvider:

from tau_coding.extensions import (
    DynamicProvider,
    NoAuth,
    OpenAICompatibleTransport,
    ProviderModel,
)


def setup(tau):
    tau.register_provider(
        DynamicProvider(
            id="local-server",
            display_name="Local server",
            transport=OpenAICompatibleTransport(
                base_url="http://127.0.0.1:8080/v1",
            ),
            auth=NoAuth(),
            models=(ProviderModel(id="cached-model"),),
            refresh_models=discover_models,
        )
    )

The provider object includes:

  • a stable provider ID and display name
  • an OpenAI-compatible transport helper
  • an explicit authentication strategy
  • a structured model snapshot
  • an optional asynchronous model refresh callback
  • an optional runtime factory for future/custom transports
  • an optional refresh timeout

Providers may register with zero models. This allows an extension to expose setup, status, and diagnostic commands before a server is configured or reachable.

2. Structured model metadata

ProviderModel supports optional:

  • display name
  • context window
  • maximum output tokens
  • input modalities
  • reasoning support
  • compatibility metadata

Unknown values remain unset. Tau does not invent model limits or capabilities.

Provider and model display names are surfaced in model-picker rows while stable IDs remain visible for diagnostics and configuration.

3. Layered, source-aware provider overlays

Dynamic providers are process-local overlays over Tau's durable provider settings.

Conceptually:

provider id: openai
  built-in/durable definition
  extension A override
  extension B override  <- effective

The registry tracks each layer by extension owner and registration order:

  • multiple extensions may register the same stable provider ID
  • re-registering from the same extension atomically replaces only that layer
  • invalid replacements preserve the previous working layer
  • unregistering removes only the caller's layer
  • removing the effective layer restores the preceding extension layer
  • removing the last extension layer restores the complete durable/built-in definition
  • setup failure and /reload remove only registrations owned by the affected extension generation

Extension-only providers disappear when unloaded. Built-in providers are never accidentally deleted by an override lifecycle.

4. Process-local state remains separate from durable configuration

Dynamic-provider definitions are composed over durable provider settings in memory. They are not written into Tau's provider catalog.

The implementation keeps the durable provider base separate from the composed runtime view, including across /resume, /new, and session adoption. This is necessary so an override can always be removed and the original provider restored.

Extension model selection:

  • updates the active session and its session record
  • does not make a process-local provider Tau's durable default
  • does not persist extension thinking or scoped-model preferences
  • does not write extension provider definitions into the user catalog

Extensions can persist non-secret endpoint and cached discovery state through their extension-scoped settings. API keys and other secrets must remain in environment variables or future credential-store integrations.

5. Explicit authentication strategies

The public API provides:

  • RequiredEnvApiKey("API_KEY")
  • OptionalEnvApiKey("API_KEY")
  • NoAuth()

Behavior is explicit:

  • required auth gives actionable guidance when the variable is missing
  • optional auth uses the key when present
  • optional auth omits Authorization when the variable is absent
  • no-auth always omits Authorization

Tau does not require a fake key and does not synthesize Bearer local.

6. Provider-owned asynchronous model discovery

A dynamic provider may define:

async def refresh_models(context) -> Sequence[ProviderModel]:
    ...

The refresh context includes:

  • a cancellation signal
  • whether network access is allowed
  • the resolved non-secret endpoint
  • an authentication resolver
  • the owning extension identity

Refresh behavior is lifecycle-safe:

  1. Cached models are registered synchronously during setup(tau).
  2. Tau refreshes only an explicitly requested startup provider.
  3. Unrelated startup does not perform mandatory local-network discovery.
  4. A successful result is validated and published atomically.
  5. A failed or timed-out refresh retains the last-known snapshot.
  6. Concurrent refresh requests for the same layer share one task.
  7. Unregister, re-registration, reload, and shutdown cancel owned refresh work.
  8. An old in-flight refresh cannot overwrite a newer provider registration.
  9. Refresh failures become extension diagnostics rather than crashing the session.

If refresh removes the active model, Tau leaves the current runtime usable until the user chooses another model and reports that the model is no longer advertised. If the active model remains available, Tau safely refreshes its runtime configuration.

7. One startup path for TUI and print mode

provider_startup.prepare_provider_startup() centralizes extension/provider startup ordering for both frontends:

  1. Resolve the target cwd, including an explicitly resumed session.
  2. Load user extensions, explicit -e extensions, and explicitly enabled project extensions.
  3. Restore cached provider snapshots during synchronous extension setup.
  4. Refresh only the explicitly targeted provider.
  5. Compose effective extension overlays over durable settings.
  6. Resolve the startup provider and model.
  7. Construct and bind the coding session.

This means the following works on the first process start:

tau -e ./my-provider --provider local-server --model coder --print "Say hello"

Project extensions remain disabled by default and require explicit enablement.

8. Async-safe runtime selection and cleanup

Extensions switch through:

await tau.select_model("local-server", "coder")

Tau creates and validates the candidate runtime before changing the active session. Therefore:

  • runtime factory or authentication failure preserves the current provider/model
  • successful selection updates the active session atomically
  • replaced Tau-owned resources are closed
  • extension runtime factories are honored
  • TUI dynamic-provider selection uses this async, non-persistent path
  • durable provider choices retain their normal persistence behavior

This avoids leaving the user with a broken session after a failed local-provider switch.

9. Settings and diagnostics

Extensions retain user-scoped JSON settings APIs:

  • load_settings()
  • save_settings(...)
  • clear_settings()

Writes are atomic and user-level, including for explicitly loaded project extensions. These settings are intended only for non-secret endpoint, cached model, and extension preference data.

Provider registration, refresh, and runtime-factory failures are isolated as extension diagnostics and appear through Tau's existing diagnostic surfaces.

User experience improvements

For users, this enables provider extensions that feel integrated rather than bolted on:

  1. A provider can appear in Tau before its server is reachable.
  2. Cached models appear immediately instead of blocking startup on network discovery.
  3. Explicit --provider and --model selections work in TUI and print mode.
  4. /model can show extension providers with friendly provider/model names.
  5. Unauthenticated local servers work without fake credentials.
  6. Failed refresh keeps cached choices available.
  7. Failed model switching keeps the current session working.
  8. Reloading or removing an override restores the previous provider automatically.
  9. Resume preserves both extension overlays and the underlying durable definitions.
  10. Extension providers cannot accidentally become permanent catalog entries.

How this improves Tau's adaptability

This PR makes provider integration an extension capability rather than a sequence of provider-specific core changes.

An extension can own:

  • endpoint configuration
  • model discovery
  • provider-specific commands and diagnostics
  • cached model state
  • authentication environment-variable policy
  • an eventual custom runtime factory

Tau provides the reusable host behavior:

  • loading and trust boundaries
  • overlay precedence and restoration
  • startup ordering
  • model surfaces
  • session switching
  • diagnostics
  • cancellation and cleanup
  • durable/process-local separation

As a result, new provider integrations can evolve independently, and Tau core only needs generic lifecycle APIs that benefit multiple extensions.

Relationship to Pi

This design follows Pi's current dynamic-provider architecture where it provides a reusable lifecycle model, while deliberately differing where Tau's product and architecture require it.

Aligned with Pi

  • extensions register provider objects rather than directly mutating host settings
  • providers may start with an empty model list
  • providers own asynchronous model refresh and authentication behavior
  • extension providers are available before startup model resolution
  • later registration and refresh update active provider/model surfaces
  • unregistering an override restores the preceding definition
  • active sessions can refresh when provider definitions change
  • extension failures are isolated instead of crashing the host

Pi's llama.cpp integration demonstrates these generic capabilities through a hidden built-in extension. This PR adopts the reusable lifecycle ideas, not the llama.cpp-specific product choices.

Intentional Tau differences

  • llama.cpp remains an external, explicitly loaded extension rather than a hidden built-in extension
  • Tau keeps source ownership on every overlay layer so reload/unregister cannot remove another extension's registration
  • optional and no-auth providers omit Authorization; Tau does not synthesize Bearer local
  • unknown model limits and capabilities remain unknown
  • Tau does not require llama.cpp router mode
  • Tau does not install, start, stop, download, load, or unload local models in this foundation
  • extension setup() remains synchronous; asynchronous work runs in host-invoked refresh callbacks
  • Tau v1 ships an OpenAI-compatible transport helper while keeping the provider object runtime-factory-ready
  • process-local extension providers are not persisted as Tau defaults or catalog definitions

These differences preserve Tau's separation of concerns:

tau_ai      provider/model transport and streaming
tau_agent   portable provider-agnostic harness
tau_coding  extension registry, startup, sessions, CLI, and TUI integration

Scope and non-goals

This PR intentionally does not add:

  • llama.cpp-specific discovery or commands
  • /local
  • provider-contributed /login
  • extension package management
  • built-in llama.cpp process management
  • model download/load/unload behavior
  • a requirement for router mode
  • OAuth for dynamic providers

Those can be built later on top of the generic lifecycle after validating it with multiple provider extensions.

Issues

Closes #443.
Addresses the generic Tau core prerequisites for #415.

Manual validation

Create an external extension that registers a cached DynamicProvider with an async /v1/models refresh callback.

Run print mode:

tau -e ./test-local-extension \
  --provider local-test \
  --model <model> \
  --print "Say hello"

Run the TUI:

tau -e ./test-local-extension

Then validate:

  1. /model shows the extension provider and model display names.
  2. Selecting the provider switches the current session without changing Tau's durable default.
  3. /reload reconstructs the provider and cancels outgoing discovery work.
  4. /resume retains the extension overlay and can still restore the durable provider after unregister.
  5. An unauthenticated endpoint receives no Authorization header.
  6. An env-key-protected endpoint receives its configured key.
  7. Failed and timed-out refresh retain cached models and report diagnostics.
  8. Failed runtime creation leaves the previous provider/model active.
  9. Overriding and unregistering a built-in provider restores the complete built-in definition.
  10. Re-registering during an in-flight refresh prevents the stale result from publishing.

Verification

  • uv run pytest — 1126 passed
  • uv run ruff check .
  • uv run ruff format --check .
  • uv run mypy
  • hugo --source website --minify — passed with the existing taxonomy-layout warning
  • uv build
  • GitHub Python checks — passed
  • GitHub documentation build — passed

@alejandro-ao alejandro-ao changed the title feat: let extensions register providers feat: add Pi-aligned dynamic extension providers Jul 22, 2026
@alejandro-ao

Copy link
Copy Markdown
Collaborator Author

Review findings (request changes):

  1. [P1] Unregistering an override does not restore the durable providersrc/tau_coding/session.py:985-997

    sync_extension_providers() removes every durable provider whose name was ever extension-owned. After an extension overriding openai unregisters or reloads, the built-in definition remains excluded. Compose overlays over the complete durable provider tuple instead of filtering durable entries by owned_names.

  2. [P1] TUI model selection persists process-local providerssrc/tau_coding/tui/app.py:4864-4874

    Dynamic choices use synchronous set_model_choice(), which persists the provider into settings/catalog and bypasses the async extension selection/runtime-factory path. Route dynamic choices through awaited select_extension_provider_model() without persisting defaults; retain normal persistence for durable providers.

  3. [P2] Resume loses the durable overlay basesrc/tau_coding/session.py:1392-1418

    Resume passes composed provider_settings but omits durable_provider_settings, so the replacement treats extension overlays as durable state. Pass durable_provider_settings=self._durable_provider_settings and cover override → resume → unregister.

  4. [P2] An in-flight refresh can overwrite a newer registrationsrc/tau_coding/extensions/runtime.py:319-328,345-369

    Re-registering the same source/provider neither cancels nor versions the existing refresh. The old callback can later publish stale models into the new layer. Cancel refresh on replacement or verify the layer revision before publishing.

Validation: uv run pytest (1122 passed), targeted tests (227 passed), Ruff check/format, and mypy all passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Plan: Pi-aligned dynamic provider registration for extensions

1 participant