# Plan: Better Pi Integration for pi-local

## Current State

**pi-local** provides custom `/local-login` and `/local-model` commands that:
- Store connections in `~/.pi/pi-local-connections.json`
- Live-query models via oMLX → LM Studio → OpenAI fallback chain
- Support model load/unload for oMLX and LM Studio endpoints
- Register providers dynamically via `pi.registerProvider(baseUrl, ProviderConfig)`

**claude-local** (reference) is a bash script with an interactive model picker that supports:
- Arrow-key navigation, `l`/`u` for load/unload, Enter to select
- Rich model display with size, context window, type info
- Session cookie auth for oMLX unload operations

**pi's native integration surface:**
- `/login` — works through `Provider.auth.apiKey.login()` or `Provider.auth.oauth`
- `/model` — shows all models from registered providers via `ModelSelectorComponent`
- `ctx.ui.custom<T>()` — full-screen TUI component hosting (used by llama extension)
- `Provider.refreshModels()` — called when `/model` opens, allowing live catalog refresh

## Goals

1. Make connections discoverable via `/login <baseUrl>` so users can add/remove them through pi's native UI
2. Have models appear in `/model` selector, refreshed on-demand
3. Build a custom picker with load/unload support (like claude-local's menu, like llama's UI)
4. Keep live-query behavior — models shouldn't be stale

## Architecture Decision: One Provider per Connection vs. Aggregator

**Option A: One `Provider` per connection (recommended)**
- Each baseUrl is registered as its own provider via the full `Provider` interface
- Provider id = normalized baseUrl (e.g., `http://127.0.0.1:1234`)
- `/login http://127.0.0.1:1234` triggers `auth.apiKey.login()` for that endpoint
- Models auto-appear in `/model` when the provider is registered and refreshed
- Pros: Clean integration, each endpoint is independently login/manageable
- Cons: Provider names are URLs (less pretty in selectors), many providers shown

**Option B: Single "local" aggregator provider**
- One `Provider` with id `"local"` that wraps all connections
- `/login local` opens a custom login flow to add/manage connections
- Models from all connections are aggregated under one provider
- Pros: Single entry in selectors, cleaner UX
- Cons: Can't do per-connection auth via `/login`, requires custom sub-flows

**Recommendation: Option A** — aligns with how llama.cpp does it, better `/login` integration, each server is independently configured.

## Implementation Plan

### Phase 1: Switch to Native `Provider` Interface

Replace the current `ProviderConfig`-based registration with full `Provider<"openai-completions">` for each connection.

**File: `src/provider.ts`** — Add a new `createLocalProvider()` function:

```typescript
import type { Provider, ApiKeyCredential, AuthContext, AuthResult, RefreshModelsContext } from "@earendil-works/pi-ai";
import { streamSimpleOpenAICompletions } from "@earendil-works/pi-ai";

export function createLocalProvider(
  baseUrl: string,
  apiKeyCommand: string,
  resolveApiKey: (raw: string) => string,
): Provider<"openai-completions"> {
  // ... see detailed spec below
}
```

The Provider provides:
- `auth.apiKey.login()` — prompts for base URL and API key via `AuthInteraction.prompt()`
- `auth.apiKey.resolve()` — resolves the stored credential (apiKeyCommand → actual key)
- `auth.apiKey.check()` — checks connectivity to the endpoint
- `refreshModels()` — live-queries the endpoint and updates model catalog
- `getModels()` — returns last-known models (sync)

### Phase 2: Custom Model Picker via `ctx.ui.custom()`

Replace the simple `ctx.ui.select()` loops with a custom TUI component using `SelectList`, similar to llama's `ui.ts`.

**File: `src/model-picker.ts`** — Add new UI components (or new file `src/ui.ts`):

