From 60a06081ad6b4d0d92c77b2f03784fde8799db8e Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Thu, 25 Jun 2026 13:45:51 +0100
Subject: [PATCH] Share one file-input control between code and agent workflow
 nodes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The code node's Input files and the agent node's Documents now render the same
NodeDocumentsControl: an All/None/Choose toggle, a multiselect of upstream
document/artifact refs, and a free-text add-ref escape (a typed A1/id appears as
a removable chip). The control owns its toggle state so Choose is reachable from
an empty selection — retiring the per-node parent flag. File passing stays
agent-only; the code node keeps its inputs field name and stages the selected
refs as files in the sandbox.
---
 .../src/agents/workflow/WorkflowBuilder.tsx   | 238 ++++++------------
 .../components/NodeDocumentsControl.tsx       | 128 ++++++++++
 .../agents/workflow/documentConfig.test.ts    |  21 ++
 .../src/agents/workflow/documentConfig.ts     |  13 +
 4 files changed, 233 insertions(+), 167 deletions(-)
 create mode 100644 frontend/src/agents/workflow/components/NodeDocumentsControl.tsx

diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx
index 5f458e78..0b7b7f41 100644
--- a/frontend/src/agents/workflow/WorkflowBuilder.tsx
+++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx
@@ -68,24 +68,20 @@ import { ConditionCase, WorkflowNode } from '../types/workflow';
 import {
   createDefaultCodeConfig,
   normalizeCodeConfig,
-  parseCodeInputs,
   parseCodeJsonSchemaDraft,
   serializeCodeConfig,
-  stringifyCodeInputs,
   validateCodeJsonSchema,
 } from './codeNodeConfig';
 import MobileBlocker from './components/MobileBlocker';
+import NodeDocumentsControl from './components/NodeDocumentsControl';
 import PromptTextArea, {
   extractUpstreamVariables,
 } from './components/PromptTextArea';
 import {
-  documentsModeToInputDocuments,
   FILE_PASSING_OPTIONS,
   FilePassing,
-  getDocumentsMode,
   normalizeFilePassing,
   toDocumentVariableOptions,
-  withChosenDocumentOptions,
 } from './documentConfig';
 import {
   AgentNode,
@@ -359,9 +355,6 @@ function WorkflowBuilderInner() {
   const reactFlowWrapper = useRef<HTMLDivElement>(null);
 
   const [selectedNode, setSelectedNode] = useState<Node | null>(null);
-  // Node id for which the author explicitly entered "Choose…" documents mode,
-  // so the picker shows even before any document is selected (empty reads as None).
-  const [docChooseNodeId, setDocChooseNodeId] = useState<string | null>(null);
   const [workflowName, setWorkflowName] = useState('New Workflow');
   const [workflowDescription, setWorkflowDescription] = useState('');
   const [showWorkflowSettings, setShowWorkflowSettings] = useState(false);
@@ -1478,6 +1471,13 @@ function WorkflowBuilderInner() {
     );
   }, [selectedNode, nodes, edges]);
 
