fix(sdk): harden runtime lifecycle and JSON-RPC

Keep DeepSeekHarness.run() reusable, but make ownership of its lazy
runtime process explicit. Document the context-manager/close contract and
update every construction example to use a context manager so repeated runs
remain valid without encouraging leaked subprocesses.

Contain notification predicate failures at the subscription boundary. Remove
only the subscriber whose callback raised, deliver that exception through its
queue, and continue dispatching to healthy subscribers so arbitrary callback
code cannot terminate the shared reader thread or strand later requests.

Enforce one in-flight prompt per server session with an atomic activePrompt
guard. Route overlap through the existing -32603 handler-error response and
clear the guard in finally, preserving parallel prompts across sessions and
sequential reuse without changing JSON-RPC request or notification shapes.

Use StringDecoder for line framing so a UTF-8 code point split across Buffer
chunks is not corrupted. Add a queued-write flush barrier, and make memoized
shutdown await it before disposal and exit while retaining exactly-once
cleanup when shutdown calls race or flushing fails.

Cover callback isolation, same-session exclusion, cross-session concurrency,
split multibyte input, delayed writes, racing shutdown, and flush failure with
deterministic tests.

packages/ui/jsonrpc/tests/plugin-apply.spec.ts

@@ -28,6 +28,7 @@ import * as jsonrpc from '../src/index.ts'
 /** One ordered observation on the plugin's outward-facing seams: a JSON-RPC frame written to `output`, or an `exit(code)` call. */
 type WireEvent =
   | { kind: 'frame'; frame: Record<string, unknown> }
