Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
7.8 kB · 244 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245import type { Duplex } from "node:stream";
const FRAME_PREFIX_BYTES = 8;
export interface PrivateFrameLimits { maxHeaderBytes: number; maxPayloadBytes: number;}
export const DEFAULT_PRIVATE_FRAME_LIMITS: PrivateFrameLimits = { maxHeaderBytes: 1024 * 1024, maxPayloadBytes: 1024 * 1024 * 1024,};
export interface PrivateFrame<THeader extends object> { header: THeader; payload: Buffer;}
export type PrivateFrameHeaderValidator<THeader extends object> = (value: unknown) => value is THeader;
function assertFrameLength(name: string, value: number, maximum: number): void { if (!Number.isSafeInteger(value) || value < 0 || value > maximum) { throw new Error(`Invalid private frame ${name}: ${value}`); }}
function isObjectHeader(value: unknown): value is object { return typeof value === "object" && value !== null && !Array.isArray(value);}
export function encodePrivateFrame<THeader extends object>( header: THeader, payload: Uint8Array = Buffer.alloc(0), limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS,): Buffer { const headerBuffer = Buffer.from(JSON.stringify(header), "utf8"); const payloadBuffer = Buffer.isBuffer(payload) ? payload : Buffer.from(payload); assertFrameLength("header length", headerBuffer.length, limits.maxHeaderBytes); assertFrameLength("payload length", payloadBuffer.length, limits.maxPayloadBytes); if (headerBuffer.length === 0) { throw new Error("Private frame header cannot be empty"); }
const frame = Buffer.allocUnsafe(FRAME_PREFIX_BYTES + headerBuffer.length + payloadBuffer.length); frame.writeUInt32BE(headerBuffer.length, 0); frame.writeUInt32BE(payloadBuffer.length, 4); headerBuffer.copy(frame, FRAME_PREFIX_BYTES); payloadBuffer.copy(frame, FRAME_PREFIX_BYTES + headerBuffer.length); return frame;}
export class PrivateFrameDecoder<THeader extends object> { // Incoming chunks are kept as segments and only joined when a whole frame has // arrived. Concatenating into one accumulator on every chunk is O(n^2) in the // frame size, which a multi-megabyte response split across socket reads feels // as seconds: the session tree of a long transcript, for example. private segments: Buffer<ArrayBufferLike>[] = []; private pendingBytes = 0;
constructor( private readonly validateHeader: PrivateFrameHeaderValidator<THeader>, private readonly limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS, ) {}
get bufferedBytes(): number { return this.pendingBytes; }
push(chunk: Uint8Array): PrivateFrame<THeader>[] { if (chunk.length > 0) { const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); this.segments.push(buffer); this.pendingBytes += buffer.length; }
const frames: PrivateFrame<THeader>[] = []; while (this.pendingBytes >= FRAME_PREFIX_BYTES) { const prefix = this.peek(FRAME_PREFIX_BYTES); const headerLength = prefix.readUInt32BE(0); const payloadLength = prefix.readUInt32BE(4); assertFrameLength("header length", headerLength, this.limits.maxHeaderBytes); assertFrameLength("payload length", payloadLength, this.limits.maxPayloadBytes); if (headerLength === 0) { throw new Error("Private frame header cannot be empty"); }
const frameLength = FRAME_PREFIX_BYTES + headerLength + payloadLength; if (this.pendingBytes < frameLength) { break; }
// Join the frame's segments exactly once (a whole frame that arrived in // one chunk is taken as-is). const frame = this.take(frameLength); const headerStart = FRAME_PREFIX_BYTES; const payloadStart = headerStart + headerLength; let decoded: unknown; try { decoded = JSON.parse(frame.toString("utf8", headerStart, payloadStart)); } catch (error) { throw new Error( `Invalid private frame header JSON: ${error instanceof Error ? error.message : String(error)}`, ); } if (!isObjectHeader(decoded) || !this.validateHeader(decoded)) { throw new Error("Invalid private frame routing header"); }
frames.push({ header: decoded, payload: Buffer.from(frame.subarray(payloadStart, payloadStart + payloadLength)), }); } return frames; }
/** Copy of the next `count` buffered bytes, without consuming them. */ private peek(count: number): Buffer { const first = this.segments[0]; if (first !== undefined && first.length >= count) { return first.subarray(0, count); } const head: Buffer[] = []; let remaining = count; for (const segment of this.segments) { const take = Math.min(segment.length, remaining); head.push(take === segment.length ? segment : segment.subarray(0, take)); remaining -= take; if (remaining === 0) break; } return Buffer.concat(head, count); }
/** Remove and return the next `count` buffered bytes as one contiguous buffer. */ private take(count: number): Buffer { const head: Buffer[] = []; let remaining = count; let consumed = 0; while (remaining > 0) { const segment = this.segments[consumed]!; if (segment.length <= remaining) { head.push(segment); remaining -= segment.length; consumed++; } else { head.push(segment.subarray(0, remaining)); // Keep the unconsumed tail as a copy so the caller's chunk is released. this.segments[consumed] = Buffer.from(segment.subarray(remaining)); remaining = 0; } } this.segments = this.segments.slice(consumed); this.pendingBytes -= count; return head.length === 1 ? head[0]! : Buffer.concat(head, count); }
finish(): void { if (this.pendingBytes !== 0) { throw new Error(`Private frame channel ended with ${this.pendingBytes} incomplete bytes`); } }}
export type PrivateFrameListener<THeader extends object> = (frame: PrivateFrame<THeader>) => void;
export class PrivateFramedChannel<THeader extends object> { private readonly decoder: PrivateFrameDecoder<THeader>; private readonly listeners = new Set<PrivateFrameListener<THeader>>(); private closed = false;
constructor( private readonly stream: Duplex, validateHeader: PrivateFrameHeaderValidator<THeader>, private readonly limits: PrivateFrameLimits = DEFAULT_PRIVATE_FRAME_LIMITS, ) { this.decoder = new PrivateFrameDecoder(validateHeader, limits); stream.on("data", this.handleData); stream.on("end", this.handleEnd); stream.on("close", this.handleClose); }
onFrame(listener: PrivateFrameListener<THeader>): () => void { this.listeners.add(listener); return () => this.listeners.delete(listener); }
async send(header: THeader, payload?: Uint8Array): Promise<void> { if (this.closed || this.stream.destroyed) { throw new Error("Private frame channel is closed"); } const frame = encodePrivateFrame(header, payload, this.limits); await new Promise<void>((resolve, reject) => { this.stream.write(frame, (error?: Error | null) => { if (error) { reject(error); } else { resolve(); } }); }); }
close(): void { if (this.closed) { return; } this.closed = true; this.detach(); this.stream.end(); }
private readonly handleData = (chunk: Buffer): void => { try { for (const frame of this.decoder.push(chunk)) { for (const listener of this.listeners) { listener(frame); } } } catch (error) { this.stream.destroy(error instanceof Error ? error : new Error(String(error))); } };
private readonly handleEnd = (): void => { try { this.decoder.finish(); } catch (error) { this.stream.destroy(error instanceof Error ? error : new Error(String(error))); } };
private readonly handleClose = (): void => { this.closed = true; this.detach(); };
private detach(): void { this.stream.off("data", this.handleData); this.stream.off("end", this.handleEnd); this.stream.off("close", this.handleClose); this.listeners.clear(); }}