deepseek-harness/packages/ui/jsonrpc/src/server.ts

/**
 * `HarnessSdkServer`: the JSON-RPC method surface the `dsh-jsonrpc` plugin
 * serves to out-of-process SDK clients (e.g. the Python `deepseek_harness`
 * package). Requests: `initialize` → `session/prompt`* → `shutdown`.
 * Notifications pushed to the host: `session.event` (every durable session
 * event, verbatim), `session.finished` (per prompt turn settle),
 * `subagent.started` / `subagent.finished` (child-session lineage and run
 * outcomes). The server owns only the SDK-facing session map — the harness
 * itself is the context the plugin mounts in; plugins, persistence, and
 * the LLM adapter set all come from the external `cordis.yml`.
 *
 * @module @deepseek-ai/dsh-jsonrpc/server
 */

import type { Context } from 'cordis'
import { resolve } from 'node:path'
import type { ContentBlock } from '@deepseek-ai/dsh-llm'
import type { AgentHandle } from '@deepseek-ai/dsh-agent'
import { SessionId, type TurnEndReason } from '@deepseek-ai/dsh-session'
import type { SubagentRunEndInfo, SubagentRunInfo } from '@deepseek-ai/dsh-subagent'
import * as LlmDeepSeek from '@deepseek-ai/dsh-llm-deepseek'
import type { JsonRpcTransportPeer } from './transport.ts'

/** Parameters of the `initialize` request (once per process, before any prompt). */
export interface InitializeParams {
  /** Working directory recorded on every SDK-created session's header. */
  cwd: string
  /** Model name every SDK-created agent runs on (see {@link HarnessSdkServer.initialize} for adapter fallback). */
  model: string
}

/** Result of the `initialize` request: the server's identity for the SDK handshake. */
export interface InitializeResult {
  /** Wire-stable server identity (`deepseek-harness-sdk-runtime`) and version. */
  serverInfo: { name: string; version: string }
}

/**
 * Parameters of a `session/prompt` request: one user turn on one SDK session,
 * with at most one in flight per session.
 */
export interface SessionPromptParams {
  /** The SDK-side session id; an unknown id lazily creates the agent+session pair. */
  sessionId: string
  /** The prompt content blocks, sent verbatim as the user message. */
  contentBlocks: ContentBlock[]
}

/** Result of a `session/prompt` request: the prompt ran to turn settle (outcome rides on `session.finished`). */
export interface SessionPromptResult {
  /** Always `true`; the turn outcome is the paired `session.finished` notification. */
  accepted: true
}

interface SessionRecord {
  handle: AgentHandle
  lastTurnEnd: TurnEndReason | undefined
  activePrompt: boolean
}

/** Runtime-local agent identity plus optional durable fork lineage. */
interface LocalAgentRecord {
  parentSessionId?: SessionId
}

/** Pending local runs that share one provider/id correlation key. */
interface PendingLocalRuns {
  count: number
  parentSessionId?: SessionId
  parentAmbiguous: boolean
}

/**
 * The SDK server over a booted harness context. Constructing it subscribes to
 * session, agent, and subagent lifecycle events, forwarding durable session
 * events and SDK-facing completion notifications while retaining local-run
 * identity across child disposal. The subscriptions live until
 * {@link shutdown}. One instance serves one transport peer for the process
 * lifetime — there is no re-`initialize`.
 */
export class HarnessSdkServer {
  private cwd = process.cwd()
  private model = 'deepseek'
  private llmFiber: { dispose(): Promise<void> } | undefined
  private readonly sessions = new Map<string, SessionRecord>()
  private readonly sessionCreations = new Map<string, Promise<SessionRecord>>()
  private readonly localAgents = new Map<SessionId, LocalAgentRecord>()
  private readonly localRuns = new Map<string, Map<SessionId, PendingLocalRuns>>()
  private readonly disposers: (() => void)[] = []
  private shutdownTask: Promise<Record<string, never>> | undefined
  private shuttingDown = false

