@ct-agents/worker 0.1.4 → 0.1.6

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ct-agents/worker",
-  "version": "0.1.4",
+  "version": "0.1.6",
   "type": "module",
   "files": [
     "src"
@@ -12,10 +12,10 @@
   },
   "dependencies": {
     "zod": "4.4.3",
-    "@ct-agents/memory": "0.1.4",
-    "@ct-agents/prompts": "0.1.4",
-    "@ct-agents/protocol": "0.1.4",
-    "@ct-agents/tools": "0.1.4"
+    "@ct-agents/memory": "0.1.6",
+    "@ct-agents/prompts": "0.1.6",
+    "@ct-agents/protocol": "0.1.6",
+    "@ct-agents/tools": "0.1.6"
   },
   "optionalDependencies": {
     "dockerode": "4.0.9"

package/src/index.ts

@@ -21,12 +21,18 @@ import type {
   ToolWorkItem,
   WorkQueue,
   WorkStats,
+  CompleteWorkItemScope,
+  RenewWorkItemLeaseInput,
 } from '@ct-agents/protocol';
 import type { MemoryResource, SkillResource } from '@ct-agents/protocol';
 import { createLoadMemoryToolHandler, createUpdateMemoryToolHandler } from '@ct-agents/memory';
 import { createLoadSkillToolHandler } from '@ct-agents/prompts';
 import { createSandboxTools } from '@ct-agents/tools';
-import { createAskUserToolHandler, createWebSearchToolHandler } from '@ct-agents/tools/builtin-tools';
+import {
+  createAskUserToolHandler,
+  createTodoWriteToolHandler,
+  createWebSearchToolHandler,
+} from '@ct-agents/tools/builtin-tools';
 import { z } from 'zod';
 import {
   createPlatformDatabaseProxyResource,
@@ -89,7 +95,11 @@ export type EnvironmentWorkerInput = {
   resourceImpls?: EnvironmentWorkerResourceImplementations;
   platformResourceProxy?: EnvironmentWorkerPlatformResourceProxy;
   abortSignal?: AbortSignal;
-  maxBatchSize?: number;
+  /** 单进程持续允许的最大 active Tool work 数。 */
+  maxConcurrency?: number;
+  /** active Tool lease 的续租周期，必须小于 loop 的 reclaimOlderThanMs。 */
+  leaseRenewIntervalMs?: number;
+  onError?: (error: unknown) => void | Promise<void>;
   harness?: {
     id: string;
     version: number;
@@ -136,6 +146,7 @@ export type RunContextResolverWorkerLoopInput = ContextResolverWorkerInput & {
 export const DEFAULT_RECLAIM_OLDER_THAN_MS = 300_000;
 export const DEFAULT_RECLAIM_INTERVAL_MS = 60_000;
 export const DEFAULT_CLAIM_BLOCK_MS = 15_000;
+export const DEFAULT_WORK_LEASE_RENEW_INTERVAL_MS = 60_000;
 
 const PLATFORM_PROXY_RESOURCE_SLOTS = new Set(['database', 'skills', 'memory']);
 
@@ -146,6 +157,7 @@ export function createPlatformBuiltinToolHandlers(): ToolHandler[] {
     createLoadMemoryToolHandler(),
     createUpdateMemoryToolHandler(),
     createWebSearchToolHandler(),
+    createTodoWriteToolHandler(),
     ...createSandboxTools(),
   ];
 }
@@ -161,6 +173,13 @@ function normalizeNonNegativeInteger(value: number | undefined, fallback: number
   return Math.max(0, Math.trunc(value));
 }
 
+function normalizePositiveInteger(value: number | undefined, fallback: number): number {
+  if (value === undefined || !Number.isFinite(value)) {
+    return fallback;
+  }
+  return Math.max(1, Math.trunc(value));
+}
+
 class PlatformBuiltinToolRegistry implements ToolRegistry {
   private readonly builtinHandlers = new Map<string, ToolHandler<unknown, unknown>>();
 
@@ -212,45 +231,97 @@ export class EnvironmentWorker {
   }
 
   public async drainOnce(options: DrainOnceOptions = {}): Promise<number> {
+    const items = await this.claim(
+      normalizePositiveInteger(this.input.maxConcurrency, 1),
+      options,
+    );
+
+    await Promise.all(items.map((item) => this.handleItem(item)));
+    return items.length;
+  }
+
+  public async claim(max: number, options: DrainOnceOptions = {}): Promise<ToolWorkItem[]> {
     const executionMode = this.input.executionMode ?? 'self_hosted';
-    const items = await this.input.workQueue.claim({
+    return this.input.workQueue.claim({
       appId: this.input.appId,
       environmentId: this.input.environmentId,
       executionMode,
-      max: this.input.maxBatchSize ?? 1,
+      max: normalizePositiveInteger(max, 1),
       blockMs: options.blockMs,
     });
-
-    await Promise.all(items.map((item) => this.handleItem(item)));
-    return items.length;
   }
 
-  private async handleItem(item: ToolWorkItem): Promise<void> {
-    let result: ToolExecutionResult;
-    const disposers: Array<() => Promise<void>> = [];
+  public async handleItem(item: ToolWorkItem): Promise<void> {
+    if (!item.leaseId) {
+      throw new Error(`claimed work item 缺少 leaseId：${item.workId}`);
+    }
+    const renewController = new AbortController();
+    const renewal = this.renewLeaseUntilStopped(item, renewController.signal);
     try {
-      const handler = this.input.toolRegistry.getTool(item.toolName);
-      const execution = await this.executeHandler(handler, item, disposers);
-      result = toToolExecutionResult(item.workId, handler, execution);
-    } catch (error) {
-      const accepted = await this.input.workQueue.complete({
-        workId: item.workId,
-        status: 'failed',
-        error: toWorkerErrorPayload(error),
-      }, this.buildCompleteScope(item));
+      let result: ToolExecutionResult;
+      const disposers: Array<() => Promise<void>> = [];
+      try {
+        const handler = this.input.toolRegistry.getTool(item.toolName);
+        const execution = await this.executeHandler(handler, item, disposers);
+        result = toToolExecutionResult(item.workId, handler, execution);
+      } catch (error) {
+        const accepted = await this.input.workQueue.complete({
+          workId: item.workId,
+          status: 'failed',
+          error: toWorkerErrorPayload(error),
+        }, this.buildCompleteScope(item));
+        if (!accepted) {
+          throw new Error(`work complete 未被接受：${item.workId}`);
+        }
+        return;
+      } finally {
+        await Promise.all(disposers.map(async (dispose) => {
+          await dispose().catch(() => undefined);
+        }));
+      }
+
+      const accepted = await this.input.workQueue.complete(result, this.buildCompleteScope(item));
       if (!accepted) {
         throw new Error(`work complete 未被接受：${item.workId}`);
       }
-      return;
     } finally {
-      await Promise.all(disposers.map(async (dispose) => {
-        await dispose().catch(() => undefined);
-      }));
+      renewController.abort();
+      await renewal;
     }
+  }
 
-    const accepted = await this.input.workQueue.complete(result, this.buildCompleteScope(item));
-    if (!accepted) {
-      throw new Error(`work complete 未被接受：${item.workId}`);
+  private async renewLeaseUntilStopped(item: ToolWorkItem, signal: AbortSignal): Promise<void> {
+    const leaseId = item.leaseId;
+    if (!leaseId) {
+      return;
+    }
+    const intervalMs = normalizePositiveInteger(
+      this.input.leaseRenewIntervalMs,
+      DEFAULT_WORK_LEASE_RENEW_INTERVAL_MS,
+    );
+    while (!signal.aborted) {
+      await delay(intervalMs, signal);
+      if (signal.aborted) {
+        return;
+      }
+      try {
+        const renewed = await this.input.workQueue.renew({
+          workId: item.workId,
+          leaseId,
+          appId: this.input.appId,
+          environmentId: this.input.environmentId,
+          executionMode: this.input.executionMode ?? item.executionMode ?? 'self_hosted',
+        });
+        if (!renewed) {
+          return;
+        }
+      } catch (error) {
+        try {
+          await this.input.onError?.(error);
+        } catch {
+          // 观测回调失败不能终止仍持有 lease 的业务 handler；下一周期继续 renew。
+        }
+      }
     }
   }
 
@@ -393,11 +464,15 @@ export class EnvironmentWorker {
   }
 
   private buildCompleteScope(item: ToolWorkItem) {
+    if (!item.leaseId) {
+      throw new Error(`claimed work item 缺少 leaseId：${item.workId}`);
+    }
     return {
       appId: this.input.appId,
       environmentId: this.input.environmentId,
       executionMode: this.input.executionMode ?? item.executionMode ?? 'self_hosted',
-    };
+      leaseId: item.leaseId,
+    } satisfies CompleteWorkItemScope;
   }
 }
 
@@ -498,7 +573,28 @@ export async function runEnvironmentWorkerLoop(input: RunEnvironmentWorkerLoopIn
   const reclaimOlderThanMs = input.reclaimOlderThanMs ?? DEFAULT_RECLAIM_OLDER_THAN_MS;
   const reclaimIntervalMs = normalizeNonNegativeInteger(input.reclaimIntervalMs, DEFAULT_RECLAIM_INTERVAL_MS);
   const blockMs = normalizeNonNegativeInteger(input.blockMs, DEFAULT_CLAIM_BLOCK_MS);
+  const maxConcurrency = normalizePositiveInteger(input.maxConcurrency, 1);
+  const leaseRenewIntervalMs = normalizePositiveInteger(
+    input.leaseRenewIntervalMs,
+    DEFAULT_WORK_LEASE_RENEW_INTERVAL_MS,
+  );
+  if (leaseRenewIntervalMs >= reclaimOlderThanMs) {
+    throw new Error('leaseRenewIntervalMs 必须小于 reclaimOlderThanMs');
+  }
   let nextReclaimAtMs = 0;
+  const active = new Set<Promise<void>>();
+
+  const startWork = (item: ToolWorkItem) => {
+    let tracked: Promise<void>;
+    tracked = worker.handleItem(item)
+      .catch(async (error) => {
+        await input.onError?.(error);
+      })
+      .finally(() => {
+        active.delete(tracked);
+      });
+    active.add(tracked);
+  };
 
   try {
     while (!input.abortSignal?.aborted) {
@@ -517,9 +613,22 @@ export async function runEnvironmentWorkerLoop(input: RunEnvironmentWorkerLoopIn
           await input.sandboxManager.sweepIdle(input.sandboxIdleTtlMs);
         }
 
-        const drained = await worker.drainOnce({ blockMs });
-        if (drained === 0) {
-          await delay(idleDelayMs, input.abortSignal);
+        const availableSlots = maxConcurrency - active.size;
+        if (availableSlots <= 0) {
+          await waitForActiveWorkOrDeadline(active, nextReclaimAtMs, input.abortSignal);
+          continue;
+        }
+
+        const items = await worker.claim(availableSlots, { blockMs });
+        for (const item of items) {
+          startWork(item);
+        }
+        if (items.length === 0) {
+          if (active.size > 0) {
+            await waitForActiveWorkOrDelay(active, idleDelayMs, input.abortSignal);
+          } else {
+            await delay(idleDelayMs, input.abortSignal);
+          }
         }
       } catch (error) {
         await input.onError?.(error);
@@ -527,6 +636,7 @@ export async function runEnvironmentWorkerLoop(input: RunEnvironmentWorkerLoopIn
       }
     }
   } finally {
+    await Promise.allSettled(active);
     await input.sandboxManager?.disposeAll();
   }
 }
@@ -621,10 +731,25 @@ export class HttpWorkQueueClient implements WorkQueue {
     return items.map(parseToolWorkItem);
   }
 
-  public async complete(result: ToolExecutionResult): Promise<boolean> {
+  public async renew(input: RenewWorkItemLeaseInput): Promise<boolean> {
+    this.assertScope(input.appId, input.environmentId);
+    const data = await this.requestJson('/work/renew', {
+      method: 'POST',
+      signal: this.input.signal,
+      body: JSON.stringify({ workId: input.workId, leaseId: input.leaseId }),
+    });
+    const record = expectRecord(data, 'renew 响应缺少 data');
+    return record.renewed === true;
+  }
+
+  public async complete(result: ToolExecutionResult, scope?: CompleteWorkItemScope): Promise<boolean> {
+    if (!scope?.leaseId) {
+      throw new Error(`work complete 缺少 leaseId：${result.workId}`);
+    }
+    this.assertScope(scope.appId, scope.environmentId);
     const data = await this.requestJson('/work/complete', {
       method: 'POST',
-      body: JSON.stringify(result),
+      body: JSON.stringify({ ...result, leaseId: scope.leaseId }),
     });
     const record = expectRecord(data, 'complete 响应缺少 data');
     if (record.accepted !== true) {
@@ -1096,10 +1221,35 @@ function delay(ms: number, signal?: AbortSignal): Promise<void> {
   }
 
   return new Promise((resolve) => {
-    const timer = setTimeout(resolve, ms);
-    signal?.addEventListener('abort', () => {
+    const finish = () => {
       clearTimeout(timer);
+      signal?.removeEventListener('abort', finish);
       resolve();
-    }, { once: true });
+    };
+    const timer = setTimeout(finish, ms);
+    signal?.addEventListener('abort', finish, { once: true });
   });
 }
+
+function waitForActiveWorkOrDeadline(
+  active: Set<Promise<void>>,
+  deadlineMs: number,
+  signal?: AbortSignal,
+): Promise<void> {
+  const waitMs = Math.max(1, deadlineMs - Date.now());
+  return Promise.race([
+    ...active,
+    delay(waitMs, signal),
+  ]).then(() => undefined);
+}
+
+function waitForActiveWorkOrDelay(
+  active: Set<Promise<void>>,
+  delayMs: number,
+  signal?: AbortSignal,
+): Promise<void> {
+  return Promise.race([
+    ...active,
+    delay(delayMs, signal),
+  ]).then(() => undefined);
+}