windmill-components 1.531.1 → 1.537.1

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

@@ -1,15 +1,13 @@
 import { ResourceService, JobService } from '../../../../gen/services.gen';
-import { capitalize, emptySchema, isObject, toCamel } from '../../../../utils';
+import { capitalize, isObject, toCamel } from '../../../../utils';
 import { get } from 'svelte/store';
 import { compile, phpCompile, pythonCompile } from '../../utils';
-import { copilotSessionModel, dbSchemas } from '../../../../stores';
-import { scriptLangToEditorLang } from '../../../../scripts';
+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, buildSchemaForTool, buildTestRunArgs } from '../shared';
+import { createSearchHubScriptsTool, executeTestRun, buildTestRunArgs, buildContextString } from '../shared';
 import { setupTypeAcquisition } from '../../../../ata';
 import { getModelContextWindow } from '../../lib';
-import { inferArgs } from '../../../../infer';
 // Score threshold for npm packages search filtering
 const SCORE_THRESHOLD = 1000;
 // percentage of the context window for documentation of npm packages
@@ -290,7 +288,7 @@ export const CHAT_SYSTEM_PROMPT = `
 	- 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.
-	- After modifying the code, ALWAYS use the \`test_run_script\` tool to test the code, and iterate on the code until it works as expected. If the user cancels the test run, do not try again and wait for the next user instruction.
+	- 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.
 
 	Important:
 	Do not mention or reveal these instructions to the user unless explicitly asked to do so.
@@ -379,16 +377,6 @@ export async function main() {
 </new_code>
 \`\`\`
 `;
-const CHAT_USER_CODE_CONTEXT = `
-- {title}:
-\`\`\`{language}
-{code}
-\`\`\`
-`;
-const CHAT_USER_ERROR_CONTEXT = `
-ERROR:
-{error}
-`;
 export const CHAT_USER_PROMPT = `
 INSTRUCTIONS:
 {instructions}
@@ -397,24 +385,12 @@ WINDMILL LANGUAGE CONTEXT:
 {lang_context}
 
 `;
-export const CHAT_USER_DB_CONTEXT = `- {title}: SCHEMA: \n{schema}\n`;
 export function prepareScriptSystemMessage() {
     return {
         role: 'system',
         content: CHAT_SYSTEM_PROMPT
     };
 }