  constructor(
    private readonly ctx: Context,
    private readonly transport: JsonRpcTransportPeer,
  ) {
    this.disposers.push(ctx.on('session/event', (session, event) => {
      if (event.type === 'turn/end') {
        const rec = this.sessions.get(String(session.id))
        if (rec) rec.lastTurnEnd = event.data.reason
      }
      this.transport.notify('session.event', { sessionId: String(session.id), event })
    }))
    this.disposers.push(ctx.on('session/created', (session) => {
      const parentSession = session.header.parentSession
      if (parentSession === undefined) return
      this.transport.notify('subagent.started', {
        parentSessionId: String(parentSession),
        childSessionId: String(session.id),
      })
    }))
    // Cache runtime-local identity and optional lineage for each agent lifetime.
    // Parent lineage is not required by the provider contract, so an empty
    // record remains a load-bearing locality marker.
    this.disposers.push(ctx.on('agent/created', (agent) => {
      const parentSessionId = agent.session.header.parentSession
      this.localAgents.set(agent.id, parentSessionId === undefined ? {} : { parentSessionId })
    }))
    this.disposers.push(ctx.on('agent/disposed', (agent) => {
      this.localAgents.delete(agent.id)
    }))
    // Snapshot locality per provider/id run key. A provider may settle one run,
    // continue the same live child in another run, and dispose that child before
    // the later result settles. Counts preserve every completion without
    // assuming settlement order. If id reuse produces disagreeing lineage, the
    // optional parent is omitted until that pending group drains rather than
    // attributed to the wrong completion.
    this.disposers.push(ctx.on('subagent/start', (info: SubagentRunInfo) => {
      const agent = this.ctx.agents.get(info.id)
      const cachedLocalAgent = this.localAgents.get(info.id)
      const localAgent = cachedLocalAgent ?? (agent === undefined
        ? undefined
        : agent.session.header.parentSession === undefined
          ? {}
          : { parentSessionId: agent.session.header.parentSession })
      if (localAgent === undefined) return
      const providerRuns = this.localRuns.get(info.provider) ?? new Map<SessionId, PendingLocalRuns>()
      const pending = providerRuns.get(info.id)
      if (pending === undefined) {
        providerRuns.set(info.id, localAgent.parentSessionId === undefined
          ? { count: 1, parentAmbiguous: false }
          : { count: 1, parentSessionId: localAgent.parentSessionId, parentAmbiguous: false })
      } else {
        pending.count += 1
        if (pending.parentSessionId !== localAgent.parentSessionId) pending.parentAmbiguous = true
      }
      this.localRuns.set(info.provider, providerRuns)
    }))
    this.disposers.push(ctx.on('subagent/end', (info: SubagentRunEndInfo) => {
      const agent = this.ctx.agents.get(info.id)
      const providerRuns = this.localRuns.get(info.provider)
      const pending = providerRuns?.get(info.id)
      if (pending !== undefined) {
        pending.count -= 1
        if (pending.count === 0) providerRuns?.delete(info.id)
        if (providerRuns?.size === 0) this.localRuns.delete(info.provider)
      }
      // This protocol reports LOCAL child sessions. A lineage-bearing child
      // has the session/created-driven start notification above; a parentless
      // local provider still gets its terminal notification. A remote provider
      // has neither a cached creation nor a live local agent and is ignored.
      if (pending === undefined && agent === undefined) return
      const parentSessionId = pending === undefined
        ? agent?.session.header.parentSession
        : pending.parentAmbiguous ? undefined : pending.parentSessionId
      this.transport.notify('subagent.finished', {
        provider: info.provider,
        agentId: String(info.id),
        ...(parentSessionId === undefined ? {} : { parentSessionId: String(parentSessionId) }),
        childSessionId: String(info.id),
        status: info.stopReason === 'completed' ? 'ok' : 'error',
        stopReason: info.stopReason,
        ...(info.lastAssistantMessage === undefined ? {} : { lastAssistantMessage: info.lastAssistantMessage }),
      })
    }))
  }

  /**
   * Handle `initialize`: record the SDK deployment facts (cwd, model) and, when
   * no registered adapter serves `params.model`, mount the DeepSeek adapter for
   * it (credentials from `$DEEPSEEK_API_KEY`/`$DEEPSEEK_BASE_URL`) — a config
   * that already registered an adapter for the model wins.
   * @param params - the SDK handshake parameters.
   * @returns the server identity for the handshake.
   */
  async initialize(params: InitializeParams): Promise<InitializeResult> {
    this.cwd = resolve(params.cwd)
    this.model = params.model
    if (!this.llmFiber && !this.hasAdapterFor(this.model)) {
      this.llmFiber = await this.ctx.plugin(LlmDeepSeek, { models: [this.model] })
    }
    return { serverInfo: { name: 'deepseek-harness-sdk-runtime', version: '0.0.1' } }
  }

