Improve scenario recording

Author: koesie10

extensions/ql-vscode/src/common/mock-gh-api/gh-api-request.ts

@@ -74,23 +74,28 @@ export interface GetVariantAnalysisRepoResultRequest {
   };
 }
 
+interface CodeSearchResponse {
+  total_count: number;
+  items: Array<{
+    repository: Repository;
+  }>;
+}
+
 interface CodeSearchRequest {
   request: {
     kind: RequestKind.CodeSearch;
     query: string;
   };
   response: {
     status: number;
-    body?: {
-      total_count?: number;
-      items?: Array<{
-        repository: Repository;
-      }>;
-    };
-    message?: string;
+    body?: CodeSearchResponse | BasicErorResponse;
   };
 }
 
+interface AutoModelResponse {
+  models: string;
+}
+
 interface AutoModelRequest {
   request: {
     kind: RequestKind.AutoModel;
@@ -100,10 +105,7 @@ interface AutoModelRequest {
   };
   response: {
     status: number;
-    body?: {
-      models: string;
-    };
-    message?: string;
+    body?: AutoModelResponse | BasicErorResponse;
   };
 }
 

extensions/ql-vscode/src/common/mock-gh-api/mock-gh-api-server.ts

@@ -12,18 +12,31 @@ import { getDirectoryNamesInsidePath } from "../files";
  * Enables mocking of the GitHub API server via HTTP interception, using msw.
  */
 export class MockGitHubApiServer extends DisposableObject {
+  private _isListening: boolean;
+
   private readonly server: SetupServer;
   private readonly recorder: Recorder;
 
   constructor() {
     super();
+    this._isListening = false;
 
     this.server = setupServer();
     this.recorder = this.push(new Recorder(this.server));
   }
 
+  public startServer(): void {
+    if (this._isListening) {
+      return;
+    }
+
+    this.server.listen({ onUnhandledRequest: "bypass" });
+    this._isListening = true;
+  }
+
   public stopServer(): void {
     this.server.close();
+    this._isListening = false;
   }
 
   public async loadScenario(
@@ -42,7 +55,6 @@ export class MockGitHubApiServer extends DisposableObject {
     const handlers = await createRequestHandlers(scenarioPath);
     this.server.resetHandlers();
     this.server.use(...handlers);
-    this.server.listen({ onUnhandledRequest: "bypass" });
   }
 
   public async saveScenario(
@@ -99,6 +111,10 @@ export class MockGitHubApiServer extends DisposableObject {
     return await getDirectoryNamesInsidePath(scenariosPath);
   }
 
+  public get isListening(): boolean {
+    return this._isListening;
+  }
+
   public get isRecording(): boolean {
     return this.recorder.isRecording;
   }

extensions/ql-vscode/src/common/mock-gh-api/recorder.ts

@@ -3,8 +3,6 @@ import { join } from "path";
 
 import { SetupServer } from "msw/node";
 
-import fetch from "node-fetch";
-
 import { DisposableObject } from "../disposable-object";
 
 import {
@@ -14,14 +12,12 @@ import {
 } from "./gh-api-request";
 
 export class Recorder extends DisposableObject {
-  private readonly allRequests = new Map<string, Request>();
   private currentRecordedScenario: GitHubApiRequest[] = [];
 
   private _isRecording = false;
 
   constructor(private readonly server: SetupServer) {
     super();
-    this.onRequestStart = this.onRequestStart.bind(this);
     this.onResponseBypass = this.onResponseBypass.bind(this);
   }
 
@@ -42,7 +38,6 @@ export class Recorder extends DisposableObject {
 
     this.clear();
 
-    this.server.events.on("request:start", this.onRequestStart);
     this.server.events.on("response:bypass", this.onResponseBypass);
   }
 
@@ -53,13 +48,11 @@ export class Recorder extends DisposableObject {
 
     this._isRecording = false;
 
-    this.server.events.removeListener("request:start", this.onRequestStart);
     this.server.events.removeListener("response:bypass", this.onResponseBypass);
   }
 
   public clear() {
     this.currentRecordedScenario = [];
-    this.allRequests.clear();
   }
 
   public async save(scenariosPath: string, name: string): Promise<string> {
@@ -109,34 +102,14 @@ export class Recorder extends DisposableObject {
     return scenarioDirectory;
   }
 
-  private onRequestStart(request: Request, requestId: string): void {
-    if (request.headers.has("x-vscode-codeql-msw-bypass")) {
-      return;
-    }
-
-    this.allRequests.set(requestId, request);
-  }
-
   private async onResponseBypass(
     response: Response,
-    _: Request,
-    requestId: string,
+    request: Request,
+    _requestId: string,
   ): Promise<void> {
-    const request = this.allRequests.get(requestId);
-    this.allRequests.delete(requestId);
-    if (!request) {
-      return;
-    }
-
-    if (response.body === undefined) {
-      return;
-    }
-
     const gitHubApiRequest = await createGitHubApiRequest(
-      request.url.toString(),
-      response.status,
-      response.body?.toString() || "",
-      response.headers,
+      request.url,
+      response,
     );
     if (!gitHubApiRequest) {
       return;
@@ -148,22 +121,23 @@ export class Recorder extends DisposableObject {
 
 async function createGitHubApiRequest(
   url: string,
-  status: number,
-  body: string,
-  headers: globalThis.Headers,
+  response: Response,
 ): Promise<GitHubApiRequest | undefined> {
   if (!url) {
     return undefined;
   }
 
+  const status = response.status;
+  const headers = response.headers;
+
   if (url.match(/\/repos\/[a-zA-Z0-9-_.]+\/[a-zA-Z0-9-_.]+$/)) {
     return {
       request: {
         kind: RequestKind.GetRepo,
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }
@@ -177,7 +151,7 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }
@@ -193,7 +167,7 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }
@@ -209,7 +183,7 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }
@@ -219,25 +193,14 @@ async function createGitHubApiRequest(
     /objects-origin\.githubusercontent\.com\/codeql-query-console\/codeql-variant-analysis-repo-tasks\/\d+\/(?<repositoryId>\d+)/,
   );
   if (repoDownloadMatch?.groups?.repositoryId) {
-    // msw currently doesn't support binary response bodies, so we need to download this separately
-    // see https://github.com/mswjs/interceptors/blob/15eafa6215a328219999403e3ff110e71699b016/src/interceptors/ClientRequest/utils/getIncomingMessageBody.ts#L24-L33
-    // Essentially, mws is trying to decode a ZIP file as UTF-8 which changes the bytes and corrupts the file.
-    const response = await fetch(url, {
-      headers: {
-        // We need to ensure we don't end up in an infinite loop, since this request will also be intercepted
-        "x-vscode-codeql-msw-bypass": "true",
-      },
-    });
-    const responseBuffer = await response.buffer();
-
     return {
       request: {
         kind: RequestKind.GetVariantAnalysisRepoResult,
         repositoryId: parseInt(repoDownloadMatch.groups.repositoryId, 10),
       },
       response: {
         status,
-        body: responseBuffer,
+        body: Buffer.from(await response.arrayBuffer()),
         contentType: headers.get("content-type") ?? "application/octet-stream",
       },
     };
@@ -252,7 +215,7 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }
@@ -267,7 +230,7 @@ async function createGitHubApiRequest(
       },
       response: {
         status,
-        body: JSON.parse(body),
+        body: await response.json(),
       },
     };
   }

extensions/ql-vscode/src/common/mock-gh-api/vscode/vscode-mock-gh-api-server.ts

@@ -42,6 +42,10 @@ export class VSCodeMockGitHubApiServer extends DisposableObject {
     };
   }
 
+  public async startServer(): Promise<void> {
+    this.server.startServer();
+  }
+
   public async stopServer(): Promise<void> {
     this.server.stopServer();
 
@@ -252,7 +256,9 @@ export class VSCodeMockGitHubApiServer extends DisposableObject {
   }
 
   private async onConfigChange(): Promise<void> {
-    if (!this.config.mockServerEnabled) {
+    if (this.config.mockServerEnabled && !this.server.isListening) {
+      await this.startServer();
+    } else if (!this.config.mockServerEnabled && this.server.isListening) {
       await this.stopServer();
     }
   }

extensions/ql-vscode/test/unit-tests/variant-analysis/gh-api/gh-api-client.test.ts

@@ -13,6 +13,7 @@ import { response as variantAnalysisRepoJson_response } from "../../../../src/co
 import { testCredentialsWithRealOctokit } from "../../../factories/authentication";
 
 const mockServer = new MockGitHubApiServer();
+beforeAll(() => mockServer.startServer());
 afterEach(() => mockServer.unloadScenario());
 afterAll(() => mockServer.stopServer());
 

extensions/ql-vscode/test/vscode-tests/cli-integration/variant-analysis/variant-analysis-submission-integration.test.ts

@@ -16,6 +16,7 @@ import { createVSCodeCommandManager } from "../../../../src/common/vscode/comman
 import { AllCommands } from "../../../../src/common/commands";
 
 const mockServer = new MockGitHubApiServer();
+beforeAll(() => mockServer.startServer());
 afterEach(() => mockServer.unloadScenario());
 afterAll(() => mockServer.stopServer());