@theia/ai-core 1.66.0-next.41 → 1.66.0-next.67

package/src/common/prompt-variable-contribution.spec.ts

@@ -0,0 +1,236 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource and others.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+
+import { enableJSDOM } from '@theia/core/lib/browser/test/jsdom';
+
+let disableJSDOM = enableJSDOM();
+
+import 'reflect-metadata';
+
+import { expect } from 'chai';
+import * as sinon from 'sinon';
+import { Container } from 'inversify';
+import { CommandService, ILogger, Logger } from '@theia/core';
+import { PromptVariableContribution, PROMPT_VARIABLE } from './prompt-variable-contribution';
+import { PromptService, PromptServiceImpl } from './prompt-service';
+import { DefaultAIVariableService, AIVariableService } from './variable-service';
+import { MockLogger } from '@theia/core/lib/common/test/mock-logger';
+
+disableJSDOM();
+
+describe('PromptVariableContribution', () => {
+    before(() => disableJSDOM = enableJSDOM());
+    after(() => disableJSDOM());
+    let contribution: PromptVariableContribution;
+    let promptService: PromptService;
+    let container: Container;
+
+    beforeEach(() => {
+        container = new Container();
+
+        // Set up PromptService
+        container.bind<PromptService>(PromptService).to(PromptServiceImpl).inSingletonScope();
+        const logger = sinon.createStubInstance(Logger);
+        const variableService = new DefaultAIVariableService({ getContributions: () => [] }, logger);
+        container.bind<AIVariableService>(AIVariableService).toConstantValue(variableService);
+        container.bind<ILogger>(ILogger).toConstantValue(new MockLogger);
+
+        // Set up CommandService stub (needed for PromptVariableContribution but not used in these tests)
+        const commandService = sinon.createStubInstance(Logger); // Using Logger as a simple mock
+        container.bind<CommandService>(CommandService).toConstantValue(commandService as unknown as CommandService);
+
+        // Bind PromptVariableContribution with proper DI
+        container.bind<PromptVariableContribution>(PromptVariableContribution).toSelf().inSingletonScope();
+
+        // Get instances
+        promptService = container.get<PromptService>(PromptService);
+        contribution = container.get<PromptVariableContribution>(PromptVariableContribution);
+    });
+
+    describe('Command Argument Substitution', () => {
+        it('substitutes $ARGUMENTS with full argument string', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Process: $ARGUMENTS',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|arg1 arg2 arg3' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Process: arg1 arg2 arg3');
+        });
+
+        it('substitutes $0 with command name', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Command $0 was called',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|args' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Command test-cmd was called');
+        });
+
+        it('substitutes $1, $2, ... with individual arguments', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'compare-cmd',
+                template: 'Compare $1 with $2',
+                isCommand: true,
+                commandName: 'compare'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'compare-cmd|item1 item2' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Compare item1 with item2');
+        });
+
+        it('handles quoted arguments in $1, $2', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'First: $1, Second: $2',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|"arg with spaces" other' },
+                {}
+            );
+
+            expect(result?.value).to.equal('First: arg with spaces, Second: other');
+        });
+
+        it('handles escaped quotes in arguments', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Arg: $1',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|"value with \\"quote\\""' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Arg: value with "quote"');
+        });
+
+        it('handles 10+ arguments correctly', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Args: $1 $10 $11',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|a b c d e f g h i j k' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Args: a j k');
+        });
+
+        it('handles command without arguments', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'hello-cmd',
+                template: 'Hello, world!',
+                isCommand: true,
+                commandName: 'hello'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'hello-cmd' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Hello, world!');
+        });
+
+        it('handles non-command prompts without substitution', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'normal-prompt',
+                template: 'This has $1 and $ARGUMENTS but is not a command'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'normal-prompt' },
+                {}
+            );
+
+            // No substitution should occur for non-commands
+            expect(result?.value).to.equal('This has $1 and $ARGUMENTS but is not a command');
+        });
+
+        it('handles missing argument placeholders gracefully', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Args: $1 $2 $3',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|only-one' },
+                {}
+            );
+
+            // Missing arguments should remain as placeholders
+            expect(result?.value).to.equal('Args: only-one $2 $3');
+        });
+    });
+
+    describe('Command Resolution', () => {
+        it('resolves command fragments correctly', async () => {
+            promptService.addBuiltInPromptFragment({
+                id: 'test-cmd',
+                template: 'Do something with $ARGUMENTS',
+                isCommand: true,
+                commandName: 'test'
+            });
+
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'test-cmd|input' },
+                {}
+            );
+
+            expect(result?.value).to.equal('Do something with input');
+            expect(result?.variable).to.deep.equal(PROMPT_VARIABLE);
+        });
+
+        it('returns empty string for non-existent prompts', async () => {
+            const result = await contribution.resolve(
+                { variable: PROMPT_VARIABLE, arg: 'non-existent|args' },
+                {}
+            );
+
+            expect(result?.value).to.equal('');
+        });
+    });
+});

package/src/common/prompt-variable-contribution.ts

