windmill-components 1.677.1 → 1.687.0

package/package/components/copilot/chat/flow/core.js

@@ -1,11 +1,12 @@
 import { ScriptService, JobService } from '../../../../gen';
 import { z } from 'zod';
 import { createDbSchemaTool, getFormattedResourceTypes, getLangContext, SUPPORTED_CHAT_SCRIPT_LANGUAGES } from '../script/core';
-import { createSearchHubScriptsTool, createToolDef, executeTestRun, buildSchemaForTool, buildTestRunArgs, buildContextString, applyCodePiecesToFlowModules, findModuleById, SPECIAL_MODULE_IDS, formatScriptLintResult, createSearchWorkspaceTool, createGetRunnableDetailsTool } from '../shared';
-import { inlineScriptStore, extractAndReplaceInlineScripts } from './inlineScriptsUtils';
-import { flowModulesSchema } from './openFlowZod';
-import { collectAllModuleIdsFromArray } from './utils';
-import { getFlowPrompt } from '../../../../system_prompts';
+import { createSearchHubScriptsTool, createToolDef, executeTestRun, buildSchemaForTool, buildTestRunArgs, buildContextString, applyCodePiecesToFlowModules, SPECIAL_MODULE_IDS, formatScriptLintResult, createSearchWorkspaceTool, createGetRunnableDetailsTool } from '../shared';
+import { findModuleInFlow, findModuleInModules } from '../../../flows/flowTree';
+import { createInlineScriptSession } from './inlineScriptsUtils';
+import { flowModuleSchema, flowModulesSchema } from './openFlowZod';
+import { collectAllFlowModuleIdsFromModules } from '../../../flows/flowTree';
+import { FLOW_CHAT_SPECIAL_MODULES, getFlowPrompt } from '../../../../system_prompts';
 /**
  * Navigate to a schema at a given path, handling arrays, objects, unions, and wrappers.
  * Uses Zod 4 internal structure.
@@ -181,6 +182,174 @@ function getExpectedFormat(schema) {
     }
     return null;
 }
+function countExactMatches(content, search) {
+    if (search.length === 0) {
+        return 0;
+    }
+    let count = 0;
+    let index = 0;
+    while ((index = content.indexOf(search, index)) !== -1) {
+        count++;
+        index += search.length;
+    }
+    return count;
+}
+function replaceFirstExactMatch(content, search, replace) {
+    const index = content.indexOf(search);
+    if (index === -1) {
+        return content;
+    }
+    return content.slice(0, index) + replace + content.slice(index + search.length);
+}
+function formatEmptyInlineScriptWarning({ emptyInlineScriptModuleIds }) {
+    if (emptyInlineScriptModuleIds.length === 0) {
+        return '';
+    }
+    const moduleList = emptyInlineScriptModuleIds.map((id) => `'${id}'`).join(', ');
+    return ` Warning: inline scripts ${moduleList} are empty for now. Use set_module_code to fill them in.`;
+}
+function validateFlowModules(rawModules) {
+    if (!Array.isArray(rawModules)) {
+        throw new Error('Flow modules must be an array');
+    }
+    const parsedModules = rawModules;
+    const result = flowModulesSchema.safeParse(parsedModules);
+    if (!result.success) {
+        const errors = result.error.issues.slice(0, 5).map((e) => {
+            const path = e.path;
+            // Try to find module id for better context
+            const moduleIndex = typeof path[0] === 'number' ? path[0] : undefined;
+            const moduleId = moduleIndex !== undefined ? parsedModules[moduleIndex]?.id : undefined;
+            const fieldPath = path.slice(1).join('.');
+            let message = e.message;
+            if (e.code === 'invalid_type') {
+                // Zod 4 message already contains "expected X, received Y"
+                // Try to extract expected format from schema, passing actual data
+                // to help resolve discriminated unions correctly
+                const targetSchema = getSchemaAtPath(flowModulesSchema, path, parsedModules);
+                if (targetSchema) {
+                    const expectedFormat = getExpectedFormat(targetSchema);
+                    if (expectedFormat) {
+                        message += `\n    Expected format: ${expectedFormat}`;
+                    }
+                }
+            }
+            if (moduleId) {
+                return `Module "${moduleId}" -> ${fieldPath}: ${message}`;
+            }
+            return `${path.join('.')}: ${message}`;
+        });
+        throw new Error(`Invalid flow modules:\n${errors.join('\n')}`);
+    }
+    const ids = collectAllFlowModuleIdsFromModules(parsedModules);
+    if (ids.length !== new Set(ids).size) {
+        throw new Error('Duplicate module IDs found in flow');
+    }
+    const reservedIds = ids.filter((id) => id === SPECIAL_MODULE_IDS.PREPROCESSOR || id === SPECIAL_MODULE_IDS.FAILURE);
+    if (reservedIds.length > 0) {
+        throw new Error('Special modules must be provided via preprocessor_module and failure_module, not inside modules');
+    }
+    return parsedModules;
+}
+function validateFlowSchema(rawSchema) {
+    if (rawSchema == null) {
+        return null;
+    }
+    if (typeof rawSchema !== 'object' || Array.isArray(rawSchema)) {
+        throw new Error('Flow schema must be an object or null');
+    }
+    return rawSchema;
+}
+function validateOptionalFlowModule(rawModule, fieldName) {
+    if (rawModule == null) {
+        return null;
+    }
+    const result = flowModuleSchema.safeParse(rawModule);
+    if (!result.success) {
+        const error = result.error.issues[0];
+        throw new Error(`Invalid ${fieldName}: ${error?.message ?? 'unknown error'}`);
+    }
+    return result.data;
+}
+function validateEditableFlowJson(rawFlow) {
+    if (!rawFlow || typeof rawFlow !== 'object' || Array.isArray(rawFlow)) {
+        throw new Error('Flow JSON must be an object');
+    }
+    const flow = rawFlow;
+    const modules = validateFlowModules(flow.modules);
+    const schema = validateFlowSchema(flow.schema);
+    const preprocessorModule = validateOptionalFlowModule(flow.preprocessor_module, 'preprocessor_module');
+    const failureModule = validateOptionalFlowModule(flow.failure_module, 'failure_module');
+    if (preprocessorModule) {
+        if (preprocessorModule.id !== SPECIAL_MODULE_IDS.PREPROCESSOR) {
+            throw new Error(`Invalid preprocessor_module: id must be "${SPECIAL_MODULE_IDS.PREPROCESSOR}"`);
+        }
+        if (preprocessorModule.value.type !== 'rawscript' && preprocessorModule.value.type !== 'script') {
+            throw new Error('Invalid preprocessor_module: only "rawscript" and "script" modules are supported');
+        }
+    }
+    if (failureModule) {
+        if (failureModule.id !== SPECIAL_MODULE_IDS.FAILURE) {
+            throw new Error(`Invalid failure_module: id must be "${SPECIAL_MODULE_IDS.FAILURE}"`);
+        }
+        if (failureModule.value.type !== 'rawscript' && failureModule.value.type !== 'script') {
+            throw new Error('Invalid failure_module: only "rawscript" and "script" modules are supported');
+        }
+    }
+    const ids = new Set(collectAllFlowModuleIdsFromModules(modules));
+    if (preprocessorModule) {
+        if (ids.has(preprocessorModule.id)) {
+            throw new Error(`Duplicate module ID found in preprocessor_module: ${preprocessorModule.id}`);
+        }
+        ids.add(preprocessorModule.id);
+    }
+    if (failureModule && ids.has(failureModule.id)) {
+        throw new Error(`Duplicate module ID found in failure_module: ${failureModule.id}`);
+    }
+    return {
+        modules,
+        schema,
+        preprocessor_module: preprocessorModule,
+        failure_module: failureModule
+    };
+}
+function buildEditableFlowJson(flow, inlineScriptSession, selectedContext = []) {
+    const codePieces = selectedContext.filter((c) => c.type === 'flow_module_code_piece');
+    const optimizedModules = inlineScriptSession
+        ? inlineScriptSession.extractAndReplaceInlineScripts(flow.value.modules)
+        : flow.value.modules;
+    const modules = applyCodePiecesToFlowModules(codePieces, optimizedModules);
+    let preprocessorModule = flow.value.preprocessor_module;
+    if (preprocessorModule?.value?.type === 'rawscript' &&
+        preprocessorModule.value.content &&
+        inlineScriptSession) {
+        inlineScriptSession.set(preprocessorModule.id, preprocessorModule.value.content);
+        preprocessorModule = {
+            ...preprocessorModule,
+            value: {
+                ...preprocessorModule.value,
+                content: `inline_script.${preprocessorModule.id}`
+            }
+        };
+    }
+    let failureModule = flow.value.failure_module;
+    if (failureModule?.value?.type === 'rawscript' && failureModule.value.content && inlineScriptSession) {
+        inlineScriptSession.set(failureModule.id, failureModule.value.content);
+        failureModule = {
+            ...failureModule,
+            value: {
+                ...failureModule.value,
+                content: `inline_script.${failureModule.id}`
+            }
+        };
+    }
+    return {
+        modules,
+        schema: flow.schema ?? null,
+        preprocessor_module: preprocessorModule ?? null,
+        failure_module: failureModule ?? null
+    };
+}
 const langSchema = z.enum(SUPPORTED_CHAT_SCRIPT_LANGUAGES);
 const resourceTypeToolSchema = z.object({
     query: z.string().describe('The query to search for, e.g. stripe, google, etc..'),
@@ -191,12 +360,85 @@ const getInstructionsForCodeGenerationToolSchema = z.object({
     language: langSchema.describe('The programming language the code will be written in')
 });
 const getInstructionsForCodeGenerationToolDef = createToolDef(getInstructionsForCodeGenerationToolSchema, 'get_instructions_for_code_generation', 'Get instructions for code generation for a raw script step');
+const specialModuleToolArgSchema = z
+    .string()
+    .nullable()
+    .describe('JSON string containing the special module object. Use null to remove the special module.');
 // Using string for modules and schema because Gemini-2.5-flash performs better with strings (MALFORMED_FUNCTION_CALL errors happens more often with objects)
 const setFlowJsonToolSchema = z.object({
     modules: z.string().optional().nullable().describe('JSON string containing the flow modules'),
-    schema: z.string().optional().nullable().describe('JSON string containing the flow input schema')
+    schema: z.string().optional().nullable().describe('JSON string containing the flow input schema'),
+    preprocessor_module: z
+        .string()
+        .optional()
+        .nullable()
+        .describe('JSON string containing the optional preprocessor module'),
+    failure_module: z
+        .string()
+        .optional()
+        .nullable()
+        .describe('JSON string containing the optional failure module')
+});
+const setFlowJsonToolDef = createToolDef(setFlowJsonToolSchema, 'set_flow_json', 'Set the complete flow modules array and optionally the flow input schema, preprocessor module, and failure module.', { strict: false });
+const setPreprocessorModuleToolSchema = z.object({
+    module: specialModuleToolArgSchema
+});
+const setPreprocessorModuleToolDef = createToolDef(setPreprocessorModuleToolSchema, 'set_preprocessor_module', 'Set or replace the flow preprocessor module. Use this when the flow needs logic that runs before the main modules.');
+const setFailureModuleToolSchema = z.object({
+    module: specialModuleToolArgSchema
+});
+const setFailureModuleToolDef = createToolDef(setFailureModuleToolSchema, 'set_failure_module', 'Set or replace the flow failure module. Use this when the flow needs a dedicated error handler.');
+const specialFlowModuleFields = {
+    preprocessor_module: SPECIAL_MODULE_IDS.PREPROCESSOR,
+    failure_module: SPECIAL_MODULE_IDS.FAILURE
+};
+function parseOptionalJsonArg(value, field) {
+    if (value === undefined || value === null) {
+        return value;
+    }
+    try {
+        return typeof value === 'string' ? JSON.parse(value) : value;
+    }
+    catch (e) {
+        const errorMessage = e instanceof Error ? e.message : String(e);
+        throw new Error(`Invalid JSON for ${field}: ${errorMessage}`);
+    }
+}
+function validateSpecialFlowModule(module, field) {
+    if (module === undefined || module === null) {
+        return module;
+    }
+    const result = flowModuleSchema.safeParse(module);
+    if (!result.success) {
+        const errors = result.error.issues.slice(0, 5).map((issue) => {
+            const path = issue.path.length > 0 ? issue.path.join('.') : field;
+            return `${path}: ${issue.message}`;
+        });
+        throw new Error(`Invalid ${field}:\n${errors.join('\n')}`);
+    }
+    const parsedModule = result.data;
+    const expectedId = specialFlowModuleFields[field];
+    if (parsedModule.id !== expectedId) {
+        throw new Error(`Invalid ${field}: id must be "${expectedId}"`);
+    }
+    if (parsedModule.value.type !== 'rawscript' && parsedModule.value.type !== 'script') {
+        throw new Error(`Invalid ${field}: only "rawscript" and "script" modules are supported`);
+    }
+    return parsedModule;
+}
+const patchFlowJsonSchema = z.object({
+    old_string: z
+        .string()
+        .min(1)
+        .describe('Exact text to find in the current compact flow JSON'),
+    new_string: z.string().describe('Replacement JSON text'),
+    replace_all: z
+        .boolean()
+        .optional()
+        .default(false)
+        .describe('When true, replace every exact match. When false, the search text must match exactly once.')
 });
-const setFlowJsonToolDef = createToolDef(setFlowJsonToolSchema, 'set_flow_json', 'Set the entire flow by providing the complete flow object. This replaces all existing modules and schema.', { strict: false });
+const patchFlowJsonToolDef = createToolDef(patchFlowJsonSchema, 'patch_flow_json', 'Make a quick exact text edit in the current compact flow JSON. Prefer this for small localized changes; use set_flow_json for larger structural rewrites.');
 // Will be overridden by setSchema
 const testRunFlowSchema = z.object({
     args: z
@@ -312,15 +554,14 @@ export const flowTools = [
             }
             const stepId = args.stepId;
             const stepArgs = args.args || {};
-            // Find the step in the flow
-            const modules = helpers.getModules();
-            let targetModule = findModuleById(modules, stepId);
+            // Find the step in the flow (includes preprocessor/failure modules)
+            let targetModule = findModuleInFlow(flow.value, stepId) ?? undefined;
             if (!targetModule) {
                 toolCallbacks.setToolStatus(toolId, {
                     content: `Step '${stepId}' not found in flow`,
                     error: `Step with id '${stepId}' does not exist in the current flow`
                 });
-                throw new Error(`Step with id '${stepId}' not found in flow. Available steps: ${modules.map((m) => m.id).join(', ')}`);
+                throw new Error(`Step with id '${stepId}' not found in flow. Available steps: ${(flow.value.modules ?? []).map((m) => m.id).join(', ')}`);
             }
             const module = targetModule;
             const moduleValue = module.value;
@@ -402,13 +643,13 @@ export const flowTools = [
     },
     {
         def: inspectInlineScriptToolDef,
-        fn: async ({ args, toolCallbacks, toolId }) => {
+        fn: async ({ args, helpers, toolCallbacks, toolId }) => {
             const parsedArgs = inspectInlineScriptSchema.parse(args);
             const moduleId = parsedArgs.moduleId;
             toolCallbacks.setToolStatus(toolId, {
                 content: `Retrieving inline script content for module '${moduleId}'...`
             });
-            const content = inlineScriptStore.get(moduleId);
+            const content = helpers.inlineScriptSession.get(moduleId);
             if (content === undefined) {
                 toolCallbacks.setToolStatus(toolId, {
                     content: `Module '${moduleId}' not found in inline script store`,
@@ -435,8 +676,6 @@ export const flowTools = [
             const parsedArgs = setModuleCodeSchema.parse(args);
             const { moduleId, code } = parsedArgs;
             toolCallbacks.setToolStatus(toolId, { content: `Setting code for module '${moduleId}'...` });
-            // Update store to keep it coherent (for subsequent set_flow_json calls with references)
-            inlineScriptStore.set(moduleId, code);
             // Update the flow directly via helper
             await helpers.setCode(moduleId, code);
             toolCallbacks.setToolStatus(toolId, {
@@ -446,78 +685,185 @@ export const flowTools = [
             return `Code for module '${moduleId}' has been updated successfully.`;
         }
     },
+    {
+        def: patchFlowJsonToolDef,
+        streamArguments: true,
+        showDetails: true,
+        showFade: true,
+        fn: async ({ args, helpers, toolId, toolCallbacks }) => {
+            const parsedArgs = patchFlowJsonSchema.parse(args);
+            const { old_string: oldString, new_string: newString, replace_all: replaceAll } = parsedArgs;
+            const { flow, selectedId } = helpers.getFlowAndSelectedId();
+            // Snapshot the current flow with a fresh session so the compact JSON matches what the model saw,
+            // then copy extracted inline scripts back into the helper session before applying the patch.
+            const inlineScriptSession = createInlineScriptSession();
+            const currentFlowJson = JSON.stringify(buildEditableFlowJson(flow, inlineScriptSession));
+            const matchCount = countExactMatches(currentFlowJson, oldString);
+            if (matchCount === 0) {
+                throw new Error('old_string was not found in the current flow JSON.');
+            }
+            if (!replaceAll && matchCount !== 1) {
+                throw new Error(`old_string matched ${matchCount} locations. Make it more specific or set replace_all to true.`);
+            }
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'Applying JSON patch...'
+            });
+            const updatedFlowJson = replaceAll
+                ? currentFlowJson.split(oldString).join(newString)
+                : replaceFirstExactMatch(currentFlowJson, oldString, newString);
+            let parsedFlow;
+            try {
+                parsedFlow = validateEditableFlowJson(JSON.parse(updatedFlowJson));
+            }
+            catch (error) {
+                const message = error instanceof Error ? error.message : String(error);
+                throw new Error(`Invalid JSON after replacement: ${message}`);
+            }
+            for (const [moduleId, content] of Object.entries(inlineScriptSession.getAll())) {
+                helpers.inlineScriptSession.set(moduleId, content);
+            }
+            const updateResult = await helpers.setFlowJson({
+                modules: parsedFlow.modules,
+                schema: parsedFlow.schema,
+                preprocessorModule: parsedFlow.preprocessor_module,
+                failureModule: parsedFlow.failure_module
+            });
+            const warning = formatEmptyInlineScriptWarning(updateResult);
+            const selectedModule = findModuleInFlow(parsedFlow, selectedId) ?? undefined;
+            if (selectedModule &&
+                'input_transforms' in selectedModule.value &&
+                selectedModule.value.input_transforms) {
+                helpers.updateExprsToSet(selectedId, selectedModule.value.input_transforms);
+            }
+            toolCallbacks.setToolStatus(toolId, {
+                content: `Updated flow JSON`,
+                result: 'Success'
+            });
+            return `Flow JSON updated.${warning}`;
+        }
+    },
+    {
+        def: setPreprocessorModuleToolDef,
+        streamArguments: true,
+        showDetails: true,
+        showFade: true,
+        fn: async ({ args, helpers, toolId, toolCallbacks }) => {
+            const parsedArgs = setPreprocessorModuleToolSchema.parse(args);
+            const parsedModule = validateSpecialFlowModule(parseOptionalJsonArg(parsedArgs.module, 'module'), 'preprocessor_module');
+            toolCallbacks.setToolStatus(toolId, {
+                content: parsedModule === null ? 'Removing preprocessor module...' : 'Setting preprocessor module...'
+            });
+            const updateResult = await helpers.setFlowJson({ preprocessorModule: parsedModule });
+            const warning = formatEmptyInlineScriptWarning(updateResult);
+            if (parsedModule &&
+                helpers.getFlowAndSelectedId().selectedId === SPECIAL_MODULE_IDS.PREPROCESSOR &&
+                'input_transforms' in parsedModule.value &&
+                parsedModule.value.input_transforms) {
+                helpers.updateExprsToSet(parsedModule.id, parsedModule.value.input_transforms);
+            }
+            toolCallbacks.setToolStatus(toolId, {
+                content: parsedModule === null ? 'Preprocessor module removed' : 'Preprocessor module updated',
+                result: 'Success'
+            });
+            return parsedModule === null
+                ? 'Preprocessor module removed'
+                : `Preprocessor module updated successfully.${warning}`;
+        }
+    },
+    {
+        def: setFailureModuleToolDef,
+        streamArguments: true,
+        showDetails: true,
+        showFade: true,
+        fn: async ({ args, helpers, toolId, toolCallbacks }) => {
+            const parsedArgs = setFailureModuleToolSchema.parse(args);
+            const parsedModule = validateSpecialFlowModule(parseOptionalJsonArg(parsedArgs.module, 'module'), 'failure_module');
+            toolCallbacks.setToolStatus(toolId, {
+                content: parsedModule === null ? 'Removing failure module...' : 'Setting failure module...'
+            });
+            const updateResult = await helpers.setFlowJson({ failureModule: parsedModule });
+            const warning = formatEmptyInlineScriptWarning(updateResult);
+            if (parsedModule &&
+                helpers.getFlowAndSelectedId().selectedId === SPECIAL_MODULE_IDS.FAILURE &&
+                'input_transforms' in parsedModule.value &&
+                parsedModule.value.input_transforms) {
+                helpers.updateExprsToSet(parsedModule.id, parsedModule.value.input_transforms);
+            }
+            toolCallbacks.setToolStatus(toolId, {
+                content: parsedModule === null ? 'Failure module removed' : 'Failure module updated',
+                result: 'Success'
+            });
+            return parsedModule === null
+                ? 'Failure module removed'
+                : `Failure module updated successfully.${warning}`;
+        }
+    },
     {
         def: setFlowJsonToolDef,
         streamArguments: true,
         showDetails: true,
         showFade: true,
         fn: async ({ args, helpers, toolId, toolCallbacks }) => {
-            const { modules, schema } = args;
+            const { modules, schema, preprocessor_module, failure_module } = args;
             let parsedModules;
             let parsedSchema;
+            let parsedPreprocessorModule;
+            let parsedFailureModule;
             // Parse JSON strings
-            try {
-                parsedModules = modules
-                    ? typeof modules === 'string'
-                        ? JSON.parse(modules)
-                        : modules
-                    : undefined;
-                parsedSchema = schema
-                    ? typeof schema === 'string'
-                        ? JSON.parse(schema)
-                        : schema
-                    : undefined;
-            }
-            catch (e) {
-                const errorMessage = e instanceof Error ? e.message : String(e);
-                throw new Error(`Invalid JSON: ${errorMessage}`);
-            }
-            // Validate modules against OpenFlow schema
-            if (parsedModules) {
-                const result = flowModulesSchema.safeParse(parsedModules);
-                if (!result.success) {
-                    const errors = result.error.issues.slice(0, 5).map((e) => {
-                        const path = e.path;
-                        // Try to find module id for better context
-                        const moduleIndex = typeof path[0] === 'number' ? path[0] : undefined;
-                        const moduleId = moduleIndex !== undefined ? parsedModules[moduleIndex]?.id : undefined;
-                        const fieldPath = path.slice(1).join('.');
-                        let message = e.message;
-                        if (e.code === 'invalid_type') {
-                            // Zod 4 message already contains "expected X, received Y"
-                            // Try to extract expected format from schema, passing actual data
-                            // to help resolve discriminated unions correctly
-                            const targetSchema = getSchemaAtPath(flowModulesSchema, path, parsedModules);
-                            if (targetSchema) {
-                                const expectedFormat = getExpectedFormat(targetSchema);
-                                if (expectedFormat) {
-                                    message += `\n    Expected format: ${expectedFormat}`;
-                                }
-                            }
-                        }
-                        if (moduleId) {
-                            return `Module "${moduleId}" -> ${fieldPath}: ${message}`;
-                        }
-                        return `${path.join('.')}: ${message}`;
-                    });
-                    throw new Error(`Invalid flow modules:\n${errors.join('\n')}`);
-                }
-                else {
-                    // check for duplicate ids
-                    const ids = collectAllModuleIdsFromArray(parsedModules);
-                    if (ids.length !== new Set(ids).size) {
-                        throw new Error('Duplicate module IDs found in flow');
-                    }
+            parsedModules = parseOptionalJsonArg(modules, 'modules');
+            parsedSchema = parseOptionalJsonArg(schema, 'schema');
+            parsedPreprocessorModule = parseOptionalJsonArg(preprocessor_module, 'preprocessor_module');
+            parsedFailureModule = parseOptionalJsonArg(failure_module, 'failure_module');
+            if (parsedModules === null) {
+                parsedModules = undefined;
+            }
+            if (parsedSchema === null) {
+                parsedSchema = undefined;
+            }
+            if (parsedModules !== undefined) {
+                parsedModules = validateFlowModules(parsedModules);
+                const reservedIds = collectAllFlowModuleIdsFromModules(parsedModules).filter((id) => id === SPECIAL_MODULE_IDS.PREPROCESSOR || id === SPECIAL_MODULE_IDS.FAILURE);
+                if (reservedIds.length > 0) {
+                    throw new Error('Special modules must be provided via preprocessor_module and failure_module, not inside modules');
                 }
             }
+            if (parsedSchema !== undefined) {
+                parsedSchema = validateFlowSchema(parsedSchema);
+            }
+            parsedPreprocessorModule = validateSpecialFlowModule(parsedPreprocessorModule, 'preprocessor_module');
+            parsedFailureModule = validateSpecialFlowModule(parsedFailureModule, 'failure_module');
+            const ids = [
+                ...(parsedModules ? collectAllFlowModuleIdsFromModules(parsedModules) : []),
+                ...([parsedPreprocessorModule, parsedFailureModule].filter((module) => module !== undefined && module !== null)
+                    .map((module) => module.id))
+            ];
+            if (ids.length !== new Set(ids).size) {
+                throw new Error('Duplicate module IDs found in flow');
+            }
             toolCallbacks.setToolStatus(toolId, {
                 content: `Setting flow...`
             });
-            await helpers.setFlowJson(parsedModules, parsedSchema);
+            const updateResult = await helpers.setFlowJson({
+                ...(parsedModules !== undefined ? { modules: parsedModules } : {}),
+                ...(parsedSchema !== undefined ? { schema: parsedSchema } : {}),
+                ...(parsedPreprocessorModule !== undefined
+                    ? { preprocessorModule: parsedPreprocessorModule }
+                    : {}),
+                ...(parsedFailureModule !== undefined ? { failureModule: parsedFailureModule } : {})
+            });
+            const warning = formatEmptyInlineScriptWarning(updateResult);
             // Update exprsToSet if the selected module has input_transforms
-            if (parsedModules) {
+            if (parsedModules !== undefined ||
+                parsedPreprocessorModule !== undefined ||
+                parsedFailureModule !== undefined) {
                 const { selectedId } = helpers.getFlowAndSelectedId();
-                const selectedModule = findModuleById(parsedModules, selectedId);
+                const selectedModule = selectedId === SPECIAL_MODULE_IDS.PREPROCESSOR
+                    ? parsedPreprocessorModule ?? undefined
+                    : selectedId === SPECIAL_MODULE_IDS.FAILURE
+                        ? parsedFailureModule ?? undefined
+                        : parsedModules
+                            ? findModuleInModules(parsedModules, selectedId)
+                            : undefined;
                 if (selectedModule &&
                     'input_transforms' in selectedModule.value &&
                     selectedModule.value.input_transforms) {
@@ -528,7 +874,7 @@ export const flowTools = [
                 content: `Flow updated`,
                 result: 'Success'
             });
-            return `Flow updated`;
+            return `Flow updated.${warning}`;
         }
     },
     {
@@ -558,7 +904,10 @@ export function prepareFlowSystemMessage(customPrompt) {
 ## Tool Selection Guide
 
 **Flow Modification:**
-- **Create or modify the entire flow** → \`set_flow_json\` (provide complete modules array and optional schema)
+- **Quick exact edits to current flow JSON** → \`patch_flow_json\` (provide \`old_string\` and \`new_string\`; default is one exact match)
+- **Update only the preprocessor** → \`set_preprocessor_module\`
+- **Update only the failure handler** → \`set_failure_module\`
+- **Create or replace the full flow** → \`set_flow_json\`
 
 **Code & Scripts:**
 - **View existing inline script code** → \`inspect_inline_script\`
@@ -580,13 +929,39 @@ export function prepareFlowSystemMessage(customPrompt) {
 - **Search resource types** → \`resource_type\`
 - **Get database schema** → \`get_db_schema\`
 
+## Quick Edits with patch_flow_json
+
+Use \`patch_flow_json\` for small, localized changes when you can target an exact snippet from the \`CURRENT FLOW JSON COMPACT\` block below.
+
+Always copy the exact search text from the \`CURRENT FLOW JSON COMPACT\` block below.
+The compact JSON is a single object with \`modules\`, \`schema\`, \`preprocessor_module\`, and \`failure_module\` keys.
+
+**Parameters:**
+- \`old_string\`: Exact JSON text to find
+- \`new_string\`: Replacement JSON text
+- \`replace_all\`: Optional boolean. Leave false unless you intentionally want to replace every exact match.
+
+**Example - Rename a referenced result:**
+\`\`\`javascript
+patch_flow_json({
+  old_string: "\"expr\":\"results.fetch_data\"",
+  new_string: "\"expr\":\"results.load_data\""
+})
+\`\`\`
+
+Use \`set_flow_json\` instead when you need to do a larger rewrite, add many new modules, or change the flow schema.
+
+${FLOW_CHAT_SPECIAL_MODULES}
+
 ## Flow Modification with set_flow_json
 
-Use the \`set_flow_json\` tool to set the entire flow structure at once. Provide the complete modules array and optionally the flow input schema.
+Use the \`set_flow_json\` tool to set the entire flow structure at once. Provide the complete modules array and optionally the flow input schema, \`preprocessor_module\`, and \`failure_module\`.
 
 **Parameters:**
 - \`modules\`: Array of flow modules (required)
 - \`schema\`: Flow input schema in JSON Schema format (optional)
+- \`preprocessor_module\`: Special module that runs before \`modules\` (optional, separate from \`modules\`)
+- \`failure_module\`: Special module that runs on failure (optional, separate from \`modules\`)
 
 **Example - Simple flow:**
 \`\`\`javascript
@@ -763,6 +1138,10 @@ To reduce token usage, rawscript content in the flow you receive is replaced wit
 **To inspect existing code:**
 - Use \`inspect_inline_script\` tool to view the current code: \`inspect_inline_script({ moduleId: "step_a" })\`
 
+**If a flow update tool warns that inline scripts are empty:**
+- The module structure was created successfully, but the code is still empty
+- Immediately call \`set_module_code\` for each warned module ID
+
 ### Writing Code for Modules
 
 **IMPORTANT: Before writing any code for a rawscript module, you MUST call the \`get_instructions_for_code_generation\` tool with the target language.** This tool provides essential language-specific instructions.
@@ -771,14 +1150,15 @@ Example: Before writing TypeScript/Bun code, call \`get_instructions_for_code_ge
 
 ### Creating Flows
 
-1. **Search for existing scripts first** (unless user explicitly asks to write from scratch):
-   - First: \`search_workspace\` to find workspace scripts and flows
-   - Use \`get_runnable_details\` to inspect a specific script or flow (inputs, description, code)
+1. **Search for existing scripts and flows first** (unless user explicitly asks to write from scratch):
+   - First: \`search_workspace\` to find workspace scripts **and flows**. Existing flows can be reused as subflow steps — prefer this over rebuilding equivalent logic inline.
+   - Use \`get_runnable_details\` to inspect a specific script or flow (inputs, description, code) so you know how to wire its \`input_transforms\`
    - Then: \`search_hub_scripts\` (only consider highly relevant results)
-   - Only create raw scripts if no suitable script is found
+   - Only create raw scripts if no suitable script or flow is found
 
 2. **Build the complete flow using \`set_flow_json\`:**
-   - If using existing script: use \`type: "script"\` with \`path\`
+   - If using an existing script: use \`type: "script"\` with \`path\`
+   - If using an existing flow as a subflow step: use \`type: "flow"\` with \`path\` (e.g. \`{ type: "flow", path: "f/flows/process_user", input_transforms: { ... } }\`). The step's \`input_transforms\` must cover the subflow's inputs.
    - If creating rawscript: use \`type: "rawscript"\` with \`language\` and \`content\`
    - **First call \`get_instructions_for_code_generation\` to get the correct code format**
    - Always define \`input_transforms\` to connect parameters to flow inputs or previous step results
@@ -839,7 +1219,7 @@ You have access to the following contexts:
         content
     };
 }
-export function prepareFlowUserMessage(instructions, flowAndSelectedId, selectedContext = []) {
+export function prepareFlowUserMessage(instructions, flowAndSelectedId, selectedContext = [], inlineScriptSession) {
     const flow = flowAndSelectedId?.flow;
     const selectedId = flowAndSelectedId?.selectedId;
     // Handle context elements
@@ -852,43 +1232,14 @@ ${instructions}`;
             content: userMessage
         };
     }
-    const codePieces = selectedContext.filter((c) => c.type === 'flow_module_code_piece');
+    const scriptSession = inlineScriptSession;
     // Clear the inline script store and extract inline scripts for token optimization
-    inlineScriptStore.clear();
-    const optimizedModules = extractAndReplaceInlineScripts(flow.value.modules);
-    // Apply code pieces to the optimized modules (returns YAML string)
-    const flowModulesYaml = applyCodePiecesToFlowModules(codePieces, optimizedModules);
-    // Handle preprocessor and failure modules
-    let optimizedPreprocessor = flow.value.preprocessor_module;
-    if (optimizedPreprocessor?.value?.type === 'rawscript' && optimizedPreprocessor.value.content) {
-        inlineScriptStore.set(optimizedPreprocessor.id, optimizedPreprocessor.value.content);
-        optimizedPreprocessor = {
-            ...optimizedPreprocessor,
-            value: {
-                ...optimizedPreprocessor.value,
-                content: `inline_script.${optimizedPreprocessor.id}`
-            }
-        };
-    }
-    let optimizedFailure = flow.value.failure_module;
-    if (optimizedFailure?.value?.type === 'rawscript' && optimizedFailure.value.content) {
-        inlineScriptStore.set(optimizedFailure.id, optimizedFailure.value.content);
-        optimizedFailure = {
-            ...optimizedFailure,
-            value: {
-                ...optimizedFailure.value,
-                content: `inline_script.${optimizedFailure.id}`
-            }
-        };
-    }
-    const finalFlow = {
-        schema: flow.schema,
-        modules: flowModulesYaml,
-        preprocessor_module: optimizedPreprocessor,
-        failure_module: optimizedFailure
-    };
-    let flowContent = `## CURRENT FLOW JSON:
-${JSON.stringify(finalFlow, null, 2)}
+    scriptSession?.clear();
+    const editableFlowJson = buildEditableFlowJson(flow, scriptSession, selectedContext);
+    let flowContent = `## CURRENT FLOW JSON COMPACT:
+\`\`\`json
+${JSON.stringify(editableFlowJson)}
+\`\`\`
 
 currently selected step:
 ${selectedId}`;