From 51b0dd8d455a7cc82ba1430e6bfa8111bb5b1869 Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Thu, 25 Jun 2026 11:43:03 +0100
Subject: [PATCH] Add a Documents control to workflow agent nodes

The agent-node panel gains a Documents section (All input docs / None / Choose
upstream document variables) and a file-passing selector (auto / native /
extract), writing the input_documents and file_passing config the engine
consumes. Selections persist as bare state-variable names so they resolve
against run state; the toggle, chosen docs, and file-passing mode re-derive from
the stored config on reload. Reuses the existing multiselect, select, and
upstream-variable enumeration.

Known minor: an input_documents entry whose upstream variable was renamed or
deleted is preserved in config but not shown as a removable chip in Choose mode.
---
 .../src/agents/workflow/WorkflowBuilder.tsx   | 146 +++++++++++++++++-
 .../workflow/components/PromptTextArea.tsx    |   4 +-
 .../agents/workflow/documentConfig.test.ts    | 126 +++++++++++++++
 .../src/agents/workflow/documentConfig.ts     |  84 ++++++++++
 4 files changed, 357 insertions(+), 3 deletions(-)
 create mode 100644 frontend/src/agents/workflow/documentConfig.test.ts
 create mode 100644 frontend/src/agents/workflow/documentConfig.ts

diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx
index b5649d2c..dfa4e532 100644
--- a/frontend/src/agents/workflow/WorkflowBuilder.tsx
+++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx
@@ -75,7 +75,17 @@ import {
   validateCodeJsonSchema,
 } from './codeNodeConfig';
 import MobileBlocker from './components/MobileBlocker';
-import PromptTextArea from './components/PromptTextArea';
+import PromptTextArea, {
+  extractUpstreamVariables,
+} from './components/PromptTextArea';
+import {
+  documentsModeToInputDocuments,
+  FILE_PASSING_OPTIONS,
+  FilePassing,
+  getDocumentsMode,
+  normalizeFilePassing,
+  toDocumentVariableOptions,
+} from './documentConfig';
 import {
   AgentNode,
   CodeNode,
@@ -104,6 +114,8 @@ interface AgentNodeConfig {
   chunks?: string;
   retriever?: string;
   json_schema?: Record<string, unknown>;
+  input_documents?: string[];
+  file_passing?: FilePassing;
 }
 
 interface UserTool {
@@ -1455,6 +1467,13 @@ function WorkflowBuilderInner() {
     return selectedModel.supports_structured_output;
   }, [selectedNode, availableModels]);
 
+  const selectedAgentDocumentOptions = useMemo(() => {
+    if (!selectedNode || selectedNode.type !== 'agent') return [];
+    return toDocumentVariableOptions(
+      extractUpstreamVariables(nodes, edges, selectedNode.id),
+    );
+  }, [selectedNode, nodes, edges]);
+
   const selectedCodeJsonSchemaText = useMemo(() => {
     if (!selectedNode || selectedNode.type !== 'code') return '';
 
@@ -2191,6 +2210,131 @@ function WorkflowBuilderInner() {
                                     emptyText="No sources available"
                                   />
                                 </div>
+                                {(() => {
+                                  const inputDocuments =
+                                    selectedNode.data.config?.input_documents;
+                                  const documentsMode =
+                                    getDocumentsMode(inputDocuments);
+                                  const chosenDocuments =
+                                    documentsMode === 'choose'
+                                      ? (inputDocuments ?? [])
+                                      : [];
+                                  const filePassing = normalizeFilePassing(
+                                    selectedNode.data.config?.file_passing,
+                                  );
+                                  const setInputDocuments = (
+                                    nextInputDocuments: string[],
+                                  ) =>
+                                    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={() =>
+                                                setInputDocuments(
+                                                  documentsModeToInputDocuments(
+                                                    mode,
+                                                    chosenDocuments,
+                                                  ),
+                                                )
+                                              }
+                                              className={`h-auto flex-1 rounded-lg px-2 py-1.5 text-xs font-medium ${
+                                                documentsMode === mode
+                                                  ? 'bg-primary text-white'
+                                                  : 'text-gray-600 dark:text-gray-300'
+                                              }`}
+                                            >
+                                              {modeLabel}
+                                            </Button>
+                                          ))}
+                                        </div>
+                                        {documentsMode === 'choose' && (
+                                          <div className="mt-2">
+                                            <MultiSelect
+                                              options={
+                                                selectedAgentDocumentOptions
+                                              }
+                                              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,
+                                              },
+                                            })
+                                          }
+                                        >
+                                          <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>
+                                    </>
+                                  );
+                                })()}
                                 <div>
                                   <label className="mb-2 block text-sm font-medium text-gray-700 dark:text-gray-300">
                                     Structured Output (JSON Schema)