@@ -69,13 +69,56 @@ export class PromptVariableContribution implements AIVariableContribution, AIVar
         resolveDependency?: (variable: AIVariableArg) => Promise<ResolvedAIVariable | undefined>
     ): Promise<ResolvedAIVariable | undefined> {
         if (request.variable.name === PROMPT_VARIABLE.name) {
-            const promptId = request.arg?.trim();
-            if (promptId) {
-                const resolvedPrompt = await this.promptService.getResolvedPromptFragmentWithoutFunctions(promptId, undefined, context, resolveDependency);
+            const arg = request.arg?.trim();
+            if (arg) {
+                // Check if this is a command-style reference (contains | separator)
+                const pipeIndex = arg.indexOf('|');
+                const promptIdOrCommandName = pipeIndex >= 0 ? arg.substring(0, pipeIndex) : arg;
+                const commandArgs = pipeIndex >= 0 ? arg.substring(pipeIndex + 1) : '';
+
+                // Determine the actual fragment ID
+                // If this is a command invocation (has args), try to find by command name first
+                let fragment = commandArgs
+                    ? this.promptService.getPromptFragmentByCommandName(promptIdOrCommandName)
+                    : undefined;
+
+                // Fall back to looking up by fragment ID if not found by command name
+                if (!fragment) {
+                    fragment = this.promptService.getRawPromptFragment(promptIdOrCommandName);
+                }
+
+                // If we still don't have a fragment, we can't resolve
+                if (!fragment) {
+                    this.logger.debug(`Could not find prompt fragment or command '${promptIdOrCommandName}'`);
+                    return {
+                        variable: request.variable,
+                        value: '',
+                        allResolvedDependencies: []
+                    };
+                }
+
+                const fragmentId = fragment.id;
+
+                // Resolve the prompt fragment normally (this handles {{variables}} and ~{functions})
+                const resolvedPrompt = await this.promptService.getResolvedPromptFragmentWithoutFunctions(
+                    fragmentId,
+                    undefined,
+                    context,
+                    resolveDependency
+                );
+
                 if (resolvedPrompt) {
+                    // If command args were provided, substitute them in the resolved text
+                    // This happens AFTER variable/function resolution, so $ARGUMENTS can be part of the template
+                    // alongside {{variables}} which get resolved first
+                    const isCommand = fragment?.isCommand === true;
+                    const finalText = isCommand && commandArgs
+                        ? this.substituteCommandArguments(resolvedPrompt.text, promptIdOrCommandName, commandArgs)
+                        : resolvedPrompt.text;
+
                     return {
                         variable: request.variable,
-                        value: resolvedPrompt.text,
+                        value: finalText,
                         allResolvedDependencies: resolvedPrompt.variables
                     };
                 }
@@ -89,6 +132,68 @@ export class PromptVariableContribution implements AIVariableContribution, AIVar
         };
     }
 
+    private substituteCommandArguments(template: string, commandName: string, commandArgs: string): string {
+        // Parse arguments (respecting quotes)
+        const args = this.parseCommandArguments(commandArgs);
+
+        // Substitute $ARGUMENTS with full arg string
+        let result = template.replace(/\$ARGUMENTS/g, commandArgs);
+
+        // Substitute $0 with command name
+        result = result.replace(/\$0/g, commandName);
+
+        // Substitute numbered arguments in reverse order to avoid collision
+        // (e.g., $10 before $1 to prevent $1 from matching the "1" in "$10")
+        for (let i = args.length; i > 0; i--) {
+            const regex = new RegExp(`\\$${i}\\b`, 'g');
+            result = result.replace(regex, args[i - 1]);
+        }
+
+        return result;
+    }
+
+    private parseCommandArguments(commandArgs: string): string[] {
+        const args: string[] = [];
+        let current = '';
+        let inQuotes = false;
+        let quoteChar = '';
+
+        for (let i = 0; i < commandArgs.length; i++) {
+            const char = commandArgs[i];
+
+            // Handle escape sequences within quotes
+            if (char === '\\' && i + 1 < commandArgs.length && inQuotes) {
+                const nextChar = commandArgs[i + 1];
+                if (nextChar === '"' || nextChar === "'" || nextChar === '\\') {
+                    current += nextChar;
+                    i++; // Skip the next character
+                    continue;
+                }
+            }
+
+            if ((char === '"' || char === "'") && !inQuotes) {
+                inQuotes = true;
+                quoteChar = char;
+            } else if (char === quoteChar && inQuotes) {
+                inQuotes = false;
+                quoteChar = '';
+            } else if (char === ' ' && !inQuotes) {
+                if (current.trim()) {
+                    args.push(current.trim());
+                    current = '';
+                }
+            } else {
+                current += char;
+            }
+        }
+
+        if (current.trim()) {
+            args.push(current.trim());
+        }
+
+        return args;
+    }
+
     protected async triggerArgumentPicker(): Promise<string | undefined> {
         // Trigger the suggestion command to show argument completions
         this.commandService.executeCommand('editor.action.triggerSuggest');