deepseek-harness/packages/subagent/subagent-control/src/index.ts

/**
 * Continuable-subagent control service (`ctx.subagentControl`): stable child
 * ids, descriptor persistence and lookup by known child id, Task-backed
 * activation, and steer-or-resume message routing. The low-level
 * `ctx.subagents` seam stays collection-, Task-, and persistence-agnostic;
 * this service owns the policy that binds one durable child session to a
 * series of disposable Task-backed activations.
 *
 * Every continuable activation — initial or resumed, parent- or human-started
 * — has exactly one Task and one result. Task settlement awaits the child
 * result, disposes the run, and only then records the outcome, so a terminal
 * Task leaves the durable child session but no live child Agent. Cancellation
 * targets the whole activation: parent and human messages that joined one
 * turn share its result and its `killed` outcome.
 *
 * @module @deepseek-ai/dsh-subagent-control
 */

import { randomUUID } from 'node:crypto'
import { Context, Service } from 'cordis'
import type { Agent } from '@deepseek-ai/dsh-agent'
import { HarnessError } from '@deepseek-ai/dsh-llm'
import type { ContentBlock } from '@deepseek-ai/dsh-llm'
import { SessionId } from '@deepseek-ai/dsh-session'
import type { SessionPersistence } from '@deepseek-ai/dsh-session-persistence'
import { foldSubagentDescriptor, snapshotSubagentDescriptor } from '@deepseek-ai/dsh-subagent'
import type { SubagentResult, SubagentRun, SubagentStartRequest } from '@deepseek-ai/dsh-subagent'
import type { TaskHooks, TaskId, TaskOutcome } from '@deepseek-ai/dsh-tasks'

declare module 'cordis' {
  interface Context {
    subagentControl: SubagentControlService
  }
}

/** Typed error for control-service routing, authorization, and delivery failures. */
export class SubagentControlError extends HarnessError {
  constructor(message: string, code: string, options?: ErrorOptions) {
    super(message, code, options)
    this.name = 'SubagentControlError'
  }
}

/** What a caller asks for when starting a continuable background child. */
export interface ContinuableStartSpec {
  /** The `ctx.subagents` provider to establish the child on. */
  readonly provider: string
  /** One-line model-facing Task label (the delegation description). */
  readonly label: string
  /**
   * The delegation request. The service resolves the stable child id and the
   * durable descriptor, then supplies the Task-owned cancellation signal and
   * `continuation` itself.
   */
  readonly request: Omit<SubagentStartRequest, 'signal' | 'continuation'>
}

/** Identities returned by {@link SubagentControlService.startContinuable}. */
export interface ContinuableStart {
  /** The durable child session id, stable across activations. */
  readonly childId: SessionId
  /** The initial activation's Task id. */
  readonly taskId: TaskId
}

/**
 * How {@link SubagentControlService.sendMessage} delivered a message:
 * `steered` joined the running activation's existing Task without creating a
 * Task of its own; `started` created a fresh Task that cold-resumes the
 * durable child with the message. Failure is an exception, never a result —
 * an undelivered message throws.
 */
export type SendMessageResult =
  | { readonly route: 'steered'; readonly taskId: TaskId }
  | { readonly route: 'started'; readonly taskId: TaskId }

/**
 * One child's current process-local activation: its Task and, after provider
 * publication, its run. Installed before any provider or persistence await
 * and removed only after run disposal and Task terminal publication. This
 * exists solely so parent and human senders can find the same activation — it
 * is not a durable catalog, admission reservation, or run-state machine.
 */
interface ActiveActivation {
  /** Assigned in the same synchronous frame as the install, when the Task registers. */
  taskId: TaskId | undefined
  /** Filled when the provider publishes; `undefined` while starting or resuming. */
  run: SubagentRun | undefined
  /** Resolved by the completion listener when the Task's terminal snapshot is recorded. */
  readonly terminal: PromiseWithResolvers<void>
}

/**
 * Map a child result to the task outcome: completed carries final text,
 * aborted is killed, and every other reason is failed without partial output.
 * @param result - child terminal result.
 * @returns outcome for the `ctx.tasks` registration.
 */