diff --git a/frontend/src/agents/workflow/components/PromptTextArea.tsx b/frontend/src/agents/workflow/components/PromptTextArea.tsx
index cb363e0b..c85a7221 100644
--- a/frontend/src/agents/workflow/components/PromptTextArea.tsx
+++ b/frontend/src/agents/workflow/components/PromptTextArea.tsx
@@ -10,7 +10,7 @@ import {
   PopoverTrigger,
 } from '@/components/ui/popover';
 
-interface WorkflowVariable {
+export interface WorkflowVariable {
   label: string;
   templatePath: string;
   section: string;
@@ -98,7 +98,7 @@ function getUpstreamNodeIds(nodeId: string, edges: Edge[]): Set<string> {
   return upstream;
 }
 
-function extractUpstreamVariables(
+export function extractUpstreamVariables(
   nodes: Node[],
   edges: Edge[],
   selectedNodeId: string,
diff --git a/frontend/src/agents/workflow/documentConfig.test.ts b/frontend/src/agents/workflow/documentConfig.test.ts
new file mode 100644
index 00000000..4ded092c
--- /dev/null
+++ b/frontend/src/agents/workflow/documentConfig.test.ts
@@ -0,0 +1,126 @@
+import { describe, expect, it } from 'vitest';
+
+import type { WorkflowVariable } from './components/PromptTextArea';
+import {
+  ALL_INPUT_DOCUMENTS_TOKEN,
+  DEFAULT_FILE_PASSING,
+  documentsModeToInputDocuments,
+  getDocumentsMode,
+  normalizeFilePassing,
+  stripAgentPrefix,
+  toDocumentVariableOptions,
+} from './documentConfig';
+
+describe('getDocumentsMode', () => {
+  it('treats missing/empty as none', () => {
+    expect(getDocumentsMode(undefined)).toBe('none');
+    expect(getDocumentsMode([])).toBe('none');
+  });
+
+  it('treats the wildcard token as all', () => {
+    expect(getDocumentsMode([ALL_INPUT_DOCUMENTS_TOKEN])).toBe('all');
+  });
+
+  it('treats a bare-name list as choose', () => {
+    expect(getDocumentsMode(['wire_doc'])).toBe('choose');
+    expect(getDocumentsMode(['wire_doc', 'A1'])).toBe('choose');
+  });
+});
+
+describe('documentsModeToInputDocuments', () => {
+  it('maps all to the wildcard token regardless of chosen names', () => {
+    expect(documentsModeToInputDocuments('all', ['wire_doc'])).toEqual([
+      ALL_INPUT_DOCUMENTS_TOKEN,
+    ]);
+  });
+
+  it('maps none to an empty list', () => {
+    expect(documentsModeToInputDocuments('none', ['wire_doc'])).toEqual([]);
+  });
+
+  it('maps choose to the chosen bare names, dropping blanks and the wildcard', () => {
+    expect(
+      documentsModeToInputDocuments('choose', [
+        'wire_doc',
+        '  ',
+        ALL_INPUT_DOCUMENTS_TOKEN,
+        'A1',
+      ]),
+    ).toEqual(['wire_doc', 'A1']);
+  });
+
+  it('round-trips a chosen selection through getDocumentsMode', () => {
+    const next = documentsModeToInputDocuments('choose', ['wire_doc']);
+    expect(getDocumentsMode(next)).toBe('choose');
+    expect(next).toEqual(['wire_doc']);
+  });
+});
+
+describe('normalizeFilePassing', () => {
+  it('passes through supported values', () => {
+    expect(normalizeFilePassing('auto')).toBe('auto');
+    expect(normalizeFilePassing('native')).toBe('native');
+    expect(normalizeFilePassing('extract')).toBe('extract');
+  });
+
+  it('falls back to the default for unknown values', () => {
+    expect(normalizeFilePassing(undefined)).toBe(DEFAULT_FILE_PASSING);
+    expect(normalizeFilePassing('bogus')).toBe(DEFAULT_FILE_PASSING);
+  });
+});
+
+describe('stripAgentPrefix', () => {
+  it('strips a dotted agent prefix to a bare name', () => {
+    expect(stripAgentPrefix('agent.wire_doc')).toBe('wire_doc');
+    expect(stripAgentPrefix('agent.node_abc_output')).toBe('node_abc_output');
+  });
+
+  it('strips a bracketed agent prefix and unescapes quotes', () => {
+    expect(stripAgentPrefix("agent['weird name']")).toBe('weird name');
+    expect(stripAgentPrefix("agent['it\\'s']")).toBe("it's");
+  });
+
+  it('leaves bare names untouched', () => {
+    expect(stripAgentPrefix('A1')).toBe('A1');
+  });
+});
+
+describe('toDocumentVariableOptions', () => {
+  const make = (templatePath: string): WorkflowVariable => ({
+    label: templatePath,
+    templatePath,
+    section: 'x',
+  });
+
+  it('keeps input_documents and upstream outputs as bare names', () => {
+    const options = toDocumentVariableOptions([
+      make('agent.input_documents'),
+      make('agent.wire_doc'),
+      make("agent['weird name']"),
+    ]);
+    expect(options).toEqual([
+      { value: 'input_documents', label: 'input_documents' },
+      { value: 'wire_doc', label: 'wire_doc' },
+      { value: 'weird name', label: 'weird name' },
+    ]);
+  });
+
+  it('excludes query, chat_history and global context variables', () => {
+    const options = toDocumentVariableOptions([
+      make('agent.query'),
+      make('agent.chat_history'),
+      make('source.content'),
+      make('system.date'),
+      make('artifacts.artifact(id)'),
+    ]);
+    expect(options).toEqual([]);
+  });
+
+  it('deduplicates by bare name', () => {
+    const options = toDocumentVariableOptions([
+      make('agent.wire_doc'),
+      make('agent.wire_doc'),
+    ]);
+    expect(options).toEqual([{ value: 'wire_doc', label: 'wire_doc' }]);
+  });
+});
diff --git a/frontend/src/agents/workflow/documentConfig.ts b/frontend/src/agents/workflow/documentConfig.ts
new file mode 100644
index 00000000..daf57971
--- /dev/null
+++ b/frontend/src/agents/workflow/documentConfig.ts
@@ -0,0 +1,84 @@
+import type { WorkflowVariable } from './components/PromptTextArea';
+
+/** Token meaning "all uploaded input documents" in the input_documents list. */
+export const ALL_INPUT_DOCUMENTS_TOKEN = '*';
+
+/** How an agent node receives selected documents. */
+export type FilePassing = 'auto' | 'native' | 'extract';
+
+export const DEFAULT_FILE_PASSING: FilePassing = 'auto';
+
+export const FILE_PASSING_OPTIONS: { value: FilePassing; label: string }[] = [
+  { value: 'auto', label: 'Auto' },
+  { value: 'native', label: 'Native' },
+  { value: 'extract', label: 'Extract' },
+];
+
+/** Three-way Documents selector position derived from input_documents. */
+export type DocumentsMode = 'all' | 'none' | 'choose';
+
+/** Normalize an unknown file_passing value to a supported option. */
+export function normalizeFilePassing(value: unknown): FilePassing {
+  return value === 'native' || value === 'extract' || value === 'auto'
+    ? value
+    : DEFAULT_FILE_PASSING;
+}
+
+/** Derive the toggle position from the stored input_documents list. */
+export function getDocumentsMode(
+  inputDocuments: string[] | undefined,
+): DocumentsMode {
+  const docs = inputDocuments ?? [];
+  if (docs.length === 0) return 'none';
+  if (docs.length === 1 && docs[0] === ALL_INPUT_DOCUMENTS_TOKEN) return 'all';
+  return 'choose';
+}
+
+/** Build the input_documents list for a chosen toggle position. */
+export function documentsModeToInputDocuments(
+  mode: DocumentsMode,
+  chosen: string[],
+): string[] {
+  if (mode === 'all') return [ALL_INPUT_DOCUMENTS_TOKEN];
+  if (mode === 'none') return [];
+  return chosen.filter(
+    (name) => name.trim() !== '' && name !== ALL_INPUT_DOCUMENTS_TOKEN,
+  );
+}
+
+/** Strip a leading `agent.` / `agent['...']` wrapper to a bare variable name. */
+export function stripAgentPrefix(templatePath: string): string {
+  const trimmed = templatePath.trim();
+  const bracketMatch = trimmed.match(/^agent\[(?:'|")(.*)(?:'|")\]$/);
+  if (bracketMatch) {
+    return bracketMatch[1].replace(/\\(['"\\])/g, '$1');
+  }
+  if (trimmed.startsWith('agent.')) {
+    return trimmed.slice('agent.'.length);
+  }
+  return trimmed;
+}
+
+/** Multiselect options of upstream document variables, stored as bare names. */
+export function toDocumentVariableOptions(
+  variables: WorkflowVariable[],
+): { value: string; label: string }[] {
+  const options: { value: string; label: string }[] = [];
+  const seen = new Set<string>();
+  for (const variable of variables) {
+    const path = variable.templatePath;
+    const isInputDocuments = path === 'agent.input_documents';
+    const isAgentOutput =
+      path.startsWith('agent.') &&
+      path !== 'agent.query' &&
+      path !== 'agent.chat_history';
+    const isAgentBracket = path.startsWith('agent[');
+    if (!isInputDocuments && !isAgentOutput && !isAgentBracket) continue;
+
+    const bareName = stripAgentPrefix(path);
+    if (!bareName || seen.has(bareName)) continue;
+    seen.add(bareName);
+    options.push({ value: bareName, label: bareName });
+  }
+  return options;
+}