From e54c632fd6a3a1d62d65eacafd790e0dfcdee8f7 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 22:56:18 -0400 Subject: [PATCH] feat(06-01): add session management service with rotation, thresholds, and delta tracking - Create createSessionManagementService factory with accumulateTokenUsage, evaluateRotation (configurable thresholds per D-01), generateHandoff, rotateSession, getHandoffContext, snapshotUsage, getUsageDelta - Export DEFAULT_THRESHOLDS (maxRuns: 20, maxTokens: 500k, maxAgeHours: 72) and RotationThresholds type - Adapter nativeContextManagement exemption skips rotation for Claude/Codex (D-04) - Handoff markdown capped at 2000 chars with structured sections (D-06) - Add 14 unit tests for service: rotation thresholds, adapter exemption, configurable thresholds, token accumulation, handoff generation, session rotation, usage delta Co-Authored-By: Claude Opus 4.6 (1M context) --- src/main/session-management-service.ts | 207 +++++++++++ tests/unit/session-management.test.ts | 453 ++++++++++++++++++++++++- 2 files changed, 659 insertions(+), 1 deletion(-) create mode 100644 src/main/session-management-service.ts diff --git a/src/main/session-management-service.ts b/src/main/session-management-service.ts new file mode 100644 index 00000000..bec12491 --- /dev/null +++ b/src/main/session-management-service.ts @@ -0,0 +1,207 @@ +import type { AppDatabase } from "./database"; +import { getConnectorDefinition } from "./connectors"; +import type { ConnectorId, TaskSessionRecord } from "@shared/types"; + +interface SessionManagementDependencies { + db: AppDatabase; + logger: { info: (...args: unknown[]) => void; warn: (...args: unknown[]) => void; error: (...args: unknown[]) => void }; +} + +// D-01/D-02: Exported for testing and external override. Conservative defaults. +export const DEFAULT_THRESHOLDS = { + maxRuns: 20, + maxTokens: 500_000, + maxAgeHours: 72, +}; + +// D-01: Type for per-adapter or company-level threshold overrides +export type RotationThresholds = Partial; + +// D-04: Handoff size cap to prevent prompt bloat +const MAX_HANDOFF_LENGTH = 2000; +const MAX_HANDOFF_RUNS = 5; + +export function createSessionManagementService({ db, logger }: SessionManagementDependencies) { + + // D-08: Accumulate token usage on task_sessions row (called for ALL adapters per D-05) + function accumulateTokenUsage(sessionId: string, inputTokens: number, outputTokens: number, runId: string): void { + db.incrementSessionUsage(sessionId, inputTokens, outputTokens, runId); + } + + // D-01: Any-of thresholds with optional overrides. D-04: Check nativeContextManagement. + // Per D-01: "Configurable per-adapter with company-level overrides" -- accepts optional + // thresholds parameter that merges with defaults. Callers can pass per-adapter or + // company-level overrides; omitted fields fall back to DEFAULT_THRESHOLDS. + function evaluateRotation( + session: TaskSessionRecord, + connectorId: string, + thresholdsOverride?: RotationThresholds, + ): { shouldRotate: boolean; reason: string | null } { + // D-04: Check adapter exemption + try { + const connector = getConnectorDefinition(connectorId as ConnectorId); + if (connector.module.nativeContextManagement === true) { + return { shouldRotate: false, reason: null }; + } + } catch { + // Unknown connector - proceed with rotation check + } + + // Merge overrides with defaults (per D-01 configurability) + const thresholds = { ...DEFAULT_THRESHOLDS, ...thresholdsOverride }; + + // D-01: Any-of thresholds - first trigger wins + if (session.runCount >= thresholds.maxRuns) { + return { shouldRotate: true, reason: "max_runs" }; + } + + const totalTokens = session.totalInputTokens + session.totalOutputTokens; + if (totalTokens >= thresholds.maxTokens) { + return { shouldRotate: true, reason: "max_tokens" }; + } + + const ageMs = Date.now() - new Date(session.createdAt).getTime(); + const ageHours = ageMs / (1000 * 60 * 60); + if (ageHours >= thresholds.maxAgeHours) { + return { shouldRotate: true, reason: "max_age_hours" }; + } + + return { shouldRotate: false, reason: null }; + } + + // D-06: Build handoff markdown from structured data, no AI call + function generateHandoff(taskId: string, session: TaskSessionRecord): string { + let task: { title: string; description: string } | null = null; + try { + task = db.getTask(taskId); + } catch { + task = null; + } + + const recentRuns = db.getRecentRunsForTask(taskId, MAX_HANDOFF_RUNS); + const errorRuns = recentRuns.filter(r => r.errorMessage); + const activityEntries = db.getRecentActivityForTask(taskId, 10); + const decisionEntries = activityEntries.filter(a => + a.action.includes("decision") || a.action.includes("approval") || a.action.includes("resolved") + ); + + const sections: string[] = []; + + sections.push("# Session Handoff"); + sections.push(""); + + if (task) { + sections.push("## Task"); + sections.push(`**${task.title}**`); + const descSnippet = task.description.length > 200 ? task.description.slice(0, 200) + "..." : task.description; + sections.push(descSnippet); + sections.push(""); + } + + if (recentRuns.length > 0) { + sections.push("## Recent Run History"); + for (const run of recentRuns) { + const summaryText = run.summary ? run.summary.slice(0, 100) : "(no summary)"; + sections.push(`- [${run.status}] ${summaryText} (${run.finishedAt ?? "in progress"})`); + } + sections.push(""); + } + + if (errorRuns.length > 0) { + sections.push("## Errors Encountered"); + for (const run of errorRuns) { + sections.push(`- ${run.errorMessage!.slice(0, 150)} (${run.finishedAt ?? "unknown"})`); + } + sections.push(""); + } + + if (decisionEntries.length > 0) { + sections.push("## Key Decisions"); + for (const entry of decisionEntries.slice(0, 5)) { + sections.push(`- ${entry.detail.slice(0, 120)} (${entry.createdAt})`); + } + sections.push(""); + } + + sections.push(`*Session stats: ${session.runCount} runs, ${session.totalInputTokens + session.totalOutputTokens} total tokens*`); + + let result = sections.join("\n"); + // Cap handoff size per research pitfall #5 + if (result.length > MAX_HANDOFF_LENGTH) { + result = result.slice(0, MAX_HANDOFF_LENGTH - 50) + "\n\n... (earlier content omitted)"; + } + return result; + } + + // Archive current session, generate handoff, create new session, update task pointer + function rotateSession(session: TaskSessionRecord): string { + const reason = "threshold_exceeded"; + const handoff = generateHandoff(session.taskId, session); + + logger.info(`[session] Rotating session ${session.id} for task ${session.taskId}: ${reason}`); + + // 1. Archive old session + db.archiveSession(session.id, reason, handoff); + + // 2. Create new session row with fresh counters, carry forward session_params + const newSessionId = db.createNewSession({ + taskId: session.taskId, + companyId: session.companyId, + agentId: session.agentId, + adapterType: session.adapterType, + sessionParamsJson: session.sessionParamsJson, + sessionDisplayId: session.sessionDisplayId, + }); + + logger.info(`[session] New session ${newSessionId} created for task ${session.taskId}`); + return newSessionId; + } + + // D-07: Read handoff from most recently rotated session for a task + function getHandoffContext(taskId: string): string | null { + // Find the most recent archived session for this task (the one just before current) + const row = db.queryAll<{ handoffMarkdown: string | null }>( + `SELECT handoff_markdown as handoffMarkdown FROM task_sessions + WHERE task_id = ? AND rotated_at IS NOT NULL + ORDER BY rotated_at DESC LIMIT 1`, + taskId + ); + return row[0]?.handoffMarkdown ?? null; + } + + // D-09: Snapshot token totals for delta tracking + function snapshotUsage(session: TaskSessionRecord): { inputTokens: number; outputTokens: number } { + return { + inputTokens: session.totalInputTokens, + outputTokens: session.totalOutputTokens, + }; + } + + // D-09: Compute delta from snapshot + function getUsageDelta( + sessionId: string, + snapshot: { inputTokens: number; outputTokens: number } + ): { inputTokensDelta: number; outputTokensDelta: number } { + // Re-read current totals from DB + const current = db.queryAll<{ totalInputTokens: number; totalOutputTokens: number }>( + `SELECT total_input_tokens as totalInputTokens, total_output_tokens as totalOutputTokens + FROM task_sessions WHERE id = ?`, + sessionId + ); + if (!current[0]) return { inputTokensDelta: 0, outputTokensDelta: 0 }; + return { + inputTokensDelta: current[0].totalInputTokens - snapshot.inputTokens, + outputTokensDelta: current[0].totalOutputTokens - snapshot.outputTokens, + }; + } + + return { + accumulateTokenUsage, + evaluateRotation, + generateHandoff, + rotateSession, + getHandoffContext, + snapshotUsage, + getUsageDelta, + }; +} diff --git a/tests/unit/session-management.test.ts b/tests/unit/session-management.test.ts index ab835479..ea9276c3 100644 --- a/tests/unit/session-management.test.ts +++ b/tests/unit/session-management.test.ts @@ -14,7 +14,8 @@ vi.mock("electron", () => ({ import { AppDatabase } from "../../src/main/database"; import { getConnectorDefinition } from "../../src/main/connectors"; -import type { ConnectorId } from "@shared/types"; +import { createSessionManagementService, DEFAULT_THRESHOLDS } from "../../src/main/session-management-service"; +import type { ConnectorId, TaskSessionRecord } from "@shared/types"; // --------------------------------------------------------------------------- // Helpers — create a minimal company + agent + task scaffold @@ -259,3 +260,453 @@ describe("adapter exemption", () => { expect(def.module.nativeContextManagement).toBe(false); }); }); + +// =========================================================================== +// Session management service tests +// =========================================================================== +describe("session management service", () => { + let tempDir: string; + let db: AppDatabase; + const mockLogger = { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + + beforeEach(async () => { + tempDir = await mkdtemp(join(tmpdir(), "agentcompany-svc-")); + db = new AppDatabase(tempDir); + db.init(); + mockLogger.info.mockClear(); + mockLogger.warn.mockClear(); + mockLogger.error.mockClear(); + }); + + afterEach(async () => { + db.close(); + await rm(tempDir, { recursive: true, force: true }); + }); + + // ------------------------------------------------------------------------- + // rotation thresholds + // ------------------------------------------------------------------------- + describe("rotation thresholds", () => { + it("returns shouldRotate: false when below all thresholds", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 5, + totalInputTokens: 10000, + totalOutputTokens: 5000, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local"); + expect(result.shouldRotate).toBe(false); + expect(result.reason).toBeNull(); + }); + + it("returns shouldRotate: true with reason max_runs when run_count >= 20", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 20, + totalInputTokens: 0, + totalOutputTokens: 0, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local"); + expect(result.shouldRotate).toBe(true); + expect(result.reason).toBe("max_runs"); + }); + + it("returns shouldRotate: true with reason max_tokens when total tokens >= 500000", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 3, + totalInputTokens: 300000, + totalOutputTokens: 200000, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local"); + expect(result.shouldRotate).toBe(true); + expect(result.reason).toBe("max_tokens"); + }); + + it("returns shouldRotate: true with reason max_age_hours when session created > 72h ago", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const oldDate = new Date(Date.now() - 73 * 60 * 60 * 1000).toISOString(); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 1, + totalInputTokens: 100, + totalOutputTokens: 100, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: oldDate, + updatedAt: oldDate, + }; + const result = svc.evaluateRotation(session, "gemini_local"); + expect(result.shouldRotate).toBe(true); + expect(result.reason).toBe("max_age_hours"); + }); + }); + + // ------------------------------------------------------------------------- + // adapter exemption in service + // ------------------------------------------------------------------------- + describe("adapter exemption in evaluateRotation", () => { + it("returns shouldRotate: false for nativeContextManagement=true even when thresholds exceeded", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "claude_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 100, // Way over threshold + totalInputTokens: 1000000, + totalOutputTokens: 500000, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date(Date.now() - 200 * 60 * 60 * 1000).toISOString(), // Way old + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "claude_local"); + expect(result.shouldRotate).toBe(false); + expect(result.reason).toBeNull(); + }); + }); + + // ------------------------------------------------------------------------- + // configurable thresholds (D-01) + // ------------------------------------------------------------------------- + describe("configurable thresholds (D-01)", () => { + it("evaluateRotation with custom thresholds override { maxRuns: 5 } triggers at 5 runs", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 5, + totalInputTokens: 100, + totalOutputTokens: 100, + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local", { maxRuns: 5 }); + expect(result.shouldRotate).toBe(true); + expect(result.reason).toBe("max_runs"); + }); + + it("partial override { maxTokens: 100000 } uses override for tokens but defaults for runs and age", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 10, // Below default 20 + totalInputTokens: 60000, + totalOutputTokens: 40000, // Total 100000 = meets override threshold + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local", { maxTokens: 100000 }); + expect(result.shouldRotate).toBe(true); + expect(result.reason).toBe("max_tokens"); + }); + + it("omitting override uses DEFAULT_THRESHOLDS", () => { + const svc = createSessionManagementService({ db, logger: mockLogger }); + const session: TaskSessionRecord = { + id: "sess-1", + companyId: "co-1", + taskId: "task-1", + agentId: "agent-1", + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + lastRunId: null, + lastError: null, + runCount: 19, // Just below default 20 + totalInputTokens: 200000, + totalOutputTokens: 200000, // Total 400000, below default 500000 + rotatedAt: null, + rotationReason: null, + handoffMarkdown: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + const result = svc.evaluateRotation(session, "gemini_local"); + expect(result.shouldRotate).toBe(false); + expect(result.reason).toBeNull(); + + // Verify DEFAULT_THRESHOLDS values + expect(DEFAULT_THRESHOLDS.maxRuns).toBe(20); + expect(DEFAULT_THRESHOLDS.maxTokens).toBe(500_000); + expect(DEFAULT_THRESHOLDS.maxAgeHours).toBe(72); + }); + }); + + // ------------------------------------------------------------------------- + // token accumulation + // ------------------------------------------------------------------------- + describe("token accumulation", () => { + it("accumulateTokenUsage calls db.incrementSessionUsage with correct args", () => { + const { taskId, companyId, agentId } = scaffold(db); + const sessionId = db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "claude_local", + sessionParamsJson: null, + sessionDisplayId: null, + }); + const svc = createSessionManagementService({ db, logger: mockLogger }); + + svc.accumulateTokenUsage(sessionId, 5000, 2000, "run-abc"); + + const session = db.getActiveTaskSession(taskId); + expect(session!.runCount).toBe(1); + expect(session!.totalInputTokens).toBe(5000); + expect(session!.totalOutputTokens).toBe(2000); + expect(session!.lastRunId).toBe("run-abc"); + }); + }); + + // ------------------------------------------------------------------------- + // handoff generation + // ------------------------------------------------------------------------- + describe("handoff generation", () => { + it("generateHandoff produces markdown with task title and run history", () => { + const { taskId, companyId, agentId, workspaceId } = scaffold(db); + const sessionId = db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + }); + + // Create a completed run for history + const run = db.createRun({ companyId, taskId, agentId, workspaceId, connectorId: "gemini_local" as ConnectorId }); + db.finishRun({ + runId: run.id, + status: "completed", + summary: "Implemented initial logic", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 0.05, + unattributedCost: false, + }); + + // Increment session usage to have stats + db.incrementSessionUsage(sessionId, 10000, 5000, run.id); + + const session = db.getActiveTaskSession(taskId)!; + const svc = createSessionManagementService({ db, logger: mockLogger }); + const handoff = svc.generateHandoff(taskId, session); + + expect(handoff).toContain("# Session Handoff"); + expect(handoff).toContain("Implement session rotation"); // task title + expect(handoff).toContain("## Recent Run History"); + expect(handoff).toContain("completed"); + expect(handoff).toContain("15000 total tokens"); // 10000 + 5000 + }); + }); + + // ------------------------------------------------------------------------- + // session rotation + // ------------------------------------------------------------------------- + describe("session rotation", () => { + it("rotateSession archives old session, creates new, copies session_params", () => { + const { taskId, companyId, agentId } = scaffold(db); + const oldSessionId = db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "gemini_local", + sessionParamsJson: JSON.stringify({ resumeId: "old-abc" }), + sessionDisplayId: "gemini-old", + }); + + // Add some usage + db.incrementSessionUsage(oldSessionId, 300000, 200000, "run-10"); + + const oldSession = db.getActiveTaskSession(taskId)!; + const svc = createSessionManagementService({ db, logger: mockLogger }); + const newSessionId = svc.rotateSession(oldSession); + + expect(newSessionId).not.toBe(oldSessionId); + + // Old session should be archived + const archivedRows = db.queryAll<{ rotatedAt: string | null; rotationReason: string | null; handoffMarkdown: string | null }>( + "SELECT rotated_at as rotatedAt, rotation_reason as rotationReason, handoff_markdown as handoffMarkdown FROM task_sessions WHERE id = ?", + oldSessionId + ); + expect(archivedRows[0].rotatedAt).toBeTruthy(); + expect(archivedRows[0].rotationReason).toBe("threshold_exceeded"); + expect(archivedRows[0].handoffMarkdown).toContain("# Session Handoff"); + + // New session should be active, with fresh counters and copied params + const newSession = db.getActiveTaskSession(taskId)!; + expect(newSession.id).toBe(newSessionId); + expect(newSession.runCount).toBe(0); + expect(newSession.totalInputTokens).toBe(0); + expect(newSession.totalOutputTokens).toBe(0); + expect(newSession.sessionParamsJson).toBe(JSON.stringify({ resumeId: "old-abc" })); + expect(newSession.sessionDisplayId).toBe("gemini-old"); + }); + }); + + // ------------------------------------------------------------------------- + // handoff context retrieval + // ------------------------------------------------------------------------- + describe("handoff context retrieval", () => { + it("getHandoffContext returns handoff_markdown from most recently rotated session", () => { + const { taskId, companyId, agentId } = scaffold(db); + const svc = createSessionManagementService({ db, logger: mockLogger }); + + // Create and rotate a session + const session1Id = db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + }); + db.incrementSessionUsage(session1Id, 100000, 50000, "run-1"); + const session1 = db.getActiveTaskSession(taskId)!; + svc.rotateSession(session1); + + const handoff = svc.getHandoffContext(taskId); + expect(handoff).not.toBeNull(); + expect(handoff).toContain("# Session Handoff"); + }); + + it("getHandoffContext returns null when no rotated sessions exist", () => { + const { taskId, companyId, agentId } = scaffold(db); + db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + }); + const svc = createSessionManagementService({ db, logger: mockLogger }); + + const handoff = svc.getHandoffContext(taskId); + expect(handoff).toBeNull(); + }); + }); + + // ------------------------------------------------------------------------- + // usage delta (D-09) + // ------------------------------------------------------------------------- + describe("usage delta", () => { + it("getUsageDelta returns token difference between snapshot and current totals", () => { + const { taskId, companyId, agentId } = scaffold(db); + const sessionId = db.createNewSession({ + taskId, + companyId, + agentId, + adapterType: "gemini_local", + sessionParamsJson: null, + sessionDisplayId: null, + }); + + // Pre-populate some usage + db.incrementSessionUsage(sessionId, 1000, 500, "run-1"); + + const session = db.getActiveTaskSession(taskId)!; + const svc = createSessionManagementService({ db, logger: mockLogger }); + + // Snapshot current state + const snapshot = svc.snapshotUsage(session); + expect(snapshot.inputTokens).toBe(1000); + expect(snapshot.outputTokens).toBe(500); + + // Accumulate more tokens + svc.accumulateTokenUsage(sessionId, 3000, 1500, "run-2"); + + // Delta should reflect new tokens only + const delta = svc.getUsageDelta(sessionId, snapshot); + expect(delta.inputTokensDelta).toBe(3000); + expect(delta.outputTokensDelta).toBe(1500); + }); + }); +});