pi-smart-router 0.1.0

package/.pi/extensions/smart-router/routing-context.ts

@@ -0,0 +1,171 @@
+import { randomUUID } from 'node:crypto';
+
+import {
+  type Context,
+  type Message,
+  type SimpleStreamOptions,
+  type TextContent,
+} from '@earendil-works/pi-ai/compat';
+
+import type {
+  Message as RoutingMessage,
+  RoutingRequest,
+  TurnType,
+} from '../../../src/domain/types/index.js';
+import { LifecycleHookState } from '../../../src/index.js';
+
+const TOKEN_ESTIMATE_KEYS = [
+  'estimatedInputTokens',
+  'estimated_input_tokens',
+  'contextTokens',
+  'tokenCount',
+  'tokens',
+] as const;
+
+function readOptionalTokenEstimate(source: unknown): number | undefined {
+  if (source === null || typeof source !== 'object') {
+    return undefined;
+  }
+
+  const record = source as Record<string, unknown>;
+  for (const key of TOKEN_ESTIMATE_KEYS) {
+    const value = record[key];
+    if (typeof value === 'number' && Number.isFinite(value) && value >= 0) {
+      return Math.floor(value);
+    }
+  }
+
+  return undefined;
+}
+
+function estimateInputTokens(
+  context: Context,
+  options?: SimpleStreamOptions,
+): number {
+  const fromOptions = readOptionalTokenEstimate(options);
+  if (fromOptions !== undefined) {
+    return fromOptions;
+  }
+
+  const fromContext = readOptionalTokenEstimate(context);
+  if (fromContext !== undefined) {
+    return fromContext;
+  }
+
+  const mapped = mapContextMessages(context.messages);
+  let charCount = mapped.reduce((sum, message) => sum + message.content.length, 0);
+  if (context.systemPrompt) {
+    charCount += context.systemPrompt.length;
+  }
+
+  if (context.messages.length === 0 && !context.systemPrompt) {
+    return 0;
+  }
+
+  return Math.max(1, Math.ceil(charCount / 4));
+}
+
+function messageContentToString(content: string | readonly (TextContent | { type: string })[]): string {
+  if (typeof content === 'string') {
+    return content;
+  }
+
+  return content
+    .filter((block): block is TextContent => block.type === 'text')
+    .map((block) => block.text)
+    .join('\n');
+}
+
+export function extractPromptText(messages: readonly Message[]): string {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const message = messages[i];
+    if (message?.role === 'user') {
+      const text = messageContentToString(message.content);
+      if (text.trim()) {
+        return text;
+      }
+    }
+  }
+  return '';
+}
+
+export function deriveTurnType(messages: readonly Message[]): TurnType {
+  if (messages.length === 0) {
+    return 'unknown';
+  }
+
+  const lastMessage = messages[messages.length - 1];
+  if (!lastMessage) {
+    return 'unknown';
+  }
+
+  if (lastMessage.role === 'toolResult') {
+    return 'tool_result';
+  }
+
+  if (lastMessage.role === 'user') {
+    const text = messageContentToString(lastMessage.content).toLowerCase();
+    if (
+      text.includes('plan') ||
+      text.includes('architect') ||
+      text.includes('design')
+    ) {
+      return 'planning';
+    }
+  }
+
+  return 'main_loop';
+}
+
+export function mapContextMessages(messages: readonly Message[]): RoutingMessage[] {
+  return messages.map((message) => {
+    if (message.role === 'user') {
+      return {
+        role: message.role,
+        content: messageContentToString(message.content),
+      };
+    }
+
+    if (message.role === 'assistant') {
+      const content = message.content
+        .map((block) => {
+          if (block.type === 'text') {
+            return block.text;
+          }
+          if (block.type === 'thinking') {
+            return block.thinking;
+          }
+          return '';
+        })
+        .filter(Boolean)
+        .join('\n');
+
+      return { role: message.role, content };
+    }
+
+    return {
+      role: 'tool',
+      content: messageContentToString(message.content),
+      tool_blocks: [],
+    };
+  });
+}
+
+export function buildRoutingRequest(
+  context: Context,
+  options: SimpleStreamOptions | undefined,
+  lifecycleHookState?: LifecycleHookState,
+): RoutingRequest {
+  const sessionId = options?.sessionId ?? randomUUID();
+  const lifecycleFlags = lifecycleHookState?.consume(sessionId) ?? {};
+
+  return {
+    request_id: randomUUID(),
+    session_id: sessionId,
+    prompt_text: extractPromptText(context.messages),
+    messages: mapContextMessages(context.messages),
+    turn_type: deriveTurnType(context.messages),
+    estimated_input_tokens: estimateInputTokens(context, options),
+    ...lifecycleFlags,
+  };
+}

