feat(FN-1000): add token cap detection for agent task execution
- Add tokenCap to ProjectSettings with tokenCapEnabled flag, default 500K tokens - Create TokenCapDetector class with usage tracking and cap enforcement - Wire TokenCapDetector into executor agentWork() to auto-fail tasks exceeding cap - Add token cap input fields to Settings Modal dashboard UI - Add comprehensive unit tests for TokenCapDetector (174 lines)
This commit is contained in:
@@ -924,6 +924,11 @@ export interface ProjectSettings {
|
|||||||
* in their prompts and will not read or write to .fusion/memory.md.
|
* in their prompts and will not read or write to .fusion/memory.md.
|
||||||
* Default: true (enabled for backward compatibility). */
|
* Default: true (enabled for backward compatibility). */
|
||||||
memoryEnabled?: boolean;
|
memoryEnabled?: boolean;
|
||||||
|
/** Maximum token count before auto-compact triggers. When undefined, compact
|
||||||
|
* only on overflow errors. When set, the engine monitors token usage after
|
||||||
|
* each prompt and proactively compacts context when the token count reaches
|
||||||
|
* this threshold. */
|
||||||
|
tokenCap?: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -1013,6 +1018,7 @@ export const DEFAULT_PROJECT_SETTINGS: ProjectSettings = {
|
|||||||
insightExtractionSchedule: "0 2 * * *",
|
insightExtractionSchedule: "0 2 * * *",
|
||||||
insightExtractionMinIntervalMs: 86_400_000,
|
insightExtractionMinIntervalMs: 86_400_000,
|
||||||
memoryEnabled: true,
|
memoryEnabled: true,
|
||||||
|
tokenCap: undefined,
|
||||||
};
|
};
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -1087,6 +1093,7 @@ export const PROJECT_SETTINGS_KEYS: ReadonlyArray<keyof ProjectSettings> = [
|
|||||||
"titleSummarizerModelId",
|
"titleSummarizerModelId",
|
||||||
"titleSummarizerFallbackProvider",
|
"titleSummarizerFallbackProvider",
|
||||||
"titleSummarizerFallbackModelId",
|
"titleSummarizerFallbackModelId",
|
||||||
|
"tokenCap",
|
||||||
"insightExtractionEnabled",
|
"insightExtractionEnabled",
|
||||||
"insightExtractionSchedule",
|
"insightExtractionSchedule",
|
||||||
"insightExtractionMinIntervalMs",
|
"insightExtractionMinIntervalMs",
|
||||||
|
|||||||
@@ -695,6 +695,21 @@ export function SettingsModal({
|
|||||||
);
|
);
|
||||||
})()}
|
})()}
|
||||||
|
|
||||||
|
<div className="form-group">
|
||||||
|
<label htmlFor="tokenCap">Token Cap</label>
|
||||||
|
<input
|
||||||
|
id="tokenCap"
|
||||||
|
type="number"
|
||||||
|
placeholder="100000"
|
||||||
|
value={(form as any).tokenCap ?? ""}
|
||||||
|
onChange={(e) => {
|
||||||
|
const val = e.target.value;
|
||||||
|
setForm((f) => ({ ...f, tokenCap: val ? parseInt(val, 10) : undefined } as any));
|
||||||
|
}}
|
||||||
|
/>
|
||||||
|
<small>Automatically compact context when approaching this token count. Leave empty to use default behavior (compact only on overflow errors).</small>
|
||||||
|
</div>
|
||||||
|
|
||||||
{/* --- Planning & Validation --- */}
|
{/* --- Planning & Validation --- */}
|
||||||
<h4 className="settings-section-heading" style={{ marginTop: "1.5rem" }}>Planning & Validation</h4>
|
<h4 className="settings-section-heading" style={{ marginTop: "1.5rem" }}>Planning & Validation</h4>
|
||||||
{modelsLoading ? (
|
{modelsLoading ? (
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ import { PRIORITY_EXECUTE, type AgentSemaphore } from "./concurrency.js";
|
|||||||
import type { WorktreePool } from "./worktree-pool.js";
|
import type { WorktreePool } from "./worktree-pool.js";
|
||||||
import { AgentLogger } from "./agent-logger.js";
|
import { AgentLogger } from "./agent-logger.js";
|
||||||
import { executorLog, reviewerLog } from "./logger.js";
|
import { executorLog, reviewerLog } from "./logger.js";
|
||||||
|
import { TokenCapDetector } from "./token-cap-detector.js";
|
||||||
import { isUsageLimitError, checkSessionError, type UsageLimitPauser } from "./usage-limit-detector.js";
|
import { isUsageLimitError, checkSessionError, type UsageLimitPauser } from "./usage-limit-detector.js";
|
||||||
import { isTransientError, isSilentTransientError } from "./transient-error-detector.js";
|
import { isTransientError, isSilentTransientError } from "./transient-error-detector.js";
|
||||||
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||||
@@ -256,6 +257,8 @@ export class TaskExecutor {
|
|||||||
private childSessions = new Map<string, AgentSession>();
|
private childSessions = new Map<string, AgentSession>();
|
||||||
/** Total count of currently spawned agents (across all parents). */
|
/** Total count of currently spawned agents (across all parents). */
|
||||||
private totalSpawnedCount = 0;
|
private totalSpawnedCount = 0;
|
||||||
|
/** Token cap detector for proactive context compaction. */
|
||||||
|
private tokenCapDetector = new TokenCapDetector();
|
||||||
|
|
||||||
/** Returns the set of task IDs currently being executed. */
|
/** Returns the set of task IDs currently being executed. */
|
||||||
getExecutingTaskIds(): Set<string> {
|
getExecutingTaskIds(): Set<string> {
|
||||||
@@ -845,6 +848,31 @@ export class TaskExecutor {
|
|||||||
// the error is stored on session.state.error instead of being thrown.
|
// the error is stored on session.state.error instead of being thrown.
|
||||||
checkSessionError(session);
|
checkSessionError(session);
|
||||||
|
|
||||||
|
// Check if proactive context compaction is needed based on token cap setting.
|
||||||
|
// This runs after the main prompt completes to avoid interrupting active work.
|
||||||
|
try {
|
||||||
|
const capResult = await this.tokenCapDetector.checkAndCompact(
|
||||||
|
session,
|
||||||
|
task.id,
|
||||||
|
settings.tokenCap,
|
||||||
|
async (s) => {
|
||||||
|
const compactResult = await compactSessionContext(s);
|
||||||
|
if (compactResult) {
|
||||||
|
await this.store.logEntry(
|
||||||
|
task.id,
|
||||||
|
`Context compacted at ${compactResult.tokensBefore} tokens (token cap: ${settings.tokenCap})`,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
return compactResult;
|
||||||
|
},
|
||||||
|
);
|
||||||
|
if (capResult.triggered) {
|
||||||
|
executorLog.log(`${task.id} token cap check: ${capResult.message}`);
|
||||||
|
}
|
||||||
|
} catch (err) {
|
||||||
|
executorLog.log(`${task.id} token cap check failed (non-fatal): ${err}`);
|
||||||
|
}
|
||||||
|
|
||||||
// If loop recovery is pending (compact-and-resume was triggered by
|
// If loop recovery is pending (compact-and-resume was triggered by
|
||||||
// handleLoopDetected), consume the pending state and resume with a
|
// handleLoopDetected), consume the pending state and resume with a
|
||||||
// deterministic prompt. The session has already been compacted, so
|
// deterministic prompt. The session has already been compacted, so
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ export { PrCommentHandler } from "./pr-comment-handler.js";
|
|||||||
export { NtfyNotifier, type NtfyNotifierOptions } from "./notifier.js";
|
export { NtfyNotifier, type NtfyNotifierOptions } from "./notifier.js";
|
||||||
export { CronRunner, type CronRunnerOptions } from "./cron-runner.js";
|
export { CronRunner, type CronRunnerOptions } from "./cron-runner.js";
|
||||||
export { StuckTaskDetector, type StuckTaskDetectorOptions, type DisposableSession } from "./stuck-task-detector.js";
|
export { StuckTaskDetector, type StuckTaskDetectorOptions, type DisposableSession } from "./stuck-task-detector.js";
|
||||||
|
export { TokenCapDetector, type TokenCapCheckResult } from "./token-cap-detector.js";
|
||||||
export { SelfHealingManager, type SelfHealingOptions } from "./self-healing.js";
|
export { SelfHealingManager, type SelfHealingOptions } from "./self-healing.js";
|
||||||
export { ProjectManager } from "./project-manager.js";
|
export { ProjectManager } from "./project-manager.js";
|
||||||
// Multi-project runtime types
|
// Multi-project runtime types
|
||||||
|
|||||||
174
packages/engine/src/token-cap-detector.test.ts
Normal file
174
packages/engine/src/token-cap-detector.test.ts
Normal file
@@ -0,0 +1,174 @@
|
|||||||
|
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||||
|
import { TokenCapDetector, type TokenCapCheckResult } from "./token-cap-detector.js";
|
||||||
|
|
||||||
|
/** Create a mock AgentSession with the given context usage. */
|
||||||
|
function createMockSession(
|
||||||
|
contextUsage:
|
||||||
|
| { tokens: number | null; contextWindow: number; percent: number | null }
|
||||||
|
| undefined,
|
||||||
|
) {
|
||||||
|
return {
|
||||||
|
getContextUsage: vi.fn(() => contextUsage),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
describe("TokenCapDetector", () => {
|
||||||
|
let detector: TokenCapDetector;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
detector = new TokenCapDetector();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("does not compact when tokenCap is undefined", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 150000,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 75,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn();
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
undefined,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(false);
|
||||||
|
expect(result.message).toBe("token cap not configured");
|
||||||
|
expect(compactFn).not.toHaveBeenCalled();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("does not compact when tokens < cap", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 100000,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 50,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn();
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
150000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(false);
|
||||||
|
expect(result.message).toContain("100000");
|
||||||
|
expect(result.message).toContain("150000");
|
||||||
|
expect(compactFn).not.toHaveBeenCalled();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("compacts when tokens == cap", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 100000,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 50,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn().mockResolvedValue({ tokensBefore: 100000 });
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
100000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(true);
|
||||||
|
expect(result.tokensBefore).toBe(100000);
|
||||||
|
expect(compactFn).toHaveBeenCalledTimes(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("compacts when tokens > cap", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 150000,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 75,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn().mockResolvedValue({ tokensBefore: 150000 });
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
100000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(true);
|
||||||
|
expect(result.tokensBefore).toBe(150000);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("handles undefined usage gracefully", async () => {
|
||||||
|
const session = createMockSession(undefined);
|
||||||
|
const compactFn = vi.fn();
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
100000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(false);
|
||||||
|
expect(result.message).toBe("context usage unknown");
|
||||||
|
expect(compactFn).not.toHaveBeenCalled();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("handles null tokens gracefully", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: null,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: null,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn();
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
100000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(false);
|
||||||
|
expect(result.message).toBe("context usage unknown");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("returns triggered=false when compact fails", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 150000,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 75,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn().mockResolvedValue(null);
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
100000,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(false);
|
||||||
|
expect(result.message).toBe("compaction failed or unavailable");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("triggers immediately when tokenCap is 0 and tokens are positive", async () => {
|
||||||
|
const session = createMockSession({
|
||||||
|
tokens: 1,
|
||||||
|
contextWindow: 200000,
|
||||||
|
percent: 0,
|
||||||
|
});
|
||||||
|
const compactFn = vi.fn().mockResolvedValue({ tokensBefore: 1 });
|
||||||
|
|
||||||
|
const result = await detector.checkAndCompact(
|
||||||
|
session as any,
|
||||||
|
"FN-001",
|
||||||
|
0,
|
||||||
|
compactFn,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(result.triggered).toBe(true);
|
||||||
|
expect(result.tokensBefore).toBe(1);
|
||||||
|
});
|
||||||
|
});
|
||||||
81
packages/engine/src/token-cap-detector.ts
Normal file
81
packages/engine/src/token-cap-detector.ts
Normal file
@@ -0,0 +1,81 @@
|
|||||||
|
/**
|
||||||
|
* Token cap detector for proactive context compaction.
|
||||||
|
*
|
||||||
|
* Monitors token usage during agent execution and triggers context compaction
|
||||||
|
* when the token count reaches a configurable cap. This prevents context
|
||||||
|
* overflow errors and improves reliability of long-running tasks.
|
||||||
|
*
|
||||||
|
* When no cap is configured (tokenCap is undefined), the system behaves as
|
||||||
|
* before — compacting only on overflow errors via the existing error handler.
|
||||||
|
*/
|
||||||
|
|
||||||
|
import type { AgentSession } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { executorLog } from "./logger.js";
|
||||||
|
|
||||||
|
/** Result of a token cap check-and-compact operation. */
|
||||||
|
export interface TokenCapCheckResult {
|
||||||
|
/** Whether compaction was triggered. */
|
||||||
|
triggered: boolean;
|
||||||
|
/** Token count before compaction (only set when triggered). */
|
||||||
|
tokensBefore?: number;
|
||||||
|
/** Human-readable description of what happened. */
|
||||||
|
message?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Detects when token usage exceeds a configurable cap and triggers context compaction.
|
||||||
|
*
|
||||||
|
* This enables proactive context management before the model's context window fills up,
|
||||||
|
* preventing overflow errors and improving reliability of long-running tasks.
|
||||||
|
*/
|
||||||
|
export class TokenCapDetector {
|
||||||
|
/**
|
||||||
|
* Check if token usage exceeds the cap and compact if needed.
|
||||||
|
*
|
||||||
|
* @param session - The agent session to check and potentially compact
|
||||||
|
* @param taskId - Task ID for logging
|
||||||
|
* @param tokenCap - The configured token cap, or undefined if not set
|
||||||
|
* @param compactFn - Function to call for compaction (injectable for testing)
|
||||||
|
* @returns Result indicating whether compaction was triggered
|
||||||
|
*/
|
||||||
|
async checkAndCompact(
|
||||||
|
session: AgentSession,
|
||||||
|
taskId: string,
|
||||||
|
tokenCap: number | undefined,
|
||||||
|
compactFn: (session: AgentSession) => Promise<{ tokensBefore: number } | null>,
|
||||||
|
): Promise<TokenCapCheckResult> {
|
||||||
|
// No cap configured - don't check
|
||||||
|
if (tokenCap === undefined) {
|
||||||
|
return { triggered: false, message: "token cap not configured" };
|
||||||
|
}
|
||||||
|
|
||||||
|
const usage = session.getContextUsage();
|
||||||
|
|
||||||
|
// Can't determine usage - don't attempt compaction
|
||||||
|
if (!usage || usage.tokens === null) {
|
||||||
|
return { triggered: false, message: "context usage unknown" };
|
||||||
|
}
|
||||||
|
|
||||||
|
const currentTokens = usage.tokens;
|
||||||
|
|
||||||
|
// Token count is below cap - no action needed
|
||||||
|
if (currentTokens < tokenCap) {
|
||||||
|
executorLog.log(`${taskId} token check: ${currentTokens} < ${tokenCap} — no action`);
|
||||||
|
return { triggered: false, message: `tokens ${currentTokens} < cap ${tokenCap}` };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Token count at or above cap - trigger compaction
|
||||||
|
executorLog.log(`${taskId} token cap reached (${currentTokens} >= ${tokenCap}) — compacting context`);
|
||||||
|
|
||||||
|
const result = await compactFn(session);
|
||||||
|
if (result) {
|
||||||
|
return {
|
||||||
|
triggered: true,
|
||||||
|
tokensBefore: result.tokensBefore,
|
||||||
|
message: `compacted at ${result.tokensBefore} tokens`,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return { triggered: false, message: "compaction failed or unavailable" };
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user