export function runOutcome(result: SubagentResult): TaskOutcome {
  switch (result.stopReason) {
    case 'completed':
      return { status: 'completed', output: finalText(result.output) }
    case 'aborted':
      return { status: 'killed' }
    case 'error':
    case 'max-tokens':
    case 'refusal':
      return { status: 'failed', detail: result.stopReason }
    // Merge-extensible reasons remain failures with their raw detail.
    default:
      return { status: 'failed', detail: String(result.stopReason) }
  }
}

/**
 * Await the child result, dispose the run, then return its task outcome. Result
 * and disposal failures become `failed`; when both fail, both details survive.
 * @param run - live run to settle and release.
 * @returns outcome after child resources are released.
 */
export async function settleRun(run: SubagentRun): Promise<TaskOutcome> {
  let outcome: TaskOutcome
  try {
    outcome = runOutcome(await run.result)
  } catch (error: unknown) {
    outcome = { status: 'failed', detail: String(error) }
  }
  try {
    await run.dispose()
  } catch (error: unknown) {
    const prefix = outcome.detail === undefined ? '' : `${outcome.detail}; `
    return { status: 'failed', detail: `${prefix}dispose failed: ${String(error)}` }
  }
  return outcome
}

/** Flatten a child's final output blocks to the task's final text. */
function finalText(blocks: ContentBlock[]): string {
  return blocks
    .filter((block): block is Extract<ContentBlock, { type: 'text' }> => block.type === 'text')
    .map(block => block.text)
    .join('')
}

/**
 * The continuable-subagent orchestration service. Tool schema and UI adapters
 * are consumers of this one contract: parent and human messages route through
 * {@link sendMessage} and share one activation result and cancellation
 * boundary, while foreground one-shot delegation keeps calling
 * `ctx.subagents.start()` directly.
 */
export class SubagentControlService extends Service {
  static inject = ['subagents', 'tasks', 'agents']

  /** Child session id → its current activation. Process-local, never durable. */
  private activations = new Map<SessionId, ActiveActivation>()

  constructor(ctx: Context) {
    super(ctx, 'subagentControl')
    // Terminal publication is one of the two removal conditions. The exact
    // Task id pins the resolution to this activation, never a later same-child one.
    ctx.tasks.onTaskDone((snapshot) => {
      for (const activation of this.activations.values()) {
        if (activation.taskId === snapshot.id) activation.terminal.resolve()
      }
    })
    ctx.effect(() => () => { this.activations.clear() }, 'subagentControl.activations()')
  }

  /**
   * Start a continuable background child: allocate its stable session id,
   * snapshot its durable descriptor, and register the initial activation's
   * Task. A synchronous validation failure (a non-JSON descriptor input,
   * missing persistence, Task preflight) throws without creating a Task; the
   * method otherwise returns both identities immediately, without waiting for
   * child publication or descriptor durability. Asynchronous startup failure
   * settles the returned Task as `failed` (or `killed` when cancelled) after
   * any published run is disposed, which can leave an unmaterialized child id
   * that later by-id operations report as unavailable.
   * @param spec - provider, Task label, and the delegation request.
   * @returns the stable child id and the initial activation's Task id.
   */
  startContinuable(spec: ContinuableStartSpec): ContinuableStart {
    this.requirePersistence()
    const childId = SessionId(randomUUID())
    const request = spec.request
    // Snapshot before Task creation: invalid descriptor JSON rejects the call
    // with no Task, and the detached value is what reaches the child log.
    const agentProvider = request.agentOptions?.provider ?? request.parent.options.provider
    const agentModel = request.agentOptions?.model ?? request.parent.options.model
    const descriptor = snapshotSubagentDescriptor({
      provider: spec.provider,
      ...agentProvider !== undefined ? { agentProvider } : {},
      ...agentModel !== undefined ? { agentModel } : {},
      ...request.persona !== undefined ? { persona: request.persona } : {},
      ...request.toolFilter !== undefined ? { toolFilter: request.toolFilter } : {},
    })
    const taskId = this.startActivation(childId, spec.label, request.parent, signal =>
      this.ctx.subagents.start(spec.provider, {
        ...request,
        signal,
        continuation: { sessionId: childId, descriptor },
      }))
    return { childId, taskId }
  }

