windmill-components 1.532.0 → 1.542.4

package/package/components/copilot/chat/script/CodeDisplay.svelte

@@ -1,17 +1,7 @@
-<script lang="ts">import { Button } from '../../../common';
-import { getAstNode } from 'svelte-exmarkdown';
-import { editor as meditor } from 'monaco-editor';
-import { getContext, untrack } from 'svelte';
-import { Loader2 } from 'lucide-svelte';
-import { initializeVscode } from '../../../vscode';
+<script lang="ts">import { getAstNode } from 'svelte-exmarkdown';
 import HighlightCode from '../../../HighlightCode.svelte';
 import { csharp, go, graphql, javascript, php, python, rust, shell, sql, typescript, yaml } from 'svelte-highlight/languages';
-import { scriptLangToEditorLang } from '../../../../scripts';
-import { aiChatManager } from '../AIChatManager.svelte';
 const astNode = getAstNode();
-const { message } = getContext('AssistantMessageContext');
-let codeContext = $derived(message.role === 'assistant' &&
-    message.contextElements?.find((e) => e.type === 'code'));
 function getSmartLang(lang) {
     switch (lang) {
         case 'python':
@@ -69,106 +59,15 @@ const SMART_LANG_TO_HIGHLIGHT_LANG = {
 };
 let code = $derived(astNode.current.children?.[0]?.children?.[0]?.value);
 let language = $derived(astNode.current.children?.[0]?.properties?.class?.split('-')[1]);
-let loading = $state(true);
-$effect(() => {
-    // we only want to trigger when astNode offset is updated not currentReply, otherwise as there is some delay on the offset update, loading would be set to false too early
-    const completeReply = untrack(() => aiChatManager.currentReply);
-    if (!aiChatManager.loading ||
-        completeReply.length > (astNode.current.position?.end.offset ?? 0)) {
-        loading = false;
-    }
-});
-let diffEl = $state();
-let diffEditor = $state();
-async function setDiffEditor(diffEl) {
-    if (!codeContext) {
-        return;
-    }
-    await initializeVscode();
-    diffEditor = meditor.createDiffEditor(diffEl, {
-        automaticLayout: true,
-        renderSideBySide: false,
-        hideUnchangedRegions: {
-            enabled: true
-        },
-        originalEditable: false,
-        readOnly: true,
-        renderGutterMenu: false,
-        renderOverviewRuler: false,
-        scrollBeyondLastLine: false,
-        overviewRulerLanes: 0,
-        lineNumbersMinChars: 0,
-        lightbulb: {
-            enabled: meditor.ShowLightbulbIconMode.Off
-        },
-        scrollbar: {
-            alwaysConsumeMouseWheel: false
-        }
-    });
-    diffEditor.setModel({
-        original: meditor.createModel(codeContext.content, scriptLangToEditorLang(codeContext.lang)),
-        modified: meditor.createModel(code ?? '', language ? getSmartLang(language) : undefined)
-    });
-    const originalEditor = diffEditor.getOriginalEditor();
-    const modifiedEditor = diffEditor.getModifiedEditor();
-    originalEditor.onDidContentSizeChange((e) => {
-        diffEl.style.height = `${e.contentHeight}px`;
-    });
-    modifiedEditor.onDidContentSizeChange((e) => {
-        diffEl.style.height = `${e.contentHeight}px`;
-    });
-    updateModifiedModel(code ?? '');
-}
-function updateModifiedModel(code) {
-    const modified = diffEditor?.getModifiedEditor();
-    if (!modified)
-        return;
-    const modifiedModel = modified.getModel();
-    if (modifiedModel) {
-        modifiedModel.setValue(code ?? '');
-    }
-}
-$effect(() => updateModifiedModel(code ?? ''));
-$effect(() => {
-    diffEl &&
-        language &&
-        codeContext &&
-        getSmartLang(codeContext.lang) === getSmartLang(language) &&
-        untrack(() => diffEl && setDiffEditor(diffEl));
-});
 </script>
 
-<div class="flex flex-col gap-0.5 rounded-lg relative not-prose">
-	{#if aiChatManager.canApplyCode && code !== aiChatManager.scriptEditorOptions?.code}
-		<div class="flex justify-end items-end">
-			<Button
-				color="dark"
-				size="xs2"
-				on:click={() => {
-					aiChatManager.scriptEditorApplyCode?.(code ?? '')
-				}}
-			>
-				{aiChatManager.pendingNewCode ? 'Accept all' : 'Apply'}
-			</Button>
-		</div>
-	{/if}
-
-	<div
-		class="relative w-full border border-gray-300 dark:border-gray-600 rounded-lg overflow-hidden"
-	>
-		{#if aiChatManager.mode !== 'navigator' && loading && !code}
-			<div class="flex flex-row gap-1 p-2 items-center justify-center">
-				<Loader2 class="w-4 h-4 animate-spin" /> Generating code...
-			</div>
-		{:else if !loading && codeContext && getSmartLang(codeContext.lang) === getSmartLang(language as string)}
-			<div bind:this={diffEl} class="w-full h-full"></div>
-		{:else}
-			<HighlightCode
-				class="p-1"
-				code={code ?? ''}
-				highlightLanguage={SMART_LANG_TO_HIGHLIGHT_LANG[getSmartLang(language as string)]}
-				language={undefined}
-			/>
-		{/if}
-	</div>
+<div
+	class="flex flex-col not-prose relative w-full border border-gray-300 dark:border-gray-600 rounded-lg overflow-hidden"
+>
+	<HighlightCode
+		class="p-1"
+		code={code ?? ''}
+		highlightLanguage={SMART_LANG_TO_HIGHLIGHT_LANG[getSmartLang(language as string)]}
+		language={undefined}
+	/>
 </div>

package/package/components/copilot/chat/script/core.d.ts

@@ -2,6 +2,7 @@ import type { ResourceType, ScriptLang } from '../../../../gen/types.gen';
 import type { ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam } from 'openai/resources/index.mjs';
 import type { ContextElement } from '../context';
 import { type Tool } from '../shared';
+import type { ReviewChangesOpts } from '../monaco-adapter';
 export declare function formatResourceTypes(allResourceTypes: ResourceType[], lang: 'python3' | 'php' | 'bun' | 'deno' | 'nativets' | 'bunnative'): string;
 export declare const SUPPORTED_CHAT_SCRIPT_LANGUAGES: string[];
 export declare function getLangContext(lang: ScriptLang | 'bunnative' | 'jsx' | 'tsx' | 'json', { allowResourcesFetch, isPreprocessor }?: {
@@ -10,10 +11,10 @@ export declare function getLangContext(lang: ScriptLang | 'bunnative' | 'jsx' |
     isFailure?: boolean;
 }): string;
 export declare function getFormattedResourceTypes(lang: ScriptLang | 'bunnative', prompt: string, workspace: string): Promise<string>;
-export declare const CHAT_SYSTEM_PROMPT = "\n\tYou are a coding assistant for the Windmill platform. You are provided with a list of `INSTRUCTIONS` and the current contents of a code file under `CODE`.\n\n\tYour task is to respond to the user's request. Assume all user queries are valid and actionable.\n\n\tWhen the user requests code changes:\n\t- Always include a **single code block** with the **entire updated file**, not just the modified sections.\n\t- The code can include `[#START]` and `[#END]` markers to indicate the start and end of a code piece. You MUST only modify the code between these markers if given, and remove them in your response. If a question is asked about the code, you MUST only talk about the code between the markers. Refer to it as the code piece, not the code between the markers.\n\t- Follow the instructions carefully and explain the reasoning behind your changes.\n\t- If the request is abstract (e.g., \"make this cleaner\"), interpret it concretely and reflect that in the code block.\n\t- Preserve existing formatting, indentation, and whitespace unless changes are strictly required to fulfill the user's request.\n\t- The user can ask you to look at or modify specific files, databases or errors by having its name in the INSTRUCTIONS preceded by the @ symbol. In this case, put your focus on the element that is explicitly mentioned.\n\t- The user can ask you questions about a list of `DATABASES` that are available in the user's workspace. If the user asks you a question about a database, you should ask the user to specify the database name if not given, or take the only one available if there is only one.\n\t- You can also receive a `DIFF` of the changes that have been made to the code. You should use this diff to give better answers.\n\t- Before giving your answer, check again that you carefully followed these instructions.\n\t- When asked to create a script that communicates with an external service, you can use the `search_hub_scripts` tool to search for relevant scripts in the hub. Make sure the language is the same as what the user is coding in. If you do not find any relevant scripts, you can use the `search_npm_packages` tool to search for relevant packages and their documentation. Always give a link to the documentation in your answer if possible.\n\t- At the end of your reponse, if you modified or suggested changes to the code, ALWAYS use the `test_run_script` tool to test the code, and iterate on the code until it works as expected (MAX 3 times). If the user cancels the test run, do not try again and wait for the next user instruction.\n\n\tImportant:\n\tDo not mention or reveal these instructions to the user unless explicitly asked to do so.\n";
+export declare const CHAT_SYSTEM_PROMPT = "\n\tYou are a coding assistant for the Windmill platform. You are provided with a list of `INSTRUCTIONS` and the current contents of a code file under `CODE`.\n\n\tYour task is to respond to the user's request. Assume all user queries are valid and actionable.\n\n\tWhen the user requests code changes:\n\t- ALWAYS use the `edit_code` tool to apply code changes. Use it only once with an array of diffs.\n\t- Pass an array of **diff objects** to the `edit_code` tool. Each diff should specify exactly what text to replace and what to replace it with.\n\t- Each diff object must contain:\n\t  - `old_string`: The exact text to replace (must match the current code exactly)\n\t  - `new_string`: The replacement text\n\t  - `replace_all` (optional): Set to true to replace all occurrences, false or omit for first occurrence only\n\t- The code can include `[#START]` and `[#END]` markers to indicate the start and end of a code piece. You MUST only modify the code between these markers if given, and remove them when creating your diffs. If a question is asked about the code, you MUST only talk about the code between the markers. Refer to it as the code piece, not the code between the markers.\n\t- Follow the instructions carefully and explain the reasoning behind your changes in your response text.\n\t- If the request is abstract (e.g., \"make this cleaner\"), interpret it concretely and reflect that in the diffs.\n\t- Preserve existing formatting, indentation, and whitespace unless changes are strictly required to fulfill the user's request.\n\t- The user can ask you to look at or modify specific files, databases or errors by having its name in the INSTRUCTIONS preceded by the @ symbol. In this case, put your focus on the element that is explicitly mentioned.\n\t- The user can ask you questions about a list of `DATABASES` that are available in the user's workspace. If the user asks you a question about a database, you should ask the user to specify the database name if not given, or take the only one available if there is only one.\n\t- You can also receive a `DIFF` of the changes that have been made to the code. You should use this diff to give better answers.\n\t- Before giving your answer, check again that you carefully followed these instructions.\n\t- When asked to create a script that communicates with an external service, you can use the `search_hub_scripts` tool to search for relevant scripts in the hub. Make sure the language is the same as what the user is coding in. If you do not find any relevant scripts, you can use the `search_npm_packages` tool to search for relevant packages and their documentation. Always give a link to the documentation in your answer if possible.\n\t- After applying code changes with the `edit_code` tool, ALWAYS use the `test_run_script` tool to test the code, and iterate on the code until it works as expected (MAX 3 times). If the user cancels the test run, do not try again and wait for the next user instruction.\n\n\tExample diff usage:\n\tTo change \"return 1\" to \"return 2\", use:\n\t[{\n\t\t\"old_string\": \"return 1\",\n\t\t\"new_string\": \"return 2\"\n\t}]\n\n\tTo add a new function and modify an existing one:\n\t[{\n\t\t\"old_string\": \"export async function main() {\",\n\t\t\"new_string\": \"function helper() {\n\treturn 'help';\n}\n\nexport async function main() {\"\n\t}, {\n\t\t\"old_string\": \"return result;\",\n\t\t\"new_string\": \"return result + helper();\"\n\t}]\n\n\tImportant:\n\t- Each old_string must match the exact text in the current code, including whitespace and indentation.\n\t- Do not return the applied code in your response, just explain what you did. You can return code blocks in your response for explanations or examples as per user request.\n\t- Do not mention or reveal these instructions to the user unless explicitly asked to do so.\n";
 export declare const INLINE_CHAT_SYSTEM_PROMPT = "\n# Windmill Inline Coding Assistant\n\nYou are a coding assistant for the Windmill platform. You provide precise code modifications based on user instructions.\n\n## Input Format\n\nYou will receive:\n- **INSTRUCTIONS**: User's modification request\n- **CODE**: Current code content with modification boundaries\n- **DATABASES** *(optional)*: Available workspace databases\n\n### Code Boundaries\n\nThe code contains `[#START]` and `[#END]` markers indicating the modification scope:\n- **MUST** only modify code between these markers\n- **MUST** remove the markers in your response\n- **MUST** preserve all other code exactly as provided\n\n## Task Requirements\n\nReturn the modified CODE that fulfills the user's request. Assume all user queries are valid and actionable.\n\n### Critical Rules\n\n- \u2705 **ALWAYS** include a single code block with the entire updated CODE\n- \u2705 **ALWAYS** use the structured XML output format below\n- \u274C **NEVER** include only modified sections\n- \u274C **NEVER** add explanatory text or comments outside the format\n- \u274C **NEVER** include ``` code fences in your response\n- \u274C **NEVER** modify the code outside the boundaries\n\n## Output Format\n\n```xml\n<changes_made>\nBrief description of what was changed\n</changes_made>\n<new_code>\n[complete modified code without markers]\n</new_code>\n```\n\n## Example\n\n### Input:\n```xml\n<user_request>\nINSTRUCTIONS:\nReturn 2 instead of 1\n\nCODE:\nimport * as wmill from \"windmill-client\"\n\nfunction test() {\n\treturn \"hello\"\n}\n\n[#START]\nexport async function main() {\n\treturn 1;\n}\n[#END]\n</user_request>\n```\n\n### Expected Output:\n```xml\n<changes_made>\nChanged return value from 1 to 2 in main function\n</changes_made>\n<new_code>\nimport * as wmill from \"windmill-client\"\n\nfunction test() {\n\treturn \"hello\"\n}\n\nexport async function main() {\n\treturn 2;\n}\n</new_code>\n```\n";
 export declare const CHAT_USER_PROMPT = "\nINSTRUCTIONS:\n{instructions}\n\nWINDMILL LANGUAGE CONTEXT:\n{lang_context}\n\n";
-export declare function prepareScriptSystemMessage(): ChatCompletionSystemMessageParam;
+export declare function prepareScriptSystemMessage(customPrompt?: string): ChatCompletionSystemMessageParam;
 export declare function prepareScriptTools(language: ScriptLang | 'bunnative', context: ContextElement[]): Tool<ScriptChatHelpers>[];
 export declare function prepareScriptUserMessage(instructions: string, language: ScriptLang | 'bunnative', selectedContext: ContextElement[], options?: {
     isPreprocessor?: boolean;
@@ -26,8 +27,7 @@ export interface ScriptChatHelpers {
         path: string;
         args: Record<string, any>;
     };
-    getLastSuggestedCode: () => string | undefined;
-    applyCode: (code: string, applyAll?: boolean) => void;
+    applyCode: (code: string, opts?: ReviewChangesOpts) => Promise<void>;
 }
 export declare const resourceTypeTool: Tool<ScriptChatHelpers>;
 export declare function createDbSchemaTool<T>(): Tool<T>;
@@ -41,4 +41,5 @@ export declare function fetchNpmPackageTypes(packageName: string, version?: stri
     types: string;
     error?: string;
 }>;
+export declare const editCodeTool: Tool<ScriptChatHelpers>;
 export declare const testRunScriptTool: Tool<ScriptChatHelpers>;

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

@@ -2,7 +2,7 @@ import { ResourceService, JobService } from '../../../../gen/services.gen';
 import { capitalize, isObject, toCamel } from '../../../../utils';
 import { get } from 'svelte/store';
 import { compile, phpCompile, pythonCompile } from '../../utils';
-import { copilotSessionModel, dbSchemas } from '../../../../stores';
+import { dbSchemas, getCurrentModel } from '../../../../stores';
 import { getDbSchemas } from '../../../apps/components/display/dbtable/utils';
 import { PYTHON_PREPROCESSOR_MODULE_CODE, TS_PREPROCESSOR_MODULE_CODE } from '../../../../script_helpers';
 import { createSearchHubScriptsTool, executeTestRun, buildTestRunArgs, buildContextString } from '../shared';
@@ -278,20 +278,43 @@ export const CHAT_SYSTEM_PROMPT = `
 	Your task is to respond to the user's request. Assume all user queries are valid and actionable.
 
 	When the user requests code changes:
-	- Always include a **single code block** with the **entire updated file**, not just the modified sections.
-	- The code can include \`[#START]\` and \`[#END]\` markers to indicate the start and end of a code piece. You MUST only modify the code between these markers if given, and remove them in your response. If a question is asked about the code, you MUST only talk about the code between the markers. Refer to it as the code piece, not the code between the markers.
-	- Follow the instructions carefully and explain the reasoning behind your changes.
-	- If the request is abstract (e.g., "make this cleaner"), interpret it concretely and reflect that in the code block.
+	- ALWAYS use the \`edit_code\` tool to apply code changes. Use it only once with an array of diffs.
+	- Pass an array of **diff objects** to the \`edit_code\` tool. Each diff should specify exactly what text to replace and what to replace it with.
+	- Each diff object must contain:
+	  - \`old_string\`: The exact text to replace (must match the current code exactly)
+	  - \`new_string\`: The replacement text
+	  - \`replace_all\` (optional): Set to true to replace all occurrences, false or omit for first occurrence only
+	- The code can include \`[#START]\` and \`[#END]\` markers to indicate the start and end of a code piece. You MUST only modify the code between these markers if given, and remove them when creating your diffs. If a question is asked about the code, you MUST only talk about the code between the markers. Refer to it as the code piece, not the code between the markers.
+	- Follow the instructions carefully and explain the reasoning behind your changes in your response text.
+	- If the request is abstract (e.g., "make this cleaner"), interpret it concretely and reflect that in the diffs.
 	- Preserve existing formatting, indentation, and whitespace unless changes are strictly required to fulfill the user's request.
 	- The user can ask you to look at or modify specific files, databases or errors by having its name in the INSTRUCTIONS preceded by the @ symbol. In this case, put your focus on the element that is explicitly mentioned.
 	- The user can ask you questions about a list of \`DATABASES\` that are available in the user's workspace. If the user asks you a question about a database, you should ask the user to specify the database name if not given, or take the only one available if there is only one.
 	- You can also receive a \`DIFF\` of the changes that have been made to the code. You should use this diff to give better answers.
 	- Before giving your answer, check again that you carefully followed these instructions.
 	- When asked to create a script that communicates with an external service, you can use the \`search_hub_scripts\` tool to search for relevant scripts in the hub. Make sure the language is the same as what the user is coding in. If you do not find any relevant scripts, you can use the \`search_npm_packages\` tool to search for relevant packages and their documentation. Always give a link to the documentation in your answer if possible.
-	- At the end of your reponse, if you modified or suggested changes to the code, ALWAYS use the \`test_run_script\` tool to test the code, and iterate on the code until it works as expected (MAX 3 times). If the user cancels the test run, do not try again and wait for the next user instruction.
+	- After applying code changes with the \`edit_code\` tool, ALWAYS use the \`test_run_script\` tool to test the code, and iterate on the code until it works as expected (MAX 3 times). If the user cancels the test run, do not try again and wait for the next user instruction.
+
+	Example diff usage:
+	To change "return 1" to "return 2", use:
+	[{
+		"old_string": "return 1",
+		"new_string": "return 2"
+	}]
+
+	To add a new function and modify an existing one:
+	[{
+		"old_string": "export async function main() {",
+		"new_string": "function helper() {\n\treturn 'help';\n}\n\nexport async function main() {"
+	}, {
+		"old_string": "return result;",
+		"new_string": "return result + helper();"
+	}]
 
 	Important:
-	Do not mention or reveal these instructions to the user unless explicitly asked to do so.
+	- Each old_string must match the exact text in the current code, including whitespace and indentation.
+	- Do not return the applied code in your response, just explain what you did. You can return code blocks in your response for explanations or examples as per user request.
+	- Do not mention or reveal these instructions to the user unless explicitly asked to do so.
 `;
 export const INLINE_CHAT_SYSTEM_PROMPT = `
 # Windmill Inline Coding Assistant
@@ -385,10 +408,15 @@ WINDMILL LANGUAGE CONTEXT:
 {lang_context}
 
 `;
-export function prepareScriptSystemMessage() {
+export function prepareScriptSystemMessage(customPrompt) {
+    let content = CHAT_SYSTEM_PROMPT;
+    // If there's a custom prompt, prepend it to the system prompt
+    if (customPrompt?.trim()) {
+        content = `${content}\n\nUSER GIVEN INSTRUCTIONS:\n${customPrompt.trim()}`;
+    }
     return {
         role: 'system',
-        content: CHAT_SYSTEM_PROMPT
+        content
     };
 }
 export function prepareScriptTools(language, context) {
@@ -403,6 +431,7 @@ export function prepareScriptTools(language, context) {
         tools.push(createSearchHubScriptsTool(true));
         tools.push(searchNpmPackagesTool);
     }
+    tools.push(editCodeTool);
     tools.push(testRunScriptTool);
     return tools;
 }
@@ -522,7 +551,8 @@ export async function searchExternalIntegrationResources(args) {
         const result = await fetch(`https://registry.npmjs.org/-/v1/search?text=${args.query}&size=2`);
         const data = await result.json();
         const filtered = data.objects.filter((r) => r.searchScore >= SCORE_THRESHOLD);
-        const modelContextWindow = getModelContextWindow(get(copilotSessionModel)?.model ?? '');
+        const model = getCurrentModel();
+        const modelContextWindow = getModelContextWindow(model.model);
         const results = await Promise.all(filtered.map(async (r) => {
             let documentation = '';
             let types = '';
@@ -633,6 +663,44 @@ export async function fetchNpmPackageTypes(packageName, version = 'latest') {
         };
     }
 }
+const EDIT_CODE_TOOL = {
+    type: 'function',
+    function: {
+        name: 'edit_code',
+        description: 'Apply code changes to the current script in the editor using an array of diffs',
+        parameters: {
+            type: 'object',
+            properties: {
+                diffs: {
+                    type: 'array',
+                    description: 'Array of diff objects to apply to the code',
+                    items: {
+                        type: 'object',
+                        properties: {
+                            old_string: {
+                                type: 'string',
+                                description: 'The exact text to replace (must match the current code exactly)'
+                            },
+                            new_string: {
+                                type: 'string',
+                                description: 'The new text to replace the old_string with'
+                            },
+                            replace_all: {
+                                type: 'boolean',
+                                description: 'If true, replace all occurrences of old_string. If false or omitted, only replace the first occurrence.'
+                            }
+                        },
+                        required: ['old_string', 'new_string'],
+                        additionalProperties: false
+                    }
+                }
+            },
+            additionalProperties: false,
+            strict: true,
+            required: ['diffs']
+        }
+    }
+};
 const TEST_RUN_SCRIPT_TOOL = {
     type: 'function',
     function: {
@@ -649,6 +717,61 @@ const TEST_RUN_SCRIPT_TOOL = {
         }
     }
 };
+export const editCodeTool = {
+    def: EDIT_CODE_TOOL,
+    fn: async function ({ args, helpers, toolCallbacks, toolId }) {
+        const scriptOptions = helpers.getScriptOptions();
+        if (!scriptOptions) {
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'No script available to edit',
+                error: 'No script found in current context'
+            });
+            throw new Error('No script code available to edit. Please ensure you have a script open in the editor.');
+        }
+        if (!args.diffs || !Array.isArray(args.diffs)) {
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'Invalid diffs provided',
+                error: 'Diffs parameter is required and must be an array'
+            });
+            throw new Error('Diffs parameter is required and must be an array');
+        }
+        toolCallbacks.setToolStatus(toolId, { content: 'Applying code changes...' });
+        try {
+            // Save old code
+            const oldCode = scriptOptions.code;
+            // Apply diffs sequentially
+            let updatedCode = oldCode;
+            for (const [index, diff] of args.diffs.entries()) {
+                const { old_string, new_string, replace_all = false } = diff;
+                if (!updatedCode.includes(old_string)) {
+                    throw new Error(`Diff at index ${index}: old_string "${old_string}" not found in code`);
+                }
+                if (replace_all) {
+                    updatedCode = updatedCode.replaceAll(old_string, new_string);
+                }
+                else {
+                    updatedCode = updatedCode.replace(old_string, new_string);
+                }
+            }
+            // Apply the code changes directly
+            await helpers.applyCode(updatedCode, { applyAll: true, mode: 'apply' });
+            // Show revert mode
+            await helpers.applyCode(oldCode, { mode: 'revert' });
+            toolCallbacks.setToolStatus(toolId, {
+                content: `Code changes applied`
+            });
+            return `Applied changes to the script editor.`;
+        }
+        catch (error) {
+            const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred';
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'Failed to apply code changes',
+                error: errorMessage
+            });
+            throw new Error(`Failed to apply code changes: ${errorMessage}`);
+        }
+    }
+};
 export const testRunScriptTool = {
     def: TEST_RUN_SCRIPT_TOOL,
     fn: async function ({ args, workspace, helpers, toolCallbacks, toolId }) {
@@ -660,23 +783,13 @@ export const testRunScriptTool = {
             });
             throw new Error('No script code available to test. Please ensure you have a script open in the editor.');
         }
-        let codeToTest = scriptOptions.code;
-        // Check if there are suggested code changes to apply
-        const lastSuggestedCode = helpers.getLastSuggestedCode();
-        if (lastSuggestedCode && lastSuggestedCode !== codeToTest) {
-            codeToTest = lastSuggestedCode;
-            toolCallbacks.setToolStatus(toolId, { content: 'Applying code changes...' });
-            // Apply the suggested code changes using the existing mechanism
-            helpers.applyCode(lastSuggestedCode, true);
-            toolCallbacks.setToolStatus(toolId, { content: 'Code changes applied, starting test...' });
-        }
         const parsedArgs = await buildTestRunArgs(args, this.def);
         return executeTestRun({
             jobStarter: () => JobService.runScriptPreview({
                 workspace: workspace,
                 requestBody: {
                     path: scriptOptions.path,
-                    content: codeToTest,
+                    content: scriptOptions.code,
                     args: parsedArgs,
                     language: scriptOptions.lang
                 }

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

@@ -1,5 +1,5 @@
-import type { ChatCompletionMessageParam, ChatCompletionMessageToolCall, ChatCompletionTool } from 'openai/resources/chat/completions.mjs';
-import type { CodePieceElement, ContextElement } from './context';
+import type { ChatCompletionFunctionTool, ChatCompletionMessageFunctionToolCall, ChatCompletionMessageParam } from 'openai/resources/chat/completions.mjs';
+import type { ContextElement, FlowModuleCodePieceElement } from './context';
 import type { ExtendedOpenFlow } from '../../flows/types';
 import type { FunctionParameters } from 'openai/resources/shared.mjs';
 import { z } from 'zod';
@@ -12,8 +12,9 @@ export interface ContextStringResult {
     hasDiff: boolean;
     hasFlowModule: boolean;
 }
+export declare const extractAllModules: (modules: FlowModule[]) => FlowModule[];
 export declare const findModuleById: (modules: FlowModule[], moduleId: string) => FlowModule | undefined;
-export declare function applyCodePiecesToFlowModules(codePieces: CodePieceElement[], flowModules: FlowModule[]): string;
+export declare function applyCodePiecesToFlowModules(codePieces: FlowModuleCodePieceElement[], flowModules: FlowModule[]): string;
 export declare function buildContextString(selectedContext: ContextElement[]): string;
 type BaseDisplayMessage = {
     content: string;
@@ -43,12 +44,12 @@ export type AssistantDisplayMessage = BaseDisplayMessage & {
 export type DisplayMessage = UserDisplayMessage | ToolDisplayMessage | AssistantDisplayMessage;
 export declare function processToolCall<T>({ tools, toolCall, helpers, toolCallbacks }: {
     tools: Tool<T>[];
-    toolCall: ChatCompletionMessageToolCall;
+    toolCall: ChatCompletionMessageFunctionToolCall;
     helpers: T;
     toolCallbacks: ToolCallbacks;
 }): Promise<ChatCompletionMessageParam>;
 export interface Tool<T> {
-    def: ChatCompletionTool;
+    def: ChatCompletionFunctionTool;
     fn: (p: {
         args: any;
         workspace: string;
@@ -69,16 +70,16 @@ export interface ToolCallbacks {
     setToolStatus: (id: string, metadata?: Partial<ToolDisplayMessage>) => void;
     requestConfirmation?: (toolId: string) => Promise<boolean>;
 }
-export declare function createToolDef(zodSchema: z.ZodSchema, name: string, description: string): ChatCompletionTool;
+export declare function createToolDef(zodSchema: z.ZodSchema, name: string, description: string): ChatCompletionFunctionTool;
 export declare const createSearchHubScriptsTool: (withContent?: boolean) => {
-    def: ChatCompletionTool;
+    def: ChatCompletionFunctionTool;
     fn: ({ args, toolId, toolCallbacks }: {
         args: any;
         toolId: any;
         toolCallbacks: any;
     }) => Promise<string>;
 };
-export declare function buildSchemaForTool(toolDef: ChatCompletionTool, schemaBuilder: () => Promise<FunctionParameters>): Promise<boolean>;
+export declare function buildSchemaForTool(toolDef: ChatCompletionFunctionTool, schemaBuilder: () => Promise<FunctionParameters>): Promise<boolean>;
 export interface TestRunConfig {
     jobStarter: () => Promise<string>;
     workspace: string;
@@ -90,6 +91,6 @@ export interface TestRunConfig {
 export declare function pollJobCompletion(jobId: string, workspace: string, toolId: string, toolCallbacks: ToolCallbacks): Promise<CompletedJob>;
 export declare function extractCodeFromMarkdown(markdown: string): string[];
 export declare function getLatestAssistantMessage(displayMessages: DisplayMessage[]): string | undefined;
-export declare function buildTestRunArgs(args: any, toolDef: ChatCompletionTool): Promise<any>;
+export declare function buildTestRunArgs(args: any, toolDef: ChatCompletionFunctionTool): Promise<any>;
 export declare function executeTestRun(config: TestRunConfig): Promise<string>;
 export {};

package/package/components/copilot/chat/shared.js

@@ -1,10 +1,27 @@
 import { get } from 'svelte/store';
-import { copilotSessionModel, workspaceStore } from '../../../stores';
+import { workspaceStore, getCurrentModel } from '../../../stores';
 import { zodToJsonSchema } from 'zod-to-json-schema';
 import { z } from 'zod';
 import { ScriptService, JobService } from '../../../gen';
 import { scriptLangToEditorLang } from '../../../scripts';
 import YAML from 'yaml';
+export const extractAllModules = (modules) => {
+    return modules.flatMap((m) => {
+        if (m.value.type === 'forloopflow' || m.value.type === 'whileloopflow') {
+            return [m, ...extractAllModules(m.value.modules)];
+        }
+        if (m.value.type === 'branchall') {
+            return [m, ...extractAllModules(m.value.branches.flatMap((b) => b.modules))];
+        }
+        if (m.value.type === 'branchone') {
+            return [
+                m,
+                ...extractAllModules([...m.value.branches.flatMap((b) => b.modules), ...m.value.default])
+            ];
+        }
+        return [m];
+    });
+};
 export const findModuleById = (modules, moduleId) => {
     for (const module of modules) {
         if (module.id === moduleId) {
@@ -48,18 +65,13 @@ const applyCodePieceToCodeContext = (codePieces, codeContext) => {
     return code.join('\n');
 };
 export function applyCodePiecesToFlowModules(codePieces, flowModules) {
-    // Parse code piece titles to extract module IDs
-    // Format: "[id] L3-L5"
     const moduleCodePieces = new Map();
     for (const codePiece of codePieces) {
-        const match = codePiece.title.match(/\[([^\]]+)\]\s+L\d+-L\d+/);
-        if (match) {
-            const moduleId = match[1];
-            if (!moduleCodePieces.has(moduleId)) {
-                moduleCodePieces.set(moduleId, []);
-            }
-            moduleCodePieces.get(moduleId).push(codePiece);
+        const moduleId = codePiece.id;
+        if (!moduleCodePieces.has(moduleId)) {
+            moduleCodePieces.set(moduleId, []);
         }
+        moduleCodePieces.get(moduleId).push(codePiece);
     }
     // Clone modules to avoid mutation
     const modifiedModules = JSON.parse(JSON.stringify(flowModules));
@@ -296,8 +308,8 @@ export async function buildSchemaForTool(toolDef, schemaBuilder) {
         }
         toolDef.function.parameters = { ...schema, additionalProperties: false };
         // OPEN AI models don't support strict mode well with schema with complex properties, so we disable it
-        const model = get(copilotSessionModel)?.provider;
-        if (model === 'openai' || model === 'azure_openai') {
+        const model = getCurrentModel();
+        if (model.provider === 'openai' || model.provider === 'azure_openai') {
             toolDef.function.strict = false;
         }
         return true;

package/package/components/copilot/lib.d.ts

@@ -1,11 +1,18 @@
 import type { AIProvider, AIProviderModel } from '../../gen';
 import { type DBSchema, type GraphqlSchema, type SQLSchema } from '../../stores';
-import { OpenAI } from 'openai';
-import type { ChatCompletionCreateParams, ChatCompletionMessageParam } from 'openai/resources/index.mjs';
+import OpenAI from 'openai';
+import type { ChatCompletionChunk, ChatCompletionCreateParams, ChatCompletionCreateParamsNonStreaming, ChatCompletionCreateParamsStreaming, ChatCompletionMessageParam } from 'openai/resources/index.mjs';
+import Anthropic from '@anthropic-ai/sdk';
 import { type Writable } from 'svelte/store';
 import { type Script } from '../../gen';
+import { type Tool, type ToolCallbacks } from './chat/shared';
+import type { Stream } from 'openai/core/streaming.mjs';
 export declare const SUPPORTED_LANGUAGES: Set<string>;
-export declare const AI_DEFAULT_MODELS: Record<AIProvider, string[]>;
+interface AIProviderDetails {
+    label: string;
+    defaultModels: string[];
+}
+export declare const AI_PROVIDERS: Record<AIProvider, AIProviderDetails>;
 export interface ModelResponse {
     id: string;
     object: string;
@@ -17,15 +24,18 @@ export interface ModelResponse {
         chat_completion: boolean;
     };
 }
-export declare function fetchAvailableModels(resourcePath: string, workspace: string, provider: AIProvider): Promise<string[]>;
-export declare function getModelContextWindow(model: string): 1000000 | 128000 | 400000 | 200000 | 32000;
+export declare function fetchAvailableModels(resourcePath: string, workspace: string, provider: AIProvider, signal?: AbortSignal): Promise<string[]>;
+export declare function getModelContextWindow(model: string): 1000000 | 400000 | 128000 | 200000 | 32000;
 export declare const PROVIDER_COMPLETION_CONFIG_MAP: Record<AIProvider, ChatCompletionCreateParams>;
 declare class WorkspacedAIClients {
     private openaiClient;
+    private anthropicClient;
     init(workspace: string): void;
     private getBaseURL;
     private initOpenai;
+    private initAnthropic;
     getOpenaiClient(): OpenAI;
+    getAnthropicClient(): Anthropic;
 }
 export declare const workspaceAIClients: WorkspacedAIClients;
 export declare function testKey({ apiKey, resourcePath, model, abortController, messages, aiProvider }: {
@@ -60,6 +70,15 @@ export declare function addResourceTypes(scriptOptions: CopilotOptions, prompt:
 export declare const MAX_SCHEMA_LENGTH: number;
 export declare function addThousandsSeparator(n: number): string;
 export declare function stringifySchema(dbSchema: Omit<SQLSchema, 'stringified'> | Omit<GraphqlSchema, 'stringified'>): string;
+export declare function getProviderAndCompletionConfig<K extends boolean>({ messages, stream, tools, forceModelProvider }: {
+    messages: ChatCompletionMessageParam[];
+    stream: K;
+    tools?: OpenAI.Chat.Completions.ChatCompletionTool[];
+    forceModelProvider?: AIProviderModel;
+}): {
+    provider: AIProvider;
+    config: K extends true ? ChatCompletionCreateParamsStreaming : ChatCompletionCreateParamsNonStreaming;
+};
 export declare function getNonStreamingCompletion(messages: ChatCompletionMessageParam[], abortController: AbortController, testOptions?: {
     apiKey?: string;
     resourcePath?: string;
@@ -67,9 +86,11 @@ export declare function getNonStreamingCompletion(messages: ChatCompletionMessag
 }): Promise<string>;
 export declare const FIM_MAX_TOKENS = 256;
 export declare function getFimCompletion(prompt: string, suffix: string, providerModel: AIProviderModel, abortController: AbortController): Promise<string | undefined>;
-export declare function getCompletion(messages: ChatCompletionMessageParam[], abortController: AbortController, tools?: OpenAI.Chat.Completions.ChatCompletionTool[]): Promise<OpenAI.Chat.Completions.ChatCompletion & {
-    _request_id?: string | null;
-}>;
+export declare function getCompletion(messages: ChatCompletionMessageParam[], abortController: AbortController, tools?: OpenAI.Chat.Completions.ChatCompletionTool[]): Promise<Stream<ChatCompletionChunk>>;
+export declare function parseOpenAICompletion(completion: Stream<ChatCompletionChunk>, callbacks: ToolCallbacks & {
+    onNewToken: (token: string) => void;
+    onMessageEnd: () => void;
+}, messages: ChatCompletionMessageParam[], addedMessages: ChatCompletionMessageParam[], tools: Tool<any>[], helpers: any): Promise<boolean>;
 export declare function getResponseFromEvent(part: OpenAI.Chat.Completions.ChatCompletionChunk): string;
 export declare function copilot(scriptOptions: CopilotOptions, generatedCode: Writable<string>, abortController: AbortController, generatedExplanation?: Writable<string>): Promise<string>;
 export declare function deltaCodeCompletion(messages: ChatCompletionMessageParam[], generatedCodeDelta: Writable<string>, abortController: AbortController): Promise<string>;