deepseek-harness/packages/ptc-runtime/ptc-runtime-node/src/persistent.ts

/** Node PTC in-memory execution with retained processes scoped to caller-owned objects. */
import { AsyncLocalStorage } from 'node:async_hooks'
import { stripTypeScriptTypes } from 'node:module'
import { isAbsolute } from 'node:path'
import { Context } from '@deepseek-ai/cordis'
import type { PtcRunRequest, PtcRunSpec, PtcRunResult, PtcRunSandbox, PtcRunFailure, PtcBindingNamespace, PtcJsonValue } from '@deepseek-ai/dsh-ptc-runtime'
import { clampTimeout } from '@deepseek-ai/dsh-timeout'
import { SandboxUnavailableError, classifyRunnerFailure, isRunnerSpawnFailure } from '@deepseek-ai/dsh-sandbox'
import type { ConfinedArgv, SandboxExecutionPolicy } from '@deepseek-ai/dsh-sandbox'
import type { SubprocessHandle, SubprocessOutcome } from '@deepseek-ai/dsh-subprocess'
import type { BirpcReturn } from 'birpc'
import type { PolicyProcesses } from './policy-processes.ts'
import type { ResolvedConfig } from './config.ts'
import { invokeBinding, validateBindings } from './bindings.ts'
import { createProcessRpc } from './rpc.ts'
import { bootstrapArgs } from './launch.ts'
import { STARTUP_ENVIRONMENT_NAMES } from './environment.ts'
import { decodePtcJsonWire } from './json-wire.ts'
import { OutputLedger } from './output-ledger.ts'
import { drainOutput, terminateProcess } from './output-stream.ts'
import { readProgramResult, type ProgramBootData, type ProgramMethods, type HostMethods } from './protocol.ts'

const messageOf = (error: unknown): string => error instanceof Error ? error.message : String(error)

function persistentOutput(maxBytes: number): OutputLedger {
  return new OutputLedger(maxBytes, `retained variables were cleared; outer output exceeded ${maxBytes} bytes`)
}

/**
 * Compare requested extra grants without resolving paths in another execution world.
 * @param left - one policy's additional roots.
 * @param right - the other policy's additional roots.
 * @returns whether both contain the same paths, ignoring order and duplicates.
 */
function sameAdditionalRoots(left: readonly string[] | undefined, right: readonly string[] | undefined): boolean {
  const effective = (roots: readonly string[] | undefined): string[] => [...new Set(roots ?? [])].sort()
  const a = effective(left)
  const b = effective(right)
  return a.length === b.length && a.every((root, index) => root === b[index])
}

interface Evaluation {
  data: ProgramBootData
  spec: PtcRunSpec
  bindings: Map<string, PtcBindingNamespace>
  invoke: typeof invokeBinding
  pending: number
  pendingBytes: number
  admit(text: string): void
  finish(error?: PtcRunFailure, value?: PtcJsonValue, discard?: boolean): void
}

/** One managed process; all mutable program state belongs to this process alone. */
class Kernel {
  private readonly unregister: (() => void) | undefined
  private confined: ConfinedArgv | undefined
  private handle?: SubprocessHandle
  private channel?: BirpcReturn<ProgramMethods, HostMethods>
  private active: Evaluation | undefined
  /** Admission remains reserved until the result includes completed process cleanup. */
  private running = false
  private cleanupPending = false
  private ready = false
  private stderr = ''
  private stopped: Promise<boolean> | undefined
  private starting?: Promise<void>
  private exitOutcome?: SubprocessOutcome
  private readonly controller = new AbortController()
  closed = false

  constructor(
    private readonly ctx: Context, private readonly config: ResolvedConfig,
    readonly policy: SandboxExecutionPolicy, processes: PolicyProcesses,
  ) {
    this.unregister = processes.register(policy, () => this.dispose())
  }

  assertIdle(): void {
    if (this.running) throw new Error('persistent PTC does not accept concurrent runs for one state owner')
    if (this.cleanupPending) throw new Error('persistent PTC environment cleanup is incomplete; release the owner before running another program')
  }

