From 5271f7ee5d0a4aba16884a4c67cda67ac54bf467 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 26 Mar 2026 23:03:23 -0400 Subject: [PATCH] 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) --- src/main/orchestrator/worker.ts | 8 ++ src/main/service-graph.ts | 15 +++ src/main/worker-service.ts | 30 +++++ tests/unit/session-management.test.ts | 175 ++++++++++++++++++++++++++ 4 files changed, 228 insertions(+) diff --git a/src/main/orchestrator/worker.ts b/src/main/orchestrator/worker.ts index 649ed567..e20ff545 100644 --- a/src/main/orchestrator/worker.ts +++ b/src/main/orchestrator/worker.ts @@ -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; } diff --git a/src/main/service-graph.ts b/src/main/service-graph.ts index f0e43376..70e87bd9 100644 --- a/src/main/service-graph.ts +++ b/src/main/service-graph.ts @@ -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 | null = null; let skillService: ReturnType | null = null; let budgetEnforcementService: ReturnType | null = null; + let sessionManagementService: ReturnType | 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, }; } diff --git a/src/main/worker-service.ts b/src/main/worker-service.ts index e14caf85..e614091f 100644 --- a/src/main/worker-service.ts +++ b/src/main/worker-service.ts @@ -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, { type: "run-log" | "run-status" }>; @@ -65,6 +68,7 @@ interface WorkerServiceDependencies { processRunCompletion: (run: RunRecord, event: z.infer & { type: "run-finished" }) => void; notify: (options: NotificationOptions) => void; budgetService?: ReturnType; + sessionService?: ReturnType; } 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", diff --git a/tests/unit/session-management.test.ts b/tests/unit/session-management.test.ts index ea9276c3..c189b143 100644 --- a/tests/unit/session-management.test.ts +++ b/tests/unit/session-management.test.ts @@ -710,3 +710,178 @@ describe("session management service", () => { }); }); }); + +// =========================================================================== +// wiring integration +// =========================================================================== +describe("wiring integration", () => { + // These tests verify the run-finished handler correctly calls session service methods. + // Uses vi.fn() mocks to isolate the wiring logic from actual DB/service implementations. + + function createMockSessionService() { + return { + accumulateTokenUsage: vi.fn(), + evaluateRotation: vi.fn().mockReturnValue({ shouldRotate: false, reason: null }), + generateHandoff: vi.fn().mockReturnValue("# Handoff"), + rotateSession: vi.fn().mockReturnValue("new-session-id"), + getHandoffContext: vi.fn().mockReturnValue(null), + snapshotUsage: vi.fn().mockReturnValue({ inputTokens: 0, outputTokens: 0 }), + getUsageDelta: vi.fn().mockReturnValue({ inputTokensDelta: 0, outputTokensDelta: 0 }), + }; + } + + // Simulate the session wiring logic extracted from worker-service run-finished handler. + // This mirrors the exact code block added to worker-service.ts, testable without starting + // an actual Electron utilityProcess or full WorkerService. + function simulateRunFinishedSessionWiring( + sessionService: ReturnType, + db: { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getActiveTaskSession: (...args: any[]) => any; + }, + task: { id: string } | null, + event: { inputTokens: number | null; outputTokens: number | null; runId: string }, + connectorId: string, + ) { + if (task && sessionService) { + const activeSession = db.getActiveTaskSession(task.id); + if (activeSession) { + sessionService.accumulateTokenUsage( + activeSession.id, + event.inputTokens ?? 0, + event.outputTokens ?? 0, + event.runId, + ); + const rotationResult = sessionService.evaluateRotation(activeSession, connectorId); + if (rotationResult.shouldRotate) { + const updatedSession = db.getActiveTaskSession(task.id); + if (updatedSession) { + sessionService.rotateSession(updatedSession); + } + } + } + } + } + + it("calls accumulateTokenUsage with correct args on run completion", () => { + const sessionService = createMockSessionService(); + const mockSession = { id: "sess-1", taskId: "task-1", runCount: 5, totalInputTokens: 1000, totalOutputTokens: 500 }; + const db = { getActiveTaskSession: vi.fn().mockReturnValue(mockSession) }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.accumulateTokenUsage).toHaveBeenCalledWith("sess-1", 200, 100, "run-1"); + }); + + it("calls evaluateRotation after accumulation", () => { + const sessionService = createMockSessionService(); + const mockSession = { id: "sess-1", taskId: "task-1", runCount: 5, totalInputTokens: 1000, totalOutputTokens: 500 }; + const db = { getActiveTaskSession: vi.fn().mockReturnValue(mockSession) }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.evaluateRotation).toHaveBeenCalledWith(mockSession, "gemini_local"); + // evaluateRotation called AFTER accumulateTokenUsage + const accOrder = sessionService.accumulateTokenUsage.mock.invocationCallOrder[0]; + const evalOrder = sessionService.evaluateRotation.mock.invocationCallOrder[0]; + expect(evalOrder).toBeGreaterThan(accOrder); + }); + + it("calls rotateSession when shouldRotate is true", () => { + const sessionService = createMockSessionService(); + sessionService.evaluateRotation.mockReturnValue({ shouldRotate: true, reason: "max_runs" }); + const mockSession = { id: "sess-1", taskId: "task-1", runCount: 20, totalInputTokens: 1000, totalOutputTokens: 500 }; + const updatedSession = { ...mockSession, runCount: 21, totalInputTokens: 1200, totalOutputTokens: 600 }; + const db = { + getActiveTaskSession: vi.fn() + .mockReturnValueOnce(mockSession) // first call: get session for accumulation + .mockReturnValueOnce(updatedSession), // second call: re-read after accumulation for rotation + }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.rotateSession).toHaveBeenCalledWith(updatedSession); + }); + + it("does NOT call rotateSession when shouldRotate is false", () => { + const sessionService = createMockSessionService(); + sessionService.evaluateRotation.mockReturnValue({ shouldRotate: false, reason: null }); + const mockSession = { id: "sess-1", taskId: "task-1", runCount: 5, totalInputTokens: 1000, totalOutputTokens: 500 }; + const db = { getActiveTaskSession: vi.fn().mockReturnValue(mockSession) }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.rotateSession).not.toHaveBeenCalled(); + }); + + it("skips session wiring when task is null", () => { + const sessionService = createMockSessionService(); + const db = { getActiveTaskSession: vi.fn() }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + null, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(db.getActiveTaskSession).not.toHaveBeenCalled(); + expect(sessionService.accumulateTokenUsage).not.toHaveBeenCalled(); + }); + + it("skips session wiring when no active session exists", () => { + const sessionService = createMockSessionService(); + const db = { getActiveTaskSession: vi.fn().mockReturnValue(null) }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: 200, outputTokens: 100, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.accumulateTokenUsage).not.toHaveBeenCalled(); + }); + + it("passes null tokens as 0 to accumulateTokenUsage", () => { + const sessionService = createMockSessionService(); + const mockSession = { id: "sess-1", taskId: "task-1", runCount: 5, totalInputTokens: 0, totalOutputTokens: 0 }; + const db = { getActiveTaskSession: vi.fn().mockReturnValue(mockSession) }; + + simulateRunFinishedSessionWiring( + sessionService, + db, + { id: "task-1" }, + { inputTokens: null, outputTokens: null, runId: "run-1" }, + "gemini_local", + ); + + expect(sessionService.accumulateTokenUsage).toHaveBeenCalledWith("sess-1", 0, 0, "run-1"); + }); +});