```
┌─────────────────────────────────────────┐
│  Local Models                           │
│  http://127.0.0.1:1234                 │
│  oMLX 0.x: 2/26 loaded, 0 loading      │
├─────────────────────────────────────────┤
│  → Qwen3.6-35B...        loaded         │
│    Llama-3.3-70B...                 78G │
│    Gemma-4-26B...                   28G │
│    ...                                 │
├─────────────────────────────────────────┤
│  Enter: select  l: load   u: unload     │
│  Esc: close                                │
```

Key behaviors:
- Arrow keys navigate, Enter selects model (sets as active)
- `l`/`u` trigger load/unload, then re-query and refresh display
- Search/filter by model name
- Status line shows current memory usage (oMLX) or loaded count
- Re-query on each picker open so model list is never stale

### Phase 3: Wire `/local-model` to Use the Custom Picker

Update `index.ts`:
- Keep `/local-login` mostly as-is but consider integrating with `/login` autocomplete completions
- Rewrite `/local-model` to:
  1. Pick connection (if multiple) via `ctx.ui.select()`
  2. Open custom picker via `ctx.ui.custom<DiscoveredModel>()`
  3. On model selection: register provider, call `pi.setModel()`

### Phase 4: Registration at Startup + Refresh on Demand

At startup (existing behavior), register the default connection's provider. On each `/local-model` invocation:
1. Call `pi.registerProvider()` with the refreshed Provider to update its model list
2. The models will then appear in `/model` for subsequent use

## Detailed File Changes

### `src/provider.ts` — New `createLocalProvider()` function

```typescript
import type { Provider, ApiKeyCredential, AuthContext, AuthResult, Model } from "@earendil-works/pi-ai";
import { streamSimpleOpenAICompletions } from "@earendil-works/pi-ai";
import type { DiscoveredModel, QueryResult } from "./model-picker";

export function createLocalProvider(
  baseUrl: string,
  storedApiKey: string,       // the raw apiKey command string
  resolveApiKey: (raw: string) => string,
  queryModels: (baseUrl: string, apiKey: string) => Promise<QueryResult>,
): Provider<"openai-completions"> {

  let models: Model<"openai-completions">[] = [];
  let queryResult: QueryResult | null = null;

  const toModel = (m: DiscoveredModel): Model<"openai-completions"> => ({
    id: m.id,
    name: m.displayName,
    api: "openai-completions",
    provider: baseUrl,
    baseUrl: `${baseUrl}/v1`,
    reasoning: m.reasoning ?? false,
    input: m.modelType?.includes("vlm") ? ["text", "image", "audio"] as const : ["text"] as const,
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: m.contextWindow ?? 128000,
    maxTokens: m.maxTokens ?? 16384,
  });

  return {
    id: baseUrl,
    name: baseUrl.replace(/^https?:\/\//, "").replace(/\/$/, ""), // display without protocol
    auth: {
      apiKey: {
        name: "Local server",
        // login() is called by /login when user picks this provider
        // For local providers, we already know the baseUrl (it IS the provider id)
        // So login just needs the API key.
        login: async (interaction): Promise<ApiKeyCredential> => {
          const apiKey = (await interaction.prompt({
            type: "secret",
            message: `API key for ${baseUrl}`,
            placeholder: storedApiKey.startsWith("!") ? "(enter key to store)" : "",
          })).trim();
          return { type: "api_key", key: apiKey || undefined };
        },
        resolve: async ({ ctx, credential }): Promise<AuthResult | undefined> => {
          const key = credential?.key ?? resolveApiKey(storedApiKey) ?? "";
          // Verify connectivity
          try {
            const result = await queryModels(baseUrl, key || "local");
            if (result.models.length >= 0) { // even empty is OK — server responded
              return { auth: { apiKey: key || "local", baseUrl: `${baseUrl}/v1` } };
            }
          } catch { /* ignore — ambient check */ }
          return undefined;
        },
      },
    },
    getModels: () => models,
    refreshModels: async (context): Promise<void> => {
      if (!context.allowNetwork || context.signal?.aborted) return;
      const key = resolveApiKey(storedApiKey) ?? "";
      try {
        queryResult = await queryModels(baseUrl, key || "local");
        models = queryResult.models.map(toModel);
      } catch {
        // Keep existing models on failure
      }
    },
    streamSimple: (model, context, options) =>
      streamSimpleOpenAICompletions(model, context, options),
  };
}
```

