memind/sse-billing.mjs

import { Transform } from 'node:stream';

function parseSseChunk(chunk) {
  const events = [];
  const blocks = chunk.split('\n\n');
  for (const block of blocks) {
    if (!block.trim()) continue;
    let data;
    for (const line of block.split('\n')) {
      if (line.startsWith('data:')) {
        data = line.slice(5).trim();
      }
    }
    if (!data) continue;
    try {
      events.push(JSON.parse(data));
    } catch {
      // ignore malformed frames
    }
  }
  return events;
}

function finishBillingKey(event) {
  const requestId = event.request_id ?? event.chat_request_id;
  if (requestId) return String(requestId);
  const state = event.token_state ?? {};
  const input =
    state.accumulatedInputTokens ?? state.accumulated_input_tokens ?? state.inputTokens ?? 0;
  const output =
    state.accumulatedOutputTokens ?? state.accumulated_output_tokens ?? state.outputTokens ?? 0;
  return `${input}:${output}`;
}

export function createSseBillingTransform({ onFinish }) {
  let buffer = '';
  const billedFinishKeys = new Set();

  const billFinishOnce = (event) => {
    if (event?.type !== 'Finish' || !event.token_state) return;
    const key = finishBillingKey(event);
    if (billedFinishKeys.has(key)) return;
    billedFinishKeys.add(key);
    void onFinish(event).catch(() => {});
  };

  return new Transform({
    transform(chunk, _encoding, callback) {
      buffer += chunk.toString('utf8');
      const parts = buffer.split('\n\n');
      buffer = parts.pop() ?? '';

      for (const block of parts) {
        if (!block.trim()) continue;
        let dataLine;
        for (const line of block.split('\n')) {
          if (line.startsWith('data:')) dataLine = line.slice(5).trim();
        }
        if (!dataLine) continue;
        try {
          billFinishOnce(JSON.parse(dataLine));
        } catch {
          // ignore
        }
      }

      callback(null, chunk);
    },
    flush(callback) {
      if (buffer.trim()) {
        for (const event of parseSseChunk(`${buffer}\n\n`)) {
          billFinishOnce(event);
        }
      }
      callback();
    },
  });
}

export function appendBalanceEvent(payload) {
  const data =
    typeof payload === 'number'
      ? { balanceCents: payload }
      : payload && typeof payload === 'object'
        ? payload
        : { balanceCents: 0 };
  return `event: balance\ndata: ${JSON.stringify(data)}\n\n`;
}