Просмотр исходного кода

fix: avoid FIFO subagent lineage attribution

Tianyi Cui 2 месяцев назад
Родитель
Сommit
d898d1faf0

+ 1 - 1
packages/ui/jsonrpc/README.md

@@ -4,7 +4,7 @@ The **SDK server plugin** (`jsonrpc`): mounting it serves a stdio JSON-RPC serve
 
 ## Wiring
 
-`inject: ['agents']` — the server creates one agent per SDK `sessionId` (get-or-create on `session/prompt`). A local subagent's shared agent/session id supplies `subagent.finished.childSessionId` directly; the server caches runtime-local identity plus optional parent lineage for the agent lifetime and snapshots it per run, because a continuation may reuse one child and the child may be disposed before a later `subagent/end`. Runs from remote providers are not reported because they create no local agent. The LLM seam is read opportunistically via `ctx.get('llm')` (not injected): when `initialize.model` has no registered adapter, the plugin mounts `dsh-llm-deepseek` for it (credentials from `$DEEPSEEK_API_KEY` / `$DEEPSEEK_BASE_URL`) — a config-registered adapter for the model wins. Everything else — persistence, the tool stacks, the adapter set — comes from the surrounding `cordis.yml`.
+`inject: ['agents']` — the server creates one agent per SDK `sessionId` (get-or-create on `session/prompt`). A local subagent's shared agent/session id supplies `subagent.finished.childSessionId` directly; the server caches runtime-local identity plus optional parent lineage for the agent lifetime and counts pending runs per provider/id, because a continuation may reuse one child and the child may be disposed before a later `subagent/end`. Settlement order is not assumed: if concurrent ID reuse makes parent lineage ambiguous, the completion remains local but omits the optional `parentSessionId` rather than attributing the wrong parent. Runs from remote providers are not reported because they create no local agent. The LLM seam is read opportunistically via `ctx.get('llm')` (not injected): when `initialize.model` has no registered adapter, the plugin mounts `dsh-llm-deepseek` for it (credentials from `$DEEPSEEK_API_KEY` / `$DEEPSEEK_BASE_URL`) — a config-registered adapter for the model wins. Everything else — persistence, the tool stacks, the adapter set — comes from the surrounding `cordis.yml`.
 
 ## Config
 

+ 36 - 13
packages/ui/jsonrpc/src/server.ts