### `src/ui.ts` — New custom picker component (new file)

Modelled on `packages/coding-agent/src/extensions/llama/ui.ts`:

```typescript
import { Container, Focusable, fuzzyFilter, Input, SelectItem, SelectList, Spacer, Text, TUI } from "@earendil-works/pi-tui";
import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import type { DiscoveredModel, QueryResult } from "./model-picker";

export type LocalAction =
  | { type: "select"; model: DiscoveredModel }
  | { type: "close" };

export interface LocalUi {
  showModels(title: string, result: QueryResult): Promise<LocalAction>;
}

class LocalView implements LocalUi, Focusable {
  // ... SelectList-based component with:
  // - Search filter input
  // - Model list with load/unload status icons
  // - l/u key bindings for load/unload
  // - Enter to select, Esc to close
}

export async function showLocalUi(
  ctx: ExtensionCommandContext,
  run: (ui: LocalUi) => Promise<void>,
): Promise<void> {
  await ctx.ui.custom<void>((tui, theme, keybindings, done) => {
    const view = new LocalView(tui, theme, keybindings);
    void run(view).then(() => done(), () => done());
    return view;
  });
}
```

### `index.ts` — Updated commands

**Startup registration:** Unchanged (still reads saved default and registers provider).

**`/local-login`:** Mostly unchanged. Could be simplified to delegate to `/login <baseUrl>` where that's applicable, but keeping the current flow is fine since it handles multi-connection management well.

**`/local-model`:** Refactored to:
1. List connections → pick one (existing)
2. Register/update Provider with `createLocalProvider()`
3. Open custom picker via `ctx.ui.custom<DiscoveredModel>()`
4. On load/unload action: call API, re-query, refresh picker display
5. On model select: `pi.setModel()` + return

### Connection Management Considerations

Current approach stores connections in `pi-local-connections.json`. This works fine and should be kept. The `Provider` interface's credential storage is separate (pi's auth store) and handles the API key. The connection file stores:
- `baseUrl` → maps to provider id
- `apiKeyCommand` → the raw key reference (!security command, $ENV, or direct)
- Last-used model metadata

## Open Questions / Trade-offs

1. **Provider name in `/model` selector** — Currently the baseUrl (e.g., `http://127.0.0.1:1234`). Not ideal for display. Could strip protocol and use host:port format, or add a "friendly name" field to connections.

2. **Multiple connections in `/model`** — Each connection shows as a separate provider. If you have 3 local servers, that's 3 providers mixed with cloud providers in the selector. This is how llama.cpp works too.

3. **`/login` for already-known connections** — If a connection is already in `pi-local-connections.json`, `/login <baseUrl>` would show pi's native API key prompt. This is fine — user can just enter the key or cancel.

4. **RefreshModels and connection failures** — If a local server is down, `refreshModels` should not crash. Keep cached models.

## Migration Path

1. Add `createLocalProvider()` to `src/provider.ts` (no breaking changes)
2. Create `src/ui.ts` with custom picker component
3. Update `index.ts` `/local-model` to use the new Provider + custom picker
4. Keep `/local-login` as-is (it works well for multi-connection management)
5. Optionally: add autocomplete completions for `/login` to suggest known local endpoints

## Estimated Effort

- Phase 1 (Provider interface): 2-3 hours
- Phase 2 (Custom picker UI): 4-6 hours (most complex, references llama/ui.ts and ModelSelectorComponent)
- Phase 3+4 (Wiring): 1-2 hours

Total: ~8-11 hours of focused work.