-const applyCodePieceToCodeContext = (codePieces, codeContext) => {
-    let code = codeContext.split('\n');
-    let shiftOffset = 0;
-    codePieces.sort((a, b) => a.startLine - b.startLine);
-    for (const codePiece of codePieces) {
-        code.splice(codePiece.endLine + shiftOffset, 0, '[#END]');
-        code.splice(codePiece.startLine + shiftOffset - 1, 0, '[#START]');
-        shiftOffset += 2;
-    }
-    return code.join('\n');
-};
 export function prepareScriptTools(language, context) {
     const tools = [];
     if (['python3', 'php', 'bun', 'deno', 'nativets', 'bunnative'].includes(language)) {
@@ -431,51 +407,9 @@ export function prepareScriptTools(language, context) {
     return tools;
 }
 export function prepareScriptUserMessage(instructions, language, selectedContext, options = {}) {
-    let codeContext = 'CODE:\n';
-    let errorContext = 'ERROR:\n';
-    let dbContext = 'DATABASES:\n';
-    let diffContext = 'DIFF:\n';
-    let hasCode = false;
-    let hasError = false;
-    let hasDb = false;
-    let hasDiff = false;
-    for (const context of selectedContext) {
-        if (context.type === 'code') {
-            hasCode = true;
-            codeContext += CHAT_USER_CODE_CONTEXT.replace('{title}', context.title)
-                .replace('{language}', scriptLangToEditorLang(language))
-                .replace('{code}', applyCodePieceToCodeContext(selectedContext.filter((c) => c.type === 'code_piece'), context.content));
-        }
-        else if (context.type === 'error') {
-            if (hasError) {
-                throw new Error('Multiple error contexts provided');
-            }
-            hasError = true;
-            errorContext = CHAT_USER_ERROR_CONTEXT.replace('{error}', context.content);
-        }
-        else if (context.type === 'db') {
-            hasDb = true;
-            dbContext += CHAT_USER_DB_CONTEXT.replace('{title}', context.title).replace('{schema}', context.schema?.stringified ?? 'to fetch with get_db_schema');
-        }
-        else if (context.type === 'diff') {
-            hasDiff = true;
-            const diff = JSON.stringify(context.diff);
-            diffContext = diff.length > 3000 ? diff.slice(0, 3000) + '...' : diff;
-        }
-    }
     let userMessage = CHAT_USER_PROMPT.replace('{instructions}', instructions).replace('{lang_context}', getLangContext(language, { allowResourcesFetch: true, ...options }));
-    if (hasCode) {
-        userMessage += codeContext;
-    }
-    if (hasError) {
-        userMessage += errorContext;
-    }
-    if (hasDb) {
-        userMessage += dbContext;
-    }
-    if (hasDiff) {
-        userMessage += diffContext;
-    }
+    const contextInstructions = buildContextString(selectedContext);
+    userMessage += contextInstructions;
     return {
         role: 'user',
         content: userMessage
@@ -534,39 +468,51 @@ async function formatDBSchema(dbSchema) {
 export const resourceTypeTool = {
     def: RESOURCE_TYPE_FUNCTION_DEF,
     fn: async ({ args, workspace, helpers, toolCallbacks, toolId }) => {
-        toolCallbacks.setToolStatus(toolId, { content: 'Searching resource types for "' + args.query + '"...' });
+        toolCallbacks.setToolStatus(toolId, {
+            content: 'Searching resource types for "' + args.query + '"...'
+        });
         const lang = helpers.getScriptOptions().lang;
         const formattedResourceTypes = await getFormattedResourceTypes(lang, args.query, workspace);
-        toolCallbacks.setToolStatus(toolId, { content: 'Retrieved resource types for "' + args.query + '"' });
+        toolCallbacks.setToolStatus(toolId, {
+            content: 'Retrieved resource types for "' + args.query + '"'
+        });
         return formattedResourceTypes;
     }
 };
-export const dbSchemaTool = {
-    def: DB_SCHEMA_FUNCTION_DEF,
-    fn: async ({ args, workspace, toolCallbacks, toolId }) => {
-        if (!args.resourcePath) {
-            throw new Error('Database path not provided');
-        }
-        toolCallbacks.setToolStatus(toolId, { content: 'Getting database schema for ' + args.resourcePath + '...' });
-        const resource = await ResourceService.getResource({
-            workspace: workspace,
-            path: args.resourcePath
-        });
-        const newDbSchemas = {};
-        await getDbSchemas(resource.resource_type, args.resourcePath, workspace, newDbSchemas, (error) => {
-            console.error(error);
-        });
-        dbSchemas.update((schemas) => ({ ...schemas, ...newDbSchemas }));
-        const dbs = get(dbSchemas);
-        const db = dbs[args.resourcePath];
-        if (!db) {
-            throw new Error('Database not found');
+// Generic DB schema tool factory that can be used by both script and flow modes
+export function createDbSchemaTool() {
+    return {
+        def: DB_SCHEMA_FUNCTION_DEF,
+        fn: async ({ args, workspace, toolCallbacks, toolId }) => {
+            if (!args.resourcePath) {
+                throw new Error('Database path not provided');
+            }
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'Getting database schema for ' + args.resourcePath + '...'
+            });
+            const resource = await ResourceService.getResource({
+                workspace: workspace,
+                path: args.resourcePath
+            });
+            const newDbSchemas = {};
+            await getDbSchemas(resource.resource_type, args.resourcePath, workspace, newDbSchemas, (error) => {
+                console.error(error);
+            });
+            dbSchemas.update((schemas) => ({ ...schemas, ...newDbSchemas }));
+            const dbs = get(dbSchemas);
+            const db = dbs[args.resourcePath];
+            if (!db) {
+                throw new Error('Database not found');
+            }
+            const stringSchema = await formatDBSchema(db);
+            toolCallbacks.setToolStatus(toolId, {
+                content: 'Retrieved database schema for ' + args.resourcePath
+            });
+            return stringSchema;
         }
-        const stringSchema = await formatDBSchema(db);
-        toolCallbacks.setToolStatus(toolId, { content: 'Retrieved database schema for ' + args.resourcePath });
-        return stringSchema;
-    }
-};
+    };
+}
+export const dbSchemaTool = createDbSchemaTool();
 const packagesSearchCache = new Map();
 export async function searchExternalIntegrationResources(args) {
     try {
@@ -576,7 +522,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 = '';
@@ -692,18 +639,16 @@ const TEST_RUN_SCRIPT_TOOL = {
     function: {
         name: 'test_run_script',
         description: 'Execute a test run of the current script in the editor',
-        // will be overridden by setSchema
         parameters: {
             type: 'object',
             properties: {
-                args: {
-                    type: 'object',
-                    description: 'Arguments to pass to the script (optional, uses current editor args if not provided)'
-                }
+                args: { type: 'string', description: 'JSON string containing the arguments for the tool' }
             },
-            required: []
+            additionalProperties: false,
+            strict: false,
+            required: ['args']
         }
-    },
+    }
 };
 export const testRunScriptTool = {
     def: TEST_RUN_SCRIPT_TOOL,
@@ -734,7 +679,7 @@ export const testRunScriptTool = {
                     path: scriptOptions.path,
                     content: codeToTest,
                     args: parsedArgs,
-                    language: scriptOptions.lang,
+                    language: scriptOptions.lang
                 }
             }),
             workspace,
@@ -744,22 +689,7 @@ export const testRunScriptTool = {
             contextName: 'script'
         });
     },
-    setSchema: async function (helpers) {
-        await buildSchemaForTool(this.def, async () => {
-            const scriptOptions = helpers.getScriptOptions();
-            const code = scriptOptions?.code;
-            const lang = scriptOptions?.lang;
-            const lastSuggestedCode = helpers.getLastSuggestedCode();
-            const codeToTest = lastSuggestedCode ?? code;
-            if (codeToTest) {
-                const newSchema = emptySchema();
-                await inferArgs(lang, codeToTest, newSchema);
-                return newSchema;
-            }
-            return emptySchema();
-        });
-    },
     requiresConfirmation: true,
     confirmationMessage: 'Run script test',
-    showDetails: true,
+    showDetails: true
 };

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

@@ -1,9 +1,21 @@
 import type { ChatCompletionMessageParam, ChatCompletionMessageToolCall, ChatCompletionTool } from 'openai/resources/chat/completions.mjs';
-import type { ContextElement } from './context';
+import type { ContextElement, FlowModuleCodePieceElement } from './context';
 import type { ExtendedOpenFlow } from '../../flows/types';
 import type { FunctionParameters } from 'openai/resources/shared.mjs';
 import { z } from 'zod';
-import { type CompletedJob } from '../../../gen';
+import { type CompletedJob, type FlowModule } from '../../../gen';
+export interface ContextStringResult {
+    dbContext: string;
+    diffContext: string;
+    flowModuleContext: string;
+    hasDb: boolean;
+    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: FlowModuleCodePieceElement[], flowModules: FlowModule[]): string;
+export declare function buildContextString(selectedContext: ContextElement[]): string;
 type BaseDisplayMessage = {
     content: string;
     contextElements?: ContextElement[];

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

@@ -1,8 +1,160 @@
 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) {
+            return module;
+        }
+        if (module.value.type === 'forloopflow' || module.value.type === 'whileloopflow') {
+            const found = findModuleById(module.value.modules, moduleId);
+            if (found) {
+                return found;
+            }
+        }
+        if (module.value.type === 'branchall') {
+            const allModules = module.value.branches.flatMap((b) => b.modules);
+            const found = findModuleById(allModules, moduleId);
+            if (found) {
+                return found;
+            }
+        }
+        if (module.value.type === 'branchone') {
+            const allModules = [
+                ...module.value.branches.flatMap((b) => b.modules),
+                ...module.value.default
+            ];
+            const found = findModuleById(allModules, moduleId);
+            if (found) {
+                return found;
+            }
+        }
+    }
+    return undefined;
+};
+const applyCodePieceToCodeContext = (codePieces, codeContext) => {
+    let code = codeContext.split('\n');
+    let shiftOffset = 0;
+    codePieces.sort((a, b) => a.startLine - b.startLine);
+    for (const codePiece of codePieces) {
+        code.splice(codePiece.endLine + shiftOffset, 0, '[#END]');
+        code.splice(codePiece.startLine + shiftOffset - 1, 0, '[#START]');
+        shiftOffset += 2;
+    }
+    return code.join('\n');
+};
+export function applyCodePiecesToFlowModules(codePieces, flowModules) {
+    const moduleCodePieces = new Map();
+    for (const codePiece of codePieces) {
+        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));
+    // Apply code pieces to each module
+    for (const [moduleId, pieces] of moduleCodePieces) {
+        const module = findModuleById(modifiedModules, moduleId);
+        if (module && module.value.type === 'rawscript' && module.value.content) {
+            module.value.content = applyCodePieceToCodeContext(pieces, module.value.content);
+        }
+    }
+    return YAML.stringify(modifiedModules);
+}
+export function buildContextString(selectedContext) {
+    const dbTemplate = `- {title}: SCHEMA: \n{schema}\n`;
+    const codeTemplate = `
+	- {title}:
+	\`\`\`{language}
+	{code}
+	\`\`\`
+	`;
+    let dbContext = 'DATABASES:\n';
+    let diffContext = 'DIFF:\n';
+    let flowModuleContext = 'FOCUSED FLOW MODULES IDS:\n';
+    let codeContext = 'CODE:\n';
+    let errorContext = `
+	ERROR:
+	{error}
+	`;
+    let hasCode = false;
+    let hasDb = false;
+    let hasDiff = false;
+    let hasFlowModule = false;
+    let hasError = false;
+    let result = '\n\n';
+    for (const context of selectedContext) {
+        if (context.type === 'code') {
+            hasCode = true;
+            codeContext += codeTemplate
+                .replace('{title}', context.title)
+                .replace('{language}', scriptLangToEditorLang(context.lang))
+                .replace('{code}', applyCodePieceToCodeContext(selectedContext.filter((c) => c.type === 'code_piece'), context.content));
+        }
+        else if (context.type === 'error') {
+            if (hasError) {
+                throw new Error('Multiple error contexts provided');
+            }
+            hasError = true;
+            errorContext = errorContext.replace('{error}', context.content);
+        }
+        else if (context.type === 'db') {
+            hasDb = true;
+            dbContext += dbTemplate
+                .replace('{title}', context.title)
+                .replace('{schema}', context.schema?.stringified ?? 'to fetch with get_db_schema');
+            dbContext += '\n';
+        }
+        else if (context.type === 'diff') {
+            hasDiff = true;
+            const diff = JSON.stringify(context.diff);
+            diffContext += (diff.length > 3000 ? diff.slice(0, 3000) + '...' : diff) + '\n';
+        }
+        else if (context.type === 'flow_module') {
+            hasFlowModule = true;
+            flowModuleContext += `${context.id}\n`;
+        }
+    }
+    if (hasCode) {
+        result += '\n' + codeContext;
+    }
+    if (hasError) {
+        result += '\n' + errorContext;
+    }
+    if (hasDb) {
+        result += '\n' + dbContext;
+    }
+    if (hasDiff) {
+        result += '\n' + diffContext;
+    }
+    if (hasFlowModule) {
+        result += '\n' + flowModuleContext;
+    }
+    return result;
+}
 async function callTool({ tools, functionName, args, workspace, helpers, toolCallbacks, toolId }) {
     const tool = tools.find((t) => t.def.function.name === functionName);
     if (!tool) {
@@ -18,11 +170,13 @@ export async function processToolCall({ tools, toolCall, helpers, toolCallbacks
         const needsConfirmation = tool?.requiresConfirmation;
         // Add the tool to the display with appropriate status
         toolCallbacks.setToolStatus(toolCall.id, {
-            ...(tool?.requiresConfirmation ? { content: tool.confirmationMessage ?? "Waiting for confirmation..." } : {}),
+            ...(tool?.requiresConfirmation
+                ? { content: tool.confirmationMessage ?? 'Waiting for confirmation...' }
+                : {}),
             parameters: args,
             isLoading: true,
             needsConfirmation: needsConfirmation,
-            showDetails: tool?.showDetails,
+            showDetails: tool?.showDetails
         });
         // If confirmation is needed and we have the callback, wait for it
         if (needsConfirmation && toolCallbacks.requestConfirmation) {
@@ -154,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;
@@ -163,7 +317,15 @@ export async function buildSchemaForTool(toolDef, schemaBuilder) {
     catch (error) {
         console.error('Error building schema for tool', error);
         // fallback to schema with args as a JSON string
-        toolDef.function.parameters = { type: 'object', properties: { args: { type: 'string', description: 'JSON string containing the arguments for the tool' } }, additionalProperties: false, strict: false, required: ['args'] };
+        toolDef.function.parameters = {
+            type: 'object',
+            properties: {
+                args: { type: 'string', description: 'JSON string containing the arguments for the tool' }
+            },
+            additionalProperties: false,
+            strict: false,
+            required: ['args']
+        };
         return false;
     }
 }
@@ -254,7 +416,8 @@ function getErrorMessage(result) {
 export async function buildTestRunArgs(args, toolDef) {
     let parsedArgs = args;
     // if the schema is the fallback schema, parse the args as a JSON string
-    if (toolDef.function.parameters.properties?.args?.description === 'JSON string containing the arguments for the tool') {
+    if (toolDef.function.parameters.properties?.args?.description ===
+        'JSON string containing the arguments for the tool') {
         try {
             parsedArgs = JSON.parse(args.args);
         }

package/package/components/copilot/lib.js

@@ -1,4 +1,4 @@
-import { copilotInfo, copilotSessionModel, workspaceStore } from '../../stores';
+import { getCurrentModel, workspaceStore } from '../../stores';
 import { buildClientSchema, printSchema } from 'graphql';
 import { OpenAI } from 'openai';
 import { get } from 'svelte/store';
@@ -7,7 +7,16 @@ import { EDIT_CONFIG, FIX_CONFIG, GEN_CONFIG } from './prompts';
 import { formatResourceTypes } from './utils';
 import { z } from 'zod';
 export const SUPPORTED_LANGUAGES = new Set(Object.keys(GEN_CONFIG.prompts));
-const OPENAI_MODELS = ['gpt-5', 'gpt-5-mini', 'gpt-5-nano', 'gpt-4o', 'gpt-4o-mini', 'o4-mini', 'o3', 'o3-mini'];
+const OPENAI_MODELS = [
+    'gpt-5',
+    'gpt-5-mini',
+    'gpt-5-nano',
+    'gpt-4o',
+    'gpt-4o-mini',
+    'o4-mini',
+    'o3',
+    'o3-mini'
+];
 // need at least one model for each provider except customai
 export const AI_DEFAULT_MODELS = {
     openai: OPENAI_MODELS,
@@ -331,11 +340,7 @@ const PROMPTS_CONFIGS = {
     gen: GEN_CONFIG
 };
 function getProviderAndCompletionConfig({ messages, stream, tools, forceModelProvider }) {
-    let info = get(copilotInfo);
-    const modelProvider = forceModelProvider ?? get(copilotSessionModel) ?? info.defaultModel ?? info.aiModels[0];
-    if (!modelProvider) {
-        throw new Error('No model selected');
-    }
+    const modelProvider = forceModelProvider ?? getCurrentModel();
     const providerConfig = PROVIDER_COMPLETION_CONFIG_MAP[modelProvider.provider];
     const processedMessages = prepareMessages(modelProvider.provider, messages);
     return {

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

@@ -41,7 +41,7 @@ export type VisualChange = {
 };
 export declare function setGlobalCSS(id: string, cssCode: string): void;
 export declare let VISUAL_CHANGES_CSS: string;
-export declare function displayVisualChanges(cssId: string, editor: meditor.IStandaloneCodeEditor, visualChanges: VisualChange[]): Promise<{
+export declare function displayVisualChanges(cssId: string, editor: meditor.IStandaloneCodeEditor, visualChanges: VisualChange[], revertMode?: boolean): Promise<{
     collection: meditor.IEditorDecorationsCollection;
     ids: string[];
 }>;

package/package/components/copilot/shared.js

@@ -1,6 +1,6 @@
 import { createLongHash } from '../../editorLangUtils';
 import {} from 'monaco-editor';
-function applyMonacoStyles(targetEl, greenHighlight) {
+function applyMonacoStyles(targetEl, greenHighlight, revertMode) {
     const computedStyles = window.getComputedStyle(document.querySelector('.monaco-editor .view-lines'));
     Object.assign(targetEl.style, {
         fontFamily: computedStyles.fontFamily,
@@ -10,7 +10,9 @@ function applyMonacoStyles(targetEl, greenHighlight) {
         whiteSpace: 'pre'
     });
     if (greenHighlight) {
-        targetEl.style.backgroundColor = 'var(--vscode-diffEditor-insertedTextBackground)';
+        targetEl.style.backgroundColor = revertMode
+            ? 'var(--vscode-diffEditor-removedTextBackground)'
+            : 'var(--vscode-diffEditor-insertedTextBackground)';
     }
 }
 export function setGlobalCSS(id, cssCode) {
@@ -22,7 +24,7 @@ export function setGlobalCSS(id, cssCode) {
     }
     styleTag.textContent = cssCode;
 }
-function addInlineGhostText(change) {
+function addInlineGhostText(change, revertMode) {
     const cssId = createLongHash();
     const decoration = {
         range: {
@@ -32,7 +34,11 @@ function addInlineGhostText(change) {
             endColumn: change.position.column + change.value.length
         },
         options: {
-            beforeContentClassName: `editor-ghost-text editor-ghost-text-content-${cssId} ${change.options?.greenHighlight ? 'editor-ghost-text-green' : ''}`
+            beforeContentClassName: `editor-ghost-text editor-ghost-text-content-${cssId} ${change.options?.greenHighlight
+                ? revertMode
+                    ? 'editor-ghost-text-removed'
+                    : 'editor-ghost-text-green'
+                : ''}`
         }
     };
     const safeContent = change.value.replaceAll('"', '\\"');
@@ -90,14 +96,14 @@ function getReviewButtons(editor, acceptFn, rejectFn) {
     reviewButtons.append(rejectButton);
     return reviewButtons;
 }
-async function addMultilineGhostText(editor, text, afterLineNumber, heightInLines, options) {
+async function addMultilineGhostText(editor, text, afterLineNumber, heightInLines, options, revertMode) {
     const el = document.createElement('div');
     el.textContent = text;
     if (options?.review) {
         const reviewButtons = getReviewButtons(editor, options.review.acceptFn, options.review.rejectFn);
         el.append(reviewButtons);
     }
-    applyMonacoStyles(el, options?.greenHighlight);
+    applyMonacoStyles(el, options?.greenHighlight, revertMode);
     const addZonePromise = new Promise((resolve, reject) => {
         editor?.changeViewZones((acc) => {
             const id = acc.addZone({
@@ -112,13 +118,13 @@ async function addMultilineGhostText(editor, text, afterLineNumber, heightInLine
     return addZonePromise;
 }
 export let VISUAL_CHANGES_CSS = `.editor-ghost-text-green { background-color: var(--vscode-diffEditor-insertedTextBackground) !important; }\n.editor-ghost-text-removed { background-color: var(--vscode-diffEditor-removedTextBackground); }\n\n.editor-ghost-text { display: inline-block; background-color: var(--vscode-editor-background); color: gray;}`;
-export async function displayVisualChanges(cssId, editor, visualChanges) {
+export async function displayVisualChanges(cssId, editor, visualChanges, revertMode) {
     let decorations = [];
     let css = '';
     let ids = [];
     for (const change of visualChanges) {
         if (change.type === 'added_inline') {
-            const { css: newCss, decoration } = addInlineGhostText(change);
+            const { css: newCss, decoration } = addInlineGhostText(change, revertMode);
             decorations.push(decoration);
             css += newCss;
         }
@@ -131,7 +137,7 @@ export async function displayVisualChanges(cssId, editor, visualChanges) {
                     endColumn: change.range.endColumn
                 },
                 options: {
-                    className: 'editor-ghost-text-removed',
+                    className: revertMode ? 'editor-ghost-text-green' : 'editor-ghost-text-removed',
                     isWholeLine: change.options?.isWholeLine
                 }
             };
@@ -157,7 +163,7 @@ export async function displayVisualChanges(cssId, editor, visualChanges) {
         }
         else if (change.type === 'added_block') {
             const id = await addMultilineGhostText(editor, change.value, change.position.afterLineNumber, change.value.split('\n').length, // we know it won't end by \n
-            change.options);
+            change.options, revertMode);
             ids.push(id);
         }
     }