@aliou/pi-ts-aperture 0.2.4 → 0.3.0

package/README.md

@@ -2,7 +2,7 @@
 
 Route Pi LLM providers through [Tailscale Aperture](https://tailscale.com/docs/features/aperture), a managed AI gateway on your tailnet.
 
-Aperture handles API key injection and request routing server-side. This extension overrides the base URL for selected providers so all LLM requests go through your Aperture instance instead of directly to provider APIs.
+Aperture handles API key injection and request routing server-side. This extension overrides selected providers so requests go through your Aperture gateway instead of directly to upstream provider APIs.
 
 ## Setup
 
@@ -16,9 +16,9 @@ Then run the setup wizard:
 /aperture:setup
 ```
 
-This will prompt you for:
-1. Your Aperture base URL (e.g. `ai.your-tailnet.ts.net`)
-2. Which providers to route through Aperture (fuzzy searchable, multi-select)
+This prompts for:
+1. Aperture base URL (for example `ai.your-tailnet.ts.net`)
+2. Providers to route through Aperture (fuzzy searchable, multi-select)
 
 Configuration is saved globally to `~/.pi/agent/extensions/aperture.json`.
 
@@ -26,19 +26,23 @@ Configuration is saved globally to `~/.pi/agent/extensions/aperture.json`.
 
 | Command | Description |
 |---|---|
-| `/aperture:setup` | Interactive wizard to configure Aperture URL and providers |
-| `/aperture:settings` | Settings UI to update base URL and provider list |
+| `/aperture:setup` | Interactive wizard to configure Aperture URL and routed providers |
+| `/aperture:settings` | Settings UI to update URL and routed provider list |
 
 ## How it works
 
 For each configured provider, the extension calls `registerProvider` with:
-- `baseUrl` set to your Aperture URL + `/v1`
-- `apiKey` set to `"-"` (Aperture ignores client keys, it injects its own)
 
-This means all requests for those providers are routed through Aperture, which handles authentication, logging, and cost tracking on its end.
+- `baseUrl` set to your Aperture URL + `/v1` (OpenAI-compatible surface used by Pi provider configs)
+- `apiKey` set to `"-"` (Aperture injects upstream credentials server-side)
+- provenance headers:
+  - `Referer: https://pi.dev`
+  - `X-Title: npm:@aliou/pi-ts-aperture`
+
+Additionally, the extension can bootstrap model IDs discovered from Aperture (`/api/providers`) for providers like OpenRouter so CLI model selection can resolve Aperture-exposed model IDs before the first prompt.
 
 ## Requirements
 
 - A Tailscale tailnet with Aperture configured
-- The device running Pi must be on the tailnet
-- Use HTTP, not HTTPS, for the Aperture URL (WireGuard handles encryption)
+- The device running Pi must be on the tailnet (or otherwise able to reach your Aperture endpoint)
+- Use the URL/scheme that matches your deployment (`http://` or `https://`)

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@aliou/pi-ts-aperture",
   "description": "Route Pi LLM providers through Tailscale Aperture",
-  "version": "0.2.4",
+  "version": "0.3.0",
   "type": "module",
   "license": "MIT",
   "private": false,

package/src/index.ts

@@ -1,105 +1,76 @@
 /**
  * Pi extension for Tailscale Aperture integration.
  *
- * Routes selected LLM providers through an Aperture gateway on your tailnet.
- * Aperture handles API key injection and request routing, so this extension
- * overrides each provider's baseUrl and sets a dummy apiKey.
+ * Keeps the entry point focused on orchestration:
+ * - load config
+ * - bootstrap provider/model visibility
+ * - register lifecycle hooks
+ * - register user commands
  */
 
 import type {
   ExtensionAPI,
   ExtensionContext,
 } from "@mariozechner/pi-coding-agent";
-import { VERSION } from "@mariozechner/pi-coding-agent";
 import { registerApertureSettings } from "./commands/settings";
 import { registerSetupCommand } from "./commands/setup";
 import { configLoader } from "./config";
+import {
+  applyAperture,
+  bootstrapProvidersFromAperture,
+  refreshActiveModel,
+  resetApertureModelsCache,
+} from "./providers/aperture";
 
-/**
- * Compute the full Aperture base URL from config, or null if not configured.
- */
-function resolveBaseUrl(): string | null {
-  const { baseUrl, providers } = configLoader.getConfig();
-  if (!baseUrl || providers.length === 0) return null;
-  return `${baseUrl.replace(/\/+$/, "")}/v1`;
-}
+function registerApertureLifecycleHook(pi: ExtensionAPI): void {
+  pi.on("before_agent_start", async (_event, ctx) => {
+    if (!ctx?.modelRegistry) return;
 
-/**
- * Override provider registrations to route through Aperture.
- * Preserves existing models so extensions that registered custom models
- * before this runs don't lose them.
- */
-function overrideProviders(
-  pi: ExtensionAPI,
-  registry: ExtensionContext["modelRegistry"],
-  providers: string[],
-  baseUrl: string,
-): void {
-  for (const provider of providers) {
-    const models = registry.getAll().filter((m) => m.provider === provider);
+    const overriddenProviders = await applyAperture(pi, ctx.modelRegistry);
+    if (!ctx.model || !overriddenProviders.includes(ctx.model.provider)) return;
 
-    pi.registerProvider(provider, {
-      baseUrl,
-      apiKey: "-",
-      ...(models.length > 0 && { api: models[0].api, models }),
-    });
-  }
+    await refreshActiveModel(pi, ctx);
+  });
 }
 
-/**
- * Apply Aperture configuration to the model registry.
- * Returns the list of providers that were overridden, or empty if no-op.
- */
-function applyAperture(
+function createConfigChangeHandler(
   pi: ExtensionAPI,
-  registry: ExtensionContext["modelRegistry"],
-): string[] {
-  const url = resolveBaseUrl();
-  if (!url) return [];
-
-  const { providers } = configLoader.getConfig();
-  overrideProviders(pi, registry, providers, url);
-  return providers;
-}
-
-export default async function (pi: ExtensionAPI): Promise<void> {
-  await configLoader.load();
-
+): (ctx: ExtensionContext) => void {
   let lastRegisteredProviders = [...configLoader.getConfig().providers];
 
-  // Apply after all extensions have registered their providers and models.
-  pi.events.on("before_agent_start", async (data) => {
-    const ctx = data as ExtensionContext;
-    if (!ctx?.modelRegistry) return;
-    applyAperture(pi, ctx.modelRegistry);
-  });
-
-  const onSetupComplete = (ctx: ExtensionContext) => {
+  return (ctx: ExtensionContext) => {
     const { providers } = configLoader.getConfig();
-    const removed = lastRegisteredProviders.filter(
-      (p) => !providers.includes(p),
+    const removedProviders = lastRegisteredProviders.filter(
+      (provider) => !providers.includes(provider),
     );
 
-    applyAperture(pi, ctx.modelRegistry);
+    resetApertureModelsCache();
+    void applyAperture(pi, ctx.modelRegistry);
     lastRegisteredProviders = [...providers];
 
-    // Re-resolve active model if it belongs to a reconfigured provider.
     if (ctx.model && providers.includes(ctx.model.provider)) {
-      const updated = ctx.modelRegistry.find(ctx.model.provider, ctx.model.id);
-      if (updated) {
+      void refreshActiveModel(pi, ctx).then((updated) => {
+        if (!updated) return;
         ctx.ui.notify(
-          `[aperture] re-routing ${ctx.model.id} through ${updated.baseUrl}`,
+          `[aperture] re-routing ${ctx.model?.id ?? "model"} through ${ctx.model?.baseUrl ?? "aperture"}`,
           "info",
         );
-        pi.setModel(updated);
-      }
+      });
     }
 
-    for (const p of removed) {
-      pi.unregisterProvider(p);
+    for (const provider of removedProviders) {
+      pi.unregisterProvider(provider);
     }
   };
+}
+
+export default async function (pi: ExtensionAPI): Promise<void> {
+  await configLoader.load();
+  await bootstrapProvidersFromAperture(pi);
+
+  registerApertureLifecycleHook(pi);
 
-  registerSetupCommand(pi, onSetupComplete);
-  registerApertureSettings(pi, onSetupComplete);
+  const onConfigChange = createConfigChangeHandler(pi);
+  registerSetupCommand(pi, onConfigChange);
+  registerApertureSettings(pi, onConfigChange);
 }

package/src/lib/aperture-api.ts

@@ -0,0 +1,32 @@
+export interface ApertureProviderInfo {
+  id?: string;
+  models?: string[];
+}
+
+/**
+ * Fetch provider -> model list mapping from Aperture and keep only selected
+ * providers configured by the user.
+ */
+export async function fetchApertureProviderModels(
+  gatewayUrl: string,
+  providers: string[],
+): Promise<Map<string, string[]>> {
+  const response = await fetch(`${gatewayUrl}/api/providers`, {
+    signal: AbortSignal.timeout(4000),
+  });
+
+  if (!response.ok) {
+    return new Map();
+  }
+
+  const data = (await response.json()) as ApertureProviderInfo[];
+  const selectedProviders = new Set(providers);
+  const modelsByProvider = new Map<string, string[]>();
+
+  for (const provider of data) {
+    if (!provider.id || !selectedProviders.has(provider.id)) continue;
+    modelsByProvider.set(provider.id, provider.models ?? []);
+  }
+
+  return modelsByProvider;
+}

package/src/providers/aperture.ts

@@ -0,0 +1,167 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  ProviderModelConfig,
+} from "@mariozechner/pi-coding-agent";
+import { configLoader } from "../config";
+import { fetchApertureProviderModels } from "../lib/aperture-api";
+import {
+  clearProviderModelsCache,
+  getProviderModelsCache,
+  setProviderModelsCache,
+} from "../state/provider-model-cache";
+import { mergeModels, toModelConfig } from "./model-config";
+
+/**
+ * Preserve provenance similarly to pi-synthetic so downstream providers can
+ * attribute traffic to Pi / this extension.
+ */
+const APERTURE_PROVENANCE_HEADERS = {
+  Referer: "https://pi.dev",
+  "X-Title": "npm:@aliou/pi-ts-aperture",
+};
+
+/**
+ * Providers for which we bootstrap models at startup (before first turn)
+ * to make CLI model selection deterministic.
+ */
+const BOOTSTRAP_DISCOVERY_PROVIDERS = new Set(["openrouter"]);
+
+/** Returns configured gateway URL without trailing slash. */
+export function resolveGatewayUrl(): string | null {
+  const { baseUrl, providers } = configLoader.getConfig();
+  if (!baseUrl || providers.length === 0) return null;
+  return baseUrl.replace(/\/+$/, "");
+}
+
+/**
+ * Returns the Aperture provider base URL used for provider registration.
+ *
+ * Aperture exposes multiple protocol paths (OpenAI, Anthropic, Gemini, ...).
+ * For this extension we route through the OpenAI-compatible `/v1` surface that
+ * Pi providers use (`openai-completions` API).
+ */
+export function resolveApertureProviderBaseUrl(): string | null {
+  const gateway = resolveGatewayUrl();
+  if (!gateway) return null;
+  return `${gateway}/v1`;
+}
+
+function resolveProviderHeaders(
+  models: ProviderModelConfig[],
+): Record<string, string> {
+  const modelHeaders = models.find((m) => m.headers)?.headers ?? {};
+  return {
+    ...APERTURE_PROVENANCE_HEADERS,
+    ...modelHeaders,
+  };
+}
+
+async function getOrLoadProviderModelsCache(
+  gatewayUrl: string,
+  providers: string[],
+): Promise<Map<string, string[]>> {
+  const current = getProviderModelsCache();
+  if (current) return current;
+
+  const loaded = await fetchApertureProviderModels(gatewayUrl, providers);
+  setProviderModelsCache(loaded);
+  return loaded;
+}
+
+export function resetApertureModelsCache(): void {
+  clearProviderModelsCache();
+}
+
+/**
+ * Apply Aperture override to configured providers:
+ * - provider baseUrl -> aperture /v1 endpoint
+ * - apiKey -> dummy token (Aperture injects real key server-side)
+ * - headers -> provenance + provider/model headers
+ */
+export async function applyAperture(
+  pi: ExtensionAPI,
+  registry: ExtensionContext["modelRegistry"],
+): Promise<string[]> {
+  const baseUrl = resolveApertureProviderBaseUrl();
+  const gatewayUrl = resolveGatewayUrl();
+  if (!baseUrl || !gatewayUrl) return [];
+
+  const { providers } = configLoader.getConfig();
+
+  let modelCache: Map<string, string[]>;
+  try {
+    modelCache = await getOrLoadProviderModelsCache(gatewayUrl, providers);
+  } catch {
+    modelCache = new Map();
+  }
+
+  for (const provider of providers) {
+    const existingModels = registry
+      .getAll()
+      .filter((m) => m.provider === provider) as ProviderModelConfig[];
+
+    const models = mergeModels(existingModels, modelCache.get(provider));
+
+    pi.registerProvider(provider, {
+      baseUrl,
+      apiKey: "-",
+      headers: resolveProviderHeaders(models),
+      ...(models.length > 0 && { api: models[0].api, models }),
+    });
+  }
+
+  return providers;
+}
+
+/**
+ * Pre-register selected providers from Aperture model discovery so CLI model
+ * resolution works even when a model is not present in Pi built-ins.
+ */
+export async function bootstrapProvidersFromAperture(
+  pi: ExtensionAPI,
+): Promise<void> {
+  const baseUrl = resolveApertureProviderBaseUrl();
+  const gatewayUrl = resolveGatewayUrl();
+  if (!baseUrl || !gatewayUrl) return;
+
+  const { providers } = configLoader.getConfig();
+
+  let modelCache: Map<string, string[]>;
+  try {
+    modelCache = await fetchApertureProviderModels(gatewayUrl, providers);
+    setProviderModelsCache(modelCache);
+  } catch {
+    return;
+  }
+
+  for (const provider of providers) {
+    if (!BOOTSTRAP_DISCOVERY_PROVIDERS.has(provider)) continue;
+
+    const modelIds = modelCache.get(provider) ?? [];
+    if (modelIds.length === 0) continue;
+
+    const models = modelIds.map((id) => toModelConfig(id));
+
+    pi.registerProvider(provider, {
+      baseUrl,
+      apiKey: "-",
+      api: "openai-completions",
+      headers: resolveProviderHeaders(models),
+      models,
+    });
+  }
+}
+
+/** Re-resolve and set current model after provider registry updates. */
+export async function refreshActiveModel(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+): Promise<boolean> {
+  if (!ctx.model) return false;
+
+  const updated = ctx.modelRegistry.find(ctx.model.provider, ctx.model.id);
+  if (!updated) return false;
+
+  return pi.setModel(updated);
+}

package/src/providers/model-config.ts

@@ -0,0 +1,54 @@
+import type { ProviderModelConfig } from "@mariozechner/pi-coding-agent";
+
+const DEFAULT_COST = {
+  input: 0,
+  output: 0,
+  cacheRead: 0,
+  cacheWrite: 0,
+};
+
+/**
+ * Build a ProviderModelConfig for Aperture-discovered model IDs.
+ *
+ * When a template model is available (same provider), preserve its API/compat
+ * shape so behavior stays consistent after rerouting.
+ */
+export function toModelConfig(
+  id: string,
+  template?: ProviderModelConfig,
+): ProviderModelConfig {
+  return {
+    id,
+    name: template?.name ?? id,
+    api: template?.api ?? "openai-completions",
+    reasoning: template?.reasoning ?? false,
+    input: template?.input ?? ["text"],
+    cost: template?.cost ?? DEFAULT_COST,
+    contextWindow: template?.contextWindow ?? 128000,
+    maxTokens: template?.maxTokens ?? 16384,
+    headers: template?.headers,
+    compat: template?.compat,
+  };
+}
+
+/**
+ * Merge known provider models with Aperture-discovered model IDs.
+ * Existing models win; missing IDs are synthesized from template defaults.
+ */
+export function mergeModels(
+  existingModels: ProviderModelConfig[],
+  apertureModelIds: string[] | undefined,
+): ProviderModelConfig[] {
+  if (!apertureModelIds || apertureModelIds.length === 0) return existingModels;
+
+  const modelsById = new Map(existingModels.map((m) => [m.id, m]));
+  const template = existingModels[0];
+
+  for (const modelId of apertureModelIds) {
+    if (!modelsById.has(modelId)) {
+      modelsById.set(modelId, toModelConfig(modelId, template));
+    }
+  }
+
+  return [...modelsById.values()];
+}

package/src/state/provider-model-cache.ts

@@ -0,0 +1,18 @@
+/**
+ * In-memory cache for Aperture provider model discovery.
+ *
+ * Ephemeral by design: reset on config changes and process restart.
+ */
+let providerModelsCache: Map<string, string[]> | null = null;
+
+export function getProviderModelsCache(): Map<string, string[]> | null {
+  return providerModelsCache;
+}
+
+export function setProviderModelsCache(models: Map<string, string[]>): void {
+  providerModelsCache = models;
+}
+
+export function clearProviderModelsCache(): void {
+  providerModelsCache = null;
+}