@ashx-j/lunr 0.2.26 → 0.2.28

package/dist/builtin-extensions/pi-subagents/src/runs/background/subagent-runner.js

@@ -34,7 +34,7 @@ import { attachPostExitStdioGuard, trySignalChild } from "../../shared/post-exit
 import { detectSubagentError, extractTextFromContent, extractToolArgsPreview, getFinalOutput, readStatus } from "../../shared/utils.js";
 import { evaluateCompletionMutationGuard } from "../shared/completion-guard.js";
 import { createMutatingFailureState, didMutatingToolFail, isMutatingTool, nextLongRunningTrigger, recordMutatingFailure, resetMutatingFailureState, resolveCurrentPath, shouldEscalateMutatingFailures, summarizeRecentMutatingFailures, } from "../shared/long-running-guard.js";
-import { parseSessionTokens } from "../../shared/session-tokens.js";
+import { parseSessionRequests, parseSessionUsage, sessionUsageDelta, tokenUsageFromUsage } from "../../shared/session-tokens.js";
 import { cleanupWorktrees, createWorktrees, diffWorktrees, findWorktreeTaskCwdConflict, formatWorktreeDiffSummary, formatWorktreeTaskCwdConflict, } from "../shared/worktree.js";
 import { resolveEffectiveThinking } from "../../shared/model-info.js";
 import { writeInitialProgressFile } from "../../shared/settings.js";