  async run(spec: PtcRunSpec, bindings: Map<string, PtcBindingNamespace>, code: string): Promise<PtcRunResult> {
    this.running = true
    const sandbox: PtcRunSandbox = { mode: this.policy.mode, denied: false }
    try {
      const result = await this.evaluate(spec, bindings, code)
      if (this.confined !== undefined) {
        sandbox.enforcement = this.confined.enforcement
        sandbox.denied = result.error?.kind !== 'sandbox-unavailable' && this.confined.denialSignatures.some(signature =>
          result.error?.message.toLowerCase().includes(signature.toLowerCase()))
      }
      return { ...result, sandbox }
    } finally { this.running = false }
  }

  private async evaluate(spec: PtcRunSpec, bindings: Map<string, PtcBindingNamespace>, code: string): Promise<PtcRunResult> {
    if (this.exitOutcome !== undefined) {
      const output = persistentOutput(this.config.maxOutputBytes)
      const failure: PtcRunFailure = spec.signal?.aborted
        ? { kind: 'abort', message: 'execution canceled; retained variables were cleared' }
        : { kind: 'worker-exit', message: `${this.exitMessage(this.exitOutcome)}; retained variables were cleared` }
      try { await this.stop() } catch (error: unknown) {
        return output.failure([], { kind: 'worker-exit', message: `managed process cleanup failed: ${messageOf(error)}` })
      }
      return output.failure([], failure)
    }
    this.stderr = ''
    const output = persistentOutput(this.config.maxOutputBytes)
    const logs: string[] = []
    const result = Promise.withResolvers<PtcRunResult>()
    const timer = spec.timeoutMs === null ? undefined : setTimeout(() => {
      evaluation.finish({ kind: 'timeout', message: `execution deadline reached (${spec.timeoutMs}ms); retained variables were cleared` }, undefined, true)
    }, spec.timeoutMs)
    const onAbort = (): void => { evaluation.finish({ kind: 'abort', message: 'execution canceled; retained variables were cleared' }, undefined, true) }
    const evaluation: Evaluation = {
      data: { code, maxOutputBytes: this.config.maxOutputBytes, namespaces: [...bindings.values()].map(binding => ({
        global: binding.global, names: Object.keys(binding.functions), ...binding.errorClass ? { errorClass: binding.errorClass } : {},
      })) },
      spec, bindings, invoke: AsyncLocalStorage.bind(invokeBinding), pending: 0, pendingBytes: 0,
      admit: (text) => {
        if (!output.admit(text, logs)) {
          logs.push(text)
          evaluation.finish({ kind: 'output-limit', message: 'output limit exceeded' }, undefined, true)
        }
      },
      finish: (error, value, discard = false) => {
        if (this.active !== evaluation) return
        clearTimeout(timer)
        spec.signal?.removeEventListener('abort', onAbort)
        this.active = undefined
        const outcome = error?.kind === 'output-limit' ? output.limit(logs)
          : error ? output.failure(logs, error) : output.success(logs, value)
        const cleanup = discard || outcome.error?.kind === 'output-limit' || spec.stateOwner === undefined ? this.stop() : Promise.resolve(true)
        void cleanup.then((drained) => {
          result.resolve(drained ? outcome : output.failure(logs, { kind: 'worker-exit', message: 'managed process exited but its output did not drain cleanly' }))
        }, (failure: unknown) => {
          result.resolve(output.failure(logs, { kind: 'worker-exit', message: `managed process cleanup failed: ${messageOf(failure)}` }))
        })
      },
    }
    this.active = evaluation
    spec.signal?.addEventListener('abort', onAbort, { once: true })
    if (spec.signal?.aborted) onAbort()
    else if (this.ready) this.dispatch(evaluation)
    else {
      this.starting = this.start(spec.cwd)
      void this.starting.catch((error: unknown) => { this.fail(error instanceof SandboxUnavailableError || (this.confined !== undefined && isRunnerSpawnFailure(error, this.confined.argv[0], spec.cwd)) ? 'sandbox-unavailable' : 'worker-exit', messageOf(error)) })
    }
    return await result.promise
  }

  async dispose(): Promise<void> {
    this.active?.finish({ kind: 'abort', message: 'PTC environment disposed; retained variables were cleared' }, undefined, true)
    await this.stop()
  }

