windmill-components 1.677.1 → 1.687.0

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

@@ -10,6 +10,44 @@ const memo = (factory) => {
     let cached;
     return () => (cached ??= factory());
 };
+function countExactMatches(content, search) {
+    if (search.length === 0) {
+        return 0;
+    }
+    let count = 0;
+    let index = 0;
+    while ((index = content.indexOf(search, index)) !== -1) {
+        count += 1;
+        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 resolveAppPatchTarget(rawPath) {
+    const trimmedPath = rawPath.trim();
+    const backendMatch = trimmedPath.match(/^\/?backend\/([^/]+)\/main\.(ts|py)$/);
+    if (backendMatch) {
+        return {
+            type: 'backend',
+            path: trimmedPath.startsWith('/') ? trimmedPath.slice(1) : trimmedPath,
+            key: backendMatch[1],
+            extension: backendMatch[2]
+        };
+    }
+    return {
+        type: 'frontend',
+        path: trimmedPath.startsWith('/') ? trimmedPath : `/${trimmedPath}`
+    };
+}
+function getBackendInlineScriptExtension(runnable) {
+    return runnable.inlineScript?.language === 'python3' ? 'py' : 'ts';
+}
 // ============= Frontend File Tools =============
 const getGetFrontendFileSchema = memo(() => z.object({
     path: z
@@ -24,6 +62,19 @@ const getSetFrontendFileSchema = memo(() => z.object({
     content: z.string().describe('The content of the file')
 }));
 const getSetFrontendFileToolDef = memo(() => createToolDef(getSetFrontendFileSchema(), 'set_frontend_file', 'Create or update a frontend file in the raw app. Returns lint diagnostics (errors and warnings).'));
+const getPatchFileSchema = memo(() => z.object({
+    path: z
+        .string()
+        .describe('Path of the file to patch. Use frontend paths like /index.tsx or inline backend paths like backend/listRecipes/main.ts.'),
+    old_string: z.string().min(1).describe('Exact text to find in the current file content'),
+    new_string: z.string().describe('Replacement text'),
+    replace_all: z
+        .boolean()
+        .optional()
+        .default(false)
+        .describe('When true, replace every exact match. When false, old_string must match exactly once.')
+}));
+const getPatchFileToolDef = memo(() => createToolDef(getPatchFileSchema(), 'patch_file', 'Make a quick exact text edit in an existing frontend file or inline backend file. Prefer this for small localized changes instead of rewriting the whole file.'));
 const getDeleteFrontendFileSchema = memo(() => z.object({
     path: z.string().describe('The path of the frontend file to delete')
 }));
@@ -245,6 +296,73 @@ export const getAppTools = memo(() => [
         showDetails: true,
         showFade: true
     },
+    {
+        def: getPatchFileToolDef(),
+        streamArguments: true,
+        showDetails: true,
+        showFade: true,
+        fn: async ({ args, helpers, toolId, toolCallbacks }) => {
+            const parsedArgs = getPatchFileSchema().parse(args);
+            const { old_string: oldString, new_string: newString, replace_all: replaceAll } = parsedArgs;
+            const target = resolveAppPatchTarget(parsedArgs.path);
+            let currentContent = '';
+            let backendRunnable;
+            if (target.type === 'frontend') {
+                if (target.path === '/wmill.d.ts') {
+                    throw new Error("'/wmill.d.ts' is generated automatically. Edit backend runnables instead.");
+                }
+                const frontendContent = helpers.getFrontendFile(target.path);
+                if (frontendContent === undefined) {
+                    throw new Error(`Frontend file '${target.path}' not found.`);
+                }
+                currentContent = frontendContent;
+            }
+            else {
+                backendRunnable = helpers.getBackendRunnable(target.key);
+                if (!backendRunnable) {
+                    throw new Error(`Backend runnable '${target.key}' not found.`);
+                }
+                if (backendRunnable.type !== 'inline' || !backendRunnable.inlineScript) {
+                    throw new Error(`'${target.path}' points to backend runnable '${target.key}', but only inline runnables can be patched as files.`);
+                }
+                const expectedExtension = getBackendInlineScriptExtension(backendRunnable);
+                if (target.extension !== expectedExtension) {
+                    throw new Error(`'${target.path}' does not match runnable '${target.key}' language. Use backend/${target.key}/main.${expectedExtension}.`);
+                }
+                currentContent = backendRunnable.inlineScript.content ?? '';
+            }
+            const matchCount = countExactMatches(currentContent, oldString);
+            if (matchCount === 0) {
+                throw new Error('old_string was not found in the current file content.');
+            }
+            if (!replaceAll && matchCount !== 1) {
+                throw new Error(`old_string matched ${matchCount} locations. Make it more specific or set replace_all to true.`);
+            }
+            const updatedContent = replaceAll
+                ? currentContent.split(oldString).join(newString)
+                : replaceFirstExactMatch(currentContent, oldString, newString);
+            toolCallbacks.setToolStatus(toolId, {
+                content: `Patching '${target.path}'...`
+            });
+            const lintResult = target.type === 'frontend'
+                ? helpers.setFrontendFile(target.path, updatedContent)
+                : await helpers.setBackendRunnable(target.key, {
+                    ...backendRunnable,
+                    inlineScript: {
+                        ...backendRunnable.inlineScript,
+                        content: updatedContent
+                    }
+                });
+            toolCallbacks.setToolStatus(toolId, {
+                content: `Patched '${target.path}'`,
+                result: 'Success'
+            });
+            return formatLintResultResponse(`Patched '${target.path}' successfully.`, lintResult);
+        },
+        preAction: ({ toolCallbacks, toolId }) => {
+            toolCallbacks.setToolStatus(toolId, { content: 'Patching file...' });
+        }
+    },
     {
         def: getDeleteFrontendFileToolDef(),
         fn: async ({ args, helpers, toolId, toolCallbacks }) => {
@@ -473,6 +591,7 @@ For inline scripts, the code must have a \`main\` function as its entrypoint.
 ### File Management
 - \`get_files()\`: Get an overview of all files (content may be truncated for large files)
 - \`get_frontend_file(path)\`: Get full content of a specific frontend file
+- \`patch_file(path, old_string, new_string, replace_all?)\`: Make a quick exact text edit to an existing file. Use frontend paths like \`/index.tsx\` and inline backend paths like \`backend/listRecipes/main.ts\`.
 - \`set_frontend_file(path, content)\`: Create or update a frontend file. Returns lint diagnostics.
 - \`delete_frontend_file(path)\`: Delete a frontend file
 - \`get_backend_runnable(key)\`: Get full configuration of a specific backend runnable
@@ -482,6 +601,15 @@ For inline scripts, the code must have a \`main\` function as its entrypoint.
 ### Linting
 - \`lint()\`: Lint all files. Returns errors/warnings grouped by frontend/backend. Use this to check for issues after making changes.
 
+## Quick Edits with patch_file
+
+Use \`patch_file\` for small, localized edits when you can copy an exact snippet from the current file contents.
+
+- \`old_string\` must be copied exactly from the current file contents
+- Leave \`replace_all\` as false unless you intentionally want to update every exact match
+- Never patch \`/wmill.d.ts\`; it is generated automatically from backend runnables
+- Use \`set_frontend_file\` or \`set_backend_runnable\` instead for large rewrites, new files, or new runnables
+
 ### Discovery
 - \`search_workspace(query, type)\`: Search workspace scripts and flows
 - \`get_runnable_details(path, type)\`: Get details (summary, description, schema, content) of a specific script or flow
@@ -620,7 +748,7 @@ When you are using the windmill-client, do not forget that as id for variables o
 
 1. Start with \`get_files()\` to get an overview of all frontend files and backend runnables (content may be truncated for large files)
 2. Use \`get_frontend_file(path)\` or \`get_backend_runnable(key)\` to get full content of specific files when needed
-3. Make changes using \`set_frontend_file\` and \`set_backend_runnable\`. These return lint diagnostics.
+3. For small localized edits, prefer \`patch_file\`. Use \`set_frontend_file\` and \`set_backend_runnable\` for larger rewrites, new files, or new runnables. These return lint diagnostics.
 4. Use \`lint()\` at the end to check for and fix any remaining errors
 
 When creating a new app, use \`search_workspace\` or \`search_hub_scripts\` to find existing scripts/flows to reuse.

package/package/components/copilot/chat/app/core.test.js

@@ -0,0 +1,192 @@
+import { describe, expect, it, vi } from 'vitest';
+import { getAppTools } from './core';
+vi.mock('../shared', () => ({
+    createToolDef: (_schema, name, description) => ({
+        type: 'function',
+        function: {
+            name,
+            description,
+            parameters: {}
+        }
+    }),
+    createSearchWorkspaceTool: () => ({
+        def: {
+            type: 'function',
+            function: {
+                name: 'search_workspace',
+                description: 'search workspace',
+                parameters: {}
+            }
+        },
+        fn: async () => ''
+    }),
+    createGetRunnableDetailsTool: () => ({
+        def: {
+            type: 'function',
+            function: {
+                name: 'get_runnable_details',
+                description: 'get runnable details',
+                parameters: {}
+            }
+        },
+        fn: async () => ''
+    }),
+    createSearchHubScriptsTool: () => ({
+        def: {
+            type: 'function',
+            function: {
+                name: 'search_hub_scripts',
+                description: 'search hub scripts',
+                parameters: {}
+            }
+        },
+        fn: async () => ''
+    })
+}));
+vi.mock('../AIChatManager.svelte', () => ({
+    aiChatManager: {
+        datatableCreationPolicy: {
+            enabled: false,
+            datatable: undefined,
+            schema: undefined
+        }
+    }
+}));
+vi.mock('$system_prompts', () => ({
+    getDatatableSdkReference: () => 'Datatable SDK reference'
+}));
+const EMPTY_LINT_RESULT = {
+    errorCount: 0,
+    warningCount: 0,
+    errors: {
+        frontend: {},
+        backend: {}
+    },
+    warnings: {
+        frontend: {},
+        backend: {}
+    }
+};
+function createHelpers(overrides = {}) {
+    return {
+        listFrontendFiles: () => [],
+        getFrontendFile: () => undefined,
+        getFrontendFiles: () => ({}),
+        setFrontendFile: () => EMPTY_LINT_RESULT,
+        deleteFrontendFile: () => undefined,
+        listBackendRunnables: () => [],
+        getBackendRunnable: () => undefined,
+        getBackendRunnables: () => ({}),
+        setBackendRunnable: async () => EMPTY_LINT_RESULT,
+        deleteBackendRunnable: () => undefined,
+        getFiles: () => ({ frontend: {}, backend: {} }),
+        getSelectedContext: () => ({ type: 'none' }),
+        snapshot: () => 1,
+        revertToSnapshot: () => undefined,
+        lint: () => EMPTY_LINT_RESULT,
+        getDatatables: async () => [],
+        getAvailableDatatableNames: () => [],
+        execDatatableSql: async () => ({ success: true }),
+        addTableToWhitelist: () => undefined,
+        ...overrides
+    };
+}
+function createToolCallbacks() {
+    return {
+        setToolStatus: vi.fn(),
+        removeToolStatus: vi.fn()
+    };
+}
+function getPatchFileTool() {
+    const tool = getAppTools().find((entry) => entry.def.function.name === 'patch_file');
+    if (!tool) {
+        throw new Error('patch_file tool not found');
+    }
+    return tool;
+}
+describe('app patch_file tool', () => {
+    it('patches frontend files with an exact replacement', async () => {
+        const setFrontendFile = vi.fn(() => EMPTY_LINT_RESULT);
+        const tool = getPatchFileTool();
+        const result = await tool.fn({
+            args: {
+                path: '/index.tsx',
+                old_string: 'Hello recipe book',
+                new_string: 'Hello cookbook'
+            },
+            workspace: 'test-workspace',
+            helpers: createHelpers({
+                getFrontendFile: () => 'export const title = "Hello recipe book"\n',
+                setFrontendFile
+            }),
+            toolCallbacks: createToolCallbacks(),
+            toolId: 'tool-1'
+        });
+        expect(setFrontendFile).toHaveBeenCalledWith('/index.tsx', 'export const title = "Hello cookbook"\n');
+        expect(result).toContain("Patched '/index.tsx' successfully.");
+    });
+    it('rejects ambiguous matches unless replace_all is set', async () => {
+        const tool = getPatchFileTool();
+        await expect(tool.fn({
+            args: {
+                path: '/index.tsx',
+                old_string: 'recipe',
+                new_string: 'meal'
+            },
+            workspace: 'test-workspace',
+            helpers: createHelpers({
+                getFrontendFile: () => 'recipe\nrecipe\n'
+            }),
+            toolCallbacks: createToolCallbacks(),
+            toolId: 'tool-2'
+        })).rejects.toThrow('old_string matched 2 locations');
+    });
+    it('patches inline backend runnables through backend/<key>/main.ts paths', async () => {
+        const runnable = {
+            name: 'Delete recipe',
+            type: 'inline',
+            inlineScript: {
+                language: 'bun',
+                content: 'export async function main() {\n\treturn "recipe"\n}\n'
+            }
+        };
+        const setBackendRunnable = vi.fn(async () => EMPTY_LINT_RESULT);
+        const tool = getPatchFileTool();
+        const result = await tool.fn({
+            args: {
+                path: 'backend/deleteRecipe/main.ts',
+                old_string: '"recipe"',
+                new_string: '"meal"'
+            },
+            workspace: 'test-workspace',
+            helpers: createHelpers({
+                getBackendRunnable: () => runnable,
+                setBackendRunnable
+            }),
+            toolCallbacks: createToolCallbacks(),
+            toolId: 'tool-3'
+        });
+        expect(setBackendRunnable).toHaveBeenCalledWith('deleteRecipe', expect.objectContaining({
+            inlineScript: expect.objectContaining({
+                content: 'export async function main() {\n\treturn "meal"\n}\n'
+            })
+        }));
+        expect(result).toContain("Patched 'backend/deleteRecipe/main.ts' successfully.");
+    });
+    it('rejects edits to generated wmill types', async () => {
+        const tool = getPatchFileTool();
+        await expect(tool.fn({
+            args: {
+                path: '/wmill.d.ts',
+                old_string: 'backend',
+                new_string: 'backendAsync'
+            },
+            workspace: 'test-workspace',
+            helpers: createHelpers({
+                getFrontendFile: () => 'generated'
+            }),
+            toolCallbacks: createToolCallbacks(),
+            toolId: 'tool-4'
+        })).rejects.toThrow('generated automatically');
+    });
+});

package/package/components/copilot/chat/chatLoop.d.ts

@@ -3,6 +3,7 @@ import Anthropic from '@anthropic-ai/sdk';
 import type { ChatCompletionMessageParam, ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam } from 'openai/resources/chat/completions.mjs';
 import type { AIProviderModel } from '../../../gen';
 import type { Tool, ToolCallbacks } from './shared';
+import { type ChatTokenUsage } from './tokenUsage';
 export interface ChatClients {
     openai: OpenAI;
     anthropic: Anthropic;
@@ -36,5 +37,7 @@ export interface ChatLoopConfig {
 }
 export interface ChatLoopResult {
     addedMessages: ChatCompletionMessageParam[];
+    tokenUsage: ChatTokenUsage;
+    hitMaxIterations: boolean;
 }
 export declare function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResult>;

package/package/components/copilot/chat/chatLoop.js

@@ -3,13 +3,17 @@ import Anthropic from '@anthropic-ai/sdk';
 import { getCompletion, parseOpenAICompletion } from '../lib';
 import { getAnthropicCompletion, parseAnthropicCompletion } from './anthropic';
 import { getOpenAIResponsesCompletion, parseOpenAIResponsesCompletion } from './openai-responses';
+import { addChatTokenUsage, emptyChatTokenUsage } from './tokenUsage';
 export async function runChatLoop(config) {
     const { messages, abortController, callbacks, clients, workspace, maxIterations, onSkipResponsesApi, getPendingUserMessage, onBeforeIteration } = config;
     let skipResponsesApi = config.skipResponsesApi ?? false;
     const addedMessages = [];
+    let tokenUsage = emptyChatTokenUsage();
     let iterations = 0;
+    let hitMaxIterations = false;
     while (true) {
         if (maxIterations !== undefined && iterations >= maxIterations) {
+            hitMaxIterations = true;
             break;
         }
         iterations++;
@@ -42,7 +46,8 @@ export async function runChatLoop(config) {
                         openaiClient: clients.openai
                     });
                     const continueCompletion = await parseOpenAIResponsesCompletion(completion, callbacks, messages, addedMessages, tools, helpers, parseOptions);
-                    if (!continueCompletion) {
+                    tokenUsage = addChatTokenUsage(tokenUsage, continueCompletion.tokenUsage);
+                    if (!continueCompletion.shouldContinue) {
                         break;
                     }
                 }
@@ -63,7 +68,8 @@ export async function runChatLoop(config) {
                     openaiClient: clients.openai
                 });
                 const continueCompletion = await parseOpenAICompletion(completion, callbacks, messages, addedMessages, tools, helpers, undefined, parseOptions);
-                if (!continueCompletion) {
+                tokenUsage = addChatTokenUsage(tokenUsage, continueCompletion.tokenUsage);
+                if (!continueCompletion.shouldContinue) {
                     break;
                 }
             }
@@ -75,7 +81,8 @@ export async function runChatLoop(config) {
             });
             if (completion) {
                 const continueCompletion = await parseAnthropicCompletion(completion, callbacks, messages, addedMessages, tools, helpers, abortController, parseOptions);
-                if (!continueCompletion) {
+                tokenUsage = addChatTokenUsage(tokenUsage, continueCompletion.tokenUsage);
+                if (!continueCompletion.shouldContinue) {
                     break;
                 }
             }
@@ -87,11 +94,12 @@ export async function runChatLoop(config) {
             });
             if (completion) {
                 const continueCompletion = await parseOpenAICompletion(completion, callbacks, messages, addedMessages, tools, helpers, undefined, parseOptions);
-                if (!continueCompletion) {
+                tokenUsage = addChatTokenUsage(tokenUsage, continueCompletion.tokenUsage);
+                if (!continueCompletion.shouldContinue) {
                     break;
                 }
             }
         }
     }
-    return { addedMessages };
+    return { addedMessages, tokenUsage, hitMaxIterations };
 }

package/package/components/copilot/chat/flow/FlowAIChat.svelte

@@ -1,29 +1,18 @@
 <script lang="ts">import FlowModuleSchemaMap from '../../../flows/map/FlowModuleSchemaMap.svelte';
 import { getContext, untrack } from 'svelte';
-import { dfs } from '../../../flows/previousResults';
-import { restoreInlineScriptReferences } from './inlineScriptsUtils';
+import { createInlineScriptSession } from './inlineScriptsUtils';
 import { loadSchemaFromModule } from '../../../flows/flowInfers';
 import { aiChatManager } from '../AIChatManager.svelte';
 import { refreshStateStore } from '../../../../svelte5Utils.svelte';
-import { getSubModules } from '../../../flows/flowExplorer';
-import { SPECIAL_MODULE_IDS } from '../shared';
+import { applyFlowJsonUpdate, updateRawScriptModuleContent } from './helperUtils';
+import { findModuleInFlow } from '../../../flows/flowTree';
 let { flowModuleSchemaMap, onTestFlow } = $props();
 const { flowStore, flowStateStore, selectionManager, currentEditor, previewArgs } = getContext('FlowEditorContext');
 const selectedId = $derived(selectionManager.getSelectedId());
 const { exprsToSet } = getContext('FlowCopilotContext') ?? {};
+const inlineScriptSession = createInlineScriptSession();
 // Get diffManager from the graph
 const diffManager = $derived(flowModuleSchemaMap?.getDiffManager());
-function getModule(id, flow = flowStore.val) {
-    if (id === SPECIAL_MODULE_IDS.PREPROCESSOR) {
-        return flow.value.preprocessor_module;
-    }
-    else if (id === SPECIAL_MODULE_IDS.FAILURE) {
-        return flow.value.failure_module;
-    }
-    else {
-        return dfs(id, flow, false)[0];
-    }
-}
 const flowHelpers = {
     // flow context
     getFlowAndSelectedId: () => {
@@ -33,16 +22,10 @@ const flowHelpers = {
             selectedId: selectedId === 'settings-metadata' ? '' : selectedId
         };
     },
-    getModules: (id) => {
-        if (id) {
-            const module = getModule(id);
-            if (!module) {
-                throw new Error('Module not found');
-            }
-            return getSubModules(module).flat();
-        }
+    getRootModules: () => {
         return flowStore.val.value.modules;
     },
+    inlineScriptSession,
     setSnapshot: (snapshot) => {
         diffManager?.setBeforeFlow(snapshot);
     },
@@ -54,7 +37,7 @@ const flowHelpers = {
         // Update current editor if needed
         const targetSnapshot = snapshot ?? diffManager.beforeFlow;
         if ($currentEditor && targetSnapshot) {
-            const module = getModule($currentEditor.stepId, targetSnapshot);
+            const module = findModuleInFlow(targetSnapshot.value, $currentEditor.stepId) ?? undefined;
             if (module) {
                 if ($currentEditor.type === 'script' && module.value.type === 'rawscript') {
                     $currentEditor.editor.setCode(module.value.content);
@@ -67,48 +50,43 @@ const flowHelpers = {
     },
     // ai chat tools
     setCode: async (id, code) => {
-        const module = getModule(id);
+        // 1. Take snapshot only if none exists (preserves baseline for cumulative changes)
+        if (!diffManager?.beforeFlow) {
+            const snapshot = $state.snapshot(flowStore).val;
+            diffManager?.setBeforeFlow(snapshot);
+            diffManager?.setEditMode(true);
+        }
+        // 2. Apply the code change
+        const module = updateRawScriptModuleContent(flowStore.val, id, code);
         if (!module) {
-            throw new Error('Module not found');
+            throw new Error('Module not found or is not a rawscript');
         }
-        if (module.value.type === 'rawscript') {
-            // 1. Take snapshot only if none exists (preserves baseline for cumulative changes)
-            if (!diffManager?.beforeFlow) {
-                const snapshot = $state.snapshot(flowStore).val;
-                diffManager?.setBeforeFlow(snapshot);
-                diffManager?.setEditMode(true);
-            }
-            // 2. Apply the code change
-            module.value.content = code;
-            const { input_transforms, schema } = await loadSchemaFromModule(module);
-            module.value.input_transforms = input_transforms;
-            refreshStateStore(flowStore);
-            // Update exprsToSet if this module is currently selected
-            if (id === selectedId && exprsToSet) {
-                exprsToSet.set(input_transforms);
-            }
-            if (flowStateStore.val[id]) {
-                flowStateStore.val[id].schema = schema;
-            }
-            else {
-                flowStateStore.val[id] = {
-                    schema
-                };
-            }
-            // 3. Manually add to moduleActions, preserving existing action types
-            // Note: currentFlow is auto-synced by FlowGraphV2's effect after refreshStateStore
-            const currentAction = diffManager?.moduleActions[id];
-            if (!currentAction) {
-                diffManager?.setModuleActions({
-                    ...diffManager?.moduleActions,
-                    [id]: { action: 'modified', pending: true }
-                });
-            }
-            // If already tracked (e.g., 'added' from setFlowJson), keep that status
+        inlineScriptSession.set(id, code);
+        const { input_transforms, schema } = await loadSchemaFromModule(module);
+        module.value.input_transforms = input_transforms;
+        refreshStateStore(flowStore);
+        // Update exprsToSet if this module is currently selected
+        if (id === selectedId && exprsToSet) {
+            exprsToSet.set(input_transforms);
+        }
+        if (flowStateStore.val[id]) {
+            flowStateStore.val[id].schema = schema;
         }
         else {
-            throw new Error('Module is not a rawscript or script');
+            flowStateStore.val[id] = {
+                schema
+            };
+        }
+        // 3. Manually add to moduleActions, preserving existing action types
+        // Note: currentFlow is auto-synced by FlowGraphV2's effect after refreshStateStore
+        const currentAction = diffManager?.moduleActions[id];
+        if (!currentAction) {
+            diffManager?.setModuleActions({
+                ...diffManager?.moduleActions,
+                [id]: { action: 'modified', pending: true }
+            });
         }
+        // If already tracked (e.g., 'added' from setFlowJson), keep that status
         if ($currentEditor && $currentEditor.type === 'script' && $currentEditor.stepId === id) {
             $currentEditor.editor.setCode(code);
         }
@@ -143,7 +121,7 @@ const flowHelpers = {
         return await onTestFlow?.(conversationId);
     },
     getLintErrors: async (moduleId) => {
-        const module = getModule(moduleId);
+        const module = findModuleInFlow(flowStore.val.value, moduleId) ?? undefined;
         if (!module || module.value.type !== 'rawscript') {
             return { errorCount: 0, warningCount: 0, errors: [], warnings: [] };
         }
@@ -164,9 +142,12 @@ const flowHelpers = {
         }
         return { errorCount: 0, warningCount: 0, errors: [], warnings: [] };
     },
-    setFlowJson: async (modules, schema) => {
+    setFlowJson: async ({ modules, schema, preprocessorModule, failureModule }) => {
         try {
-            if (modules || schema) {
+            if (modules !== undefined ||
+                schema !== undefined ||
+                preprocessorModule !== undefined ||
+                failureModule !== undefined) {
                 // Take snapshot of current flowStore and set as beforeFlow
                 if (!diffManager?.hasPendingChanges) {
                     const snapshot = $state.snapshot(flowStore).val;
@@ -174,21 +155,17 @@ const flowHelpers = {
                     diffManager?.setEditMode(true);
                 }
             }
-            if (modules) {
-                // Restore inline script references back to full content
-                const restoredModules = restoreInlineScriptReferences(modules);
-                // Directly modify flowStore (immediate effect)
-                flowStore.val.value.modules = restoredModules;
-            }
-            // Update schema if provided
-            if (schema !== undefined) {
-                flowStore.val.schema = schema;
-            }
+            const result = applyFlowJsonUpdate(flowStore.val, inlineScriptSession, {
+                modules,
+                schema,
+                preprocessorModule,
+                failureModule
+            });
             // Refresh the state store to update UI
             refreshStateStore(flowStore);
+            return result;
         }
         catch (error) {
-            console.error('setFlowJson error:', error);
             throw new Error(`Failed to parse or apply JSON: ${error instanceof Error ? error.message : String(error)}`);
         }
     }
@@ -198,7 +175,9 @@ $effect(() => {
         selectedId &&
         diffManager?.moduleActions[selectedId]?.pending &&
         $currentEditor.editor.getAiChatEditorHandler()) {
-        const moduleLastSnapshot = getModule(selectedId, diffManager.beforeFlow);
+        const moduleLastSnapshot = diffManager.beforeFlow
+            ? (findModuleInFlow(diffManager.beforeFlow.value, selectedId) ?? undefined)
+            : undefined;
         const content = moduleLastSnapshot?.value.type === 'rawscript' ? moduleLastSnapshot.value.content : '';
         if (content.length > 0) {
             untrack(() => $currentEditor.editor.reviewAppliedCode(content, {