Add ability to stop automodeling

Author: charisk

extensions/ql-vscode/src/common/interface-types.ts

@@ -554,6 +554,11 @@ interface GenerateExternalApiFromLlmMessage {
   modeledMethods: Record<string, ModeledMethod>;
 }
 
+interface StopGeneratingExternalApiFromLlmMessage {
+  t: "stopGeneratingExternalApiFromLlm";
+  dependencyName: string;
+}
+
 interface ModelDependencyMessage {
   t: "modelDependency";
 }
@@ -575,4 +580,5 @@ export type FromDataExtensionsEditorMessage =
   | SaveModeledMethods
   | GenerateExternalApiMessage
   | GenerateExternalApiFromLlmMessage
+  | StopGeneratingExternalApiFromLlmMessage
   | ModelDependencyMessage;

extensions/ql-vscode/src/data-extensions-editor/auto-model-codeml-queries.ts

@@ -160,6 +160,7 @@ type AutoModelQueriesOptions = {
   queryStorageDir: string;
 
   progress: ProgressCallback;
+  cancellationTokenSource: CancellationTokenSource;
 };
 
 export type AutoModelQueriesResult = {
@@ -174,12 +175,11 @@ export async function runAutoModelQueries({
   databaseItem,
   queryStorageDir,
   progress,
+  cancellationTokenSource,
 }: AutoModelQueriesOptions): Promise<AutoModelQueriesResult | undefined> {
   // maxStep for this part is 1500
   const maxStep = 1500;
 
-  const cancellationTokenSource = new CancellationTokenSource();
-
   const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
 
   // CodeQL needs to have access to the database to be able to retrieve the

extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts

@@ -15,13 +15,16 @@ import { CodeQLCliServer } from "../codeql-cli/cli";
 import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
+import { CancellationTokenSource } from "vscode";
 
 // Limit the number of candidates we send to the model in each request
 // to avoid long requests.
 // Note that the model may return fewer than this number of candidates.
 const candidateBatchSize = 20;
 
 export class AutoModeler {
+  private readonly jobs: Map<string, CancellationTokenSource>;
+
   constructor(
     private readonly app: App,
     private readonly cliServer: CodeQLCliServer,
@@ -34,27 +37,56 @@ export class AutoModeler {
     private readonly addModeledMethods: (
       modeledMethods: Record<string, ModeledMethod>,
     ) => Promise<void>,
-  ) {}
+  ) {
+    this.jobs = new Map<string, CancellationTokenSource>();
+  }
 
   public async startModeling(
     dependency: string,
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
     mode: Mode,
   ): Promise<void> {
-    await this.modelDependency(
-      dependency,
-      externalApiUsages,
-      modeledMethods,
-      mode,
-    );
+    if (this.jobs.has(dependency)) {
+      return;
+    }
+
+    const cancellationTokenSource = new CancellationTokenSource();
+    this.jobs.set(dependency, cancellationTokenSource);
+
+    try {
+      await this.modelDependency(
+        dependency,
+        externalApiUsages,
+        modeledMethods,
+        mode,
+        cancellationTokenSource,
+      );
+    } finally {
+      this.jobs.delete(dependency);
+    }
+  }
+
+  public async stopModeling(dependency: string): Promise<void> {
+    void extLogger.log(`Stopping modeling for dependency ${dependency}`);
+    const cancellationTokenSource = this.jobs.get(dependency);
+    if (cancellationTokenSource) {
+      cancellationTokenSource.cancel();
+    }
+  }
+
+  public async stopAllModeling(): Promise<void> {
+    for (const cancellationTokenSource of this.jobs.values()) {
+      cancellationTokenSource.cancel();
+    }
   }
 
   private async modelDependency(
     dependency: string,
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
     mode: Mode,
+    cancellationTokenSource: CancellationTokenSource,
   ): Promise<void> {
     void extLogger.log(`Modeling dependency ${dependency}`);
     await withProgress(async (progress) => {
@@ -79,6 +111,10 @@ export class AutoModeler {
       );
       try {
         for (let i = 0; i < batchNumber; i++) {
+          if (cancellationTokenSource.token.isCancellationRequested) {
+            break;
+          }
+
           const start = i * candidateBatchSize;
           const end = start + candidateBatchSize;
           const candidatesToProcess = allCandidateMethods.slice(start, end);
@@ -100,6 +136,7 @@ export class AutoModeler {
             mode,
             progress,
             maxStep,
+            cancellationTokenSource,
           );
         }
       } finally {
@@ -114,6 +151,7 @@ export class AutoModeler {
     mode: Mode,
     progress: ProgressCallback,
     maxStep: number,
+    cancellationTokenSource: CancellationTokenSource,
   ): Promise<void> {
     const usages = await runAutoModelQueries({
       mode,
@@ -123,6 +161,7 @@ export class AutoModeler {
       queryStorageDir: this.queryStorageDir,
       databaseItem: this.databaseItem,
       progress: (update) => progress({ ...update, maxStep }),
+      cancellationTokenSource,
     });
     if (!usages) {
       return;

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -187,6 +187,9 @@ export class DataExtensionsEditorView extends AbstractWebview<
           );
         }
         break;
+      case "stopGeneratingExternalApiFromLlm":
+        await this.autoModeler.stopModeling(msg.dependencyName);
+        break;
       case "modelDependency":
         await this.modelDependency();
         break;

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -241,6 +241,13 @@ export function DataExtensionsEditor({
     [],
   );
 
+  const onStopGenerateFromLlmClick = useCallback((dependencyName: string) => {
+    vscode.postMessage({
+      t: "stopGeneratingExternalApiFromLlm",
+      dependencyName,
+    });
+  }, []);
+
   const onOpenDatabaseClick = useCallback(() => {
     vscode.postMessage({
       t: "openDatabase",
@@ -345,6 +352,7 @@ export function DataExtensionsEditor({
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
+          onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
           onGenerateFromSourceClick={onGenerateFromSourceClick}
           onModelDependencyClick={onModelDependencyClick}
         />

extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx

@@ -89,6 +89,7 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onStopGenerateFromLlmClick: (dependencyName: string) => void;
   onGenerateFromSourceClick: () => void;
   onModelDependencyClick: () => void;
 };
@@ -105,6 +106,7 @@ export const LibraryRow = ({
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
+  onStopGenerateFromLlmClick,
   onGenerateFromSourceClick,
   onModelDependencyClick,
 }: Props) => {
@@ -127,6 +129,15 @@ export const LibraryRow = ({
     [title, externalApiUsages, modeledMethods, onGenerateFromLlmClick],
   );
 
+  const handleStopModelWithAI = useCallback(
+    async (e: React.MouseEvent) => {
+      onStopGenerateFromLlmClick(title);
+      e.stopPropagation();
+      e.preventDefault();
+    },
+    [title, onStopGenerateFromLlmClick],
+  );
+
   const handleModelFromSource = useCallback(
     async (e: React.MouseEvent) => {
       onGenerateFromSourceClick();
@@ -167,6 +178,12 @@ export const LibraryRow = ({
     );
   }, [externalApiUsages, modifiedSignatures]);
 
+  const canStopAutoModeling = useMemo(() => {
+    return externalApiUsages.some((externalApiUsage) =>
+      inProgressSignatures.has(externalApiUsage.signature),
+    );
+  }, [externalApiUsages, inProgressSignatures]);
+
   return (
     <LibraryContainer>
       <TitleContainer onClick={toggleExpanded} aria-expanded={isExpanded}>
@@ -185,12 +202,18 @@ export const LibraryRow = ({
           </ModeledPercentage>
           {hasUnsavedChanges ? <VSCodeTag>UNSAVED</VSCodeTag> : null}
         </NameContainer>
-        {viewState.showLlmButton && (
+        {viewState.showLlmButton && !canStopAutoModeling && (
           <VSCodeButton appearance="icon" onClick={handleModelWithAI}>
             <Codicon name="lightbulb-autofix" label="Model with AI" />
             &nbsp;Model with AI
           </VSCodeButton>
         )}
+        {viewState.showLlmButton && canStopAutoModeling && (
+          <VSCodeButton appearance="icon" onClick={handleStopModelWithAI}>
+            <Codicon name="debug-stop" label="Stop model with AI" />
+            &nbsp;Stop
+          </VSCodeButton>
+        )}
         {viewState.mode === Mode.Application && (
           <VSCodeButton appearance="icon" onClick={handleModelFromSource}>
             <Codicon name="code" label="Model from source" />

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx

@@ -31,6 +31,7 @@ type Props = {
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
+  onStopGenerateFromLlmClick: (dependencyName: string) => void;
   onGenerateFromSourceClick: () => void;
   onModelDependencyClick: () => void;
 };
@@ -49,6 +50,7 @@ export const ModeledMethodsList = ({
   onChange,
   onSaveModelClick,
   onGenerateFromLlmClick,
+  onStopGenerateFromLlmClick,
   onGenerateFromSourceClick,
   onModelDependencyClick,
 }: Props) => {
@@ -93,6 +95,7 @@ export const ModeledMethodsList = ({
           onChange={onChange}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
+          onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}
           onGenerateFromSourceClick={onGenerateFromSourceClick}
           onModelDependencyClick={onModelDependencyClick}
         />

extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/auto-model-codeml-queries.test.ts

@@ -19,6 +19,7 @@ import { MethodSignature } from "../../../../src/data-extensions-editor/external
 import { join } from "path";
 import { exists, readFile } from "fs-extra";
 import { load as loadYaml } from "js-yaml";
+import { CancellationTokenSource } from "vscode-jsonrpc";
 
 describe("runAutoModelQueries", () => {
   const qlpack = {
@@ -142,6 +143,7 @@ describe("runAutoModelQueries", () => {
       }),
       queryStorageDir: "/tmp/queries",
       progress: jest.fn(),
+      cancellationTokenSource: new CancellationTokenSource(),
     };
 
     const result = await runAutoModelQueries(options);