package/.pi/extensions/smart-router/routing-outcomes.ts

@@ -0,0 +1,44 @@
+import type { RoutingDecision, RoutingRequest } from '../../../src/domain/types/index.js';
+import type { SessionRoutingSnapshot } from '../../../src/infrastructure/telemetry/outcome-recorder.js';
+import type { StreamDelegationDeps } from './types.js';
+
+export function capturePreRouteOutcomes(
+  request: RoutingRequest,
+  deps: StreamDelegationDeps,
+  priorSnapshot: SessionRoutingSnapshot | undefined,
+  hadPin: boolean,
+): void {
+  if (!priorSnapshot || !deps.outcomeRecorder) {
+    return;
+  }
+
+  const sessionId = request.session_id;
+
+  if (request.compaction_flag && hadPin) {
+    deps.outcomeRecorder.recordCompactionPinBreak(priorSnapshot, sessionId);
+  }
+
+  if (request.force_model_id) {
+    deps.outcomeRecorder.recordModelOverride(
+      priorSnapshot,
+      sessionId,
+      request.force_model_id,
+    );
+  }
+}
+
+export function updateSessionRoutingSnapshot(
+  deps: StreamDelegationDeps,
+  sessionId: string | undefined,
+  request: RoutingRequest,
+  decision: RoutingDecision,
+): void {
+  if (!sessionId || !deps.sessionRouting) {
+    return;
+  }
+
+  deps.sessionRouting.set(sessionId, {
+    lastRequestId: request.request_id,
+    lastSelectedModelId: decision.selected_model_id,
+  });
+}

package/.pi/extensions/smart-router/session-lifecycle.ts