  /**
   * Handle `session/prompt`: get-or-create the session's agent, send the
   * content as the user message, await turn settle (quiescence), then notify
   * `session.finished` with the settled turn's outcome. A session accepts at
   * most one prompt at a time; an overlapping request fails immediately while
   * other sessions remain independent.
   * @param params - the target session id and prompt content.
   * @returns `{ accepted: true }` after the turn settled.
   */
  async prompt(params: SessionPromptParams): Promise<SessionPromptResult> {
    const rec = await this.getOrCreateSession(params.sessionId)
    if (rec.activePrompt) throw new Error(`session already has an active prompt: ${params.sessionId}`)
    rec.activePrompt = true
    try {
      rec.lastTurnEnd = undefined
      rec.handle.agent.send(params.contentBlocks)
      await rec.handle.agent.whenIdle()
      const status = this.finishedStatus(rec.lastTurnEnd)
      this.transport.notify('session.finished', {
        sessionId: params.sessionId,
        status,
        reason: rec.lastTurnEnd,
      })
      return { accepted: true }
    } finally {
      rec.activePrompt = false
    }
  }

  /**
   * Handle `shutdown`: dispose every SDK-created agent handle (awaiting loop
   * quiescence), unmount the adapter fiber this server mounted (if any), and
   * detach the event subscriptions. The CONTEXT stays up — the bin disposes it
   * as part of process exit.
   * @returns an empty object (the JSON-RPC result).
   */
  shutdown(): Promise<Record<string, never>> {
    this.shutdownTask ??= this.performShutdown()
    return this.shutdownTask
  }

  private async performShutdown(): Promise<Record<string, never>> {
    this.shuttingDown = true
    const pendingCreations = [...this.sessionCreations.values()]
    await Promise.allSettled(pendingCreations)
    this.sessionCreations.clear()
    const records = [...this.sessions.values()]
    this.sessions.clear()
    this.localAgents.clear()
    this.localRuns.clear()
    const failures: unknown[] = []
    while (this.disposers.length > 0) {
      try {
        this.disposers.pop()?.()
      } catch (error) {
        failures.push(error)
      }
    }
    const teardownResults = await Promise.allSettled([
      ...records.map(rec => Promise.resolve().then(() => rec.handle.dispose())),
      ...(this.llmFiber === undefined ? [] : [Promise.resolve().then(() => this.llmFiber?.dispose())]),
    ])
    this.llmFiber = undefined
    failures.push(...teardownResults
      .filter((result): result is PromiseRejectedResult => result.status === 'rejected')
      .map(result => result.reason as unknown))
    if (failures.length === 1) throw failures[0]
    if (failures.length > 1) throw new AggregateError(failures, 'SDK server teardown failed')
    return {}
  }

  /**
   * Dispatch one incoming JSON-RPC request to its typed handler. Throws (→ a
   * JSON-RPC error response) on an unknown method.
   * @param method - the JSON-RPC method name.
   * @param params - the raw params object from the wire.
   * @returns the handler's result, to be serialized as the response.
   */
  async handleRequest(method: string, params: Record<string, unknown> | undefined): Promise<unknown> {
    switch (method) {
      case 'initialize':
        return this.initialize(params as unknown as InitializeParams)
      case 'session/prompt':
        return this.prompt(params as unknown as SessionPromptParams)
      case 'shutdown':
        return this.shutdown()
      default:
        throw new Error(`unknown DeepSeek Harness SDK runtime method: ${method}`)
    }
  }

  private async getOrCreateSession(sessionId: string): Promise<SessionRecord> {
    if (this.shuttingDown) throw new Error('SDK server is shutting down')
    const existing = this.sessions.get(sessionId)
    if (existing) return existing
    const pending = this.sessionCreations.get(sessionId)
    if (pending) return pending
    const creation = this.createSession(sessionId)
    this.sessionCreations.set(sessionId, creation)
    void creation.then(
      () => { this.sessionCreations.delete(sessionId) },
      () => { this.sessionCreations.delete(sessionId) },
    )
    return creation
  }

  private async createSession(sessionId: string): Promise<SessionRecord> {
    const handle = await this.ctx.agents.create({
      sessionId: SessionId(sessionId),
      meta: { cwd: this.cwd },
      agentOptions: { model: this.model },
    })
    const rec: SessionRecord = { handle, lastTurnEnd: undefined, activePrompt: false }
    this.sessions.set(sessionId, rec)
    return rec
  }

  private finishedStatus(reason: TurnEndReason | undefined): 'ok' | 'error' {
    if (!reason) return 'error'
    return reason.kind === 'completed' ? 'ok' : 'error'
  }

  private hasAdapterFor(model: string): boolean {
    return this.ctx.get('llm')?.models().includes(model) ?? false
  }
}