@animalabs/membrane 0.5.79 → 0.5.81

package/src/utils/usage.ts

@@ -0,0 +1,184 @@
+import type { ModelPricing, ProviderResponse, UsageCacheConvention } from '../types/provider.js';
+import type { CostBreakdown, DetailedUsage, TurnRoundUsage } from '../types/response.js';
+import { addCostBreakdowns, calculateCost, warnUnpricedModel } from './cost.js';
+
+/** Adapters whose convention is undeclared warn once each, not once per call. */
+const warnedUndeclaredAdapters = new Set<string>();
+
+/**
+ * Restate a provider's usage in membrane's single convention: `inputTokens`,
+ * `cacheReadTokens` and `cacheCreationTokens` DISJOINT, each priced at its own
+ * rate. Adapters reporting `cache-inclusive` counts have the cached span
+ * subtracted out of `inputTokens`; the total prompt size stays recoverable as
+ * `inputTokens + cacheReadTokens`.
+ *
+ * Without this, one field name carried two incomparable meanings — the cache
+ * hit ratio was unbounded on cache-excluded adapters and cost double-charged
+ * the cached span on cache-inclusive ones.
+ */
+export function normalizeUsageToCacheExcluded(
+  usage: ProviderResponse['usage'],
+  adapterName: string,
+  adapterConvention: UsageCacheConvention | undefined,
+): ProviderResponse['usage'] {
+  // An adapter that declares nothing is in exactly the state `unknown` names,
+  // which is why the interface field is optional and lands here.
+  const convention = usage.cacheConvention ?? adapterConvention ?? 'unknown';
+  const cacheReadTokens = usage.cacheReadTokens ?? 0;
+
+  if (convention === 'unknown' && cacheReadTokens > 0 && !warnedUndeclaredAdapters.has(adapterName)) {
+    warnedUndeclaredAdapters.add(adapterName);
+    console.warn(
+      `[membrane:usage] adapter "${adapterName}" reported ${cacheReadTokens} cache-read tokens but`
+      + ' declares usageCacheConvention "unknown", so membrane cannot tell whether inputTokens'
+      + ` (${usage.inputTokens}) already includes them. Counts are passed through unchanged;`
+      + ' cost and cache hit ratio may be off by the cached span until the adapter declares'
+      + ' "cache-excluded" or "cache-inclusive".'
+    );
+  }
+
+  if (convention !== 'cache-inclusive' || cacheReadTokens === 0) return usage;
+
+  return { ...usage, inputTokens: Math.max(0, usage.inputTokens - cacheReadTokens) };
+}
+
+/**
+ * Share of the prompt that was served from cache. Both terms are in the
+ * cache-excluded convention, so this is a true ratio in [0, 1] regardless of
+ * which provider produced the counts.
+ */
+export function calculateCacheHitRatio(
+  usage: { inputTokens: number; cacheReadTokens?: number },
+): number {
+  const cacheReadTokens = usage.cacheReadTokens ?? 0;
+  const promptTokens = (usage.inputTokens ?? 0) + cacheReadTokens;
+  if (promptTokens === 0) return 0;
+  return cacheReadTokens / promptTokens;
+}
+
+/** Test seam: the once-per-adapter warn latch is process-wide otherwise. */
+export function resetUndeclaredConventionWarnings(): void {
+  warnedUndeclaredAdapters.clear();
+}
+
+/** Unconvertible provider item types warn once each, not once per response. */
+const warnedUnconvertibleProviderItems = new Set<string>();
+
+/**
+ * A provider content item membrane cannot normalize is preserved verbatim on a
+ * zero-width carrier rather than dropped, but the caller should know its
+ * content is invisible to normalized consumers. Once per type, not per item.
+ */
+export function warnUnconvertibleProviderItem(itemType: string): void {
+  if (warnedUnconvertibleProviderItems.has(itemType)) return;
+  warnedUnconvertibleProviderItems.add(itemType);
+  console.warn(
+    `[membrane:content] no normalized ContentBlock for provider item type "${itemType}"`
+    + ' — preserving it verbatim as a rawItem carrier so it can be replayed, but its'
+    + ' content is not visible to normalized consumers.'
+  );
+}
+
+/** Test seam: the once-per-type warn latch is process-wide otherwise. */
+export function resetUnconvertibleProviderItemWarnings(): void {
+  warnedUnconvertibleProviderItems.clear();
+}
+
+/**
+ * Accumulates one turn's usage ACROSS ROUNDS, pricing each round under the
+ * model that actually served it.
+ *
+ * A routed turn is not billed at one rate: OpenRouter re-picks a provider per
+ * call and an alias can resolve to a new snapshot mid-turn. Every streaming
+ * loop used to re-resolve pricing when the served model changed and then
+ * recompute the WHOLE accumulated usage at that latest rate, retroactively
+ * re-billing every earlier round — two 1M-token rounds at 1,000/M then
+ * 10,000/M reported 20,000 against a real 11,000.
+ *
+ * The four tool loops (callback/yielding × XML/native) each carried their own
+ * copy of the accumulate-and-price block; this is the single one they share.
+ */
+export class TurnUsageAccumulator {
+  private readonly rounds: TurnRoundUsage[] = [];
+  private readonly tokens: DetailedUsage = { inputTokens: 0, outputTokens: 0 };
+  private summedCost: CostBreakdown | undefined;
+  /** False once any round could not be priced, or two rounds disagreed on currency. */
+  private costCoversEveryRound = true;
+  private lastServed: string | undefined;
+
+  constructor(
+    private readonly requestedModel: string,
+    private readonly resolveRoundPricing: (servedModel?: string) => ModelPricing | undefined,
+  ) {}
+
+  /**
+   * Fold one provider round in, priced at its own served model, and return the
+   * turn total as it now stands (a fresh snapshot — callers hand this to
+   * `onUsage` and to stream events, and must not see it mutate underneath).
+   */
+  addRound(servedModel: string | undefined, usage: ProviderResponse['usage']): DetailedUsage {
+    if (servedModel) this.lastServed = servedModel;
+
+    this.tokens.inputTokens += usage.inputTokens;
+    this.tokens.outputTokens += usage.outputTokens;
+    if (usage.cacheCreationTokens) {
+      this.tokens.cacheCreationTokens = (this.tokens.cacheCreationTokens ?? 0) + usage.cacheCreationTokens;
+    }
+    if (usage.cacheReadTokens) {
+      this.tokens.cacheReadTokens = (this.tokens.cacheReadTokens ?? 0) + usage.cacheReadTokens;
+    }
+    // Already inside outputTokens (billed at the output rate) — summed for
+    // attribution, never added to the output total.
+    if (usage.thinkingTokens) {
+      this.tokens.thinkingTokens = (this.tokens.thinkingTokens ?? 0) + usage.thinkingTokens;
+    }
+
+    const roundModel = servedModel || this.requestedModel;
+    const roundPricing = this.resolveRoundPricing(servedModel);
+    const roundCost = roundPricing ? calculateCost(usage, roundPricing) : undefined;
+    if (!roundPricing) warnUnpricedModel(roundModel);
+
+    this.rounds.push({
+      model: roundModel,
+      usage: {
+        inputTokens: usage.inputTokens,
+        outputTokens: usage.outputTokens,
+        ...(usage.cacheCreationTokens != null ? { cacheCreationTokens: usage.cacheCreationTokens } : {}),
+        ...(usage.cacheReadTokens != null ? { cacheReadTokens: usage.cacheReadTokens } : {}),
+        ...(usage.thinkingTokens != null ? { thinkingTokens: usage.thinkingTokens } : {}),
+        ...(roundCost ? { estimatedCost: roundCost } : {}),
+      },
+    });
+
+    // A partial sum presented as the turn total is a false number, and an
+    // absent cost already means "membrane does not know". So one unpriced (or
+    // unsummable) round drops the total rather than under-reporting it; the
+    // rounds that WERE priced stay readable in the roster.
+    if (!roundCost) {
+      this.costCoversEveryRound = false;
+    } else if (this.costCoversEveryRound) {
+      this.summedCost = this.summedCost ? addCostBreakdowns(this.summedCost, roundCost) : roundCost;
+      if (!this.summedCost) this.costCoversEveryRound = false;
+    }
+
+    return this.total;
+  }
+
+  /** Turn totals: summed tokens, plus the summed per-round cost when every round had rates. */
+  get total(): DetailedUsage {
+    return {
+      ...this.tokens,
+      ...(this.costCoversEveryRound && this.summedCost ? { estimatedCost: this.summedCost } : {}),
+    };
+  }
+
+  /** The model that served the LAST round, or undefined if no round named one. */
+  get lastServedModel(): string | undefined {
+    return this.lastServed;
+  }
+
+  /** One entry per provider round, in order — the audit trail behind the summed total. */
+  get perRound(): TurnRoundUsage[] {
+    return this.rounds.map((round) => ({ model: round.model, usage: { ...round.usage } }));
+  }
+}

