feat(06-02): wire token usage tracking, session accumulation, and rotation into worker pipeline

- Add inputTokens/outputTokens fields to worker run-finished messages and workerMessageSchema
- Wire sessionService into WorkerServiceDependencies and ServiceGraph lazy slot
- Add token accumulation and rotation evaluation in run-finished handler
- Add 7 wiring integration tests verifying accumulation, rotation, and guard paths

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
ZacharyZhang-NYandClaude Opus 4.6 committed 2026-03-26 23:03:23 -04:00
1 parent e54c632fd6
commit 5271f7ee5d
4 files changed
+228

No files matched your search

+8
View File
@@ -65,6 +65,8 @@ type ChildMessage =
sessionDisplayId: string | null;
costUsd: number | null;
unattributedCost: boolean;
inputTokens: number | null;
outputTokens: number | null;
};
const parentPort = process.parentPort;
@@ -236,6 +238,8 @@ async function startRun(job: RunJob) {
sessionDisplayId: result.sessionDisplayId ?? null,
costUsd: typeof result.costUsd === "number" ? result.costUsd : null,
unattributedCost: !result.costUsd,
inputTokens: result.usage?.inputTokens ?? null,
outputTokens: result.usage?.outputTokens ?? null,
});
} catch (error) {
post({
@@ -251,6 +255,8 @@ async function startRun(job: RunJob) {
sessionDisplayId: null,
costUsd: null,
unattributedCost: true,
inputTokens: null,
outputTokens: null,
});
} finally {
cancelledRuns.delete(job.runId);
@@ -279,6 +285,8 @@ function cancelRun(runId: string) {
sessionDisplayId: null,
costUsd: null,
unattributedCost: true,
inputTokens: null,
outputTokens: null,
});
return;
}
+15
View File
@@ -14,6 +14,7 @@ import { createWorkerService } from "./worker-service";
import { createAppLifecycleService } from "./app-lifecycle-service";
import { createSkillService } from "./skill-service";
import { createBudgetEnforcementService } from "./budget-enforcement-service";
import { createSessionManagementService } from "./session-management-service";
import type { DesktopEvent, RunRecord } from "@shared/types";
import type { AutomationEventContext } from "./automation-service";
import type { NotificationOptions } from "./notification-service";
@@ -72,6 +73,7 @@ export function createServiceGraph({
let appLifecycleService: ReturnType<typeof createAppLifecycleService> | null = null;
let skillService: ReturnType<typeof createSkillService> | null = null;
let budgetEnforcementService: ReturnType<typeof createBudgetEnforcementService> | null = null;
let sessionManagementService: ReturnType<typeof createSessionManagementService> | null = null;
function getOrCreateSkillService() {
if (!skillService) {
@@ -96,6 +98,16 @@ export function createServiceGraph({
return budgetEnforcementService;
}
function getOrCreateSessionManagementService() {
if (!sessionManagementService) {
sessionManagementService = createSessionManagementService({
db,
logger,
});
}
return sessionManagementService;
}
const wakeAgentIfPossible =(agentId: string | null | undefined, companyId: string, trigger: string) => {
if (!orchestrationService) {
throw new Error("Orchestration service is not initialized.");
@@ -192,6 +204,7 @@ export function createServiceGraph({
processRunCompletion: (run, event) => processRunCompletion(run, event),
notify,
budgetService: getOrCreateBudgetEnforcementService(),
sessionService: getOrCreateSessionManagementService(),
});
orchestrationService = createOrchestrationService({
@@ -281,6 +294,7 @@ export function createServiceGraph({
appLifecycleService = null;
skillService = null;
budgetEnforcementService = null;
sessionManagementService = null;
db.close();
// Force exit after grace period — covers all platforms (macOS Dock, Windows tray, Linux)
// Worker child processes (claude/codex CLI) may keep the event loop alive indefinitely
@@ -302,5 +316,6 @@ export function createServiceGraph({
getAppLifecycleService: () => appLifecycleService,
getOrCreateSkillService,
getOrCreateBudgetEnforcementService,
getOrCreateSessionManagementService,
};
}
+30
View File
@@ -5,6 +5,7 @@ import type { AppDatabase } from "./database";
import type { RunRecord } from "@shared/types";
import type { NotificationOptions } from "./notification-service";
import type { createBudgetEnforcementService } from "./budget-enforcement-service";
import type { createSessionManagementService } from "./session-management-service";
const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
export const workerMessageSchema = z.union([
@@ -33,6 +34,8 @@ export const workerMessageSchema = z.union([
sessionDisplayId: z.string().nullable(),
costUsd: z.number().nullable(),
unattributedCost: z.boolean(),
inputTokens: z.number().nullable(),
outputTokens: z.number().nullable(),
}),
]);
type WorkerUiEvent = Extract<z.infer<typeof workerMessageSchema>, { type: "run-log" | "run-status" }>;
@@ -65,6 +68,7 @@ interface WorkerServiceDependencies {
processRunCompletion: (run: RunRecord, event: z.infer<typeof workerMessageSchema> & { type: "run-finished" }) => void;
notify: (options: NotificationOptions) => void;
budgetService?: ReturnType<typeof createBudgetEnforcementService>;
sessionService?: ReturnType<typeof createSessionManagementService>;
}
export function createWorkerService({
@@ -79,6 +83,7 @@ export function createWorkerService({
processRunCompletion,
notify,
budgetService,
sessionService,
}: WorkerServiceDependencies) {
const MAX_RETRY_ATTEMPTS = 2;
const RETRY_COOLDOWN_MS = 30_000; // 30 seconds
@@ -290,6 +295,31 @@ export function createWorkerService({
entityId: run.id,
detail: event.summary ?? event.errorMessage ?? "",
});
// SESS-01/SESS-02/SESS-03/SESS-05: Token accumulation and rotation evaluation
if (task && sessionService) {
const activeSession = db.getActiveTaskSession(task.id);
if (activeSession) {
// D-08/D-05: Accumulate token usage for ALL adapters (including exempt ones)
sessionService.accumulateTokenUsage(
activeSession.id,
event.inputTokens ?? 0,
event.outputTokens ?? 0,
event.runId,
);
// D-03: Evaluate rotation after accumulation (after each run completes)
const rotationResult = sessionService.evaluateRotation(activeSession, run.connectorId);
if (rotationResult.shouldRotate) {
// Re-read to get updated counts after accumulation
const updatedSession = db.getActiveTaskSession(task.id);
if (updatedSession) {
sessionService.rotateSession(updatedSession);
logger.info(`[session] Rotated session for task ${task.id}: ${rotationResult.reason}`);
}
}
}
}
publishDomainChanged();
emitEvent({
type: "run-status",