fix(agent): abort not break immediately

Author: ycjcl868

src/main/agent/index.ts

@@ -226,7 +226,9 @@ export class ComputerUseAgent {
         logger.info('[vlmParams_conversations]:', vlmParams.conversations);
         logger.info('[vlmParams_images_len]:', vlmParams.images.length);
 
-        const vlmRes = await vlm.invoke(vlmParams);
+        const vlmRes = await vlm.invoke(vlmParams, {
+          abortController,
+        });
 
         if (!vlmRes?.prediction) {
           continue;

src/main/agent/llm/base.ts

@@ -9,6 +9,10 @@ export interface VlmRequest {
   images: string[];
 }
 
+export interface VlmRequestOptions {
+  abortController?: AbortController | null;
+}
+
 export interface VlmResponse {
   prediction: string;
   reflections?: string[];
@@ -32,5 +36,8 @@ export abstract class VLM<
   K extends VlmResponse = VlmResponse,
 > {
   abstract get vlmModel(): string;
-  abstract invoke({ conversations, images }: T): Promise<K>;
+  abstract invoke(
+    { conversations, images }: T,
+    options?: VlmRequestOptions,
+  ): Promise<K>;
 }

src/main/agent/llm/ui-tars.ts

@@ -10,7 +10,7 @@ import { store } from '@main/store/create';
 import { preprocessResizeImage } from '@main/utils/image';
 
 import { MAX_PIXELS } from '../constant';
-import { VLM, VlmRequest, VlmResponse } from './base';
+import { VLM, VlmRequest, VlmRequestOptions, VlmResponse } from './base';
 
 export interface UITARSOptions {
   reflection: boolean;
@@ -31,7 +31,11 @@ export class UITARS implements VLM<VlmRequest, VlmResponse> {
 
   // [image, prompt]
   // [gpt, image]
-  async invoke({ conversations, images }: VlmRequest) {
+  async invoke(
+    { conversations, images }: VlmRequest,
+    options?: VlmRequestOptions,
+  ) {
+    const { abortController } = options ?? {};
     const compressedImages = await Promise.all(
       images.map((image) => preprocessResizeImage(image, MAX_PIXELS)),
     );
@@ -51,19 +55,24 @@ export class UITARS implements VLM<VlmRequest, VlmResponse> {
 
     const startTime = Date.now();
     const result = await openai.chat.completions
-      .create({
-        model: this.vlmModel,
-        max_tokens: 1000,
-        stream: false,
-        temperature: 0,
-        top_p: 0.7,
-        seed: null,
-        stop: null,
-        frequency_penalty: null,
-        presence_penalty: null,
-        // messages
-        messages,
-      })
+      .create(
+        {
+          model: this.vlmModel,
+          max_tokens: 1000,
+          stream: false,
+          temperature: 0,
+          top_p: 0.7,
+          seed: null,
+          stop: null,
+          frequency_penalty: null,
+          presence_penalty: null,
+          // messages
+          messages,
+        },
+        {
+          signal: abortController?.signal,
+        },
+      )
       .finally(() => {
         logger.info(`[vlm_invoke_time_cost]: ${Date.now() - startTime}ms`);
       });

src/main/store/create.ts

@@ -21,6 +21,8 @@ import { runAgent } from './runAgent';
 import { SettingStore } from './setting';
 import { AppState } from './types';
 
+const abortController = new AbortController();
+
 export const store = createStore<AppState>(
   (set, get) =>
     ({
@@ -33,7 +35,7 @@ export const store = createStore<AppState>(
       getSetting: (key) => SettingStore.get(key),
       ensurePermissions: {},
 
-      abortController: null,
+      abortController,
       thinking: false,
 
       // dispatch for renderer
@@ -81,7 +83,7 @@ export const store = createStore<AppState>(
       },
 
       RUN_AGENT: async () => {
-        set({ abortController: new AbortController(), thinking: true });
+        set({ thinking: true });
 
         await runAgent(set, get);
 
@@ -90,7 +92,8 @@ export const store = createStore<AppState>(
       STOP_RUN: () => {
         set({ status: StatusEnum.END, thinking: false });
         showWindow();
-        get().abortController?.abort();
+        abortController.abort();
+
         closeScreenMarker();
       },
       SET_INSTRUCTIONS: (instructions) => {