126 lines
2.9 KiB
TypeScript
126 lines
2.9 KiB
TypeScript
import type { AgentEvent } from '../shared/contracts'
|
|
|
|
type BufferedAgentEvent = Extract<
|
|
AgentEvent,
|
|
{ type: 'text' | 'reasoning' }
|
|
>
|
|
|
|
export type AgentEventBufferOptions = {
|
|
onEvent(event: AgentEvent): void
|
|
onError?(error: unknown): void
|
|
flushIntervalMs?: number
|
|
maximumBufferedBytes?: number
|
|
}
|
|
|
|
const DEFAULT_FLUSH_INTERVAL_MS = 32
|
|
const DEFAULT_MAXIMUM_BUFFERED_BYTES = 64 * 1024
|
|
|
|
export class AgentEventBuffer {
|
|
private readonly onEvent: (event: AgentEvent) => void
|
|
private readonly onError: ((error: unknown) => void) | undefined
|
|
private readonly flushIntervalMs: number
|
|
private readonly maximumBufferedBytes: number
|
|
private pending: BufferedAgentEvent | undefined
|
|
private pendingBytes = 0
|
|
private timer: ReturnType<typeof setTimeout> | undefined
|
|
private closed = false
|
|
|
|
constructor(options: AgentEventBufferOptions) {
|
|
this.onEvent = options.onEvent
|
|
this.onError = options.onError
|
|
this.flushIntervalMs =
|
|
options.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS
|
|
this.maximumBufferedBytes =
|
|
options.maximumBufferedBytes ??
|
|
DEFAULT_MAXIMUM_BUFFERED_BYTES
|
|
if (this.flushIntervalMs <= 0) {
|
|
throw new Error('flushIntervalMs must be positive')
|
|
}
|
|
if (this.maximumBufferedBytes <= 0) {
|
|
throw new Error('maximumBufferedBytes must be positive')
|
|
}
|
|
}
|
|
|
|
push(event: AgentEvent): void {
|
|
if (this.closed) {
|
|
return
|
|
}
|
|
if (event.type !== 'text' && event.type !== 'reasoning') {
|
|
this.flush()
|
|
this.onEvent(event)
|
|
return
|
|
}
|
|
|
|
const eventBytes = Buffer.byteLength(event.delta)
|
|
const matchesPending =
|
|
this.pending?.requestId === event.requestId &&
|
|
this.pending.type === event.type
|
|
if (!matchesPending) {
|
|
this.flush()
|
|
} else if (
|
|
this.pendingBytes + eventBytes >
|
|
this.maximumBufferedBytes
|
|
) {
|
|
this.flush()
|
|
}
|
|
|
|
if (eventBytes > this.maximumBufferedBytes) {
|
|
this.onEvent(event)
|
|
return
|
|
}
|
|
if (this.pending) {
|
|
this.pending = {
|
|
...this.pending,
|
|
delta: this.pending.delta + event.delta
|
|
}
|
|
this.pendingBytes += eventBytes
|
|
return
|
|
}
|
|
|
|
this.pending = { ...event }
|
|
this.pendingBytes = eventBytes
|
|
this.scheduleFlush()
|
|
}
|
|
|
|
flush(): void {
|
|
this.clearTimer()
|
|
const event = this.pending
|
|
this.pending = undefined
|
|
this.pendingBytes = 0
|
|
if (event) {
|
|
this.onEvent(event)
|
|
}
|
|
}
|
|
|
|
close(): void {
|
|
if (this.closed) {
|
|
return
|
|
}
|
|
this.closed = true
|
|
this.flush()
|
|
}
|
|
|
|
private scheduleFlush(): void {
|
|
if (this.timer) {
|
|
return
|
|
}
|
|
this.timer = setTimeout(() => {
|
|
this.timer = undefined
|
|
try {
|
|
this.flush()
|
|
} catch (error) {
|
|
this.onError?.(error)
|
|
}
|
|
}, this.flushIntervalMs)
|
|
this.timer.unref?.()
|
|
}
|
|
|
|
private clearTimer(): void {
|
|
if (!this.timer) {
|
|
return
|
|
}
|
|
clearTimeout(this.timer)
|
|
this.timer = undefined
|
|
}
|
|
}
|