  /**
   * Deliver one message to a known continuable child: steer its running
   * activation, or cold-resume the durable session into a fresh Task-backed
   * activation. The two routes are reported distinctly so timing-dependent
   * routing is observable. A throw means the message was NOT delivered — in
   * particular, losing a race with Task settlement does not fall through to
   * cold resume within the same call; a later retry after Task terminal may
   * start the next activation. The started Task owns descriptor lookup and
   * direct-parent authorization (its AbortSignal exists before that lookup),
   * so an unknown, foreign, or descriptor-less child settles the started Task
   * as `failed` with a detail reporting the id as unavailable.
   * @param parent - the live parent agent sending the message (model tool or
   *   human adapter); Task access is authorized by its session id.
   * @param childId - the stable child session id.
   * @param message - the content to deliver.
   * @returns whether the message `steered` the existing Task or `started` a new one.
   */
  sendMessage(parent: Agent, childId: SessionId, message: ContentBlock[]): SendMessageResult {
    this.assertOwnership(childId)
    const activation = this.activations.get(childId)
    if (activation !== undefined) {
      return { route: 'steered', taskId: this.steerActivation(activation, parent, childId, message) }
    }
    return { route: 'started', taskId: this.resumeActivation(parent, childId, message) }
  }

  /**
   * Synchronous ownership compare before any by-id routing: a live registry
   * Agent outside the association — or different from the associated run's
   * agent — was started by something else. Fail instead of adopting an idle
   * Agent or attaching an untracked turn.
   */
  private assertOwnership(childId: SessionId): void {
    const live = this.ctx.agents.get(childId)
    if (live === undefined) return
    const activation = this.activations.get(childId)
    if (activation === undefined) {
      throw new SubagentControlError(
        `subagent "${childId}" has a live agent outside control-service ownership; the message was not delivered`,
        'OWNERSHIP_CONFLICT',
      )
    }
    if (activation.run !== undefined && activation.run.localAgent !== live) {
      throw new SubagentControlError(
        `subagent "${childId}" registry agent is not the associated activation's agent; the message was not delivered`,
        'OWNERSHIP_CONFLICT',
      )
    }
  }

  /** Deliver to the running activation's Task through strict live steering. */
  private steerActivation(
    activation: ActiveActivation,
    parent: Agent,
    childId: SessionId,
    message: ContentBlock[],
  ): TaskId {
    const taskId = activation.taskId
    /* v8 ignore next 3 -- the install and Task registration share one synchronous frame, so an observed activation carries its Task id. */
    if (taskId === undefined) {
      throw new SubagentControlError(`subagent "${childId}" activation is starting; the message was not delivered`, 'NOT_DELIVERED')
    }
    // Owner-session authorization plus the live status for the strict check.
    const snapshot = this.ctx.tasks.get(taskId, parent)
    if (snapshot.status !== 'running') {
      throw new SubagentControlError(
        `subagent "${childId}" task ${taskId} is ${snapshot.status}; the message was not delivered `
        + '— retry after it settles to start the next activation',
        'NOT_DELIVERED',
      )
    }
    const run = activation.run
    if (run === undefined) {
      throw new SubagentControlError(`subagent "${childId}" activation is starting; the message was not delivered`, 'NOT_DELIVERED')
    }
    if (run.steer === undefined) {
      throw new SubagentControlError(
        `subagent "${childId}" provider does not accept live delivery; the message was not delivered`,
        'NOT_DELIVERED',
      )
    }
    try {
      run.steer(message)
    } catch (error: unknown) {
      // Strict steering lost the race with turn settlement. Deliberately no
      // cold-resume fallback here: that would attach the message to a turn the
      // caller did not observe.
      throw new SubagentControlError(
        `subagent "${childId}" stopped before delivery; the message was not delivered`,
        'NOT_DELIVERED',
        { cause: error },
      )
    }
    return taskId
  }