  private stop(): Promise<boolean> {
    if (this.stopped) return this.stopped
    this.closed = true
    this.cleanupPending = true
    this.controller.abort()
    this.channel?.$close()
    this.stopped = (async () => {
      await this.starting?.catch(() => {})
      const handle = this.handle
      return handle === undefined || await terminateProcess(handle, this.config.graceMs)
    })().then((drained) => {
      this.unregister?.()
      this.cleanupPending = false
      if (!drained) this.ctx.logger.warn('persistent Node process exited but its output did not drain cleanly')
      return drained
    }, (error: unknown) => { this.stopped = undefined; throw error })
    return this.stopped
  }

  private fail(kind: PtcRunFailure['kind'], message: string): void {
    if (this.active) this.active.finish({ kind, message: `${message}; retained variables were cleared` }, undefined, true)
    else {
      if (!this.closed) this.ctx.logger.warn(message)
      void this.stop().catch((error: unknown) => { this.ctx.logger.warn(messageOf(error)) })
    }
  }

  private dispatch(evaluation: Evaluation): void {
    void this.channel?.run(evaluation.spec.cwd, evaluation.data).then((raw) => {
      const { error, value } = readProgramResult(raw)
      // Raw stdio can be queued beside the RPC result on separate pipes.
      setImmediate(() => { evaluation.finish(error, value) })
    }).catch((error: unknown) => { if (this.active === evaluation && !this.channel?.$closed) this.fail('protocol', messageOf(error)) })
  }

  private exitMessage(outcome: SubprocessOutcome): string {
    return `Node process exited (${String(outcome.exitCode)}, signal ${String(outcome.signal)})${this.stderr ? `: ${this.stderr}` : ''}`
  }

  private async start(cwd: string): Promise<void> {
    const executable = await this.ctx.subprocess.resolveExecutable(this.config.launch.executable, undefined, this.controller.signal)
    if (this.closed) return
    const argv = [executable, `--max-old-space-size=${this.config.maxOldGenerationSizeMb}`, ...bootstrapArgs(this.ctx.fs, this.config.launch, this.config.maxMessageBytes, true)]
    const policy = this.policy
    this.confined = policy.mode === 'danger-full-access' ? undefined : await this.ctx.sandbox.confine(argv, { ...policy, mode: policy.mode }, this.controller.signal)
    // oxlint-disable-next-line typescript/no-unnecessary-condition -- Cancellation can close the kernel during confinement.
    if (this.closed) return
    const env: NodeJS.ProcessEnv = Object.fromEntries(Object.keys(process.env)
      .filter(key => !STARTUP_ENVIRONMENT_NAMES.has(key.toUpperCase()) && key.toUpperCase() !== 'ELECTRON_RUN_AS_NODE')
      .map(key => [key, undefined]))
    const handle = this.ctx.subprocess.spawn({ argv: this.confined?.argv ?? argv, cwd, env, stdio: { stdin: 'ignore', stdout: 'pipe', stderr: 'pipe', control: 'pipe' }, graceMs: this.config.graceMs })
    this.handle = handle
    if (!handle.control || !handle.stdout || !handle.stderr) throw new Error('subprocess provider did not supply persistent PTC pipes')
    for (const stream of [handle.stdout, handle.stderr]) {
      stream.setEncoding('utf8')
      stream.on('data', (text: string) => {
        if (stream === handle.stderr) this.stderr = (this.stderr + text).slice(-this.config.maxOutputBytes)
        this.active?.admit(text)
      })
      stream.on('error', (error: Error) => { this.fail('worker-exit', error.message) })
    }
    // Report a protocol failure and abandon the handler; the closed endpoint sends no response.
    const refuse = (message: string): never => { this.fail('protocol', message); throw new Error(message) }
    this.channel = createProcessRpc<ProgramMethods, HostMethods>(handle.control, this.config.maxMessageBytes, {
      ready: () => {
        if (this.ready || !this.active) { this.fail('protocol', 'unexpected persistent PTC readiness'); return }
        this.ready = true
        this.dispatch(this.active)
      },
      log: (text: unknown) => {
        if (this.active === undefined || typeof text !== 'string') this.fail('protocol', 'invalid persistent PTC log')
        else this.active.admit(text)
      },
      outputLimit: () => {
        if (this.active === undefined) this.fail('protocol', 'idle persistent PTC output limit')
        else this.active.finish({ kind: 'output-limit', message: 'output limit exceeded' }, undefined, true)
      },
      invoke: async (global, name, wire) => {
        const evaluation = this.active
        if (evaluation === undefined) return refuse('idle persistent PTC binding call')
        const args = decodePtcJsonWire(wire)
        if (args === undefined) return refuse('binding arguments must be lossless JSON')
        const bytes = Buffer.byteLength(JSON.stringify(wire))
        if (evaluation.pending >= this.config.maxPendingCalls || evaluation.pendingBytes + bytes > this.config.maxMessageBytes) {
          return refuse('pending binding calls exceed configured limits')
        }
        // A retained function can name a binding this evaluation no longer declares; invokeBinding rejects it.
        const functions = evaluation.bindings.get(global)?.functions
        const fn = functions !== undefined && Object.hasOwn(functions, name) ? functions[name] : undefined
        evaluation.pending += 1
        evaluation.pendingBytes += bytes
        try {
          const value = await evaluation.invoke(fn, args)
          // The child already rejected this wait; a late result must not fill the pipe shared with the next evaluation.
          if (this.active !== evaluation) throw new Error('PTC evaluation ended before the binding completed')
          return value
        } finally {
          evaluation.pending -= 1
          evaluation.pendingBytes -= bytes
        }
      },
    }, [], (error, kind) => {
      if (kind === 'protocol') this.fail('protocol', error.message)
      else setImmediate(() => { if (!this.closed && this.exitOutcome === undefined) handle.terminate() })
    })
    void handle.done.then(async (outcome) => {
      this.exitOutcome = outcome
      await Promise.all([drainOutput(handle.stdout, this.config.graceMs), drainOutput(handle.stderr, this.config.graceMs)])
      // The run result can already be queued when process exit arrives.
      setImmediate(() => {
        if (!this.closed) this.fail(this.confined !== undefined && classifyRunnerFailure(outcome.exitCode, this.stderr, this.confined.runnerFailureRules) !== undefined ? 'sandbox-unavailable' : 'worker-exit', this.exitMessage(outcome))
      })
    }, (error: unknown) => { this.fail('worker-exit', messageOf(error)) })
  }
}

