deepseek-harness/packages/core/tools/src/code-mode.ts

/**
 * Code Mode `run_code` transport. Programs call the registry's agent-visible
 * tools through nested executions scheduled under the native concurrency
 * contract; each sub-dispatch is logged for reconstruction, while only the
 * outer curated result enters model history.
 * @module @deepseek-ai/dsh-tools/src/code-mode
 */

import { CallId, HarnessError } from '@deepseek-ai/dsh-llm'
import type { ContentBlock } from '@deepseek-ai/dsh-llm'
import type { CodeBindingFunction, CodeRunResult, CodeRuntime } from '@deepseek-ai/dsh-code-runtime'
import { snapshotJsonValue } from '@deepseek-ai/dsh-session'
import type { JsonValue } from '@deepseek-ai/dsh-session'
import { defineTool } from './schema.ts'
import { TOOL_REGISTRY_SCHEDULER } from './index.ts'
import type { ToolDefinition, ToolExecutionResult, ToolRegistry, ToolRunContext } from './index.ts'

declare module '@deepseek-ai/dsh-session' {
  interface SessionEventMap {
    /**
     * One sub-dispatch STARTING inside a `run_code` program: the parent
     * `run_code` call id, the deterministic sub-call id (`<parent>:code:<n>`,
     * numbered in submission order), and the tool `name` with its
     * JSON-normalized `arguments` — the exact value dispatched, normalized
     * BEFORE dispatch, so this append can never fail on payload shape.
     * Appended when the scheduler actually starts the call (not at
     * submission), so a start means the tool body pipeline was entered; a
     * call abandoned in the queue logs nothing. Log-only: `deriveMessages()`
     * ignores it; UIs use it for live per-sub-call running state and pair it
     * with `tool/code-dispatch` by `subCallId` (timing = the two events'
     * `time` fields).
     */
    'tool/code-dispatch-start': { parentCallId: CallId; subCallId: CallId; name: string; arguments: unknown }
    /**
     * One bridged sub-dispatch SETTLING: the pairing ids (matching the
     * `tool/code-dispatch-start` with the same `subCallId`), the tool `name`
     * with the same JSON-normalized `arguments`, and the sub-call's complete
     * model-facing outcome in `tool/result`'s own vocabulary
     * (`content` + `isError`), so UIs render a sub-call through the exact
     * code path that renders a native call. Every started sub-call settles
     * with exactly one of these (abort included: the aborted pipeline result
     * is an `isError` outcome).
     * Log-only: `deriveMessages()` ignores it, so sub-calls never re-enter
     * model context; persistence and UIs get every call. Appended inside the
     * parent `run_code`'s execution (the bridge drains in-flight dispatches
     * before returning), so the turn-enclosure invariant holds by
     * construction.
     */
    'tool/code-dispatch': { parentCallId: CallId; subCallId: CallId; name: string; arguments: unknown; isError: boolean; content: ContentBlock[] }
  }
}

/** The model-facing name of the Code Mode tool. */
export const RUN_CODE_NAME = 'run_code'

/** The `tools:sdk` section order: inside the 100–199 tool-guidance band, after per-tool guidance sections. */
export const SDK_SECTION_ORDER = 150

/**
 * Thrown by `run_code` when the program run itself failed — a program
 * exception, a budget expiry, an abort, or substrate death. Extends
 * {@link HarnessError} (`code: 'CODE_RUN_FAILED'`); the registry's execution
 * pipeline converts it into a structured `isError` result whose text carries
 * the failure kind plus the captured logs, so the model can self-correct.
 */
export class CodeRunFailedError extends HarnessError {
  constructor(message: string) {
    super(message, 'CODE_RUN_FAILED')
    this.name = 'CodeRunFailedError'
  }
}

/**
 * Snapshot one binding call's argument as lossless JSON, then snapshot that
 * detached value again so dispatch and logging stay independent without
 * reintroducing structured-clone's platform-specific nesting limit.
 */
