- Replaced O(n) `Array.shift()` calls with O(1) head-index tracking in `RawSseDebugBuffer`. - Implemented lazy compaction to reclaim the dead-prefix memory only when the leading window grows sufficiently large. - Added comprehensive tests to verify ordering, dropped count accuracy, and buffer limits under heavy load.
295 lines
10 KiB
TypeScript
295 lines
10 KiB
TypeScript
import type { Model, ProviderResponseMetadata, RawSseEvent } from "@oh-my-pi/pi-ai";
|
||
|
||
const MAX_RAW_SSE_EVENTS = 1_000;
|
||
const MAX_RAW_SSE_CHARS = 512_000;
|
||
const MAX_RAW_SSE_EVENT_CHARS = 64_000;
|
||
|
||
export type RawSseDebugRecord =
|
||
| {
|
||
kind: "response";
|
||
sequence: number;
|
||
timestamp: number;
|
||
provider?: string;
|
||
model?: string;
|
||
api?: string;
|
||
status: number;
|
||
requestId?: string | null;
|
||
transport?: string;
|
||
}
|
||
| {
|
||
kind: "event";
|
||
sequence: number;
|
||
timestamp: number;
|
||
provider?: string;
|
||
model?: string;
|
||
api?: string;
|
||
event: string | null;
|
||
raw: string[];
|
||
truncated: boolean;
|
||
originalChars: number;
|
||
};
|
||
|
||
export interface RawSseDebugSnapshot {
|
||
records: readonly RawSseDebugRecord[];
|
||
droppedRecords: number;
|
||
droppedChars: number;
|
||
totalEvents: number;
|
||
lastUpdatedAt?: number;
|
||
}
|
||
|
||
// Per-record char counts are stored in a parallel array (`#recordChars`) on
|
||
// the buffer rather than stamped onto each record via a symbol property.
|
||
// Stamping triggered hidden-class transitions in V8/JSC — the previous
|
||
// revision saw `trimRawLines` regress 4× (0.5s → 2.0s in a 50s profile)
|
||
// because every event-record allocation went through the slow dictionary
|
||
// path. The parallel array keeps records as plain monomorphic objects.
|
||
type TrimResult = { raw: string[]; truncated: boolean; originalChars: number; chars: number };
|
||
|
||
// Single-pass trim. Returns the final `chars` count using the historical
|
||
// formula `reduce(line.length + 1, init = 1)` so the new accounting matches
|
||
// the previous `countRecordChars` byte-for-byte (the trailing +1 covers the
|
||
// record-level newline that `rawRecordText` appends in `toRawText`).
|
||
//
|
||
// When the event fits within budget the input `raw` array is returned
|
||
// **by reference** — see the ownership contract documented at
|
||
// `RawSseDebugBuffer.recordEvent` below.
|
||
function trimRawLines(raw: string[]): TrimResult {
|
||
let originalChars = 0;
|
||
for (let i = 0; i < raw.length; i++) originalChars += raw[i].length + 1;
|
||
|
||
if (originalChars <= MAX_RAW_SSE_EVENT_CHARS) {
|
||
return { raw, truncated: false, originalChars, chars: originalChars + 1 };
|
||
}
|
||
|
||
const trimmed: string[] = [];
|
||
let remaining = MAX_RAW_SSE_EVENT_CHARS;
|
||
let chars = 1; // matches reduce(.., init = 1)
|
||
for (const line of raw) {
|
||
if (remaining <= 0) break;
|
||
if (line.length + 1 <= remaining) {
|
||
trimmed.push(line);
|
||
chars += line.length + 1;
|
||
remaining -= line.length + 1;
|
||
continue;
|
||
}
|
||
const slice = line.slice(0, Math.max(0, remaining));
|
||
trimmed.push(slice);
|
||
chars += slice.length + 1;
|
||
remaining = 0;
|
||
}
|
||
const tail = `: omp-debug-truncated originalChars=${originalChars}`;
|
||
trimmed.push(tail);
|
||
chars += tail.length + 1;
|
||
return { raw: trimmed, truncated: true, originalChars, chars };
|
||
}
|
||
|
||
export function formatRawSseIsoTime(timestamp: number): string {
|
||
return new Date(timestamp).toISOString();
|
||
}
|
||
|
||
export function formatRawSseResponseComment(record: Extract<RawSseDebugRecord, { kind: "response" }>): string {
|
||
const fields = [
|
||
"omp-response",
|
||
`ts=${formatRawSseIsoTime(record.timestamp)}`,
|
||
`status=${record.status}`,
|
||
record.provider ? `provider=${record.provider}` : undefined,
|
||
record.model ? `model=${record.model}` : undefined,
|
||
record.api ? `api=${record.api}` : undefined,
|
||
record.requestId ? `requestId=${record.requestId}` : undefined,
|
||
record.transport ? `transport=${record.transport}` : undefined,
|
||
].filter((field): field is string => field !== undefined);
|
||
return `: ${fields.join(" ")}`;
|
||
}
|
||
|
||
export function rawSseRecordLines(record: RawSseDebugRecord): string[] {
|
||
if (record.kind === "response") return [formatRawSseResponseComment(record)];
|
||
return record.raw;
|
||
}
|
||
|
||
function rawRecordText(record: RawSseDebugRecord): string {
|
||
return `${rawSseRecordLines(record).join("\n")}\n`;
|
||
}
|
||
|
||
function metadataTransport(response: ProviderResponseMetadata): string | undefined {
|
||
const value = response.metadata?.lastTransport;
|
||
return typeof value === "string" ? value : undefined;
|
||
}
|
||
|
||
export class RawSseDebugBuffer {
|
||
#records: RawSseDebugRecord[] = [];
|
||
// Parallel to `#records`: `#recordChars[i]` is the precomputed char count
|
||
// for `#records[i]`. Kept in lockstep by `#append` (push both) and
|
||
// `#enforceLimits` (advance `#head` to evict, then `slice` both together
|
||
// when compacting). See the comment above the class for why this is a
|
||
// sidecar array instead of a per-record property.
|
||
#recordChars: number[] = [];
|
||
// Head-index ring over `#records`/`#recordChars`: index of the oldest live
|
||
// record. Eviction advances `#head` (amortized O(1)) rather than an O(n)
|
||
// front `shift()`; the dead `[0, #head)` prefix is reclaimed lazily by
|
||
// `#enforceLimits`. Live count is `#records.length - #head`; the live
|
||
// records are `#records[#head ..]`.
|
||
#head = 0;
|
||
#totalChars = 0;
|
||
#droppedRecords = 0;
|
||
#droppedChars = 0;
|
||
#totalEvents = 0;
|
||
#lastUpdatedAt: number | undefined;
|
||
#nextSequence = 1;
|
||
#listeners = new Set<() => void>();
|
||
#emitScheduled = false;
|
||
|
||
subscribe(listener: () => void): () => void {
|
||
this.#listeners.add(listener);
|
||
return () => this.#listeners.delete(listener);
|
||
}
|
||
|
||
recordResponse(response: ProviderResponseMetadata, model?: Model): void {
|
||
const record: RawSseDebugRecord = {
|
||
kind: "response",
|
||
sequence: this.#nextSequence++,
|
||
timestamp: Date.now(),
|
||
provider: model?.provider,
|
||
model: model?.id,
|
||
api: model?.api,
|
||
status: response.status,
|
||
requestId: response.requestId,
|
||
transport: metadataTransport(response),
|
||
};
|
||
this.#append(record, formatRawSseResponseComment(record).length + 1);
|
||
}
|
||
|
||
// Ownership contract for `event.raw`:
|
||
// The caller (`notifyRawSseEvent` in `packages/ai/src/utils/sse-debug.ts`)
|
||
// hands us a freshly-allocated `string[]` per event and never retains,
|
||
// mutates, or re-dispatches it.
|
||
// That lets `trimRawLines` keep the array by reference instead of
|
||
// cloning on every chunk — a measurable savings on the streaming hot
|
||
// path. If a future observer-chain mutates the array, restore the
|
||
// `raw.slice()` defensive copy inside `trimRawLines`.
|
||
recordEvent(event: RawSseEvent, model?: Model): void {
|
||
const trimmed = trimRawLines(event.raw);
|
||
this.#totalEvents += 1;
|
||
this.#append(
|
||
{
|
||
kind: "event",
|
||
sequence: this.#nextSequence++,
|
||
timestamp: Date.now(),
|
||
provider: model?.provider,
|
||
model: model?.id,
|
||
api: model?.api,
|
||
event: event.event,
|
||
raw: trimmed.raw,
|
||
truncated: trimmed.truncated,
|
||
originalChars: trimmed.originalChars,
|
||
},
|
||
trimmed.chars,
|
||
);
|
||
}
|
||
|
||
snapshot(): RawSseDebugSnapshot {
|
||
return {
|
||
records: this.#records.slice(this.#head),
|
||
droppedRecords: this.#droppedRecords,
|
||
droppedChars: this.#droppedChars,
|
||
totalEvents: this.#totalEvents,
|
||
lastUpdatedAt: this.#lastUpdatedAt,
|
||
};
|
||
}
|
||
|
||
toRawText(): string {
|
||
// Reads the live window directly: `rawRecordText` only computes a string
|
||
// from each record, so no caller-visible mutation is possible. With a
|
||
// non-empty dead prefix we map a slice past `#head`; `#head === 0` (the
|
||
// common case) maps `#records` in place with no extra copy.
|
||
const live = this.#head === 0 ? this.#records : this.#records.slice(this.#head);
|
||
const body = live.map(rawRecordText).join("\n");
|
||
if (this.#droppedRecords === 0) return body;
|
||
const dropped = `: omp-debug-dropped records=${this.#droppedRecords} chars=${this.#droppedChars}\n\n`;
|
||
return body.length > 0 ? `${dropped}${body}` : dropped;
|
||
}
|
||
|
||
#append(record: RawSseDebugRecord, chars: number): void {
|
||
this.#records.push(record);
|
||
this.#recordChars.push(chars);
|
||
this.#totalChars += chars;
|
||
this.#lastUpdatedAt = record.timestamp;
|
||
this.#enforceLimits();
|
||
this.#emit();
|
||
}
|
||
|
||
#enforceLimits(): void {
|
||
while (this.#records.length - this.#head > MAX_RAW_SSE_EVENTS || this.#totalChars > MAX_RAW_SSE_CHARS) {
|
||
if (this.#records.length - this.#head === 0) break;
|
||
const chars = this.#recordChars[this.#head] ?? 0;
|
||
this.#head += 1;
|
||
this.#totalChars = Math.max(0, this.#totalChars - chars);
|
||
this.#droppedRecords += 1;
|
||
this.#droppedChars += chars;
|
||
}
|
||
// Reclaim the consumed `[0, #head)` prefix once it grows large: one O(n)
|
||
// memmove amortized over many O(1) evictions, bounding the backing arrays
|
||
// to ~2x the live window. `#head >= MAX_RAW_SSE_EVENTS` covers the
|
||
// full-record-count steady state; `#head > liveCount` covers a small live
|
||
// window held by a few large records under the char budget.
|
||
const liveCount = this.#records.length - this.#head;
|
||
if (this.#head >= MAX_RAW_SSE_EVENTS || this.#head > liveCount) {
|
||
this.#records = this.#records.slice(this.#head);
|
||
this.#recordChars = this.#recordChars.slice(this.#head);
|
||
this.#head = 0;
|
||
}
|
||
}
|
||
|
||
#emit(): void {
|
||
const count = this.#listeners.size;
|
||
if (count === 0) return;
|
||
// With a single listener (the common case — RawSse debug viewer is the
|
||
// only subscriber), keep eager emit so per-event semantics are
|
||
// preserved. With multiple listeners, coalesce bursts of events into
|
||
// one microtask-deferred fan-out to avoid N×M listener invocations
|
||
// during a streaming response.
|
||
if (count === 1) {
|
||
this.#fanOut();
|
||
return;
|
||
}
|
||
if (this.#emitScheduled) return;
|
||
this.#emitScheduled = true;
|
||
queueMicrotask(() => {
|
||
this.#emitScheduled = false;
|
||
this.#fanOut();
|
||
});
|
||
}
|
||
|
||
#fanOut(): void {
|
||
for (const listener of this.#listeners) {
|
||
try {
|
||
listener();
|
||
} catch {
|
||
// Debug viewers must not be able to break stream capture.
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
const globalFallbackBuffer = new RawSseDebugBuffer();
|
||
const kRawSseDebugBuffer = Symbol("debug.rawSseBuffer");
|
||
type OwnerWithBuffer = object & { rawSseDebugBuffer?: unknown; [kRawSseDebugBuffer]?: RawSseDebugBuffer };
|
||
|
||
export function resolveRawSseDebugBuffer(owner?: object): RawSseDebugBuffer {
|
||
if (!owner) return globalFallbackBuffer;
|
||
|
||
const tagged = owner as OwnerWithBuffer;
|
||
const declared = tagged.rawSseDebugBuffer;
|
||
if (declared instanceof RawSseDebugBuffer) return declared;
|
||
|
||
const existing = tagged[kRawSseDebugBuffer];
|
||
if (existing) return existing;
|
||
|
||
const buffer = new RawSseDebugBuffer();
|
||
try {
|
||
tagged[kRawSseDebugBuffer] = buffer;
|
||
} catch {
|
||
// Non-extensible owner: caller gets a fresh buffer on each call.
|
||
}
|
||
return buffer;
|
||
}
|