@@ -0,0 +1,140 @@
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
+
+import { SessionPinner } from '../../../src/domain/pinning/session-pinner.js';
+import {
+  bindSharedModelRegistry,
+  ensureFleetFresh,
+  formatLmuStatus,
+  rebuildFleet,
+} from './fleet-bootstrap.js';
+import { notifyPricingStalenessIfNeeded } from './pricing-lifecycle.js';
+import type { FleetMode, SmartRouterRuntime } from './types.js';
+
+export const FLEET_MODE_ENTRY_TYPE = 'smart-router-fleet-mode' as const;
+
+const SMART_ROUTER_PROVIDER = 'smart-router' as const;
+const SMART_ROUTER_AUTO_ID = 'auto' as const;
+
+export function isSmartRouterActive(model: { provider: string; id: string }): boolean {
+  return model.provider === SMART_ROUTER_PROVIDER && model.id === SMART_ROUTER_AUTO_ID;
+}
+
+function parseFleetModeEntry(data: unknown): FleetMode | undefined {
+  if (
+    typeof data === 'object' &&
+    data !== null &&
+    'mode' in data &&
+    (data.mode === 'scoped' || data.mode === 'all')
+  ) {
+    return data.mode;
+  }
+  return undefined;
+}
+
+export function restoreFleetModeFromSession(ctx: ExtensionContext): FleetMode | undefined {
+  const entries = ctx.sessionManager.getEntries();
+  for (let i = entries.length - 1; i >= 0; i--) {
+    const entry = entries[i];
+    if (entry?.type === 'custom' && entry.customType === FLEET_MODE_ENTRY_TYPE) {
+      return parseFleetModeEntry(entry.data);
+    }
+  }
+  return undefined;
+}
+
+function restoreLmuFromLedger(runtime: SmartRouterRuntime, sessionId: string): void {
+  const lastExec = runtime.executionLedger.getLastExecution(sessionId);
+  if (lastExec) {
+    runtime.setLmuStatus?.(lastExec.id);
+  } else if (runtime.lastDecision) {
+    runtime.setLmuStatus?.(runtime.lastDecision.selected_model_id);
+  }
+}
+
+function wireLmuStatusHandlers(
+  runtime: SmartRouterRuntime,
+  ctx: ExtensionContext,
+  getActiveModel: () => { provider: string; id: string },
+): void {
+  runtime.setLmuStatus = (modelId) => {
+    if (!isSmartRouterActive(getActiveModel())) {
+      return;
+    }
+    ctx.ui.setStatus(
+      'smart-router-lmu',
+      formatLmuStatus(modelId, ctx.ui.theme as { fg: (color: string, text: string) => string }),
+    );
+  };
+  runtime.clearLmuStatus = () => {
+    ctx.ui.setStatus('smart-router-lmu', undefined);
+  };
+}
+
+export function setupSessionHooks(
+  pi: ExtensionAPI,
+  runtime: SmartRouterRuntime,
+  sessionPinner: SessionPinner,
+  datasetNotify: { fn: ((message: string) => void) | undefined },
+): void {
+  let activeModel: { provider: string; id: string } | undefined;
+
+  pi.on('session_start', async (_event, ctx) => {
+    activeModel = ctx.model;
+    bindSharedModelRegistry(runtime, ctx.modelRegistry);
+    runtime.sessionCwd = ctx.cwd;
+    runtime.streamDeps.ensureFleetFresh = async () => {
+      if (runtime.sessionCwd === undefined) {
+        return;
+      }
+      await ensureFleetFresh(runtime, pi, runtime.sessionCwd);
+    };
+
+    const restoredMode = restoreFleetModeFromSession(ctx);
+    if (restoredMode) {
+      runtime.fleetMode = restoredMode;
+    }
+    await rebuildFleet(runtime, pi, ctx.cwd);
+
+    notifyPricingStalenessIfNeeded(runtime, (message, level) => {
+      ctx.ui.notify(message, level);
+    });
+
+    wireLmuStatusHandlers(runtime, ctx, () => {
+      const model = activeModel ?? ctx.model;
+      return model ?? { provider: '', id: '' };
+    });
+    runtime.notifyDatasetEnabled = (message) => {
+      ctx.ui.notify(message, 'info');
+    };
+    datasetNotify.fn = runtime.notifyDatasetEnabled;
+
+    const sessionId = ctx.sessionManager.getSessionId();
+    await sessionPinner.restoreSessionPin(sessionId);
+
+    if (ctx.model !== undefined && isSmartRouterActive(ctx.model)) {
+      restoreLmuFromLedger(runtime, sessionId);
+    } else {
+      runtime.clearLmuStatus?.();
+    }
+  });
+
+  pi.on('model_select', (event, ctx) => {
+    activeModel = event.model;
+    if (isSmartRouterActive(event.model)) {
+      restoreLmuFromLedger(runtime, ctx.sessionManager.getSessionId());
+    } else {
+      runtime.clearLmuStatus?.();
+    }
+  });
+
+  pi.on('session_shutdown', async (_event, ctx) => {
+    activeModel = undefined;
+    delete runtime.setLmuStatus;
+    delete runtime.clearLmuStatus;
+    delete runtime.notifyDatasetEnabled;
+    delete runtime.sessionCwd;
+    delete runtime.streamDeps.ensureFleetFresh;
+    datasetNotify.fn = undefined;
+    ctx.ui.setStatus('smart-router-lmu', undefined);
+  });
+}

package/.pi/extensions/smart-router/stream-delegation.ts