function jsonNormalizeArgs(value: unknown): { dispatched: unknown; logged: unknown } {
  let snapshot: JsonValue | undefined
  try {
    snapshot = snapshotJsonValue(value) as JsonValue | undefined
  } catch (error: unknown) {
    throw new Error(`tool arguments must be lossless JSON: ${error instanceof Error ? error.message : String(error)}`)
  }
  if (snapshot === undefined) {
    throw new Error('tool arguments must be lossless JSON (call the tool with an arguments object, e.g. `{}`)')
  }
  const logged = snapshotJsonValue(snapshot)
  /* v8 ignore next -- snapshot is already a detached lossless JSON value. */
  if (logged === undefined) {
    throw new Error('tool arguments could not be detached for durable logging')
  }
  return { dispatched: snapshot, logged }
}

/** Two-space JSON presentation, matching the existing shallow `run_code` text contract. */
const JSON_INDENT = '  '

/**
 * ECMAScript caps `JSON.stringify`'s `space` string at ten characters. The
 * renderer also caps TOTAL indentation there, compacting deeper subtrees, so
 * formatted output remains linear in the canonical JSON size.
 */
const MAX_JSON_INDENT_CHARS = 10

/** A pending fragment in the iterative JSON presentation traversal. */
type JsonRenderTask =
  | { kind: 'text'; text: string }
  | { kind: 'value'; value: JsonValue; depth: number; compact: boolean }

/** Render one non-string JSON root without recursive traversal or unbounded indentation growth. */
function renderJsonValue(value: Exclude<JsonValue, string>): string {
  const chunks: string[] = []
  const tasks: JsonRenderTask[] = [{ kind: 'value', value, depth: 0, compact: false }]
  for (let task = tasks.pop(); task !== undefined; task = tasks.pop()) {
    if (task.kind === 'text') {
      chunks.push(task.text)
      continue
    }

    const current = task.value
    if (current === null || typeof current === 'boolean' || typeof current === 'number') {
      chunks.push(String(current))
      continue
    }
    if (typeof current === 'string') {
      chunks.push(JSON.stringify(current))
      continue
    }

    const compact = task.compact || (task.depth + 1) * JSON_INDENT.length > MAX_JSON_INDENT_CHARS
    const childDepth = task.depth + 1
    if (Array.isArray(current)) {
      chunks.push('[')
      if (current.length === 0) {
        chunks.push(']')
        continue
      }
      tasks.push({ kind: 'text', text: compact ? ']' : `\n${JSON_INDENT.repeat(task.depth)}]` })
      for (let index = current.length - 1; index >= 0; index--) {
        const item = current[index]
        /* v8 ignore next -- canonical JsonValue arrays are dense. */
        if (item === undefined) throw new Error('cannot render a sparse JSON array')
        tasks.push({ kind: 'value', value: item, depth: childDepth, compact })
        tasks.push({
          kind: 'text',
          text: compact
            ? index === 0 ? '' : ','
            : `${index === 0 ? '\n' : ',\n'}${JSON_INDENT.repeat(childDepth)}`,
        })
      }
      continue
    }

    const keys = Object.keys(current)
    chunks.push('{')
    if (keys.length === 0) {
      chunks.push('}')
      continue
    }
    tasks.push({ kind: 'text', text: compact ? '}' : `\n${JSON_INDENT.repeat(task.depth)}}` })
    for (let index = keys.length - 1; index >= 0; index--) {
      const key = keys[index]
      /* v8 ignore next -- the loop is bounded by the captured key count. */
      if (key === undefined) throw new Error('cannot render a missing JSON object key')
      const item = current[key]
      /* v8 ignore next -- canonical JsonValue records contain no undefined properties. */
      if (item === undefined) throw new Error('cannot render an undefined JSON object property')
      tasks.push({ kind: 'value', value: item, depth: childDepth, compact })
      tasks.push({
        kind: 'text',
        text: compact
          ? `${index === 0 ? '' : ','}${JSON.stringify(key)}:`
          : `${index === 0 ? '\n' : ',\n'}${JSON_INDENT.repeat(childDepth)}${JSON.stringify(key)}: `,
      })
    }
  }
  return chunks.join('')
}

/** Render one present program completion value for the model-facing result text. */
function renderValue(value: JsonValue): string {
  return typeof value === 'string' ? value : renderJsonValue(value)
}