  /**
   * Cold-resume a persisted child into a fresh Task-backed activation. The
   * Task owns its `AbortController` before descriptor lookup: the load,
   * direct-parent authorization, and descriptor fold run inside the
   * activation, with cancellation rechecked after the un-signalled
   * persistence await so an early `task_kill` prevents any later child work.
   */
  private resumeActivation(parent: Agent, childId: SessionId, message: ContentBlock[]): TaskId {
    const persistence = this.requirePersistence()
    return this.startActivation(childId, resumeLabel(message), parent, async (signal) => {
      let loaded: Awaited<ReturnType<typeof persistence.load>>
      try {
        loaded = await persistence.load(childId)
      } catch (error: unknown) {
        throw new SubagentControlError(
          `subagent "${childId}" is unavailable`,
          'NOT_RESUMABLE',
          { cause: error },
        )
      }
      // The persistence seam takes no signal; recheck before any child work.
      if (signal.aborted) throw new SubagentControlError('subagent resume was cancelled during lookup', 'CANCELLED')
      // Authorize the persisted header before folding: only the direct parent
      // recorded at creation may continue this child.
      if (loaded.meta.parentSession !== parent.id) {
        throw new SubagentControlError(
          `subagent "${childId}" belongs to another parent session`,
          'UNAUTHORIZED',
        )
      }
      // Fold only the child's own suffix: a fork seed replays the parent's
      // log, which may carry an ANCESTOR's descriptor when the parent is
      // itself a continuable child.
      const descriptor = foldSubagentDescriptor(loaded.events.slice(loaded.meta.seedLength ?? 0))
      if (descriptor === undefined) {
        throw new SubagentControlError(
          `subagent "${childId}" has no supported continuation descriptor`,
          'NOT_RESUMABLE',
        )
      }
      return this.ctx.subagents.resume(descriptor.provider, {
        sessionId: childId,
        prompt: message,
        parent,
        signal,
        descriptor,
      })
    })
  }

  /**
   * Install the activation association, register its Task, and bind the two
   * removal conditions. The association is installed before any persistence
   * or provider await — the producer body runs synchronously up to its first
   * await — and removed only after run disposal (the producer settled) and
   * Task terminal publication. This synchronous install admits one activation
   * per child in this process; a competing untracked publication still loses
   * at the Agent registry collision boundary inside the provider.
   */
  private startActivation(
    childId: SessionId,
    label: string,
    owner: Agent,
    begin: (signal: AbortSignal) => Promise<SubagentRun>,
  ): TaskId {
    const activation: ActiveActivation = {
      taskId: undefined,
      run: undefined,
      terminal: Promise.withResolvers<void>(),
    }
    this.activations.set(childId, activation)
    let taskId: TaskId
    try {
      taskId = this.ctx.tasks.start({
        kind: 'subagent',
        label,
        owner,
        run: (): TaskHooks => {
          const controller = new AbortController()
          const done = (async (): Promise<TaskOutcome> => {
            try {
              const run = await begin(controller.signal)
              activation.run = run
              return await settleRun(run)
            } catch (error: unknown) {
              // A pre-publication abort rejects only after the provider's
              // creation transaction rolled back to quiescence, so recording
              // `killed` here honors the settlement-after-rollback contract.
              return controller.signal.aborted
                ? { status: 'killed' }
                : { status: 'failed', detail: String(error) }
            }
          })()
          void Promise.allSettled([done, activation.terminal.promise]).then(() => {
            /* v8 ignore else -- service teardown clears the map while a producer is still settling. */
            if (this.activations.get(childId) === activation) this.activations.delete(childId)
          })
          return {
            cancel: (reason?: string) => {
              // Cancellation targets the whole activation: every message that
              // joined this turn shares the `killed` outcome.
              controller.abort(reason ?? 'subagent activation killed')
            },
            done,
            // No readOutput: the child session owns intermediate detail.
          }
        },
      })
    } catch (error: unknown) {
      // Task preflight failed; nothing started, so the install rolls back.
      this.activations.delete(childId)
      throw error
    }
    // Same synchronous frame as the install: an observer that can run at all
    // runs after this assignment.
    activation.taskId = taskId
    return taskId
  }

  /** Resolve the persistence service continuable children require, or fail loud. */
  private requirePersistence(): SessionPersistence {
    const persistence = this.ctx.get('sessionPersistence')
    if (persistence === undefined) {
      throw new SubagentControlError(
        'continuable subagents require session persistence (load a dsh-session-persistence backend)',
        'PERSISTENCE_UNAVAILABLE',
      )
    }
    return persistence
  }
}

/** Derive a resumed activation's Task label from its message. */
function resumeLabel(message: ContentBlock[]): string {
  const text = finalText(message).trim().replace(/\s+/g, ' ')
  if (text.length === 0) return 'subagent follow-up'
  return text.length > 80 ? `${text.slice(0, 79)}…` : text
}

export default SubagentControlService