@@ -0,0 +1,75 @@
+import {
+  type Api,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+  createAssistantMessageEventStream,
+} from '@earendil-works/pi-ai/compat';
+
+import type { RoutingDecision, RoutingFeatureSidecar } from '../../../src/domain/types/index.js';
+import { createErrorMessage } from './delegation-runtime.js';
+import { routeAndDelegate } from './route-and-delegate.js';
+import type { StreamDelegationDeps } from './types.js';
+
+function isRoutingLogEnabled(): boolean {
+  return process.env.SMART_ROUTER_LOG_ROUTING === '1';
+}
+
+export function logRoutingDecision(
+  decision: RoutingDecision,
+  delegate?: { provider: string; modelId: string; api: Api },
+): void {
+  if (!isRoutingLogEnabled()) {
+    return;
+  }
+
+  console.warn(
+    '[smart-router] routing decision',
+    JSON.stringify({
+      request_id: decision.request_id,
+      selected_model_id: decision.selected_model_id,
+      tier: decision.tier,
+      stage: decision.stage,
+      reason_code: decision.reason_code,
+      routing_latency_ms: decision.routing_latency_ms,
+      features: decision.features ?? null,
+      delegate,
+    }),
+  );
+}
+
+/** Read dataset feature sidecar from a routing decision (SP-057). */
+export function getRoutingFeatureSidecar(
+  decision: RoutingDecision,
+): RoutingFeatureSidecar | undefined {
+  return decision.features;
+}
+
+export function createStreamSimple(deps: StreamDelegationDeps) {
+  return function streamSimple(
+    model: Model<Api>,
+    context: Context,
+    options?: SimpleStreamOptions,
+  ): AssistantMessageEventStream {
+    const stream = createAssistantMessageEventStream();
+
+    void (async () => {
+      try {
+        await routeAndDelegate(context, options, deps, stream);
+      } catch (error) {
+        stream.push({
+          type: 'error',
+          reason: options?.signal?.aborted ? 'aborted' : 'error',
+          error: createErrorMessage(model, options, error),
+        });
+        stream.end();
+      }
+    })();
+
+    return stream;
+  };
+}
+
+export { resolveDelegationOptions, buildDelegationContext } from './delegation-runtime.js';
+export { resolveTargetModel } from './route-and-delegate.js';

package/.pi/extensions/smart-router/types.ts

@@ -0,0 +1,90 @@
+import type {
+  Api,
+  AssistantMessageEventStream,
+  Context,
+  Model,
+  SimpleStreamOptions,
+} from '@earendil-works/pi-ai/compat';
+import type { ModelRegistry } from '@earendil-works/pi-coding-agent';
+
+import type { HydraMatcher } from '../../../src/domain/matching/hydra-matcher.js';
+import { ExecutionLedger } from '../../../src/domain/delegation/execution-ledger.js';
+import { SessionPinner } from '../../../src/domain/pinning/session-pinner.js';
+import type {
+  ModelProfile,
+  PriceCatalog,
+  RoutingDecision,
+} from '../../../src/domain/types/index.js';
+import type { StorePort } from '../../../src/domain/types/store-port.js';
+import {
+  DatasetRecorder,
+} from '../../../src/infrastructure/telemetry/dataset-recorder.js';
+import {
+  OutcomeRecorder,
+  type SessionRoutingSnapshot,
+} from '../../../src/infrastructure/telemetry/outcome-recorder.js';
+import {
+  LifecycleHookState,
+  type RouterHandle,
+} from '../../../src/index.js';
+
+export type FleetMode = 'scoped' | 'all';
+
+export type SmartRouterCommand =
+  | { command: 'status' }
+  | { command: 'history'; limit: number }
+  | { command: 'mode'; mode: FleetMode }
+  | { command: 'pricing'; subcommand: 'refresh' }
+  | { command: 'export'; subcommand: 'dataset'; limit: number }
+  | { command: 'export'; subcommand: 'telemetry-contrib'; limit: number }
+  | { command: 'feedback'; rating: 'good' | 'bad' }
+  | { command: 'unpin' };
+
+/** Provider stream delegate; defaults to pi-ai streamSimple when omitted. */
+export type DelegateStreamFn = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+export interface StreamDelegationDeps {
+  router: RouterHandle;
+  modelRegistry: ModelRegistry;
+  fleet: ModelProfile[];
+  /** Cheap scope fingerprint check before each routed turn. */
+  ensureFleetFresh?: () => Promise<void>;
+  readonly executionLedger: ExecutionLedger;
+  /** Injectable for tests; production uses pi-ai streamSimple. */
+  delegateStream?: DelegateStreamFn;
+  readonly lifecycleHookState?: LifecycleHookState;
+  readonly datasetRecorder?: DatasetRecorder;
+  readonly outcomeRecorder?: OutcomeRecorder;
+  readonly sessionPinner?: SessionPinner;
+  readonly sessionRouting?: Map<string, SessionRoutingSnapshot>;
+  onRoutingDecision?: (decision: RoutingDecision) => void;
+  /** Fired when a delegated provider stream completes successfully. */
+  onDelegatedModel?: (model: { readonly provider: string; readonly id: string }) => void;
+}
+
+export interface SmartRouterRuntime {
+  fleetMode: FleetMode;
+  lastDecision: RoutingDecision | undefined;
+  priceCatalog: PriceCatalog | null;
+  /** Cached scope fingerprint; rebuild when this changes. */
+  fleetScopeFingerprint?: string;
+  /** Session cwd for ensureFleetFresh before routed turns. */
+  sessionCwd?: string;
+  modelRegistry: ModelRegistry;
+  readonly store: StorePort;
+  readonly sessionPinner: SessionPinner;
+  readonly executionLedger: ExecutionLedger;
+  readonly lifecycleHookState: LifecycleHookState;
+  readonly datasetRecorder?: DatasetRecorder;
+  readonly outcomeRecorder?: OutcomeRecorder;
+  readonly sessionRouting: Map<string, SessionRoutingSnapshot>;
+  streamDeps: StreamDelegationDeps;
+  hydraMatcher: HydraMatcher | undefined;
+  setLmuStatus?: (modelId: string) => void;
+  clearLmuStatus?: () => void;
+  notifyDatasetEnabled?: (message: string) => void;
+}

