pi-long-task 0.3.11 → 0.3.13

package/src/worker_session.ts

@@ -1,5 +1,13 @@
+import path from "node:path";
+
 import { coverageGoalAction, coverageGoalVerification, parseCoverageGoal } from "./coverage_goal.ts";
-import { hasTaskResult, hasTaskResultStatus, isDoneStatus, parseReportedStatus } from "./result_writer.ts";
+import {
+  hasCompleteTaskResult,
+  hasTaskResult,
+  isDoneStatus,
+  parseCompleteTaskResult,
+  parseReportedStatus,
+} from "./result_writer.ts";
 import type { Task } from "./todo_parser.ts";
 
 export interface WorkerTaskPromptOptions {
@@ -243,7 +251,10 @@ export interface CreateWorkerSessionOptions {
   model?: unknown;
   modelName?: string;
   thinkingLevel?: string;
+  modelRuntime?: unknown;
+  /** @deprecated Compatibility injection for Pi SDK versions before 0.80.8. */
   authStorage?: unknown;
+  /** @deprecated Compatibility injection for Pi SDK versions before 0.80.8. */
   modelRegistry?: unknown;
   settingsManager?: unknown;
   resourceLoader?: unknown;
@@ -306,14 +317,72 @@ remaining:
 - <remaining item or "none">`;
 }
 
+interface WorkerPiSdkModelExports {
+  ModelRuntime?: {
+    create(options?: { authPath?: string; modelsPath?: string | null }): Promise<unknown>;
+  };
+  AuthStorage?: {
+    create(authPath?: string): unknown;
+  };
+  ModelRegistry?: {
+    create?(authStorage: unknown, modelsPath?: string): unknown;
+  };
+}
+
+export interface WorkerModelContext {
+  resolver: unknown;
+  sessionOptions: Record<string, unknown>;
+  api: "modelRuntime" | "legacy";
+}
+
+export async function createWorkerModelContext(
+  sdk: WorkerPiSdkModelExports,
+  options: Pick<CreateWorkerSessionOptions, "modelRuntime" | "authStorage" | "modelRegistry">,
+  agentDir: string,
+): Promise<WorkerModelContext> {
+  if (options.modelRuntime) {
+    return {
+      resolver: options.modelRuntime,
+      sessionOptions: { modelRuntime: options.modelRuntime },
+      api: "modelRuntime",
+    };
+  }
+
+  if (sdk.ModelRuntime?.create) {
+    const modelRuntime = await sdk.ModelRuntime.create({
+      authPath: path.join(agentDir, "auth.json"),
+      modelsPath: path.join(agentDir, "models.json"),
+    });
+    return {
+      resolver: modelRuntime,
+      sessionOptions: { modelRuntime },
+      api: "modelRuntime",
+    };
+  }
+
+  const authStorage = options.authStorage ?? sdk.AuthStorage?.create(path.join(agentDir, "auth.json"));
+  const modelRegistry =
+    options.modelRegistry ?? sdk.ModelRegistry?.create?.(authStorage, path.join(agentDir, "models.json"));
+  if (!authStorage || !modelRegistry) {
+    throw new Error(
+      "Pi Long Task cannot initialize worker model services: this Pi SDK exposes neither ModelRuntime.create() nor the legacy AuthStorage.create()/ModelRegistry.create() API.",
+    );
+  }
+
+  return {
+    resolver: modelRegistry,
+    sessionOptions: { authStorage, modelRegistry },
+    api: "legacy",
+  };
+}
+
 export async function createIsolatedWorkerSession(
   options: CreateWorkerSessionOptions,
 ): Promise<WorkerSessionFactoryResult> {
   const pi = await import("@earendil-works/pi-coding-agent");
   const cwd = options.cwd;
   const agentDir = options.agentDir ?? pi.getAgentDir();
-  const authStorage = options.authStorage ?? pi.AuthStorage.create();
-  const modelRegistry = options.modelRegistry ?? pi.ModelRegistry.create(authStorage as never);
+  const modelContext = await createWorkerModelContext(pi as unknown as WorkerPiSdkModelExports, options, agentDir);
   const settingsManager = options.settingsManager ?? pi.SettingsManager.create(cwd, agentDir);
 
   applyWorkerSettingsDefaults(settingsManager);
@@ -330,12 +399,12 @@ export async function createIsolatedWorkerSession(
   await resourceLoader.reload();
 
   const model =
-    options.model ?? (options.modelName ? await resolveWorkerModel(modelRegistry, options.modelName) : undefined);
+    options.model ??
+    (options.modelName ? await resolveWorkerModel(modelContext.resolver, options.modelName) : undefined);
   const createOptions: Record<string, unknown> = {
     cwd,
     agentDir,
-    authStorage,
-    modelRegistry,
+    ...modelContext.sessionOptions,
     settingsManager,
     resourceLoader,
     tools: [...(options.tools ?? DEFAULT_WORKER_TOOLS)],
@@ -373,6 +442,7 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
   let messageUsageCostTotal = 0;
   let hasMessageUsageCost = false;
   let sessionStatsCostTotal: number | undefined;
+  let resolvePromptWait: (() => void) | undefined;
 
   const prompt = buildTaskPrompt(options);
   const taskTimeoutSeconds = options.taskTimeoutSeconds ?? DEFAULT_TASK_TIMEOUT_SECONDS;
@@ -420,38 +490,86 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     timers.add(timer);
   };
 
-  const requestGracefulTaskResult = async (message: string, options: { shutdown?: boolean } = {}) => {
+  const requestGracefulTaskResult = (message: string, options: { shutdown?: boolean } = {}) => {
     if (!session || finished || aborted) {
       return;
     }
     if (options.shutdown) {
       shutdownRequested = true;
     }
-    if (session.isBashRunning && session.abortBash) {
-      session.abortBash();
-      compactionEvents.push("aborted running bash before graceful shutdown request");
-    }
-    if ((session.isStreaming || session.isBashRunning) && session.steer) {
-      await session.steer(message);
-    } else {
-      await session.prompt(message);
+    try {
+      if (session.isBashRunning && session.abortBash) {
+        session.abortBash();
+        compactionEvents.push("aborted running bash before graceful shutdown request");
+      }
+      const request =
+        (session.isStreaming || session.isBashRunning) && session.steer
+          ? session.steer(message)
+          : session.followUp
+            ? session.followUp(message)
+            : session.steer
+              ? session.steer(message)
+              : undefined;
+      if (!request) {
+        compactionEvents.push("graceful TASK_RESULT request skipped: session does not support steer/followUp");
+        return;
+      }
+      void request.catch((exc: unknown) => {
+        compactionEvents.push(`graceful TASK_RESULT request failed: ${errorMessage(exc)}`);
+      });
+    } catch (exc) {
+      compactionEvents.push(`graceful TASK_RESULT request failed: ${errorMessage(exc)}`);
     }
   };
 
-  const abortSession = async (reason: string) => {
+  const abortSession = (reason: string) => {
     if (!session || finished || aborted) {
       return;
     }
     aborted = true;
     shutdownRequested = true;
     error = error ?? reason;
-    await session.abort?.();
+    try {
+      const abortResult = session.abort?.();
+      void Promise.resolve(abortResult).catch((exc: unknown) => {
+        compactionEvents.push(`session abort failed: ${errorMessage(exc)}`);
+      });
+    } catch (exc) {
+      compactionEvents.push(`session abort failed: ${errorMessage(exc)}`);
+    }
+    resolvePromptWait?.();
+    resolvePromptWait = undefined;
   };
 
-  const abortListener = () => {
-    void abortSession("worker session aborted by outer signal").catch((exc: unknown) => {
-      compactionEvents.push(`abort by outer signal failed: ${errorMessage(exc)}`);
+  const waitForPrompt = async (text: string): Promise<void> => {
+    if (!session) {
+      return;
+    }
+    let settled = false;
+    const completed = new Promise<void>((resolve) => {
+      resolvePromptWait = resolve;
     });
+    void session.prompt(text).then(
+      () => {
+        settled = true;
+        resolvePromptWait?.();
+        resolvePromptWait = undefined;
+      },
+      (exc: unknown) => {
+        settled = true;
+        error = error ?? errorMessage(exc);
+        resolvePromptWait?.();
+        resolvePromptWait = undefined;
+      },
+    );
+    await completed;
+    if (!settled && !aborted) {
+      error = error ?? "worker prompt ended without settling";
+    }
+  };
+
+  const abortListener = () => {
+    abortSession("worker session aborted by outer signal");
   };
 
   try {
@@ -549,12 +667,12 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     options.abortSignal?.addEventListener("abort", abortListener, { once: true });
 
     if (taskTimeoutSeconds > 0) {
-      schedule(async () => {
-        if (finished || hasTaskResultStatus(assistantText)) {
+      schedule(() => {
+        if (finished || hasCompleteTaskResult(assistantText)) {
           return;
         }
         timedOut = true;
-        await requestGracefulTaskResult(buildTimeLimitMessage(taskTimeoutSeconds), { shutdown: true });
+        requestGracefulTaskResult(buildTimeLimitMessage(taskTimeoutSeconds), { shutdown: true });
         if (gracefulShutdownSeconds > 0) {
           schedule(
             () => abortSession(`task exceeded ${taskTimeoutSeconds.toFixed(0)}s timeout`),
@@ -564,12 +682,14 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
       }, taskTimeoutSeconds * 1000);
     }
 
-    await session.prompt(prompt);
+    await waitForPrompt(prompt);
     assistantText = latestAssistantText(session, assistantText);
 
-    if (!hasTaskResultStatus(assistantText) && !aborted && !options.abortSignal?.aborted) {
-      contextObservations.push("missing TASK_RESULT status after initial prompt; requested required block once");
-      await requestGracefulTaskResult(buildMissingTaskResultMessage());
+    if (!hasCompleteTaskResult(assistantText) && !error && !aborted && !timedOut && !options.abortSignal?.aborted) {
+      contextObservations.push(
+        "missing TASK_RESULT status after initial prompt, or required fields were incomplete; requested required block once",
+      );
+      await waitForPrompt(buildMissingTaskResultMessage());
       assistantText = latestAssistantText(session, assistantText);
     }
   } catch (exc) {
@@ -584,7 +704,11 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
       sessionFile = session.sessionFile ?? sessionFile;
       sessionId = session.sessionId ?? sessionId;
       sessionStatsCostTotal = await workerUsageCostFromSessionStats(session);
-      session.dispose?.();
+      try {
+        await Promise.resolve(session.dispose?.());
+      } catch (exc) {
+        compactionEvents.push(`session dispose failed: ${errorMessage(exc)}`);
+      }
     }
   }
 
@@ -592,7 +716,8 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     assistantText = buildLongTaskFailureTaskResult(error ?? (timedOut ? "task timed out" : "worker session aborted"));
   }
 
-  const reportedStatus = parseReportedStatus(assistantText);
+  const parsedResult = parseCompleteTaskResult(assistantText);
+  const reportedStatus = parsedResult?.status ?? parseReportedStatus(assistantText);
   const capturedWorkerCost = selectWorkerCostTotal({
     messageCostTotal: hasMessageUsageCost ? messageUsageCostTotal : undefined,
     statsCostTotal: sessionStatsCostTotal,
@@ -603,7 +728,7 @@ export async function runWorkerTask(options: RunWorkerTaskOptions): Promise<Sess
     startedAt,
     endedAt: now().toISOString(),
     reportedStatus,
-    done: isDoneStatus(reportedStatus),
+    done: Boolean(parsedResult && isDoneStatus(reportedStatus) && !error && !aborted && !timedOut),
     assistantText,
     sessionFile,
     sessionId,
@@ -672,24 +797,18 @@ function applyWorkerSettingsDefaults(settingsManager: unknown): void {
   }
 }
 
-async function resolveWorkerModel(modelRegistry: unknown, modelName: string): Promise<unknown> {
+async function resolveWorkerModel(modelResolver: unknown, modelName: string): Promise<unknown> {
   const [provider, ...modelIdParts] = modelName.split("/");
   const modelId = modelIdParts.join("/");
-  if (provider && modelId && isRecord(modelRegistry) && typeof modelRegistry.find === "function") {
-    const registryModel = modelRegistry.find(provider, modelId);
-    if (registryModel) {
-      return registryModel;
-    }
+  if (!provider || !modelId || !isRecord(modelResolver)) {
+    return undefined;
   }
 
-  try {
-    const ai = await import("@earendil-works/pi-ai");
-    if (typeof ai.getModel === "function" && provider && modelId) {
-      const getModel = ai.getModel as (providerName: string, modelId: string) => unknown;
-      return getModel(provider, modelId);
-    }
-  } catch {
-    // Optional peer resolution can fail in tests that inject a session factory.
+  if (typeof modelResolver.getModel === "function") {
+    return modelResolver.getModel(provider, modelId);
+  }
+  if (typeof modelResolver.find === "function") {
+    return modelResolver.find(provider, modelId);
   }
   return undefined;
 }