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.ts | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

(limited to 'packages/core/src/code_assist/server.ts')

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