package/.pi/extensions/smart-router/utils.ts

@@ -0,0 +1,58 @@
+import { join } from 'node:path';
+
+import {
+  createResilientStore,
+  SqliteStore,
+  SqliteStoreError,
+} from '../../../src/infrastructure/persistence/sqlite-store.js';
+import type { StorePort } from '../../../src/domain/types/store-port.js';
+import type { RateLimitPort } from '../../../src/infrastructure/gateway/gateway-dispatch.js';
+
+export const DEFAULT_ROUTER_STATE_DB_PATH = '.pi-smart-router/state.db';
+export const DEFAULT_RATE_LIMIT_MAX_TOKENS = 60;
+export const DEFAULT_RATE_LIMIT_REFILL_RATE = 1;
+
+export function getRouterStateDbPath(cwd: string): string {
+  const configured = process.env.ROUTER_STATE_DB_PATH?.trim();
+  if (configured && configured.length > 0) {
+    return configured;
+  }
+  return join(cwd, DEFAULT_ROUTER_STATE_DB_PATH);
+}
+
+export function createExtensionStore(cwd: string): StorePort {
+  return createResilientStore({
+    dbPath: getRouterStateDbPath(cwd),
+    models: [],
+  }).store;
+}
+
+export function createSqliteRateLimiter(sqliteStore: SqliteStore): RateLimitPort {
+  return {
+    consumeToken(key: string, cost = 1) {
+      try {
+        return sqliteStore.consumeToken(key, cost);
+      } catch (error) {
+        if (
+          error instanceof SqliteStoreError &&
+          error.message.includes('Token bucket not found')
+        ) {
+          sqliteStore.initBucket(
+            key,
+            DEFAULT_RATE_LIMIT_MAX_TOKENS,
+            DEFAULT_RATE_LIMIT_REFILL_RATE,
+          );
+          return sqliteStore.consumeToken(key, cost);
+        }
+        throw error;
+      }
+    },
+  };
+}
+
+export function resolveRateLimiter(store: StorePort): RateLimitPort | undefined {
+  if (!(store instanceof SqliteStore)) {
+    return undefined;
+  }
+  return createSqliteRateLimiter(store);
+}