deepseek-harness/packages/code-runtime/code-runtime-node/src/channel.ts

/** Length-prefixed JSON transport with bounded input and queued writes. */
import { Buffer } from 'node:buffer'
import type { Duplex } from 'node:stream'
import type { CodeJsonValue } from '@deepseek-ai/dsh-code-runtime'
import { jsonValueBytesUpTo } from './output-json.ts'

const stringify = JSON.stringify
const parse = JSON.parse

/** One co-shipped process channel; the consumer owns frame validation and terminal outcomes. */
export class JsonChannel {
  private readonly stream: Duplex
  private readonly maxBytes: number
  private readonly receive: (message: unknown, bytes: number) => void
  private readonly failure: (error: Error, kind: 'io' | 'protocol') => void
  private readonly header = Buffer.alloc(4)
  private headerBytes = 0
  private payload: Buffer | undefined
  private payloadBytes = 0
  private queuedBytes = 0
  private closed = false
  private tail: Promise<void> = Promise.resolve()

  constructor(
    stream: Duplex,
    maxBytes: number,
    receive: (message: unknown, bytes: number) => void,
    failure: (error: Error, kind: 'io' | 'protocol') => void,
  ) {
    this.stream = stream
    this.maxBytes = maxBytes
    this.receive = receive
    this.failure = failure
    stream.on('data', this.onData)
    stream.on('error', this.onError)
    stream.on('end', this.onEnd)
  }

  private readonly onError = (error: Error): void => { if (!this.closed) this.failure(error, 'io') }
  private readonly onEnd = (): void => { if (!this.closed) this.failure(new Error('control channel ended before the program settled'), 'io') }
  private readonly onData = (chunk: Buffer): void => {
    if (this.closed) return
    try {
      let offset = 0
      while (offset < chunk.length && !this.closed) {
        if (this.payload === undefined) {
          const bytes = Math.min(4 - this.headerBytes, chunk.length - offset)
          chunk.copy(this.header, this.headerBytes, offset, offset + bytes)
          this.headerBytes += bytes
          offset += bytes
          if (this.headerBytes !== 4) continue
          const length = this.header.readUInt32BE(0)
          if (length === 0 || length > this.maxBytes) throw new Error(`control frame exceeds ${this.maxBytes} bytes or is empty`)
          this.payload = Buffer.allocUnsafe(length)
          this.payloadBytes = 0
          this.headerBytes = 0
        }
        const bytes = Math.min(this.payload.length - this.payloadBytes, chunk.length - offset)
        chunk.copy(this.payload, this.payloadBytes, offset, offset + bytes)
        this.payloadBytes += bytes
        offset += bytes
        if (this.payloadBytes !== this.payload.length) continue
        const frame = this.payload
        this.payload = undefined
        this.payloadBytes = 0
        this.receive(parse(new TextDecoder('utf-8', { fatal: true }).decode(frame)), frame.length)
      }
    } catch (error: unknown) {
      this.failure(error instanceof Error ? error : new Error(String(error)), 'protocol')
    }
  }

  /**
   * Queue a bounded frame and honor stream backpressure.
   * @param message - JSON-only co-shipped protocol value.
   * @returns Resolves when this frame has been written, or rejects after transport failure.
   */
  send(message: unknown): Promise<void> {
    if (this.closed) return Promise.reject(new Error('control channel is closed'))
    const size = jsonValueBytesUpTo(message as CodeJsonValue, this.maxBytes)
    if (size === undefined || this.queuedBytes + size > this.maxBytes) {
      return Promise.reject(new Error(`control output exceeds ${this.maxBytes} queued bytes`))
    }
    const body = Buffer.from(stringify(message), 'utf8')
    const header = Buffer.alloc(4)
    header.writeUInt32BE(body.length)
    this.queuedBytes += body.length
    const write = async (): Promise<void> => {
      if (this.closed) throw new Error('control channel is closed')
      await new Promise<void>((resolve, reject) => {
        const stop = (error?: Error): void => {
          this.stream.off('error', onError)
          this.stream.off('close', onClose)
          if (error) reject(error)
          else resolve()
        }
        const onError = (error: Error): void => { stop(error) }
        const onClose = (): void => { stop(new Error('control channel closed during a write')) }
        this.stream.once('error', onError)
        this.stream.once('close', onClose)
        this.stream.cork()
        this.stream.write(header)
        this.stream.write(body, (error) => { stop(error ?? undefined) })
        this.stream.uncork()
      })
    }
    const result = this.tail.then(write).finally(() => { this.queuedBytes -= body.length })
    this.tail = result.catch(() => {})
    return result
  }

  /** Stop reads and close the owned endpoint; pending writes reject on closure. */
  close(): void {
    if (this.closed) return
    this.closed = true
    this.payload = undefined
    this.stream.off('data', this.onData)
    this.stream.off('end', this.onEnd)
    this.stream.destroy()
  }

  /** Wait for accepted writes to finish or fail. */
  async drain(): Promise<void> { await this.tail }
}