From 11f524c1251e1f604c9b64a215b81512e23cd396 Mon Sep 17 00:00:00 2001
From: Tommaso Sciortino <sciortino@gmail.com>
Date: Mon, 16 Jun 2025 13:24:42 -0700
Subject: Propagate abort signal to ccpa generateContent. (#1106)

---
 packages/core/src/code_assist/server.test.ts |  2 ++
 packages/core/src/code_assist/server.ts      | 11 ++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

(limited to 'packages/core/src/code_assist')

diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
index bb033163..922d20fb 100644
--- a/packages/core/src/code_assist/server.test.ts
+++ b/packages/core/src/code_assist/server.test.ts
@@ -45,6 +45,7 @@ describe('CodeAssistServer', () => {
     expect(server.callEndpoint).toHaveBeenCalledWith(
       'generateContent',
       expect.any(Object),
+      undefined,
     );
     expect(response.candidates?.[0]?.content?.parts?.[0]?.text).toBe(
       'response',
@@ -82,6 +83,7 @@ describe('CodeAssistServer', () => {
       expect(server.streamEndpoint).toHaveBeenCalledWith(
         'streamGenerateContent',
         expect.any(Object),
+        undefined,
       );
       expect(res.candidates?.[0]?.content?.parts?.[0]?.text).toBe('response');
     }
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index fc3a0ce2..2c5d6db9 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/packages/core/src/code_assist/server.ts
@@ -53,6 +53,7 @@ export class CodeAssistServer implements ContentGenerator {
     const resps = await this.streamEndpoint<CodeAssistResponse>(
       'streamGenerateContent',
       toCodeAssistRequest(req, this.projectId),
+      req.config?.abortSignal,
     );
     return (async function* (): AsyncGenerator<GenerateContentResponse> {
       for await (const resp of resps) {
@@ -67,6 +68,7 @@ export class CodeAssistServer implements ContentGenerator {
     const resp = await this.callEndpoint<CodeAssistResponse>(
       'generateContent',
       toCodeAssistRequest(req, this.projectId),
+      req.config?.abortSignal,
     );
     return fromCodeAsistResponse(resp);
   }
@@ -99,7 +101,11 @@ export class CodeAssistServer implements ContentGenerator {
     throw Error();
   }
 
-  async callEndpoint<T>(method: string, req: object): Promise<T> {
+  async callEndpoint<T>(
+    method: string,
+    req: object,
+    signal?: AbortSignal,
+  ): Promise<T> {
     const res = await this.auth.request({
       url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`,
       method: 'POST',
@@ -109,6 +115,7 @@ export class CodeAssistServer implements ContentGenerator {
       },
       responseType: 'json',
       body: JSON.stringify(req),
+      signal,
     });
     return res.data as T;
   }
@@ -116,6 +123,7 @@ export class CodeAssistServer implements ContentGenerator {
   async streamEndpoint<T>(
     method: string,
     req: object,
+    signal?: AbortSignal,
   ): Promise<AsyncGenerator<T>> {
     const res = await this.auth.request({
       url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`,
@@ -129,6 +137,7 @@ export class CodeAssistServer implements ContentGenerator {
       },
       responseType: 'stream',
       body: JSON.stringify(req),
+      signal,
     });
 
     return (async function* (): AsyncGenerator<T> {
-- 
cgit v1.2.3