@tokenbuddy/tokenbuddy 1.0.4 → 1.0.6

package/src/daemon.ts

@@ -9,8 +9,22 @@ import {
   applyProviderInstall,
   detectProviders,
   previewProviderInstall,
-  rollbackProviderInstall
+  rollbackProviderInstall,
+  type ClaudeCodeModelMappingConfig,
 } from "./provider-install.js";
+import {
+  discoverSellerBackedModels,
+  fetchSellerManifest,
+  fetchSellerRegistry,
+  manifestModelIds,
+  manifestPaymentMethods,
+  manifestProtocols,
+  normalizeSellerUrl,
+  type RegistrySeller,
+  type SellerManifest,
+  type SellerRegistryDocument,
+  type SellerRoutingPreference,
+} from "./seller-catalog.js";
 
 const logger = createModuleLogger("tb-proxyd");
 const PROXY_JSON_BODY_LIMIT = "10mb";
@@ -21,30 +35,7 @@ export interface DaemonConfig {
   dbPath: string;
   sellerRegistryUrl: string;
   selectionMode?: "auto" | "manual";
-}
-
-interface RegistrySeller {
-  id: string;
-  name?: string;
-  url: string;
-  supportedProtocols?: string[];
-  paymentMethods?: string[];
-}
-
-interface SellerRegistryDocument {
-  version: number;
-  defaultSeller?: string;
-  sellers: RegistrySeller[];
-}
-
-interface SellerManifest {
-  sellerId?: string;
-  seller_id?: string;
-  supportedProtocols?: string[];
-  supported_protocols?: string[];
-  paymentMethods?: string[];
-  payment_methods?: string[];
-  models?: Array<{ id: string; [key: string]: unknown }>;
+  selectedSellerId?: string;
 }
 
 interface SellerRoute {
@@ -61,6 +52,225 @@ interface UsageSummary {
   billedMicros: number;
 }
 
+interface ResponsesStreamState {
+  itemId: string;
+  text: string;
+  contentPartStarted: boolean;
+}
+
+class ResponsesStreamNormalizer {
+  private pending = "";
+  private readonly state = new Map<string, ResponsesStreamState>();
+
+  public push(chunk: string): string {
+    this.pending += chunk;
+    const blocks = this.pending.split("\n\n");
+    this.pending = blocks.pop() || "";
+    return blocks
+      .map((block) => this.normalizeBlock(block))
+      .filter((block) => block.length > 0)
+      .join("\n\n");
+  }
+
+  public finish(): string {
+    if (!this.pending.trim()) {
+      return "";
+    }
+    const block = this.normalizeBlock(this.pending);
+    this.pending = "";
+    return block;
+  }
+
+  private normalizeBlock(block: string): string {
+    if (!block.trim()) {
+      return "";
+    }
+    // Each \n\n separates an event in SSE format
+    const subBlocks = block.split("\n\n");
+    const output: string[] = [];
+
+    for (const sub of subBlocks) {
+      if (!sub.trim() || sub.trim() === "data: [DONE]") {
+        if (sub.trim()) output.push(sub);
+        continue;
+      }
+
+      const lines = sub.split("\n");
+      const eventLine = lines.find((l) => l.startsWith("event:"));
+      const dataLine = lines.find((l) => l.startsWith("data:"));
+      if (!dataLine) {
+        output.push(sub);
+        continue;
+      }
+      const rawData = dataLine.replace(/^data:\s?/, "");
+      if (rawData === "[DONE]") {
+        output.push(sub);
+        continue;
+      }
+
+      let payload: any;
+      try {
+        payload = JSON.parse(rawData);
+      } catch {
+        output.push(sub);
+        continue;
+      }
+
+      const eventName =
+        (eventLine?.replace(/^event:\s?/, "") || payload?.type) as string;
+      if (!eventName || !eventName.startsWith("response.")) {
+        output.push(sub);
+        continue;
+      }
+
+      // When upstream already sends content_part.added, record it in state
+      if (
+        eventName === "response.content_part.added" &&
+        payload?.item_id
+      ) {
+        const current = this.state.get(payload.item_id as string);
+        if (current) current.contentPartStarted = true;
+        output.push(sub);
+        continue;
+      }
+
+      // response.output_item.added: inject content_part.added only if upstream hasn't
+      if (
+        eventName === "response.output_item.added" &&
+        payload?.item?.type === "message" &&
+        payload?.item?.id
+      ) {
+        const itemId = payload.item.id as string;
+        const current = this.getState(itemId);
+        const item = { ...payload.item };
+        item.content = [{ type: "output_text", text: "", annotations: [] }];
+        output.push(this.serializeEvent(eventName, {
+          ...payload,
+          output_index: payload.output_index ?? 0,
+          item
+        }));
+        if (!current.contentPartStarted) {
+          current.contentPartStarted = true;
+          output.push(this.serializeEvent("response.content_part.added", {
+            type: "response.content_part.added",
+            item_id: itemId,
+            output_index: payload.output_index ?? 0,
+            content_index: 0,
+            part: { type: "output_text", text: "", annotations: [] }
+          }));
+        }
+        continue;
+      }
+
+      // response.output_text.delta: inject content_part.added if missing
+      if (eventName === "response.output_text.delta" && payload?.item_id) {
+        const itemId = payload.item_id as string;
+        const current = this.getState(itemId);
+        if (!current.contentPartStarted) {
+          current.contentPartStarted = true;
+          output.push(this.serializeEvent("response.content_part.added", {
+            type: "response.content_part.added",
+            item_id: itemId,
+            output_index: payload.output_index ?? 0,
+            content_index: payload.content_index ?? 0,
+            part: { type: "output_text", text: "", annotations: [] }
+          }));
+        }
+        const deltaText =
+          typeof payload.delta === "string"
+            ? payload.delta
+            : typeof payload.delta?.text === "string"
+              ? payload.delta.text
+              : "";
+        current.text += deltaText;
+        output.push(this.serializeEvent(eventName, {
+          ...payload,
+          output_index: payload.output_index ?? 0,
+          content_index: payload.content_index ?? 0
+        }));
+        continue;
+      }
+
+      // response.output_text.done: also emit content_part.done
+      if (eventName === "response.output_text.done" && payload?.item_id) {
+        const itemId = payload.item_id as string;
+        const current = this.getState(itemId);
+        output.push(this.serializeEvent(eventName, {
+          ...payload,
+          output_index: payload.output_index ?? 0,
+          content_index: payload.content_index ?? 0
+        }));
+        output.push(this.serializeEvent("response.content_part.done", {
+          type: "response.content_part.done",
+          item_id: itemId,
+          output_index: payload.output_index ?? 0,
+          content_index: payload.content_index ?? 0,
+          part: { type: "output_text", text: current.text, annotations: [] }
+        }));
+        continue;
+      }
+
+      // response.output_item.done: normalize content to output_text type
+      if (
+        eventName === "response.output_item.done" &&
+        payload?.item?.type === "message" &&
+        payload?.item?.id
+      ) {
+        const itemId = payload.item.id as string;
+        const current = this.getState(itemId);
+        const item = {
+          ...payload.item,
+          content: [{ type: "output_text", text: current.text, annotations: [] }]
+        };
+        output.push(this.serializeEvent(eventName, {
+          ...payload,
+          output_index: payload.output_index ?? 0,
+          item
+        }));
+        continue;
+      }
+
+      // response.completed: patch output if empty
+      if (eventName === "response.completed" && payload?.response) {
+        const response = { ...payload.response };
+        if (!Array.isArray(response.output) || response.output.length === 0) {
+          const first = this.state.values().next()
+            .value as ResponsesStreamState | undefined;
+          if (first) {
+            response.output = [{
+              id: first.itemId,
+              type: "message",
+              status: "completed",
+              role: "assistant",
+              content: [{ type: "output_text", text: first.text, annotations: [] }]
+            }];
+            response.output_text = first.text;
+          }
+        }
+        output.push(this.serializeEvent(eventName, { ...payload, response }));
+        continue;
+      }
+
+      // All other events: pass through unchanged
+      output.push(sub);
+    }
+
+    return output.join("\n\n");
+  }
+
+  private getState(itemId: string): ResponsesStreamState {
+    const current = this.state.get(itemId);
+    if (current) return current;
+    const created = { itemId, text: "", contentPartStarted: false };
+    this.state.set(itemId, created);
+    return created;
+  }
+
+  private serializeEvent(name: string, data: any): string {
+    return `event: ${name}\ndata: ${JSON.stringify(data)}`;
+  }
+}
+
 interface PurchaseCreateResponse {
   purchaseId?: string;
   purchase_id?: string;
@@ -91,13 +301,22 @@ export class TokenbuddyDaemon {
   private controlServer?: any;
   private proxyServer?: any;
   private selectionMode: "auto" | "manual";
+  private selectedSellerId?: string;
 
   private activePurchases = new Map<string, Promise<string>>();
 
   constructor(config: DaemonConfig) {
-    this.config = config;
-    this.selectionMode = config.selectionMode || "auto";
     this.tokenStore = new BuyerStore({ dbPath: config.dbPath });
+    const routingPreference =
+      this.tokenStore.getDaemonRuntimeConfig<SellerRoutingPreference>("routing")
+        ?.config;
+    this.config = config;
+    this.selectionMode =
+      config.selectionMode ||
+      (routingPreference?.mode === "fixed" ? "manual" : routingPreference?.mode) ||
+      "auto";
+    this.selectedSellerId =
+      config.selectedSellerId || routingPreference?.sellerId;
   }
 
   private activeControlPort(): number {
@@ -111,43 +330,25 @@ export class TokenbuddyDaemon {
   }
 
   private async fetchRegistry(): Promise<SellerRegistryDocument> {
-    const response = await fetch(this.config.sellerRegistryUrl);
-    if (!response.ok) {
-      throw new Error(`registry returned ${response.status}`);
-    }
-    const data = await response.json() as SellerRegistryDocument;
-    if (!data || !Array.isArray(data.sellers)) {
-      throw new Error("registry response missing sellers");
-    }
-    return data;
-  }
-
-  private async fetchSellerManifest(seller: RegistrySeller): Promise<SellerManifest> {
-    const baseUrl = seller.url.replace(/\/+$/, "");
-    const response = await fetch(`${baseUrl}/manifest`);
-    if (!response.ok) {
-      throw new Error(`manifest returned ${response.status}`);
-    }
-    return await response.json() as SellerManifest;
+    return await fetchSellerRegistry(this.config.sellerRegistryUrl);
   }
 
   private runtimeSummary() {
+    const sellerRoutingMode = this.selectedSellerId ? "fixed" : this.selectionMode;
     return {
       status: "running",
       pid: process.pid,
       controlPort: this.activeControlPort(),
       proxyPort: this.activeProxyPort(),
       selectionMode: this.selectionMode,
+      sellerRoutingMode,
+      selectedSellerId: this.selectedSellerId,
       dbPath: this.config.dbPath,
       sellerRegistryUrl: this.config.sellerRegistryUrl,
       store: this.tokenStore.summary()
     };
   }
 
-  private normalizeSellerUrl(seller: RegistrySeller): string {
-    return seller.url.replace(/\/+$/, "");
-  }
-
   private endpointProtocol(endpoint: string): string | undefined {
     if (endpoint === "/v1/chat/completions") {
       return "chat_completions";
@@ -173,19 +374,63 @@ export class TokenbuddyDaemon {
     return undefined;
   }
 
-  private manifestProtocols(manifest: SellerManifest, seller: RegistrySeller): string[] {
-    const protocols = manifest.supportedProtocols || manifest.supported_protocols || seller.supportedProtocols || [];
-    return protocols.includes("anthropic_messages") && !protocols.includes("messages")
-      ? [...protocols, "messages"]
-      : protocols;
+  private stripLocalClaudeOneMMarker(modelId: string): string {
+    const trimmed = modelId.trimEnd();
+    return trimmed.toLowerCase().endsWith("[1m]")
+      ? trimmed.slice(0, -4).trimEnd()
+      : trimmed;
+  }
+
+  private resolveClaudeRoleModel(modelId: string): string {
+    const runtimeConfig = this.tokenStore.getProviderRuntimeConfig<ClaudeCodeModelMappingConfig>("claude-code")?.config;
+    const stripped = this.stripLocalClaudeOneMMarker(modelId);
+    if (!runtimeConfig || runtimeConfig.selectionKind !== "claude-role-mapping") {
+      return stripped;
+    }
+
+    const lowered = stripped.toLowerCase();
+    if (lowered === "haiku" || lowered.startsWith("claude-haiku")) {
+      return runtimeConfig.roles.haiku?.upstreamModel || runtimeConfig.fallbackModel || stripped;
+    }
+    if (lowered === "sonnet" || lowered.startsWith("claude-sonnet")) {
+      return runtimeConfig.roles.sonnet?.upstreamModel || runtimeConfig.fallbackModel || stripped;
+    }
+    if (lowered === "opus" || lowered.startsWith("claude-opus")) {
+      return runtimeConfig.roles.opus?.upstreamModel || runtimeConfig.fallbackModel || stripped;
+    }
+    return runtimeConfig.fallbackModel || stripped;
   }
 
-  private manifestPaymentMethods(manifest: SellerManifest, seller: RegistrySeller): string[] {
-    return manifest.paymentMethods || manifest.payment_methods || seller.paymentMethods || [];
+  private resolveRouteModelId(endpoint: string, body: unknown): { requestedModelId?: string; resolvedModelId?: string } {
+    const requestedModelId = this.extractModelId(endpoint, body);
+    if (!requestedModelId) {
+      return {};
+    }
+    if (this.endpointProtocol(endpoint) === "messages") {
+      return {
+        requestedModelId,
+        resolvedModelId: this.resolveClaudeRoleModel(requestedModelId)
+      };
+    }
+    return {
+      requestedModelId,
+      resolvedModelId: requestedModelId
+    };
   }
 
-  private manifestModelIds(manifest: SellerManifest): string[] {
-    return (manifest.models || []).map((model) => model.id).filter((id) => typeof id === "string" && id.length > 0);
+  private applyResolvedModelToBody(
+    endpoint: string,
+    body: Record<string, unknown>,
+    resolvedModelId: string
+  ): Record<string, unknown> {
+    const nextBody: Record<string, unknown> = { ...body };
+    if ("model" in nextBody) {
+      nextBody.model = resolvedModelId;
+    }
+    if (endpoint === "/v1/responses" && "model_id" in nextBody) {
+      nextBody.model_id = resolvedModelId;
+    }
+    return nextBody;
   }
 
   private defaultPaymentMethod(): string | undefined {
@@ -206,13 +451,16 @@ export class TokenbuddyDaemon {
     const registry = await this.fetchRegistry();
     const defaultSellers = registry.sellers.filter((seller) => seller.id === registry.defaultSeller);
     const backupSellers = registry.sellers.filter((seller) => seller.id !== registry.defaultSeller);
-    const sellers = this.selectionMode === "manual" ? defaultSellers : [...defaultSellers, ...backupSellers];
+    const manualSellers = this.selectedSellerId
+      ? registry.sellers.filter((seller) => seller.id === this.selectedSellerId)
+      : defaultSellers;
+    const sellers = this.selectionMode === "manual" ? manualSellers : [...defaultSellers, ...backupSellers];
 
     const routes: SellerRoute[] = [];
     for (const seller of sellers) {
       let manifest: SellerManifest;
       try {
-        manifest = await this.fetchSellerManifest(seller);
+        manifest = await fetchSellerManifest(seller);
       } catch (error: unknown) {
         logger.warn("route.manifest.failed", "seller manifest unavailable during route selection", {
           sellerKey: seller.id,
@@ -223,9 +471,9 @@ export class TokenbuddyDaemon {
         continue;
       }
 
-      const protocols = this.manifestProtocols(manifest, seller);
-      const paymentMethods = this.manifestPaymentMethods(manifest, seller);
-      const modelIds = this.manifestModelIds(manifest);
+      const protocols = manifestProtocols(manifest, seller);
+      const paymentMethods = manifestPaymentMethods(manifest, seller);
+      const modelIds = manifestModelIds(manifest);
       if (!protocols.includes(protocol) || !paymentMethods.includes(paymentMethod) || !modelIds.includes(modelId)) {
         continue;
       }
@@ -303,51 +551,10 @@ export class TokenbuddyDaemon {
     models: Array<{ id: string; sellerId: string; sellerName?: string; sellerUrl: string; supportedProtocols: string[]; paymentMethods: string[] }>;
     sellers: Array<{ id: string; name?: string; url: string; status: string; manifestSellerId?: string; errorMessage?: string }>;
   }> {
-    const registry = await this.fetchRegistry();
-    const sellerResults = await Promise.all(registry.sellers.map(async (seller) => {
-      try {
-        const manifest = await this.fetchSellerManifest(seller);
-        const protocols = this.manifestProtocols(manifest, seller);
-        const paymentMethods = this.manifestPaymentMethods(manifest, seller);
-        const models = (manifest.models || []).map((model) => ({
-          id: model.id,
-          sellerId: seller.id,
-          sellerName: seller.name,
-          sellerUrl: seller.url,
-          supportedProtocols: protocols,
-          paymentMethods
-        }));
-        return {
-          seller: {
-            id: seller.id,
-            name: seller.name,
-            url: seller.url,
-            status: "ok",
-            manifestSellerId: manifest.sellerId || manifest.seller_id || seller.id
-          },
-          models
-        };
-      } catch (error: unknown) {
-        const errorMessage = error instanceof Error ? error.message : String(error);
-        logger.warn("models.refresh.seller_failed", "seller manifest refresh failed", {
-          sellerId: seller.id,
-          errorMessage
-        });
-        return {
-          seller: {
-            id: seller.id,
-            name: seller.name,
-            url: seller.url,
-            status: "failed",
-            errorMessage
-          },
-          models: []
-        };
-      }
-    }));
+    const catalog = await discoverSellerBackedModels(this.config.sellerRegistryUrl);
     return {
-      models: sellerResults.flatMap((result) => result.models),
-      sellers: sellerResults.map((result) => result.seller)
+      models: catalog.models,
+      sellers: catalog.sellers
     };
   }
 
@@ -409,7 +616,7 @@ export class TokenbuddyDaemon {
 
   private async getOrPurchaseToken(route: SellerRoute): Promise<string> {
     const sellerKey = route.seller.id;
-    const sellerUrl = this.normalizeSellerUrl(route.seller);
+    const sellerUrl = normalizeSellerUrl(route.seller);
     const { modelId, paymentMethod } = route;
     const cached = this.tokenStore.getToken(sellerKey);
     const rebuyMinBalanceMicros = this.tokenRebuyMinBalanceMicros();
@@ -600,7 +807,7 @@ export class TokenbuddyDaemon {
     const timeoutMs = this.clawtipProofTimeoutMs();
     const payload = JSON.stringify({
       sellerKey: route.seller.id,
-      sellerUrl: this.normalizeSellerUrl(route.seller),
+      sellerUrl: normalizeSellerUrl(route.seller),
       modelId: route.modelId,
       purchase: createData,
       paymentInstructions: createData.paymentInstructions || createData.payment_instructions,
@@ -686,7 +893,7 @@ export class TokenbuddyDaemon {
   private copyUpstreamHeaders(upstreamResponse: globalThis.Response, res: Response): void {
     upstreamResponse.headers.forEach((value, key) => {
       const lowerKey = key.toLowerCase();
-      if (["connection", "keep-alive", "transfer-encoding", "upgrade"].includes(lowerKey)) {
+      if (["connection", "content-encoding", "content-length", "keep-alive", "transfer-encoding", "upgrade"].includes(lowerKey)) {
         return;
       }
       res.setHeader(key, value);
@@ -696,7 +903,8 @@ export class TokenbuddyDaemon {
   private async forwardProxyRequest(endpoint: string, req: Request, res: Response): Promise<void> {
     const startedAt = Date.now();
     const body = req.body || {};
-    const modelId = this.extractModelId(endpoint, body);
+    const { requestedModelId, resolvedModelId } = this.resolveRouteModelId(endpoint, body);
+    const modelId = resolvedModelId;
     const requestId = req.header("x-request-id") || (body && typeof body === "object" ? (body as { requestId?: string }).requestId : undefined) || `proxy_req_${crypto.randomBytes(8).toString("hex")}`;
     const idempotencyKey = req.header("idempotency-key") || `idem_${crypto.randomBytes(12).toString("hex")}`;
 
@@ -718,15 +926,16 @@ export class TokenbuddyDaemon {
             requestId,
             sellerKey,
             model: modelId,
+            requestedModel: requestedModelId,
             endpoint,
             stream: Boolean((body as { stream?: unknown }).stream)
           });
           const token = await this.getOrPurchaseToken(route);
-          const sellerUrl = this.normalizeSellerUrl(route.seller);
-          const upstreamBody = {
+          const sellerUrl = normalizeSellerUrl(route.seller);
+          const upstreamBody = this.applyResolvedModelToBody(endpoint, {
             ...(body as Record<string, unknown>),
             requestId
-          };
+          }, modelId);
 
           logger.info("proxy.upstream_fetch.started", "proxy upstream fetch started", {
             requestId,
@@ -786,13 +995,29 @@ export class TokenbuddyDaemon {
               return;
             }
             let bytes = 0;
+            const decoder = new TextDecoder();
+            const responsesStreamNormalizer = new ResponsesStreamNormalizer();
             while (true) {
               const { done, value } = await reader.read();
               if (done) {
                 break;
               }
               bytes += value.byteLength;
-              res.write(Buffer.from(value));
+              if (endpoint === "/v1/responses") {
+                const chunk = decoder.decode(value, { stream: true });
+                const normalized = responsesStreamNormalizer.push(chunk);
+                if (normalized.length > 0) {
+                  res.write(`${normalized}\n\n`);
+                }
+              } else {
+                res.write(Buffer.from(value));
+              }
+            }
+            if (endpoint === "/v1/responses") {
+              const trailing = responsesStreamNormalizer.finish();
+              if (trailing.length > 0) {
+                res.write(`${trailing}\n\n`);
+              }
             }
             res.end();
             const billedMicros = Math.max(1, bytes);
@@ -1026,7 +1251,9 @@ export class TokenbuddyDaemon {
         const changes = previewProviderInstall({
           providers: Array.isArray(req.body?.providers) ? req.body.providers : [],
           proxyUrl: String(req.body?.proxyUrl || ""),
-          model: String(req.body?.model || ""),
+          model: typeof req.body?.model === "string" ? req.body.model : undefined,
+          providerSelections: req.body?.providerSelections,
+          sellerRouting: req.body?.sellerRouting,
           home: typeof req.body?.home === "string" ? req.body.home : undefined
         });
         logger.info("provider.install.previewed", "provider install previewed", {
@@ -1053,7 +1280,9 @@ export class TokenbuddyDaemon {
         const applied = applyProviderInstall({
           providers: Array.isArray(req.body?.providers) ? req.body.providers : [],
           proxyUrl: String(req.body?.proxyUrl || ""),
-          model: String(req.body?.model || ""),
+          model: typeof req.body?.model === "string" ? req.body.model : undefined,
+          providerSelections: req.body?.providerSelections,
+          sellerRouting: req.body?.sellerRouting,
           home: typeof req.body?.home === "string" ? req.body.home : undefined
         }, this.tokenStore);
         logger.info("provider.install.applied", "provider install applied", {