feat: add app management, local inference controls, and task statistics

src/main/knowledge/embedding-inference-broker.ts

@@ -1,4 +1,6 @@
 import { randomUUID } from 'node:crypto'
+import { localInferenceService } from '../local-inference-service'
+import type { InferenceResources } from '../../shared/local-inference-contracts'
 import {
   EMBEDDING_INFERENCE_MAX_IN_FLIGHT,
   embeddingInferenceProtocolBase,
@@ -15,13 +17,13 @@ const MAX_TIMEOUT_MS = 300_000
 const DEFAULT_MAX_RESTARTS = 2
 const MAX_RESTARTS = 10
 const SHUTDOWN_GRACE_MS = 1_000
-const MAX_RETIRED_REQUESTS = 256
 
 export interface EmbeddingInferenceTransport {
+  getResources?(): InferenceResources
   postMessage(message: EmbeddingInferenceRequest): void
   onMessage(listener: (message: unknown) => void): () => void
   onClose(listener: () => void): () => void
-  close(): void
+  close(): void | Promise<void>
 }
 
 export type EmbeddingInferenceTransportFactory = () =>
@@ -42,6 +44,8 @@ type PendingRequest = {
   readonly expectedCount: number
   readonly resolve: (vectors: number[][]) => void
   readonly reject: (error: Error) => void
+  readonly rejectCaller: (error: Error) => void
+  callerError?: Error
   readonly timeout: ReturnType<typeof setTimeout>
   removeAbortListener?: () => void
 }
@@ -62,9 +66,10 @@ export class EmbeddingInferenceBrokerError extends Error {
       | 'START_FAILURE'
       | 'ENGINE_FAILURE'
       | 'INVALID_RESULT'
-      | 'OVERLOADED'
+      | 'OVERLOADED',
+    detail?: string
   ) {
-    super(`Embedding inference failed (${code})`)
+    super(`Embedding inference failed (${code})${detail ? `: ${detail}` : ''}`)
     this.name = 'EmbeddingInferenceBrokerError'
   }
 }