/** Canonical value returned by the outer Code Mode transport. */
type RunCodeOutput = { logs: string[]; result?: JsonValue }

/**
 * Build the `run_code` {@link ToolDefinition}: required `code` and
 * `description` parameters, executed through the dispatch bridge described
 * above. The
 * registry reserves it as presentation infrastructure under non-native modes,
 * outside the filterable global/scoped capability layers.
 * @param registry - the owning registry (sub-calls go through its `execute`,
 *   bindings cover its registered tools).
 * @param requireRuntime - resolves `ctx.codeRuntime` or throws the loud
 *   misconfiguration error (shared with the registry's assembly-time checks).
 * @param maxParallel - the run's overlap cap for parallel-classified
 *   sub-calls (the registry passes its validated `maxParallelSubCalls`).
 * @returns the registry-ready definition.
 */
export function createRunCodeTool(registry: ToolRegistry, requireRuntime: () => CodeRuntime, maxParallel: number): ToolDefinition {
  return defineTool({
    name: RUN_CODE_NAME,
    description:
      'Execute a TypeScript program against the available tools. Write the BODY of an '
      + 'async function (erasable syntax only; top-level `await` and `return` work) and '
      + 'call tools as `await tools.name(args)` per the declarations in the system prompt. '
      + 'Only what you print or return comes back — curate it.',
    parameters: {
      code: { type: 'string', required: true, description: 'The program: the body of an async TypeScript function.' },
      description: {
        type: 'string',
        required: true,
        description: 'Clear, concise description of what this program does in active voice, '
          + '5-10 words (shown in the UI). Examples: "Count TODO markers across packages"; '
          + '"Read failing test and its fixture"; "Rename config key in every cordis.yml".',
      },
    },
    output: {
      schema: {
        type: 'object',
        additionalProperties: false,
        properties: {
          logs: { type: 'array', required: true, items: { type: 'string' } },
          result: { type: 'json' },
        },
      },
      render: (_args, value) => {
        const rendered = value.result === undefined ? '' : renderValue(value.result)
        const parts = [value.logs.join('\n'), rendered].filter(part => part.length > 0)
        return [{ type: 'text', text: parts.length > 0 ? parts.join('\n') : '(run_code completed with no output)' }]
      },
    },
    async execute(args, exec): Promise<RunCodeOutput> {
      if (args.description.trim().length === 0) {
        throw new Error('invalid description: expected a non-empty string')
      }
      const runtime = requireRuntime()

      // The run-scoped abort: follows the outer signal in, and fires when the
      // run settles for ANY reason, so an in-flight sub-dispatch is aborted
      // (its executor kills on this signal) instead of orphaned, and
      // queued-unstarted dispatches are abandoned.
      const runController = new AbortController()
      const onOuterAbort = (): void => { runController.abort(exec.signal.reason) }
      exec.signal.addEventListener('abort', onOuterAbort, { once: true })

      let dispatches = 0
      // The per-run scheduler, reusing the NATIVE concurrency contract through
      // the registry's staged view (the loop scheduler's own seam): submitted
      // calls START strictly in submission order; only the around-dispatch/body
      // stage overlaps — ordered pre-execute runs at start time and ordered
      // post-execute/context commitment runs in submission order through the
      // commit cursor below, so stateful policy listeners observe submission
      // order exactly as they do under the native loop. Consecutive
      // parallel-classified calls overlap up to maxParallel; an exclusive call
      // waits for the pool to drain, runs alone, and bars later calls.
      // Classification is re-read via executionMode() immediately before each
      // start (a registry mutation while queued can flip a call exclusive),
      // matching the native scheduler's lazy reclassification.
      interface PendingDispatch {
        /** Ordered stage: append the start event, prepare, dispatch (body overlaps), park for commit. */
        start(): Promise<void>
        classify(): 'parallel' | 'exclusive'
        abandon(): void
        /** Ordered stage: post-execute + context deferral + settle event, in submission order. */
        commit(): Promise<void>
        /** Set once the dispatch stage settles; commit() runs after this resolves. */
        dispatched?: Promise<void>
      }
      const pendingQueue: PendingDispatch[] = []
      const inFlight = new Set<Promise<void>>()
      const commitQueue: PendingDispatch[] = []
      let committing = false
      let exclusiveActive = false
      let pumping = false
      /** Ordered commit cursor: drain the head-of-line settled dispatches one at a time. */
      const commitReady = async (): Promise<void> => {
        if (committing) return
        committing = true
        try {
          while (commitQueue.length > 0) {
            const head = commitQueue[0]
            /* v8 ignore next -- the loop condition bounds the index. */
            if (head === undefined) break
            if (head.dispatched === undefined) break
            await head.dispatched
            commitQueue.shift()
            await head.commit()
          }
        } finally {
          committing = false
        }
      }
      const pump = (): void => {
        // The finally-driven re-entry below would otherwise recurse.
        if (pumping) return
        pumping = true
        try {
          for (;;) {
            const head = pendingQueue[0]
            if (head === undefined) return
            if (runController.signal.aborted) {
              pendingQueue.shift()
              head.abandon()
              continue
            }
            // Reclassify at start time (fail-closed on registry changes).
            const mode = head.classify()
            if (exclusiveActive || inFlight.size >= (mode === 'exclusive' ? 1 : maxParallel)) return
            if (mode === 'exclusive') {
              if (inFlight.size > 0) return
              exclusiveActive = true
            }
            pendingQueue.shift()
            commitQueue.push(head)
            const flight = head.start().finally(() => {
              inFlight.delete(flight)
              if (mode === 'exclusive') exclusiveActive = false
              // Commit ordering and slot refill are independent: the cursor
              // may wait head-of-line on an earlier dispatch while later
              // slots keep starting.
              void commitReady()
              pump()
            })
            inFlight.add(flight)
          }
        } finally {
          pumping = false
        }
      }
      /** Every in-flight dispatch settled AND committed; nothing can start (the run is aborted at call time). */
      const drainDispatches = async (): Promise<void> => {
        // Abandon queued-unstarted tasks first, then await the live set until quiescent.
        pump()
        while (inFlight.size > 0) await Promise.allSettled([...inFlight])
        await commitReady()
      }

      // Read through a call, not a bare property: the abort state genuinely
      // changes across awaits, and a direct `.aborted` re-check after one
      // would be narrowed away by control flow analysis.
      const runOver = (): boolean => runController.signal.aborted

      const binding = (name: string): CodeBindingFunction => async (rawArgs: unknown): Promise<JsonValue> => {
        if (runOver()) {
          throw new Error(`run_code run is over (${String(runController.signal.reason)}); ${name} not dispatched`)
        }
        const normalized = jsonNormalizeArgs(rawArgs)
        const n = ++dispatches
        const subCallId = CallId(`${String(exec.callId)}:code:${n}`)
        const input = {
          callId: subCallId,
          name,
          arguments: normalized.dispatched,
          ...exec.agent ? { agent: exec.agent } : {},
          parent: exec.token,
          signal: runController.signal,
        }
        type DispatchOutcome = { isError: true; message: string } | { isError: false; value: JsonValue }
        const scheduler = registry[TOOL_REGISTRY_SCHEDULER]
        const outcome = await new Promise<DispatchOutcome>((resolve, reject) => {
          // Set by start(): what commit() finalizes in submission order.
          let parked:
            | { kind: 'post-result' | 'final-result'; exec: ToolRunContext; result: ToolExecutionResult }
            | undefined
          const settle = (result: ToolExecutionResult): void => {
            exec.agent?.session.append('tool/code-dispatch', {
              parentCallId: exec.callId,
              subCallId,
              name,
              // The SIBLING parse of the dispatched value: byte-identical JSON,
              // but a separate object — a tool mutating its args cannot desync
              // this record from what it actually received.
              arguments: normalized.logged,
              isError: result.isError,
              // The registry deep-froze this projection at result finalization;
              // append snapshots it again, so the log copy stays detached.
              content: result.content,
            })
            resolve(result.isError
              ? { isError: true, message: result.error.message }
              : { isError: false, value: result.value })
          }
          pendingQueue.push({
            // Re-read per pump pass against the same agent view the SDK
            // declared; fail-closed exclusive when undeclared/invalid.
            classify: () => registry.executionMode(input).kind,
            abandon: () => {
              reject(new Error(`run_code run is over (${String(runController.signal.reason)}); ${name} tool call abandoned`))
            },
            start(): Promise<void> {
              exec.agent?.session.append('tool/code-dispatch-start', {
                parentCallId: exec.callId,
                subCallId,
                name,
                arguments: normalized.logged,
              })
              // Ordered prepare (pre-execute/guards) runs here — starts are
              // strictly submission-ordered; only dispatch overlaps.
              this.dispatched = (async () => {
                const prepared = await scheduler.prepare(input)
                if (prepared.kind === 'dispatch') {
                  const dispatchOutcome = await scheduler.dispatch(prepared.exec)
                  parked = { kind: dispatchOutcome.kind, exec: prepared.exec, result: dispatchOutcome.result }
                  return
                }
                parked = { kind: prepared.kind, exec: prepared.exec, result: prepared.result }
              })()
              return this.dispatched
            },
            async commit(): Promise<void> {
              /* v8 ignore next -- commit() runs only after this.dispatched resolved, which set parked. */
              if (parked === undefined) return
              const result = parked.kind === 'post-result'
                ? await scheduler.finalize(parked.exec, parked.result)
                : scheduler.finish(parked.exec, parked.result)
              for (const context of result.additionalContexts ?? []) {
                exec.deferContext(context)
              }
              settle(result)
            },
          })
          pump()
        })
        // A budget expiry or outer cancel that lands while this call was in
        // flight already aborted the dispatch; stop the program now rather
        // than hand it a result from a run that is over.
        if (runOver()) {
          throw new Error(`run_code run is over (${String(runController.signal.reason)}); ${name} result discarded`)
        }
        // The worker turns a binding rejection into ToolCallError and adds
        // only the binding name. Native content and internal error metadata
        // stay outside the program-facing failure contract.
        if (outcome.isError) throw new Error(outcome.message)
        return outcome.value
      }

      // Null-prototype + defineProperty, mirroring the worker-side namespace
      // build: a registered tool named `__proto__` must become an ordinary
      // own key (a plain-object assignment would hit the prototype setter,
      // silently dropping the binding), and the runtime host resolves
      // binding names as own properties only.
      const functions: Record<string, CodeBindingFunction> = Object.create(null) as Record<string, CodeBindingFunction>
      // Enumerate the CALLING AGENT's visible set (scoped tools join,
      // restricted globals vanish) — the same view the SDK section declared,
      // so a program can bind exactly what its prompt promised; sub-dispatch
      // re-resolves per call through the same view (exec.agent threads down).
      for (const schema of registry.schemas(exec.agent)) {
        if (schema.name === RUN_CODE_NAME) continue
        Object.defineProperty(functions, schema.name, { enumerable: true, value: binding(schema.name) })
      }

      try {
        let result: CodeRunResult
        try {
          result = await runtime.run({
            program: args.code,
            bindings: [{
              global: 'tools',
              functions,
              errorClass: { name: 'ToolCallError', memberNameProperty: 'toolName' },
            }],
            signal: runController.signal,
          })
        } finally {
          // Abort sub-dispatches and drain every in-flight dispatch before
          // closing the turn (queued-unstarted ones are abandoned unlogged).
          // Binding failures remain observable through their individual promises.
          runController.abort('run_code settled')
          await drainDispatches()
        }

        if (result.error) {
          const logsText = result.logs.length > 0 ? `\nCaptured output:\n${result.logs.join('\n')}` : ''
          throw new CodeRunFailedError(`code run failed (${result.error.kind}): ${result.error.message}${logsText}`)
        }
        return {
          logs: result.logs,
          ...result.value !== undefined ? { result: result.value } : {},
        }
      } finally {
        exec.signal.removeEventListener('abort', onOuterAbort)
      }
    },
    // The model-authored description is the call's always-visible UI label
    // (the bash `description` precedent); the program itself rides rawInput.
    presentCall: args => ({
      card: 'generic',
      title: args.description,
      kind: 'execute',
      rawInput: args.code,
    }),
    // Deliberately no presentResult: the generic surface fallback keeps this
    // title and reads durable result content without duplicating a large raw
    // result into the host view payload.
  })
}