+  | { kind: 'write-complete'; ids: (string | number)[] }
   | { kind: 'exit'; code: number }
 
 interface ApplyHarness {
@@ -36,6 +37,7 @@ interface ApplyHarness {
   fiber: Awaited<ReturnType<Context['plugin']>>
   /** Every output frame and exit call, in observation order — ordering assertions read this. */
   events: WireEvent[]
+  outputErrors: Error[]
   send(frame: Record<string, unknown>): void
   sendRaw(text: string): void
   frames(): Record<string, unknown>[]
@@ -65,7 +67,10 @@ async function settle(): Promise<void> {
  * server.spec recipe) and mount the jsonrpc plugin on it through the real
  * namespace mount path, with in-memory seams standing in for stdio/exit.
  */
-async function mountPlugin(storageDir: string): Promise<ApplyHarness> {
+async function mountPlugin(
+  storageDir: string,
+  options: { writeDelayMs?: number; failFlush?: boolean } = {},
+): Promise<ApplyHarness> {
   const ctx = new Context()
   await ctx.plugin(agentCore)
   await ctx.plugin(SessionPersistenceJsonl, { root: storageDir })
@@ -73,22 +78,39 @@ async function mountPlugin(storageDir: string): Promise<ApplyHarness> {
 
   const input = new PassThrough()
   const events: WireEvent[] = []
+  const outputErrors: Error[] = []
   let pendingOutput = ''
-  // A hand-rolled Writable (not a PassThrough): _write records each decoded
-  // frame synchronously, so `events` preserves the true frame-vs-exit order.
+  // A hand-rolled Writable (not a PassThrough): _write records frames on
+  // admission and write-complete only when its callback fires, so a delayed
+  // output proves exit waits for the transport's flush barrier.
   const output = new Writable({
     write(chunk: Buffer, _encoding, callback) {
+      const ids: (string | number)[] = []
       pendingOutput += chunk.toString('utf8')
       for (;;) {
         const newline = pendingOutput.indexOf('\n')
         if (newline < 0) break
         const line = pendingOutput.slice(0, newline).trim()
         pendingOutput = pendingOutput.slice(newline + 1)
-        if (line) events.push({ kind: 'frame', frame: JSON.parse(line) as Record<string, unknown> })
+        if (line) {
+          const frame = JSON.parse(line) as Record<string, unknown>
+          events.push({ kind: 'frame', frame })
+          if (typeof frame.id === 'string' || typeof frame.id === 'number') ids.push(frame.id)
+        }
       }
-      callback()
+      const complete = (): void => {
+        if (options.failFlush === true && chunk.length === 0) {
+          callback(new Error('flush callback failed'))
+          return
+        }
+        events.push({ kind: 'write-complete', ids })
+        callback()
+      }
+      if ((options.writeDelayMs ?? 0) > 0) setTimeout(complete, options.writeDelayMs)
+      else complete()
     },
   })
+  output.on('error', (error: Error) => { outputErrors.push(error) })
   const exit = (code: number): void => { events.push({ kind: 'exit', code }) }
 
   const fiber = await ctx.plugin(jsonrpc, { input, output, exit })
@@ -99,6 +121,7 @@ async function mountPlugin(storageDir: string): Promise<ApplyHarness> {
     ctx,
     fiber,
     events,
+    outputErrors,
     send: (frame) => { input.write(`${JSON.stringify(frame)}\n`) },
     sendRaw: (text) => { input.write(text) },
     frames,
@@ -199,7 +222,7 @@ describe('dsh-jsonrpc plugin apply', () => {
 
   it('answers shutdown before exiting 0 exactly once, even against a racing second shutdown', async () => {
     const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-apply-shutdown-'))
-    const harness = await mountPlugin(storageDir)
+    const harness = await mountPlugin(storageDir, { writeDelayMs: 10 })
     try {
       // Two shutdown frames in ONE chunk: both are dispatched from the same
       // read-loop pass, so both setImmediate exit callbacks get scheduled and
@@ -211,15 +234,21 @@ describe('dsh-jsonrpc plugin apply', () => {
       await waitFor(() => harness.exits().length > 0 ? true : undefined, 'exit recorder call')
       expect(harness.exits()).toEqual([0])
 
-      // Response-then-exit ordering: both shutdown responses were flushed to
-      // output BEFORE exit(0) ran (the setImmediate in the request handler).
+      // Response-then-exit ordering: both response write callbacks and the
+      // empty flush barrier complete before exit(0), even on delayed output.
       const exitIndex = harness.events.findIndex(event => event.kind === 'exit')
       const firstResponse = harness.events.findIndex(event => event.kind === 'frame' && event.frame.id === 'sd-1')
       const secondResponse = harness.events.findIndex(event => event.kind === 'frame' && event.frame.id === 'sd-2')
+      const firstComplete = harness.events.findIndex(event => event.kind === 'write-complete' && event.ids.includes('sd-1'))
+      const secondComplete = harness.events.findIndex(event => event.kind === 'write-complete' && event.ids.includes('sd-2'))
+      const flushComplete = harness.events.findIndex(event => event.kind === 'write-complete' && event.ids.length === 0)
       expect(firstResponse).toBeGreaterThanOrEqual(0)
       expect(secondResponse).toBeGreaterThanOrEqual(0)
-      expect(exitIndex).toBeGreaterThan(firstResponse)
-      expect(exitIndex).toBeGreaterThan(secondResponse)
+      expect(firstComplete).toBeGreaterThan(firstResponse)
+      expect(secondComplete).toBeGreaterThan(secondResponse)
+      expect(flushComplete).toBeGreaterThan(firstComplete)
+      expect(flushComplete).toBeGreaterThan(secondComplete)
+      expect(exitIndex).toBeGreaterThan(flushComplete)
 
       // Idempotent: the racing second shutdown never produces a second exit.
       await settle()
@@ -236,6 +265,27 @@ describe('dsh-jsonrpc plugin apply', () => {
     }
   })
 
+  it('still disposes and exits once when the flush callback fails', async () => {
+    const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-apply-flush-failure-'))
+    const harness = await mountPlugin(storageDir, { failFlush: true })
+    try {
+      harness.send({ jsonrpc: '2.0', id: 'sd-fail', method: 'shutdown' })
+
+      await waitFor(() => harness.exits().length > 0 ? true : undefined, 'exit after flush failure')
+      await settle()
+      expect(harness.exits()).toEqual([0])
+      expect(harness.outputErrors.map(error => error.message)).toEqual(['flush callback failed'])
+
+      const before = harness.frames().length
+      harness.send({ jsonrpc: '2.0', id: 'after-flush-failure', method: 'initialize', params: { cwd: storageDir, model: 'x' } })
+      await settle()
+      expect(harness.frames().length).toBe(before)
+    } finally {
+      await harness.dispose()
+      await rm(storageDir, { recursive: true, force: true })
+    }
+  })
+
   it('stops serving on a bare fiber dispose (HMR-style unload) without calling exit', async () => {
     const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-apply-dispose-'))
     const harness = await mountPlugin(storageDir)

packages/ui/jsonrpc/tests/server.spec.ts

@@ -152,6 +152,57 @@ describe('HarnessSdkServer', () => {
     }
   })
 
+  it('rejects overlapping prompts for one session without serializing other sessions', async () => {
+    let releaseMain: (() => void) | undefined
+    const firstMainIdle = new Promise<void>((resolve) => { releaseMain = resolve })
+    const mainWhenIdle = vi.fn<() => Promise<void>>()
+      .mockReturnValueOnce(firstMainIdle)
+      .mockResolvedValue(undefined)
+    const mainSend = vi.fn()
+    const mainAgent = {
+      send: mainSend,
+      whenIdle: mainWhenIdle,
+    } as unknown as Agent
+    const otherSend = vi.fn()
+    const otherAgent = {
+      send: otherSend,
+      whenIdle: vi.fn(() => Promise.resolve()),
+    } as unknown as Agent
+    const mainHandle = { agent: mainAgent, dispose: vi.fn(() => Promise.resolve()) }
+    const otherHandle = { agent: otherAgent, dispose: vi.fn(() => Promise.resolve()) }
+    const create = vi.fn(async (options: { agentId: AgentId }) =>
+      String(options.agentId) === 'main' ? mainHandle : otherHandle)
+    const ctx = {
+      on: vi.fn(() => () => undefined),
+      agents: { create, get: () => undefined },
+      get: () => undefined,
+    } as unknown as Context
+    const server = new HarnessSdkServer(ctx, new FakeTransport())
+    const prompt = (sessionId: string, text: string) => server.prompt({
+      sessionId,
+      contentBlocks: [{ type: 'text', text }],
+    })
+
+    const first = prompt('main', 'first')
+    await vi.waitFor(() => { expect(mainSend).toHaveBeenCalledOnce() })
+
+    await expect(prompt('main', 'overlap')).rejects.toThrow('session already has an active prompt: main')
+    await expect(prompt('other', 'independent')).resolves.toEqual({ accepted: true })
+    releaseMain?.()
+    await expect(first).resolves.toEqual({ accepted: true })
+    await expect(prompt('main', 'sequential')).resolves.toEqual({ accepted: true })
+
+    mainWhenIdle.mockRejectedValueOnce(new Error('turn wait failed'))
+    await expect(prompt('main', 'failing')).rejects.toThrow('turn wait failed')
+    await expect(prompt('main', 'after failure')).resolves.toEqual({ accepted: true })
+
+    expect(mainSend).toHaveBeenCalledTimes(4)
+    expect(otherSend).toHaveBeenCalledOnce()
+    await server.shutdown()
+    expect(mainHandle.dispose).toHaveBeenCalledOnce()
+    expect(otherHandle.dispose).toHaveBeenCalledOnce()
+  })
+
   it('notifies the host when a child session is created with parent lineage', async () => {
     const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-subagent-'))
     const ctx = await makeHarness(storageDir)
@@ -490,11 +541,11 @@ describe('HarnessSdkServer', () => {
       get: () => undefined,
     } as unknown as Context
     const server = new HarnessSdkServer(ctx, new FakeTransport()) as unknown as {
-      sessions: Map<string, { handle: AgentHandle; lastTurnEnd: undefined }>
+      sessions: Map<string, { handle: AgentHandle; lastTurnEnd: undefined; activePrompt: boolean }>
       shutdown(): Promise<Record<string, never>>
     }
-    server.sessions.set('first', { handle: { agent: {} as Agent, dispose: firstDispose }, lastTurnEnd: undefined })
-    server.sessions.set('second', { handle: { agent: {} as Agent, dispose: secondDispose }, lastTurnEnd: undefined })
+    server.sessions.set('first', { handle: { agent: {} as Agent, dispose: firstDispose }, lastTurnEnd: undefined, activePrompt: false })
+    server.sessions.set('second', { handle: { agent: {} as Agent, dispose: secondDispose }, lastTurnEnd: undefined, activePrompt: false })
 
     await expect(server.shutdown()).rejects.toThrow('SDK server teardown failed')
     expect(firstDispose).toHaveBeenCalledOnce()

packages/ui/jsonrpc/tests/transport.spec.ts

@@ -1,5 +1,5 @@
 import { once } from 'node:events'
-import { PassThrough } from 'node:stream'
+import { PassThrough, Writable } from 'node:stream'
 import { describe, expect, it } from 'vitest'
 import { JsonRpcLineTransport } from '../src/index.ts'
 
@@ -122,6 +122,64 @@ describe('JsonRpcLineTransport', () => {
     b.close()
   })
 
+  it('preserves multibyte UTF-8 characters split across Buffer chunks', async () => {
+    const input = new PassThrough()
+    const output = new PassThrough()
+    const transport = new JsonRpcLineTransport(input, output)
+    const notifications: Record<string, unknown>[] = []
+    transport.onNotification((method, params) => { notifications.push({ method, params }) })
+    transport.start()
+
+    const frame = Buffer.from(`${JSON.stringify({ jsonrpc: '2.0', method: 'message', params: { text: '你好' } })}\n`)
+    const character = Buffer.from('你')
+    const characterStart = frame.indexOf(character)
+    expect(characterStart).toBeGreaterThanOrEqual(0)
+    input.write(frame.subarray(0, characterStart + 1))
+    input.write(frame.subarray(characterStart + 1))
+    await new Promise(resolve => setTimeout(resolve, 10))
+
+    expect(notifications).toEqual([{ method: 'message', params: { text: '你好' } }])
+    transport.close()
+  })
+
+  it('flush waits for all earlier output writes', async () => {
+    const events: string[] = []
+    const output = new Writable({
+      write(chunk: Buffer, _encoding, callback) {
+        const label = chunk.length === 0 ? 'barrier' : 'frame'
+        events.push(`start:${label}`)
+        setTimeout(() => {
+          events.push(`finish:${label}`)
+          callback()
+        }, 5)
+      },
+    })
+    const transport = new JsonRpcLineTransport(new PassThrough(), output)
+
+    transport.notify('tick')
+    await transport.flush()
+
+    expect(events).toEqual([
+      'start:frame',
+      'finish:frame',
+      'start:barrier',
+      'finish:barrier',
+    ])
+    transport.close()
+  })
+
+  it('reports an output callback failure from flush', async () => {
+    const output = {
+      write(_chunk: string, callback?: (error?: Error) => void) {
+        callback?.(new Error('flush failed'))
+        return true
+      },
+    }
+    const transport = new JsonRpcLineTransport(new PassThrough(), output as never)
+
+    await expect(transport.flush()).rejects.toThrow('flush failed')
+  })
+
   it('rejects pending requests when the input closes', async () => {
     const { aToB, b } = transportPair()
     b.start()