@@ -127,7 +127,7 @@ function findLatestSessionFile(sessionDir) {
             .map((f) => path.join(sessionDir, f));
         if (files.length === 0)
             return null;
-        files.sort((a, b) => fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs);
+        files.sort((a, b) => fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs || b.localeCompare(a));
         return files[0] ?? null;
     }
     catch {
@@ -136,19 +136,18 @@ function findLatestSessionFile(sessionDir) {
     }
 }
 function emptyUsage() {
-    return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 };
+    return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0, unknownRequests: 0, partialRequests: 0, unknownCosts: 0 };
 }
 function tokenUsageFromAttempts(attempts) {
-    if (!attempts || attempts.length === 0)
+    if (!attempts?.some((attempt) => attempt.usage))
         return null;
-    let input = 0;
-    let output = 0;
-    for (const attempt of attempts) {
-        input += attempt.usage?.input ?? 0;
-        output += attempt.usage?.output ?? 0;
-    }
-    const total = input + output;
-    return total > 0 ? { input, output, total } : null;
+    const usage = emptyUsage();
+    for (const attempt of attempts)
+        if (attempt.usage) {
+            for (const key of Object.keys(usage))
+                usage[key] += attempt.usage[key] ?? 0;
+        }
+    return tokenUsageFromUsage(usage);
 }
 function costSummaryFromAttempts(attempts) {
     if (!attempts || attempts.length === 0)
@@ -157,7 +156,7 @@ function costSummaryFromAttempts(attempts) {
     let outputTokens = 0;
     let costUsd = 0;
     for (const attempt of attempts) {
-        inputTokens += attempt.usage?.input ?? 0;
+        inputTokens += (attempt.usage?.input ?? 0) + (attempt.usage?.cacheRead ?? 0) + (attempt.usage?.cacheWrite ?? 0);
         outputTokens += attempt.usage?.output ?? 0;
         costUsd += attempt.usage?.cost ?? 0;
     }
@@ -345,6 +344,9 @@ function runPiStreaming(args, cwd, outputFile, env, piPackageRoot, piArgv1, maxS
                     usage.cacheRead += eventUsage.cacheRead ?? 0;
                     usage.cacheWrite += eventUsage.cacheWrite ?? 0;
                     usage.cost += eventUsage.cost?.total ?? 0;
+                    usage.unknownRequests += Number(eventUsage.measurement === "unknown");
+                    usage.partialRequests += Number(eventUsage.measurement === "partial");
+                    usage.unknownCosts += Number(eventUsage.costSource === "unknown");
                 }
                 if (isTerminalAssistantStop(event.message)) {
                     if (!event.message.errorMessage && extractTextFromContent(event.message.content).trim())
@@ -771,6 +773,8 @@ export async function runSingleStep(step, ctx) {
                 model: imported.model,
                 attemptedModels: imported.attemptedModels,
                 modelAttempts: imported.modelAttempts,
+                receiptUsage: imported.usage,
+                usageRequests: imported.usageRequests,
                 totalCost: imported.totalCost,
                 structuredOutput: timedOut || stopped ? undefined : imported.structuredOutput,
                 structuredOutputPath: timedOut || stopped ? undefined : imported.structuredOutputPath,
@@ -827,6 +831,7 @@ export async function runSingleStep(step, ctx) {
             : [undefined];
     const attemptedModels = [];
     const modelAttempts = [];
+    const priorRequestIds = new Set(parseSessionRequests(step.sessionFile ?? sessionDir).map((request) => request.id));
     const attemptNotes = [];
     const eventsPath = path.join(path.dirname(ctx.outputFile), "events.jsonl");
     let finalResult;
@@ -903,7 +908,11 @@ export async function runSingleStep(step, ctx) {
             waitToolEnabled: step.waitToolEnabled,
             parentPermissionMode,
         });
+        const usageBeforeAttempt = parseSessionUsage(step.sessionFile ?? sessionDir);
         const run = await runPiStreaming(args, step.cwd ?? ctx.cwd, ctx.outputFile, env, ctx.piPackageRoot, ctx.piArgv1, step.maxSubagentDepth, { eventsPath, runId: ctx.id, stepIndex: ctx.flatIndex, agent: step.agent }, ctx.registerInterrupt, ctx.onChildEvent, transcriptWriter, ctx.registerTimeout, ctx.timeoutMessage, ctx.registerStop, ctx.stopMessage, ctx.registerTurnBudgetAbort, ctx.onWriterProcess);
+        const persistedUsage = parseSessionUsage(step.sessionFile ?? sessionDir);
+        if (persistedUsage)
+            run.usage = sessionUsageDelta(persistedUsage, usageBeforeAttempt);
         if (run.turnBudget)
             turnBudget = run.turnBudget;
         else if (ctx.turnBudget) {
@@ -1093,6 +1102,7 @@ export async function runSingleStep(step, ctx) {
             }, null, 2), "utf-8");
         }
     }
+    const usageRequests = parseSessionRequests(step.sessionFile ?? sessionDir).filter((request) => Boolean(step.sessionFile) || !priorRequestIds.has(request.id));
     return {
         childId: step.childId,
         description: step.description,
@@ -1102,7 +1112,9 @@ export async function runSingleStep(step, ctx) {
         exitCode: effectiveFinalExitCode,
         error: effectiveFinalError,
         protocolError: finalResult?.protocolError,
-        sessionFile: step.sessionFile,
+        sessionFile: step.sessionFile ?? (sessionDir ? findLatestSessionFile(sessionDir) ?? undefined : undefined),
+        receiptUsage: step.sessionFile ? parseSessionUsage(step.sessionFile) ?? undefined : undefined,
+        usageRequests: usageRequests.length > 0 ? usageRequests : undefined,
         intercomTarget: ctx.childIntercomTarget,
         model: finalResult?.model,
         attemptedModels: attemptedModels.length > 0 ? attemptedModels : undefined,
@@ -1251,6 +1263,7 @@ async function runSubagent(config, onWriterProcess) {
     const stopMessage = "Subagent stopped by user.";
     const timeoutAbortController = new AbortController();
     const stopAbortController = new AbortController();
+    const launchInterruptController = new AbortController();
     let previousCumulativeTokens = { input: 0, output: 0, total: 0 };
     let latestSessionFile;
     const flatSteps = flattenSteps(steps);
@@ -2051,7 +2064,7 @@ async function runSubagent(config, onWriterProcess) {
             step.turnCount = (step.turnCount ?? 0) + 1;
             const usage = event.message.usage;
             if (usage) {
-                const input = usage.input ?? usage.inputTokens ?? 0;
+                const input = (usage.input ?? usage.inputTokens ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
                 const output = usage.output ?? usage.outputTokens ?? 0;
                 const previousInput = step.tokens?.input ?? 0;
                 const previousOutput = step.tokens?.output ?? 0;
@@ -2139,11 +2152,18 @@ async function runSubagent(config, onWriterProcess) {
         }, 1000);
         activityTimer.unref?.();
     }
+    const launchCancellation = {
+        signal: combinedAbortSignal([timeoutAbortController.signal, stopAbortController.signal, launchInterruptController.signal]),
+        onAbort: (task) => stopped
+            ? stoppedStepResult(task.agent)
+            : timedOut ? timedOutStepResult(task.agent) : pausedStepResult(task.agent),
+    };
     const interruptRunner = () => {
         consumeInterruptRequest(asyncDir);
         if (interrupted || statusPayload.state !== "running")
             return;
         interrupted = true;
+        launchInterruptController.abort();
         const now = Date.now();
         statusPayload.state = "paused";
         currentActivityState = undefined;
@@ -2619,11 +2639,15 @@ async function runSubagent(config, onWriterProcess) {
                 if (singleResult.exitCode !== 0 && failFast)
                     aborted = true;
                 return stopped || childStopped ? { ...singleResult, output: stopMessage, error: stopMessage, exitCode: 1, interrupted: false, timedOut: false, stopped: true, skipped: false } : timedOut ? { ...singleResult, output: timeoutMessage ?? "Subagent timed out.", error: timeoutMessage ?? "Subagent timed out.", exitCode: 1, interrupted: false, timedOut: true, skipped: false } : { ...singleResult, skipped: false };
-            }, globalSemaphore);
+            }, globalSemaphore, launchCancellation);
             flatIndex += dynamicSteps.length;
             for (const pr of parallelResults) {
                 results.push({
                     agent: pr.agent,
+                    childId: pr.childId,
+                    description: pr.description,
+                    receiptUsage: pr.receiptUsage,
+                    usageRequests: pr.usageRequests,
                     output: pr.output,
                     error: pr.error,
                     protocolError: pr.protocolError,
@@ -2843,7 +2867,7 @@ async function runSubagent(config, onWriterProcess) {
                         type: "subagent.step.started", ts: taskStartTime, runId: id, stepIndex: fi, agent: task.agent,
                     }));
                     const taskSessionDir = config.sessionDir
-                        ? path.join(config.sessionDir, `parallel-${taskIdx}`)
+                        ? path.join(config.sessionDir, `parallel-${stepIndex}-${taskIdx}`)
                         : undefined;
                     const { taskForRun, taskCwd } = prepareParallelTaskRun(task, cwd, worktreeSetup, taskIdx);
                     flushPendingStepSteers(fi);
@@ -2942,14 +2966,11 @@ async function runSubagent(config, onWriterProcess) {
                     if (singleResult.exitCode !== 0 && failFast)
                         aborted = true;
                     return stopped || childStopped ? { ...singleResult, output: stopMessage, error: stopMessage, exitCode: 1, interrupted: false, timedOut: false, stopped: true, skipped: false } : timedOut ? { ...singleResult, output: timeoutMessage ?? "Subagent timed out.", error: timeoutMessage ?? "Subagent timed out.", exitCode: 1, interrupted: false, timedOut: true, skipped: false } : { ...singleResult, skipped: false };
-                }, globalSemaphore);
+                }, globalSemaphore, launchCancellation);
                 flatIndex += group.parallel.length;
                 for (let t = 0; t < group.parallel.length; t++) {
                     const fi = groupStartFlatIndex + t;
-                    const sessionTokens = config.sessionDir
-                        ? parseSessionTokens(path.join(config.sessionDir, `parallel-${t}`))
-                        : null;
-                    const taskTokens = sessionTokens ?? tokenUsageFromAttempts(parallelResults[t]?.modelAttempts);
+                    const taskTokens = tokenUsageFromAttempts(parallelResults[t]?.modelAttempts);
                     if (!taskTokens)
                         continue;
                     statusPayload.steps[fi].tokens = taskTokens;
@@ -2965,6 +2986,10 @@ async function runSubagent(config, onWriterProcess) {
                 for (const pr of parallelResults) {
                     results.push({
                         agent: pr.agent,
+                        childId: pr.childId,
+                        description: pr.description,
+                        receiptUsage: pr.receiptUsage,
+                        usageRequests: pr.usageRequests,
                         output: pr.output,
                         error: pr.error,
                         protocolError: pr.protocolError,
@@ -3091,6 +3116,10 @@ async function runSubagent(config, onWriterProcess) {
             const childStopped = singleResult.stopped === true;
             results.push({
                 agent: singleResult.agent,
+                childId: singleResult.childId,
+                description: singleResult.description,
+                receiptUsage: singleResult.receiptUsage,
+                usageRequests: singleResult.usageRequests,
                 output: stopped || childStopped ? stopMessage : timedOut ? (timeoutMessage ?? "Subagent timed out.") : singleResult.output,
                 error: stopped || childStopped ? stopMessage : timedOut ? (timeoutMessage ?? "Subagent timed out.") : singleResult.error,
                 protocolError: singleResult.protocolError,
@@ -3127,26 +3156,15 @@ async function runSubagent(config, onWriterProcess) {
                 }, stepIndex);
             }
             statusPayload.outputs = outputs;
-            const cumulativeTokens = config.sessionDir ? parseSessionTokens(config.sessionDir) : null;
-            let stepTokens = cumulativeTokens
-                ? {
-                    input: cumulativeTokens.input - previousCumulativeTokens.input,
-                    output: cumulativeTokens.output - previousCumulativeTokens.output,
-                    total: cumulativeTokens.total - previousCumulativeTokens.total,
-                }
-                : null;
-            if (cumulativeTokens) {
-                previousCumulativeTokens = cumulativeTokens;
-            }
-            else {
-                stepTokens = tokenUsageFromAttempts(singleResult.modelAttempts);
-                if (stepTokens) {
-                    previousCumulativeTokens = {
-                        input: previousCumulativeTokens.input + stepTokens.input,
-                        output: previousCumulativeTokens.output + stepTokens.output,
-                        total: previousCumulativeTokens.total + stepTokens.total,
-                    };
-                }
+            const stepTokens = tokenUsageFromAttempts(singleResult.modelAttempts);
+            if (stepTokens) {
+                previousCumulativeTokens = {
+                    input: previousCumulativeTokens.input + stepTokens.input,
+                    output: previousCumulativeTokens.output + stepTokens.output,
+                    total: previousCumulativeTokens.total + stepTokens.total,
+                    cacheRead: (previousCumulativeTokens.cacheRead ?? 0) + (stepTokens.cacheRead ?? 0),
+                    cacheWrite: (previousCumulativeTokens.cacheWrite ?? 0) + (stepTokens.cacheWrite ?? 0),
+                };
             }
             const stepEndTime = Date.now();
             const childInterrupted = singleResult.interrupted === true;
@@ -3236,6 +3254,17 @@ async function runSubagent(config, onWriterProcess) {
         outputTokens: sum.outputTokens + (result.totalCost?.outputTokens ?? 0),
         costUsd: sum.costUsd + (result.totalCost?.costUsd ?? 0),
     }), { inputTokens: 0, outputTokens: 0, costUsd: 0 });
+    statusPayload.totalTokens = results.reduce((total, result) => {
+        const usage = tokenUsageFromAttempts(result.modelAttempts);
+        if (usage) {
+            total.input += usage.input;
+            total.output += usage.output;
+            total.total += usage.total;
+            total.cacheRead += usage.cacheRead ?? 0;
+            total.cacheWrite += usage.cacheWrite ?? 0;
+        }
+        return total;
+    }, { input: 0, output: 0, total: 0, cacheRead: 0, cacheWrite: 0 });
     const finalTotalCost = totalCost.inputTokens > 0 || totalCost.outputTokens > 0 || totalCost.costUsd > 0 ? totalCost : undefined;
     const finalFlatAgents = statusPayload.steps.map((step) => step.agent);
     const agentName = finalFlatAgents.length === 1
@@ -3399,6 +3428,13 @@ async function runSubagent(config, onWriterProcess) {
                 attemptedModels: r.attemptedModels,
                 modelAttempts: r.modelAttempts,
                 totalCost: r.totalCost,
+                usageRequests: r.usageRequests,
+                usage: r.receiptUsage ?? r.modelAttempts?.reduce((total, attempt) => {
+                    if (attempt.usage)
+                        for (const key of Object.keys(total))
+                            total[key] += attempt.usage[key] ?? 0;
+                    return total;
+                }, emptyUsage()),
                 artifactPaths: r.artifactPaths,
                 truncated: r.truncated,
                 transcriptPath: r.transcriptPath,