fix(multiagent): buffer graph node printer output to prevent interleaving (strands-agents#3077)

Author: chaynabors

strands-ts/src/multiagent/__tests__/graph.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it, vi } from 'vitest'
 import { Agent } from '../../agent/agent.js'
+import { AgentPrinter } from '../../agent/printer.js'
 import { MockMessageModel } from '../../__fixtures__/mock-message-model.js'
 import { MockSnapshotStorage } from '../../__fixtures__/mock-storage-provider.js'
 import { collectGenerator } from '../../__fixtures__/model-test-helpers.js'
@@ -716,6 +717,106 @@ describe('Graph', () => {
       )
     })
 
+    it('does not buffer agent printer output for sequential graphs', async () => {
+      const a = makeAgent('a', 'a-reply')
+      const b = makeAgent('b', 'b-reply')
+      const printerA = new AgentPrinter(() => {})
+      const printerB = new AgentPrinter(() => {})
+      ;(a as unknown as { _printer: AgentPrinter })._printer = printerA
+      ;(b as unknown as { _printer: AgentPrinter })._printer = printerB
+
+      const graph = new Graph({
+        nodes: [a, b],
+        edges: [['a', 'b']],
+      })
+
+      const seenA: unknown[] = []
+      const seenB: unknown[] = []
+      const gen = graph.stream('go')
+      for await (const event of gen) {
+        if (event.type === 'nodeStreamUpdateEvent') {
+          if (event.nodeId === 'a') seenA.push((a as unknown as { _printer: AgentPrinter })._printer)
+          if (event.nodeId === 'b') seenB.push((b as unknown as { _printer: AgentPrinter })._printer)
+        }
+      }
+
+      expect(seenA.length).toBeGreaterThan(0)
+      expect(seenB.length).toBeGreaterThan(0)
+      expect(seenA.every((p) => p === printerA)).toBe(true)
+      expect(seenB.every((p) => p === printerB)).toBe(true)
+    })
+
+    it('buffers agent printer output when topology allows parallel execution', async () => {
+      const a = makeAgent('a', 'a-reply')
+      const b = makeAgent('b', 'b-reply')
+      const printerA = new AgentPrinter(() => {})
+      const printerB = new AgentPrinter(() => {})
+      ;(a as unknown as { _printer: AgentPrinter })._printer = printerA
+      ;(b as unknown as { _printer: AgentPrinter })._printer = printerB
+
+      // Two sources with no edges between them — both can run concurrently.
+      const graph = new Graph({
+        nodes: [a, b],
+        edges: [],
+      })
+
+      const seenA: unknown[] = []
+      const seenB: unknown[] = []
+      const gen = graph.stream('go')
+      for await (const event of gen) {
+        if (event.type === 'nodeStreamUpdateEvent') {
+          if (event.nodeId === 'a') seenA.push((a as unknown as { _printer: AgentPrinter })._printer)
+          if (event.nodeId === 'b') seenB.push((b as unknown as { _printer: AgentPrinter })._printer)
+        }
+      }
+
+      expect(seenA.some((p) => p !== printerA)).toBe(true)
+      expect(seenB.some((p) => p !== printerB)).toBe(true)
+      expect((a as unknown as { _printer: AgentPrinter })._printer).toBe(printerA)
+      expect((b as unknown as { _printer: AgentPrinter })._printer).toBe(printerB)
+    })
+
+    it('does not interleave output from parallel nodes on a shared writer', async () => {
+      // Fan-out graph (A -> B, A -> C). B and C run in parallel after A completes.
+      // Each emits several TextBlocks so their printer writes are naturally
+      // interleavable; the buffering behavior guarantees each node's writes
+      // land as a contiguous run on the shared writer.
+      const a = makeAgent('a', 'a-reply')
+      const b = new Agent({
+        model: new MockMessageModel().addTurn(['b-1 ', 'b-2 ', 'b-3'].map((text) => new TextBlock(text))),
+        printer: false,
+        id: 'b',
+      })
+      const c = new Agent({
+        model: new MockMessageModel().addTurn(['c-1 ', 'c-2 ', 'c-3'].map((text) => new TextBlock(text))),
+        printer: false,
+        id: 'c',
+      })
+
+      const writes: string[] = []
+      ;(b as unknown as { _printer: AgentPrinter })._printer = new AgentPrinter((text) => writes.push('b:' + text))
+      ;(c as unknown as { _printer: AgentPrinter })._printer = new AgentPrinter((text) => writes.push('c:' + text))
+
+      const graph = new Graph({
+        nodes: [a, b, c],
+        edges: [
+          ['a', 'b'],
+          ['a', 'c'],
+        ],
+      })
+
+      await graph.invoke('go')
+
+      // Each node's writes form one contiguous slice — no other node's writes appear between them.
+      const bIndices = writes.map((w, i) => (w.startsWith('b:') ? i : -1)).filter((i) => i !== -1)
+      const cIndices = writes.map((w, i) => (w.startsWith('c:') ? i : -1)).filter((i) => i !== -1)
+      expect(bIndices.length).toBeGreaterThan(0)
+      expect(cIndices.length).toBeGreaterThan(0)
+      const isContiguous = (indices: number[]): boolean => indices.every((v, i) => i === 0 || v === indices[i - 1]! + 1)
+      expect(isContiguous(bIndices)).toBe(true)
+      expect(isContiguous(cIndices)).toBe(true)
+    })
+
     it('cleans up running nodes when consumer breaks mid-stream', async () => {
       const graph = new Graph({
         nodes: [makeAgent('a', 'a-reply'), makeAgent('b', 'b-reply')],

strands-ts/src/multiagent/__tests__/nodes.test.ts

@@ -1,6 +1,7 @@
 import { beforeEach, describe, expect, it } from 'vitest'
 import { z } from 'zod'
 import { Agent } from '../../agent/agent.js'
+import { AgentPrinter } from '../../agent/printer.js'
 import { BeforeInvocationEvent } from '../../hooks/events.js'
 import type { MultiAgentInput } from '../multiagent.js'
 import { MockMessageModel } from '../../__fixtures__/mock-message-model.js'
@@ -233,6 +234,40 @@ describe('AgentNode', () => {
       expect(preserveContextAgent.appState.get<{ count: number }>('count')).toBe(2)
     })
 
+    it('leaves the original printer in place when bufferOutput is not set', async () => {
+      const writes: string[] = []
+      const printer = new AgentPrinter((text) => writes.push(text))
+      ;(agent as unknown as { _printer: AgentPrinter })._printer = printer
+
+      // Capture the active printer at every stream event — confirms it never gets
+      // swapped for a buffering printer mid-stream on the sequential path.
+      const seenPrinters: unknown[] = []
+      const gen = node.stream([new TextBlock('prompt')], state)
+      for await (const _event of gen) {
+        seenPrinters.push((agent as unknown as { _printer: AgentPrinter })._printer)
+      }
+
+      expect(seenPrinters.every((p) => p === printer)).toBe(true)
+      expect(writes.join('')).toContain('reply')
+    })
+
+    it('buffers printer output when bufferOutput is true and flushes on completion', async () => {
+      const writes: string[] = []
+      const printer = new AgentPrinter((text) => writes.push(text))
+      ;(agent as unknown as { _printer: AgentPrinter })._printer = printer
+
+      // While streaming, a buffering printer is installed in place of the original.
+      const printerDuringStream: unknown[] = []
+      const gen = node.stream([new TextBlock('prompt')], state, { bufferOutput: true })
+      for await (const _event of gen) {
+        printerDuringStream.push((agent as unknown as { _printer: AgentPrinter })._printer)
+      }
+
+      expect(printerDuringStream.some((p) => p !== printer)).toBe(true)
+      expect((agent as unknown as { _printer: AgentPrinter })._printer).toBe(printer)
+      expect(writes.join('')).toContain('reply')
+    })
+
     it('passes structuredOutputSchema from options to the agent', async () => {
       const schema = z.object({ agentName: z.string().optional(), message: z.string() })
 

strands-ts/src/multiagent/graph.ts

@@ -136,6 +136,8 @@ export class Graph implements MultiAgent {
   private readonly _pluginRegistry: MultiAgentPluginRegistry
   private readonly _hookRegistry: HookRegistryImplementation
   private readonly _sources: Node[]
+  /** Whether the topology can ever run two nodes in parallel. Gates per-node printer buffering. */
+  private readonly _canRunConcurrently: boolean
   private readonly _tracer: Tracer
   readonly sessionManager?: SessionManager | undefined
   private _initialized: boolean
@@ -168,6 +170,7 @@ export class Graph implements MultiAgent {
     this.edges = this._resolveEdges(edges)
     this._sources = this._resolveSources(sources)
     this._validateSources()
+    this._canRunConcurrently = this._computeCanRunConcurrently()
 
     this.sessionManager = sessionManager
 
@@ -525,7 +528,11 @@ export class Graph implements MultiAgent {
 
     try {
       const gen = this._tracer.withSpanContext(nodeSpan, () =>
-        node.stream(input, state, { invocationState, ...(cancelSignal && { cancelSignal }) })
+        node.stream(input, state, {
+          invocationState,
+          ...(cancelSignal && { cancelSignal }),
+          ...(this._canRunConcurrently && { bufferOutput: true }),
+        })
       )
       let next = await this._tracer.withSpanContext(nodeSpan, () => gen.next())
       while (!next.done) {
@@ -682,6 +689,18 @@ export class Graph implements MultiAgent {
     return [...this.nodes.values()].filter((node) => !targetIds.has(node.id))
   }
 
+  /**
+   * Static check: can two nodes ever run in parallel? False when `maxConcurrency` is 1,
+   * there's a single source, and no node has more than one outgoing edge.
+   */
+  private _computeCanRunConcurrently(): boolean {
+    if (this.config.maxConcurrency < 2) return false
+    if (this._sources.length > 1) return true
+    const fanOut = new Map<string, number>()
+    for (const e of this.edges) fanOut.set(e.source.id, (fanOut.get(e.source.id) ?? 0) + 1)
+    return [...fanOut.values()].some((count) => count > 1)
+  }
+
   /**
    * Identifies terminus nodes and returns their combined content.
    * A terminus node is where an execution path ended: completed with no

strands-ts/src/multiagent/nodes.ts

@@ -1,4 +1,5 @@
 import { Agent } from '../agent/agent.js'
+import { AgentPrinter, type Printer } from '../agent/printer.js'
 import type { InvocationState, InvokeOptions, InvokableAgent, AgentStreamEvent } from '../types/agent.js'
 import type { MultiAgentInput } from './multiagent.js'
 import { dropStaleInterruptedResult } from './multiagent.js'
@@ -48,6 +49,12 @@ export interface NodeInputOptions {
    * orchestrators to enforce per-node timeouts or propagate external cancellation.
    */
   cancelSignal?: AbortSignal
+
+  /**
+   * Buffer the agent's printer output and flush it on completion so concurrent siblings
+   * don't interleave on stdout.
+   */
+  bufferOutput?: boolean
 }
 
 /**
@@ -271,6 +278,17 @@ export class AgentNode extends Node {
       this._agent.loadSnapshot(nodeState.interruptedSnapshot)
     }
 
+    // Swap the agent's printer for a buffer so concurrent siblings don't interleave on stdout.
+    const agentInternals =
+      options?.bufferOutput && isAgent ? (this._agent as unknown as { _printer?: Printer }) : undefined
+    const originalPrinter = agentInternals?._printer
+    let buffered = ''
+    if (agentInternals && originalPrinter) {
+      agentInternals._printer = new AgentPrinter((text) => {
+        buffered += text
+      })
+    }
+
     try {
       const invokeOptions: InvokeOptions = {
         ...(options?.structuredOutputSchema && { structuredOutputSchema: options.structuredOutputSchema }),
@@ -314,6 +332,10 @@ export class AgentNode extends Node {
       if (preRunSnapshot) {
         ;(this._agent as Agent).loadSnapshot(preRunSnapshot)
       }
+      if (agentInternals && originalPrinter) {
+        agentInternals._printer = originalPrinter
+        if (buffered.length > 0) originalPrinter.write(buffered)
+      }
     }
   }
 }