/** Sandboxed, best-effort retained state; restart, disposal, and interrupted execution clear it. */
export class InMemoryExecution {
  /**
   * Model-facing rules for retained variables, authority, output, and recovery.
   * @returns Usage guidance appended to the run_code schema.
   */
  get executionInstructions(): string {
    return 'The program contains TypeScript REPL statements; only erasable syntax is accepted, with no enum or namespaces. Type annotations are advisory. Top-level declarations and await are supported; top-level return is not supported. Emit output with console.log(...) or a final lossless JSON expression. Only printed values and that final expression become program output. Variables, functions, and imported modules remain in memory for this Agent across calls. Use await import(...) for modules. Direct filesystem access follows this execution\'s sandbox policy; process.env starts empty. Changing the session sandbox mode clears retained variables. An approved sandbox escalation runs in a separate temporary environment: it cannot access retained variables, and its declarations are discarded afterward. Restart, Agent disposal, cancellation, timeout, output overflow, or process failure clears the environment. If a variable is undefined, recreate it and reacquire its resources. State is not restored from the conversation, and failed evaluations can leave earlier mutations in memory. Redeclaring an existing let, const, or class name fails; assign to a mutable binding or use a new name. Await all tool calls; background callbacks cannot call tools after their evaluation ends. Await child-process completion before ending an evaluation; raw subprocess output belongs to whichever evaluation is active when it arrives.'
  }
  private readonly owners = new WeakMap<object, Kernel>()
  private readonly kernels = new Set<Kernel>()
  private disposed = false

  /**
   * Own retained execution resources under the existing Node provider.
   * @param ctx - Provider lifecycle and execution services.
   * @param config - Validated limits and selected executable.
   * @param processes - Provider-owned policy invalidation and cleanup.
   */
  constructor(private readonly ctx: Context, private readonly config: ResolvedConfig, private readonly processes: PolicyProcesses) {
    if (this.config.launch.kind === 'embedded') throw new Error('Persistent PTC requires a node-script launch')
    ctx.effect(() => async () => {
      this.disposed = true
      const results = await Promise.allSettled([...this.kernels].map(kernel => kernel.dispose()))
      this.kernels.clear()
      const failed = results.find(result => result.status === 'rejected')
      if (failed?.status === 'rejected') throw failed.reason
    }, 'Persistent PTC environments')
  }