+  const selectedCodeDocumentOptions = useMemo(() => {
+    if (!selectedNode || selectedNode.type !== 'code') return [];
+    return toDocumentVariableOptions(
+      extractUpstreamVariables(nodes, edges, selectedNode.id),
+    );
+  }, [selectedNode, nodes, edges]);
+
   const selectedCodeJsonSchemaText = useMemo(() => {
     if (!selectedNode || selectedNode.type !== 'code') return '';
 
@@ -2214,144 +2214,60 @@ function WorkflowBuilderInner() {
                                     emptyText="No sources available"
                                   />
                                 </div>
-                                {(() => {
-                                  const inputDocuments =
-                                    selectedNode.data.config?.input_documents;
-                                  const documentsMode =
-                                    getDocumentsMode(inputDocuments);
-                                  const showChoose =
-                                    documentsMode === 'choose' ||
-                                    docChooseNodeId === selectedNode.id;
-                                  const chosenDocuments = showChoose
-                                    ? documentsModeToInputDocuments(
-                                        'choose',
-                                        inputDocuments ?? [],
-                                      )
-                                    : [];
-                                  const filePassing = normalizeFilePassing(
-                                    selectedNode.data.config?.file_passing,
-                                  );
-                                  const setInputDocuments = (
-                                    nextInputDocuments: string[],
-                                  ) =>
+                                <NodeDocumentsControl
+                                  key={selectedNode.id}
+                                  value={
+                                    selectedNode.data.config?.input_documents ??
+                                    []
+                                  }
+                                  onChange={(nextInputDocuments) =>
                                     handleUpdateNodeData({
                                       config: {
                                         ...(selectedNode.data.config || {}),
                                         input_documents: nextInputDocuments,
                                       },
-                                    });
-                                  return (
-                                    <>
-                                      <div>
-                                        <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
-                                          Documents
-                                        </label>
-                                        <div className="border-border bg-card flex gap-1 rounded-xl border p-1">
-                                          {(
-                                            [
-                                              ['all', 'All input docs'],
-                                              ['none', 'None'],
-                                              ['choose', 'Choose…'],
-                                            ] as const
-                                          ).map(([mode, modeLabel]) => (
-                                            <Button
-                                              key={mode}
-                                              type="button"
-                                              variant="ghost"
-                                              size="sm"
-                                              onClick={() => {
-                                                setDocChooseNodeId(
-                                                  mode === 'choose'
-                                                    ? selectedNode.id
-                                                    : null,
-                                                );
-                                                setInputDocuments(
-                                                  documentsModeToInputDocuments(
-                                                    mode,
-                                                    chosenDocuments,
-                                                  ),
-                                                );
-                                              }}
-                                              className={`h-auto flex-1 rounded-lg px-2 py-1.5 text-xs font-medium ${
-                                                (showChoose
-                                                  ? 'choose'
-                                                  : documentsMode) === mode
-                                                  ? 'bg-primary text-white'
-                                                  : 'text-gray-600 dark:text-gray-300'
-                                              }`}
-                                            >
-                                              {modeLabel}
-                                            </Button>
-                                          ))}
-                                        </div>
-                                        {showChoose && (
-                                          <div className="mt-2">
-                                            <MultiSelect
-                                              options={withChosenDocumentOptions(
-                                                selectedAgentDocumentOptions,
-                                                chosenDocuments,
-                                              )}
-                                              selected={chosenDocuments}
-                                              onChange={(nextDocuments) =>
-                                                setInputDocuments(
-                                                  documentsModeToInputDocuments(
-                                                    'choose',
-                                                    nextDocuments,
-                                                  ),
-                                                )
-                                              }
-                                              placeholder="Select documents..."
-                                              searchPlaceholder="Search variables..."
-                                              emptyText="No upstream documents"
-                                            />
-                                          </div>
-                                        )}
-                                        <p className="text-muted-foreground mt-1 text-xs">
-                                          Documents passed to this agent from
-                                          uploads or upstream nodes.
-                                        </p>
-                                      </div>
-                                      <div>
-                                        <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
-                                          File passing
-                                        </label>
-                                        <Select
-                                          value={filePassing}
-                                          onValueChange={(value) =>
-                                            handleUpdateNodeData({
-                                              config: {
-                                                ...(selectedNode.data.config ||
-                                                  {}),
-                                                file_passing:
-                                                  value as FilePassing,
-                                              },
-                                            })
-                                          }
+                                    })
+                                  }
+                                  options={selectedAgentDocumentOptions}
+                                  label="Documents"
+                                  helpText="Documents passed to this agent from uploads or upstream nodes."
+                                />
+                                <div>
+                                  <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
+                                    File passing
+                                  </label>
+                                  <Select
+                                    value={normalizeFilePassing(
+                                      selectedNode.data.config?.file_passing,
+                                    )}
+                                    onValueChange={(value) =>
+                                      handleUpdateNodeData({
+                                        config: {
+                                          ...(selectedNode.data.config || {}),
+                                          file_passing: value as FilePassing,
+                                        },
+                                      })
+                                    }
+                                  >
+                                    <SelectTrigger className="w-full">
+                                      <SelectValue />
+                                    </SelectTrigger>
+                                    <SelectContent>
+                                      {FILE_PASSING_OPTIONS.map((option) => (
+                                        <SelectItem
+                                          key={option.value}
+                                          value={option.value}
                                         >
-                                          <SelectTrigger className="w-full">
-                                            <SelectValue />
-                                          </SelectTrigger>
-                                          <SelectContent>
-                                            {FILE_PASSING_OPTIONS.map(
-                                              (option) => (
-                                                <SelectItem
-                                                  key={option.value}
-                                                  value={option.value}
-                                                >
-                                                  {option.label}
-                                                </SelectItem>
-                                              ),
-                                            )}
-                                          </SelectContent>
-                                        </Select>
-                                        <p className="text-muted-foreground mt-1 text-xs">
-                                          Auto: send native when the model
-                                          supports it, otherwise extract text.
-                                        </p>
-                                      </div>
-                                    </>
-                                  );
-                                })()}
+                                          {option.label}
+                                        </SelectItem>
+                                      ))}
+                                    </SelectContent>
+                                  </Select>
+                                  <p className="text-muted-foreground mt-1 text-xs">
+                                    Auto: send native when the model supports
+                                    it, otherwise extract text.
+                                  </p>
+                                </div>
                                 <div>
                                   <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
                                     Structured Output (JSON Schema)
@@ -2911,33 +2827,21 @@ function WorkflowBuilderInner() {
                                     placeholder={'print("hello world")'}
                                   />
                                 </div>
-                                <div>
-                                  <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
-                                    Inputs
-                                  </label>
-                                  <textarea
-                                    value={stringifyCodeInputs(
-                                      selectedNode.data.config?.inputs,
-                                    )}
-                                    onChange={(e) =>
-                                      handleUpdateNodeData({
-                                        config: {
-                                          ...(selectedNode.data.config || {}),
-                                          inputs: parseCodeInputs(
-                                            e.target.value,
-                                          ),
-                                        },
-                                      })
-                                    }
-                                    className="border-border focus-visible:ring-ring/50 focus-visible:border-ring bg-card w-full rounded-xl border px-3 py-2 text-sm transition-all outline-none focus-visible:ring-2 dark:text-white"
-                                    rows={2}
-                                    placeholder="One per line: state variable or artifact ref"
-                                  />
-                                  <p className="text-muted-foreground mt-1 text-xs">
-                                    Artifact refs or state-variable names, one
-                                    per line.
-                                  </p>
-                                </div>
+                                <NodeDocumentsControl
+                                  key={selectedNode.id}
+                                  value={selectedNode.data.config?.inputs ?? []}
+                                  onChange={(nextInputs) =>
+                                    handleUpdateNodeData({
+                                      config: {
+                                        ...(selectedNode.data.config || {}),
+                                        inputs: nextInputs,
+                                      },
+                                    })
+                                  }
+                                  options={selectedCodeDocumentOptions}
+                                  label="Input files"
+                                  helpText="Artifacts/upstream refs staged as files in the sandbox (one becomes inputs/<name>)."
+                                />
                                 <div>
                                   <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
                                     Output Variable
diff --git a/frontend/src/agents/workflow/components/NodeDocumentsControl.tsx b/frontend/src/agents/workflow/components/NodeDocumentsControl.tsx
new file mode 100644
index 00000000..17bfced3
--- /dev/null
+++ b/frontend/src/agents/workflow/components/NodeDocumentsControl.tsx
@@ -0,0 +1,128 @@
+import { Plus } from 'lucide-react';
+import { useState } from 'react';
+
+import { Button } from '@/components/ui/button';
+import { Input } from '@/components/ui/input';
+import { MultiSelect } from '@/components/ui/multi-select';
+
+import {
+  appendDocumentRef,
+  documentsModeToInputDocuments,
+  DocumentsMode,
+  getDocumentsMode,
+  withChosenDocumentOptions,
+} from '../documentConfig';
+
+interface NodeDocumentsControlProps {
+  value: string[];
+  onChange: (next: string[]) => void;
+  options: { value: string; label: string }[];
+  label: string;
+  helpText?: string;
+}
+
+const MODE_OPTIONS: { mode: DocumentsMode; label: string }[] = [
+  { mode: 'all', label: 'All input docs' },
+  { mode: 'none', label: 'None' },
+  { mode: 'choose', label: 'Choose…' },
+];
+
+/** Shared All/None/Choose documents picker for agent and code workflow nodes. */
+export default function NodeDocumentsControl({
+  value,
+  onChange,
+  options,
+  label,
+  helpText,
+}: NodeDocumentsControlProps) {
+  // Track mode in component state so "Choose" stays reachable even when the
+  // chosen list is empty (an empty list otherwise reads back as "None").
+  const [mode, setMode] = useState<DocumentsMode>(() =>
+    getDocumentsMode(value),
+  );
+  const [refDraft, setRefDraft] = useState('');
+
+  const chosen = documentsModeToInputDocuments('choose', value);
+  const showChoose = mode === 'choose';
+
+  const selectMode = (next: DocumentsMode) => {
+    setMode(next);
+    onChange(documentsModeToInputDocuments(next, chosen));
+  };
+
+  const addRef = () => {
+    const next = appendDocumentRef(chosen, refDraft);
+    setRefDraft('');
+    if (next === chosen) return;
+    setMode('choose');
+    onChange(documentsModeToInputDocuments('choose', next));
+  };
+
+  return (
+    <div>
+      <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
+        {label}
+      </label>
+      <div className="border-border bg-card flex gap-1 rounded-xl border p-1">
+        {MODE_OPTIONS.map(({ mode: optionMode, label: modeLabel }) => (
+          <Button
+            key={optionMode}
+            type="button"
+            variant="ghost"
+            size="sm"
+            onClick={() => selectMode(optionMode)}
+            className={`h-auto flex-1 rounded-lg px-2 py-1.5 text-xs font-medium ${
+              mode === optionMode
+                ? 'bg-primary text-white'
+                : 'text-gray-600 dark:text-gray-300'
+            }`}
+          >
+            {modeLabel}
+          </Button>
+        ))}
+      </div>
+      {showChoose && (
+        <div className="mt-2 flex flex-col gap-2">
+          <MultiSelect
+            options={withChosenDocumentOptions(options, chosen)}
+            selected={chosen}
+            onChange={(next) =>
+              onChange(documentsModeToInputDocuments('choose', next))
+            }
+            placeholder="Select documents..."
+            searchPlaceholder="Search variables..."
+            emptyText="No upstream documents"
+          />
+          <div className="flex gap-2">
+            <Input
+              type="text"
+              value={refDraft}
+              onChange={(e) => setRefDraft(e.target.value)}
+              onKeyDown={(e) => {
+                if (e.key === 'Enter') {
+                  e.preventDefault();
+                  addRef();
+                }
+              }}
+              className="bg-card h-auto rounded-xl px-3 py-2 text-sm shadow-none"
+              placeholder="Add ref (e.g. A1)"
+            />
+            <Button
+              type="button"
+              variant="ghost"
+              size="sm"
+              onClick={addRef}
+              className="h-auto shrink-0 gap-1 rounded-lg px-3 py-2 text-sm font-medium text-gray-600 dark:text-gray-300"
+            >
+              <Plus size={14} />
+              Add
+            </Button>
+          </div>
+        </div>
+      )}
+      {helpText && (
+        <p className="text-muted-foreground mt-1 text-xs">{helpText}</p>
+      )}
+    </div>
+  );
+}
diff --git a/frontend/src/agents/workflow/documentConfig.test.ts b/frontend/src/agents/workflow/documentConfig.test.ts
index 320c645f..a4ce2ce8 100644
--- a/frontend/src/agents/workflow/documentConfig.test.ts
+++ b/frontend/src/agents/workflow/documentConfig.test.ts
@@ -3,6 +3,7 @@ import { describe, expect, it } from 'vitest';
 import type { WorkflowVariable } from './components/PromptTextArea';
 import {
   ALL_INPUT_DOCUMENTS_TOKEN,
+  appendDocumentRef,
   DEFAULT_FILE_PASSING,
   documentsModeToInputDocuments,
   getDocumentsMode,
@@ -126,6 +127,26 @@ describe('toDocumentVariableOptions', () => {
   });
 });
 
+describe('appendDocumentRef', () => {
+  it('appends a trimmed literal ref', () => {
+    expect(appendDocumentRef(['wire_doc'], '  A1 ')).toEqual([
+      'wire_doc',
+      'A1',
+    ]);
+  });
+
+  it('ignores blanks and the wildcard token', () => {
+    expect(appendDocumentRef(['wire_doc'], '   ')).toEqual(['wire_doc']);
+    expect(appendDocumentRef(['wire_doc'], ALL_INPUT_DOCUMENTS_TOKEN)).toEqual([
+      'wire_doc',
+    ]);
+  });
+
+  it('does not append an already-present ref', () => {
+    expect(appendDocumentRef(['A1'], 'A1')).toEqual(['A1']);
+  });
+});
+
 describe('withChosenDocumentOptions', () => {
   const base = [{ value: 'wire_doc', label: 'wire_doc' }];
 
diff --git a/frontend/src/agents/workflow/documentConfig.ts b/frontend/src/agents/workflow/documentConfig.ts
index 9cda0bcf..c524c07e 100644
--- a/frontend/src/agents/workflow/documentConfig.ts
+++ b/frontend/src/agents/workflow/documentConfig.ts
@@ -83,6 +83,19 @@ export function toDocumentVariableOptions(
   return options;
 }
 
+/** Append a typed literal ref to a chosen list, ignoring blanks, the wildcard and duplicates. */
+export function appendDocumentRef(chosen: string[], ref: string): string[] {
+  const value = ref.trim();
+  if (
+    value === '' ||
+    value === ALL_INPUT_DOCUMENTS_TOKEN ||
+    chosen.includes(value)
+  ) {
+    return chosen;
+  }
+  return [...chosen, value];
+}
+
 /** Append stored chosen names lacking an upstream option so they stay visible and removable. */
 export function withChosenDocumentOptions(
   options: { value: string; label: string }[],