@@ -63,6 +63,13 @@ 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
@@ -78,7 +85,7 @@ export class HarnessSdkServer {
   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<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
@@ -112,10 +119,12 @@ export class HarnessSdkServer {
     this.disposers.push(ctx.on('agent/disposed', (agent) => {
       this.localAgents.delete(agent.id)
     }))
-    // Snapshot locality per run. A provider may settle one run, continue the
-    // same live child in another run, and dispose that child before the later
-    // result settles. Consuming an agent-lifetime marker at the first end would
-    // lose the later notification; this queue pairs each start with one end.
+    // 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)
@@ -125,21 +134,35 @@ export class HarnessSdkServer {
           ? {}
           : { parentSessionId: agent.session.header.parentSession })
       if (localAgent === undefined) return
-      const runs = this.localRuns.get(info.id) ?? []
-      runs.push(localAgent)
-      this.localRuns.set(info.id, runs)
+      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 runs = this.localRuns.get(info.id)
-      const cachedLocalAgent = runs?.shift()
-      if (runs?.length === 0) this.localRuns.delete(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 (cachedLocalAgent === undefined && agent === undefined) return
-      const parentSessionId = cachedLocalAgent?.parentSessionId ?? agent?.session.header.parentSession
+      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),

+ 121 - 2
packages/ui/jsonrpc/tests/server.spec.ts

@@ -373,6 +373,100 @@ describe('HarnessSdkServer', () => {
     }
   })
 
+  it('omits ambiguous lineage when one local id is reused and runs settle out of order', async () => {
+    const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-subagent-reuse-'))
+    const ctx = await makeHarness(storageDir)
+    try {
+      const transport = new FakeTransport()
+      const server = new HarnessSdkServer(ctx, transport)
+      const oldParent = await ctx.agents.create({
+        sessionId: SessionId('old-parent'),
+        meta: { cwd: storageDir },
+        agentOptions: { model: 'deepseek' },
+      })
+      const oldChild = await ctx.agents.create({
+        sessionId: SessionId('reused-child'),
+        meta: { cwd: storageDir, parentSession: SessionId('old-parent') },
+        agentOptions: { model: 'deepseek' },
+      })
+      const first = Promise.withResolvers<SubagentResult>()
+      const sameLifetime = Promise.withResolvers<SubagentResult>()
+      const replacement = Promise.withResolvers<SubagentResult>()
+      const results = [first.promise, sameLifetime.promise, replacement.promise]
+      let starts = 0
+      const disposeProvider = ctx.subagents.registerProvider({
+        name: 'reused',
+        capabilities: { outputSchema: false, depthLimit: false, toolFilter: false, persona: false },
+        inheritsParentContext: false,
+        start() {
+          const result = results[starts]
+          starts += 1
+          if (result === undefined) throw new Error('unexpected fourth reused-id run')
+          return Promise.resolve({ id: SessionId('reused-child'), result, dispose: () => Promise.resolve() })
+        },
+      })
+
+      const firstRun = await ctx.subagents.start('reused', {
+        parent: oldParent.agent,
+        prompt: [],
+        signal: new AbortController().signal,
+      })
+      const sameLifetimeRun = await ctx.subagents.start('reused', {
+        parent: oldParent.agent,
+        prompt: [],
+        signal: new AbortController().signal,
+      })
+      sameLifetime.resolve({ output: [{ type: 'text', text: 'same lifetime' }], stopReason: 'completed' })
+      await sameLifetimeRun.result
+      await oldChild.dispose()
+      const newParent = await ctx.agents.create({
+        sessionId: SessionId('new-parent'),
+        meta: { cwd: storageDir },
+        agentOptions: { model: 'deepseek' },
+      })
+      const newChild = await ctx.agents.create({
+        sessionId: SessionId('reused-child'),
+        meta: { cwd: storageDir, parentSession: SessionId('new-parent') },
+        agentOptions: { model: 'deepseek' },
+      })
+      const secondRun = await ctx.subagents.start('reused', {
+        parent: newParent.agent,
+        prompt: [],
+        signal: new AbortController().signal,
+      })
+
+      replacement.resolve({ output: [{ type: 'text', text: 'new lifetime' }], stopReason: 'completed' })
+      await secondRun.result
+      first.resolve({ output: [{ type: 'text', text: 'old lifetime' }], stopReason: 'completed' })
+      await firstRun.result
+      await Promise.resolve()
+
+      const finished = transport.notifications.filter(notification =>
+        notification.method === 'subagent.finished'
+        && notification.params?.childSessionId === 'reused-child',
+      )
+      expect(finished.map(notification => notification.params?.lastAssistantMessage)).toEqual([
+        [{ type: 'text', text: 'same lifetime' }],
+        [{ type: 'text', text: 'new lifetime' }],
+        [{ type: 'text', text: 'old lifetime' }],
+      ])
+      expect(finished[0]?.params?.parentSessionId).toBe('old-parent')
+      expect(finished.slice(1).every(notification => !Object.hasOwn(notification.params ?? {}, 'parentSessionId'))).toBe(true)
+
+      await firstRun.dispose()
+      await sameLifetimeRun.dispose()
+      await secondRun.dispose()
+      disposeProvider()
+      await newChild.dispose()
+      await oldParent.dispose()
+      await newParent.dispose()
+      await server.shutdown()
+    } finally {
+      await ctx.fiber.dispose()
+      await rm(storageDir, { recursive: true, force: true })
+    }
+  })
+
   it('falls back to live lineage and ignores runs without a local child session', async () => {
     const storageDir = await mkdtemp(join(tmpdir(), 'dsh-jsonrpc-subagent-fallback-'))
     const ctx = await makeHarness(storageDir)
@@ -395,13 +489,38 @@ describe('HarnessSdkServer', () => {
         meta: { cwd: storageDir },
         agentOptions: { model: 'deepseek' },
       })
+      const missedStartResult = Promise.withResolvers<SubagentResult>()
+      const disposeMissedStartProvider = ctx.subagents.registerProvider({
+        name: 'fork',
+        capabilities: { outputSchema: false, depthLimit: false, toolFilter: false, persona: false },
+        inheritsParentContext: true,
+        start: () => Promise.resolve({
+          id: SessionId('fallback-child-session'),
+          result: missedStartResult.promise,
+          dispose: () => Promise.resolve(),
+        }),
+      })
+      // Start before the server subscribes, so the terminal fallback must use
+      // the still-live registry entry rather than a cached start record.
+      const missedStartRun = await ctx.subagents.start('fork', {
+        parent: parentHandle.agent,
+        prompt: [],
+        signal: new AbortController().signal,
+      })
       const transport = new FakeTransport()
       const server = new HarnessSdkServer(ctx, transport)
 
+      missedStartResult.resolve({ output: [], stopReason: 'max-tokens' })
+      await missedStartRun.result
+      await Promise.resolve()
+      await missedStartRun.dispose()
+      disposeMissedStartProvider()
+      // The server also missed this agent's creation, but observes the start;
+      // recover its lineage from the still-live registry entry.
       await settleSubagent(ctx, parentHandle.agent, {
-        provider: 'fork',
+        provider: 'fork-live-fallback',
         id: SessionId('fallback-child-session'),
-        stopReason: 'max-tokens',
+        stopReason: 'completed',
         lastAssistantMessage: [],
       })
       await settleSubagent(ctx, parentHandle.agent, {