@@ -100,13 +105,65 @@ export class EmbeddingInferenceBroker {
   private readonly timeoutMs: number
   private readonly maxRestarts: number
   private readonly pending = new Map<string, PendingRequest>()
-  private readonly retiredRequestIds = new Set<string>()
   private active?: ActiveTransport
   private starting?: Promise<ActiveTransport>
   private restarts = 0
   private everStarted = false
   private closed = false
   private shutdownComplete?: () => void
+  private paused = false
+  private stopping = false
+  private managementError?: string
+
+  getManagementState(): 'stopped' | 'stopping' | 'starting' | 'running' | 'idle' | 'error' {
+    if (this.stopping) return 'stopping'
+    if (this.paused || this.closed) return 'stopped'
+    if (this.starting) return 'starting'
+    if (this.active) return 'running'
+    return this.managementError ? 'error' : 'idle'
+  }
+
+  getManagementError(): string | undefined { return this.managementError }
+
+  getResources(): InferenceResources {
+    return this.active?.transport.getResources?.() ?? {
+      scope: 'unavailable', reason: this.active ? '当前执行端未提供进程指标' : '服务进程未运行'
+    }
+  }
+
+  async start(): Promise<void> {
+    if (this.closed) throw new EmbeddingInferenceBrokerError('CLOSED')
+    this.paused = false
+    this.restarts = 0
+    try {
+      await this.ensureTransport()
+      this.managementError = undefined
+    } catch (error) {
+      this.managementError = String(error)
+      throw error
+    }
+  }
+
+  async stop(): Promise<void> {
+    this.paused = true
+    this.stopping = true
+    const error = new Error('向量执行服务已由用户停止；任务不会自动重放')
+    this.rejectCallers(error)
+    try {
+      if (this.starting) await this.starting.catch(() => undefined)
+      const active = this.active
+      if (!active) return
+      await active.transport.close()
+      this.disposeTransport(active.transport, false)
+      this.rejectAll(error)
+    } catch (error) {
+      this.paused = false
+      this.managementError = error instanceof Error ? error.message : String(error)
+      throw error
+    } finally {
+      this.stopping = false
+    }
+  }
 
   constructor(options: EmbeddingInferenceBrokerOptions) {
     if (typeof options.createTransport !== 'function') {
@@ -131,12 +188,29 @@ export class EmbeddingInferenceBroker {
     texts: readonly string[],
     role: EmbeddingInferenceRole,
     signal?: AbortSignal
+  ): Promise<number[][]> {
+    const controller = new AbortController()
+    return new Promise<number[][]>((resolve, reject) => {
+      // Track worker execution separately from prompt caller cancellation/timeout.
+      void localInferenceService.run('embedding', role === 'query' ? '向量请求 · 查询' : '向量请求 · 文档',
+        (markCancelling) => this.embedRequest(texts, role, signal, controller.signal, (error) => {
+          markCancelling()
+          reject(error)
+        }),
+        { cancel: () => controller.abort() }).then(resolve, reject)
+    })
+  }
+
+  private async embedRequest(
+    texts: readonly string[], role: EmbeddingInferenceRole, signal: AbortSignal | undefined,
+    managementSignal: AbortSignal,
+    rejectCaller: (error: Error) => void
   ): Promise<number[][]> {
     const normalizedTexts = validateEmbeddingInferenceTexts(texts)
     if (role !== 'query' && role !== 'document') {
       throw new TypeError('role must be query or document')
     }
-    if (this.closed) {
+    if (this.closed || this.paused) {
       throw new EmbeddingInferenceBrokerError('CLOSED')
     }
     if (signal?.aborted) {
@@ -147,7 +221,7 @@ export class EmbeddingInferenceBroker {
     }
 
     const active = await this.ensureTransport()
-    if (this.closed) {
+    if (this.closed || this.paused) {
       throw new EmbeddingInferenceBrokerError('CLOSED')
     }
     if (signal?.aborted) {
@@ -166,6 +240,7 @@ export class EmbeddingInferenceBroker {
         expectedCount: normalizedTexts.length,
         resolve,
         reject,
+        rejectCaller,
         timeout
       }
       if (signal) {
@@ -178,6 +253,21 @@ export class EmbeddingInferenceBroker {
       }
       this.pending.set(requestId, pending)
 
+      const requestCancellation = (): void => {
+        // Management cancellation waits for the worker's result or CANCELLED response.
+        try {
+          active.transport.postMessage({ ...embeddingInferenceProtocolBase(), type: 'cancel', requestId })
+        } catch {
+          this.failTransport(active.transport, new EmbeddingInferenceBrokerError('PROCESS_CLOSED'))
+        }
+      }
+      const removeSourceAbort = pending.removeAbortListener
+      pending.removeAbortListener = () => {
+        removeSourceAbort?.()
+        managementSignal.removeEventListener('abort', requestCancellation)
+      }
+      managementSignal.addEventListener('abort', requestCancellation, { once: true })
+
       try {
         active.transport.postMessage({
           ...embeddingInferenceProtocolBase(),
@@ -186,6 +276,7 @@ export class EmbeddingInferenceBroker {
           role,
           texts: normalizedTexts
         })
+        if (managementSignal.aborted) requestCancellation()
       } catch {
         this.failTransport(
           active.transport,
@@ -200,10 +291,7 @@ export class EmbeddingInferenceBroker {
       return
     }
     this.closed = true
-    for (const requestId of this.pending.keys()) {
-      this.retireRequestId(requestId)
-    }
-    this.rejectAll(new EmbeddingInferenceBrokerError('CLOSED'))
+    this.rejectCallers(new EmbeddingInferenceBrokerError('CLOSED'))
 
     let active = this.active
     if (!active && this.starting) {
@@ -261,7 +349,7 @@ export class EmbeddingInferenceBroker {
       .then(() => this.createTransport())
       .then((transport) => {
         if (this.closed) {
-          transport.close()
+          void Promise.resolve(transport.close()).catch(() => undefined)
           throw new EmbeddingInferenceBrokerError('CLOSED')
         }
         const active: ActiveTransport = {
@@ -272,7 +360,8 @@ export class EmbeddingInferenceBroker {
           removeCloseListener: transport.onClose(() =>
             this.failTransport(
               transport,
-              new EmbeddingInferenceBrokerError('PROCESS_CLOSED')
+              new EmbeddingInferenceBrokerError('PROCESS_CLOSED'),
+              true
             )
           )
         }
@@ -280,6 +369,7 @@ export class EmbeddingInferenceBroker {
         return active
       })
       .catch((error: unknown) => {
+        this.managementError = error instanceof Error ? error.message : String(error)
         if (error instanceof EmbeddingInferenceBrokerError) {
           throw error
         }
@@ -317,6 +407,7 @@ export class EmbeddingInferenceBroker {
         )
         return
       }
+      this.rejectAll(new EmbeddingInferenceBrokerError('CLOSED'))
       this.shutdownComplete()
       return
     }
@@ -330,9 +421,6 @@ export class EmbeddingInferenceBroker {
 
     const pending = this.pending.get(message.requestId)
     if (!pending) {
-      if (this.retiredRequestIds.delete(message.requestId)) {
-        return
-      }
       this.failTransport(
         transport,
         new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
@@ -345,7 +433,7 @@ export class EmbeddingInferenceBroker {
         message.requestId,
         message.code === 'CANCELLED'
           ? abortError()
-          : new EmbeddingInferenceBrokerError(message.code)
+          : new EmbeddingInferenceBrokerError(message.code, message.message)
       )
       return
     }
@@ -365,11 +453,12 @@ export class EmbeddingInferenceBroker {
 
   private cancelPending(requestId: string, error: Error): void {
     const pending = this.pending.get(requestId)
-    if (!pending) {
+    if (!pending || pending.callerError) {
       return
     }
-    this.retireRequestId(requestId)
-    this.finishPending(requestId, error)
+    pending.callerError = error
+    clearTimeout(pending.timeout)
+    pending.rejectCaller(error)
     try {
       this.active?.transport.postMessage({
         ...embeddingInferenceProtocolBase(),
@@ -398,6 +487,7 @@ export class EmbeddingInferenceBroker {
     this.pending.delete(requestId)
     clearTimeout(pending.timeout)
     pending.removeAbortListener?.()
+    error = pending.callerError ?? error
     if (error) {
       pending.reject(error)
     } else {
@@ -405,26 +495,23 @@ export class EmbeddingInferenceBroker {
     }
   }
 
-  private retireRequestId(requestId: string): void {
-    this.retiredRequestIds.add(requestId)
-    if (this.retiredRequestIds.size > MAX_RETIRED_REQUESTS) {
-      const oldest = this.retiredRequestIds.values().next().value
-      if (typeof oldest === 'string') {
-        this.retiredRequestIds.delete(oldest)
-      }
-    }
-  }
-
   private failTransport(
     transport: EmbeddingInferenceTransport,
-    error: Error
+    error: Error,
+    closed = false
   ): void {
     if (this.active?.transport !== transport) {
       return
     }
+    this.rejectCallers(error)
+    if (closed) {
+      this.disposeTransport(transport, false)
+      this.rejectAll(error)
+      this.shutdownComplete?.()
+      return
+    }
+    this.managementError = error.message
     this.disposeTransport(transport, true)
-    this.rejectAll(error)
-    this.shutdownComplete?.()
   }
 
   private disposeTransport(
@@ -435,15 +522,30 @@ export class EmbeddingInferenceBroker {
     if (!active || active.transport !== transport) {
       return
     }
+    if (close) {
+      try {
+        void Promise.resolve(transport.close()).then(() => {
+          if (this.active !== active) return
+          this.disposeTransport(transport, false)
+          this.rejectAll(new EmbeddingInferenceBrokerError('PROCESS_CLOSED'))
+        }).catch((error: unknown) => {
+          this.managementError = error instanceof Error ? error.message : String(error)
+        })
+      } catch (error) {
+        this.managementError = error instanceof Error ? error.message : String(error)
+      }
+      return
+    }
     this.active = undefined
     active.removeMessageListener()
     active.removeCloseListener()
-    if (close) {
-      try {
-        transport.close()
-      } catch {
-        // The process is already unavailable.
-      }
+  }
+
+  private rejectCallers(error: Error): void {
+    for (const pending of this.pending.values()) {
+      pending.callerError ??= error
+      clearTimeout(pending.timeout)
+      pending.rejectCaller(pending.callerError)
     }
   }