package/src/yielding-stream.ts

@@ -70,6 +70,18 @@ export class YieldingStreamImpl implements YieldingStream {
   // Flag indicating the stream producer is done
   private producerDone = false;
 
+  // Set when the CONSUMER leaves — `for await … break` (iterator.return) or
+  // iterator.throw. Nobody will read this stream again, so queued events are
+  // dropped rather than accumulated for a reader that is never coming back.
+  private consumerDeparted = false;
+
+  // Held so the listener can be removed at terminal: a long-lived caller
+  // signal shared across many streams otherwise accumulates one closure per
+  // stream, none of which are ever released.
+  private readonly onExternalAbort = () => {
+    this.cancel();
+  };
+
   constructor(
     private readonly options: YieldingStreamOptions,
     private readonly runInference: (stream: YieldingStreamImpl) => Promise<void>
@@ -78,9 +90,7 @@ export class YieldingStreamImpl implements YieldingStream {
 
     // Link external signal if provided
     if (options.signal) {
-      options.signal.addEventListener('abort', () => {
-        this.cancel();
-      });
+      options.signal.addEventListener('abort', this.onExternalAbort, { once: true });
     }
   }
 
@@ -179,6 +189,7 @@ export class YieldingStreamImpl implements YieldingStream {
     this.emit({ type: 'aborted', reason: 'user' });
     this.producerDone = true;
     this.state = { status: 'done' };
+    this.releaseExternalSignal();
   }
 
   // ============================================================================
@@ -215,6 +226,27 @@ export class YieldingStreamImpl implements YieldingStream {
           await this.waitForEvent();
         }
       },
+
+      // `for await … break`, `return` out of the loop body, and a throw in
+      // the loop body all call return(). Without it the consumer walks away
+      // and the producer keeps streaming, resuming and billing. Cancelling
+      // here is the same act the consumer would perform by hand.
+      return: async (): Promise<IteratorResult<StreamEvent>> => {
+        this.departConsumer();
+        return { value: undefined as unknown as StreamEvent, done: true };
+      },
+
+      // Reached through `yield*` delegation, or by a caller injecting an
+      // error by hand. Same departure as return() — nobody is reading this
+      // stream again — and then the error continues on its way, exactly as an
+      // async generator with no handler of its own would rethrow it. Reporting
+      // `done: true` here instead swallowed it: a `yield*` delegating to this
+      // stream resumed after the delegation as if nothing had been thrown in,
+      // and a direct `.throw(e)` resolved rather than rejecting.
+      throw: async (error?: unknown): Promise<IteratorResult<StreamEvent>> => {
+        this.departConsumer();
+        throw error;
+      },
     };
   }
 
