fix(streaming): reset stale output drilldowns

Author: TheodoreSpeaks

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/chat/components/output-select/output-select.test.tsx

@@ -3,7 +3,11 @@
  */
 import { act, type ReactNode } from 'react'
 import { createRoot, type Root } from 'react-dom/client'
-import { afterEach, describe, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { outputMenuState } = vi.hoisted(() => ({
+  outputMenuState: { includeNestedWorkflow: true },
+}))
 
 vi.mock('@sim/emcn', () => ({
   cn: (...values: unknown[]) => values.flat().filter(Boolean).join(' '),
@@ -134,31 +138,37 @@ vi.mock('@/lib/workflows/streaming/nested-output-options', () => {
 
   return {
     collectReferencedWorkflowIds: () => [],
-    buildWorkflowOutputOptions: () => [rootOutput, nestedOutput],
-    buildWorkflowOutputMenu: () => [
-      {
+    buildWorkflowOutputOptions: () =>
+      outputMenuState.includeNestedWorkflow ? [rootOutput, nestedOutput] : [rootOutput],
+    buildWorkflowOutputMenu: () => {
+      const rootNode = {
         blockId: 'summary',
         blockName: 'Summarizer',
         blockType: 'agent',
         outputs: [rootOutput],
         children: [],
-      },
-      {
-        blockId: 'workflow',
-        blockName: 'Research',
-        blockType: 'workflow_input',
-        outputs: [],
-        children: [
-          {
-            blockId: 'workflow/agent',
-            blockName: 'Writer',
-            blockType: 'agent',
-            outputs: [nestedOutput],
-            children: [],
-          },
-        ],
-      },
-    ],
+      }
+      return outputMenuState.includeNestedWorkflow
+        ? [
+            rootNode,
+            {
+              blockId: 'workflow',
+              blockName: 'Research',
+              blockType: 'workflow_input',
+              outputs: [],
+              children: [
+                {
+                  blockId: 'workflow/agent',
+                  blockName: 'Writer',
+                  blockType: 'agent',
+                  outputs: [nestedOutput],
+                  children: [],
+                },
+              ],
+            },
+          ]
+        : [rootNode]
+    },
   }
 })
 
@@ -167,6 +177,10 @@ import { OutputSelect } from '@/app/workspace/[workspaceId]/w/[workflowId]/compo
 let root: Root | null = null
 let container: HTMLDivElement | null = null
 
+beforeEach(() => {
+  outputMenuState.includeNestedWorkflow = true
+})
+
 function outputSelect(
   workflowId: string,
   selectedOutputs: string[],
@@ -251,4 +265,15 @@ describe('OutputSelect nested workflow menu', () => {
     expect(document.body.textContent).toContain('Summarizer')
     expect(document.body.textContent).not.toContain('Back')
   })
+
+  it('returns to the root menu when a workflow edit invalidates the active path', () => {
+    const onOutputSelect = renderOutputSelect([])
+    clickOption('Outputs')
+
+    outputMenuState.includeNestedWorkflow = false
+    rerenderOutputSelect('root', [], onOutputSelect)
+
+    expect(document.body.textContent).toContain('Summarizer')
+    expect(document.body.textContent).not.toContain('Back')
+  })
 })

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/chat/components/output-select/output-select.tsx

@@ -57,6 +57,20 @@ interface OutputSelectProps {
   className?: string
 }
 
+interface OutputSelectMenuProps {
+  outputMenu: readonly WorkflowOutputMenuNode[]
+  workflowOutputs: readonly WorkflowOutputOption[]
+  selectedOutputs: string[]
+  onOutputSelect: (outputIds: string[]) => void
+  disabled: boolean
+  placeholder: string
+  valueMode: 'id' | 'label'
+  align: 'start' | 'end' | 'center'
+  maxHeight: number
+  size: 'sm' | 'md'
+  className?: string
+}
+
 function getOutputValue(output: WorkflowOutputOption, valueMode: 'id' | 'label'): string {
   return valueMode === 'label' && !output.blockId.includes('/') ? output.label : output.id
 }
@@ -88,7 +102,7 @@ function resolveOutputMenuNode(
  * @returns The OutputSelect component
  */
 export function OutputSelect(props: OutputSelectProps) {
-  return <OutputSelectContent key={props.workflowId ?? 'no-workflow'} {...props} />
+  return <OutputSelectContent {...props} />
 }
 
 function OutputSelectContent({
@@ -103,7 +117,6 @@ function OutputSelectContent({
   size = 'sm',
   className,
 }: OutputSelectProps) {
-  const [menuPath, setMenuPath] = useState<string[]>([])
   const blocks = useWorkflowStore((state) => state.blocks)
   const edges = useWorkflowStore((state) => state.edges)
   const { isShowingDiff, isDiffReady, hasActiveDiff, baselineWorkflow } = useWorkflowDiffStore(
@@ -184,6 +197,43 @@ function OutputSelectContent({
       })
     : []
   const outputMenu = buildWorkflowOutputMenu(workflowOutputs)
+  const outputMenuRevision = JSON.stringify([
+    workflowId,
+    ...workflowOutputs.map((output) => output.id),
+  ])
+
+  return (
+    <OutputSelectMenu
+      key={outputMenuRevision}
+      outputMenu={outputMenu}
+      workflowOutputs={workflowOutputs}
+      selectedOutputs={selectedOutputs}
+      onOutputSelect={onOutputSelect}
+      disabled={disabled}
+      placeholder={placeholder}
+      valueMode={valueMode}
+      align={align}
+      maxHeight={maxHeight}
+      size={size}
+      className={className}
+    />
+  )
+}
+
+function OutputSelectMenu({
+  outputMenu,
+  workflowOutputs,
+  selectedOutputs,
+  onOutputSelect,
+  disabled,
+  placeholder,
+  valueMode,
+  align,
+  maxHeight,
+  size,
+  className,
+}: OutputSelectMenuProps) {
+  const [menuPath, setMenuPath] = useState<string[]>([])
   const activeMenuNode = resolveOutputMenuNode(outputMenu, menuPath)
 
   const validOutputCount = selectedOutputs.filter((val) =>

apps/sim/lib/workflows/streaming/streaming.test.ts

@@ -9,6 +9,7 @@ import {
   agentStreamProtocolResponseHeaders,
   createStreamingResponse,
 } from '@/lib/workflows/streaming/streaming'
+import type { AgentStreamSink } from '@/providers/stream-events'
 
 const workflowStreamingLoggerCallIndex = loggerMock.createLogger.mock.calls.findIndex(
   ([name]) => name === 'WorkflowStreaming'
@@ -1207,7 +1208,7 @@ describe('createStreamingResponse agent-events-v1', () => {
       },
       executeFn: async ({ onStream }) => {
         let textController!: ReadableStreamDefaultController<Uint8Array>
-        let sink: { onEvent: (event: unknown) => void | Promise<void> } | undefined
+        let sink: AgentStreamSink | undefined
         const textStream = new ReadableStream<Uint8Array>({
           start(controller) {
             textController = controller
@@ -1218,7 +1219,7 @@ describe('createStreamingResponse agent-events-v1', () => {
           blockId: 'agent-1',
           stream: textStream,
           streamFormat: 'text',
-          subscribe: (nextSink: { onEvent: (event: unknown) => void | Promise<void> }) => {
+          subscribe: (nextSink: AgentStreamSink) => {
             sink = nextSink
             return () => {}
           },
@@ -1297,7 +1298,7 @@ describe('createStreamingResponse agent-events-v1', () => {
       },
       executeFn: async ({ onStream }) => {
         let textController!: ReadableStreamDefaultController<Uint8Array>
-        let sink: { onEvent: (event: unknown) => void | Promise<void> } | undefined
+        let sink: AgentStreamSink | undefined
         const textStream = new ReadableStream<Uint8Array>({
           start(controller) {
             textController = controller
@@ -1308,7 +1309,7 @@ describe('createStreamingResponse agent-events-v1', () => {
           blockId: 'agent-1',
           stream: textStream,
           streamFormat: 'text',
-          subscribe: (nextSink: { onEvent: (event: unknown) => void | Promise<void> }) => {
+          subscribe: (nextSink: AgentStreamSink) => {
             sink = nextSink
             return () => {}
           },
@@ -1543,7 +1544,7 @@ describe('createStreamingResponse agent-events-v1', () => {
       },
       executeFn: async ({ onStream }) => {
         let textController!: ReadableStreamDefaultController<Uint8Array>
-        let sink: { onEvent: (event: unknown) => void | Promise<void> } | undefined
+        let sink: AgentStreamSink | undefined
         const textStream = new ReadableStream<Uint8Array>({
           start(controller) {
             textController = controller
@@ -1554,7 +1555,7 @@ describe('createStreamingResponse agent-events-v1', () => {
           blockId: 'agent-1',
           stream: textStream,
           streamFormat: 'text',
-          subscribe: (nextSink: any) => {
+          subscribe: (nextSink: AgentStreamSink) => {
             sink = nextSink
             return () => {
               sink = undefined