  /**
   * Resolve the file policy, directory and deadline for a retained REPL.
   * @param request - Program, bindings, owner, directory and elapsed budget.
   * @returns Resolved inputs with the selected file policy.
   */
  resolve(request: PtcRunRequest): PtcRunSpec {
    if (this.disposed) throw new Error('persistent PTC: resolve after disposal')
    const sandboxPolicy = request.sandboxPolicy ?? this.ctx.sandboxPolicy.resolve()
    const cwd = request.cwd ?? sandboxPolicy.workspaceRoot
    if (!isAbsolute(cwd)) throw new Error('persistent PTC: cwd must be absolute')
    return { ...request, sandboxPolicy, cwd, timeoutMs: request.timeoutMs === null ? null : clampTimeout(request.timeoutMs, this.config.timeoutMs, this.config.maxTimeoutMs, 'persistent PTC timeoutMs') }
  }

  /**
   * Evaluate statements using this owner's process, or a temporary process without an owner.
   * @param spec - Resolved request; concurrent calls for one owner are rejected.
   * @returns Logs, a final JSON expression, or a structured execution failure.
   */
  async run(spec: PtcRunSpec): Promise<PtcRunResult> {
    if (this.disposed) throw new Error('persistent PTC: run after disposal')
    if (spec.sandboxPolicy === undefined || !isAbsolute(spec.cwd) || (spec.timeoutMs !== null && (!Number.isFinite(spec.timeoutMs) || spec.timeoutMs <= 0 || spec.timeoutMs > this.config.maxTimeoutMs))) throw new Error('persistent PTC: run requires resolved inputs with a sandbox policy')
    const bindings = validateBindings(spec)
    let kernel = spec.stateOwner === undefined ? undefined : this.owners.get(spec.stateOwner)
    kernel?.assertIdle()
    if (kernel && !kernel.closed && (kernel.policy.mode !== spec.sandboxPolicy.mode
      || kernel.policy.workspaceRoot !== spec.sandboxPolicy.workspaceRoot
      || kernel.policy.sessionId !== spec.sandboxPolicy.sessionId
      || !sameAdditionalRoots(kernel.policy.additionalWritableRoots, spec.sandboxPolicy.additionalWritableRoots))) {
      throw new Error('persistent PTC: release the retained environment before changing its sandbox policy')
    }
    let code: string
    try { code = stripTypeScriptTypes(spec.program) } catch (error: unknown) {
      const output = persistentOutput(this.config.maxOutputBytes)
      const outcome = output.failure([], { kind: 'exception', message: messageOf(error) })
      if (outcome.error?.kind === 'output-limit' && spec.stateOwner !== undefined) {
        try { await this.release(spec.stateOwner) } catch (failure: unknown) {
          return output.failure([], { kind: 'worker-exit', message: `managed process cleanup failed: ${messageOf(failure)}` })
        }
      }
      return outcome
    }
    if (!kernel || kernel.closed) {
      if (kernel) this.kernels.delete(kernel)
      // Direct callers can own mutable policy arrays; the retained process owns its grant snapshot.
      kernel = new Kernel(this.ctx, this.config, {
        ...spec.sandboxPolicy,
        ...spec.sandboxPolicy.additionalWritableRoots === undefined
          ? {}
          : { additionalWritableRoots: [...spec.sandboxPolicy.additionalWritableRoots] },
      }, this.processes)
      this.kernels.add(kernel)
      if (spec.stateOwner !== undefined) this.owners.set(spec.stateOwner, kernel)
    }
    try { return await kernel.run(spec, bindings, code) } finally {
      if (spec.stateOwner === undefined) this.kernels.delete(kernel)
    }
  }

  /**
   * Release retained state and await its process cleanup.
   * @param owner - Exact object supplied to run.
   * @returns Completion after the owned process and output streams stop.
   */
  async release(owner: object): Promise<void> {
    const kernel = this.owners.get(owner)
    if (!kernel) return
    await kernel.dispose()
    if (this.owners.get(owner) === kernel) this.owners.delete(owner)
    this.kernels.delete(kernel)
  }
}