@@ -226,6 +258,9 @@ export class YieldingStreamImpl implements YieldingStream {
    * Push an event to be yielded to the consumer.
    */
   emit(event: StreamEvent): void {
+    // The consumer has left for good; queueing here would grow without bound
+    // for a reader that cannot observe it.
+    if (this.consumerDeparted) return;
     this.eventQueue.push(event);
     this.notifyEventWaiter();
   }
@@ -255,6 +290,7 @@ export class YieldingStreamImpl implements YieldingStream {
    */
   markDone(): void {
     this.producerDone = true;
+    this.releaseExternalSignal();
     this.notifyEventWaiter();
   }
 
@@ -269,6 +305,21 @@ export class YieldingStreamImpl implements YieldingStream {
   // Private Helpers
   // ============================================================================
 
+  /**
+   * The consumer has stopped iterating for good: drop everything queued for
+   * it, refuse further queueing, and cancel the producer.
+   */
+  private departConsumer(): void {
+    this.consumerDeparted = true;
+    this.eventQueue.length = 0;
+    this.cancel();
+    this.eventQueue.length = 0;
+  }
+
+  private releaseExternalSignal(): void {
+    this.options.signal?.removeEventListener('abort', this.onExternalAbort);
+  }
+
   private startInference(): void {
     if (this.state.status !== 'idle') {
       return; // Already started