Skip to content

Commit 286f33e

Browse files
Siri-RaySiri-Ray
committed
roll over near-limit agent sessions (#5816)
(cherry picked from commit 43b6403)
1 parent 158b32d commit 286f33e

17 files changed

Lines changed: 780 additions & 4 deletions

‎apps/daemon/src/agent-session-resume.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,47 @@ export interface AgentResumeContext {
2525
isResuming: boolean;
2626
/** Hash of the stable instruction block last sent on this session, or null. */
2727
storedStablePromptHash: string | null;
28+
/** Effective provider input size recorded on the session's last turn. */
29+
storedInputTokens: number | null;
2830
/** Set when a stored session existed but was rejected; see the type. */
2931
invalidationReason: ResumeInvalidationReason | null;
3032
}
3133

34+
function readStoredSessionInputTokens(
35+
db: SqliteDb,
36+
messageId: string | null | undefined,
37+
): number | null {
38+
if (!messageId) return null;
39+
const row = db
40+
.prepare('SELECT events_json AS eventsJson FROM messages WHERE id = ?')
41+
.get(messageId) as { eventsJson?: unknown } | undefined;
42+
if (!row || typeof row.eventsJson !== 'string') return null;
43+
let events: unknown;
44+
try {
45+
events = JSON.parse(row.eventsJson);
46+
} catch {
47+
return null;
48+
}
49+
if (!Array.isArray(events)) return null;
50+
for (let index = events.length - 1; index >= 0; index -= 1) {
51+
const event = events[index];
52+
if (!event || typeof event !== 'object' || Array.isArray(event)) continue;
53+
const usage = event as {
54+
kind?: unknown;
55+
inputTokens?: unknown;
56+
inputTokensEffective?: unknown;
57+
};
58+
if (usage.kind !== 'usage') continue;
59+
const effective = typeof usage.inputTokensEffective === 'number'
60+
? usage.inputTokensEffective
61+
: usage.inputTokens;
62+
if (typeof effective === 'number' && Number.isFinite(effective) && effective > 0) {
63+
return Math.floor(effective);
64+
}
65+
}
66+
return null;
67+
}
68+
3269
export type CapturedAgentSessionResult = 'stored' | 'cleared' | 'skipped';
3370

3471
/**
@@ -112,6 +149,9 @@ export function resolveAgentResumeContext(
112149
newSessionId: randomUUID(),
113150
isResuming: resumable,
114151
storedStablePromptHash: resumable ? (record?.stablePromptHash ?? null) : null,
152+
storedInputTokens: resumable
153+
? readStoredSessionInputTokens(db, record?.lastMessageId)
154+
: null,
115155
invalidationReason,
116156
};
117157
}

‎apps/daemon/src/langfuse-bridge.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,20 @@ interface DaemonRunRecord {
9898
promptTelemetry?: PromptStackTelemetry;
9999
projectAttachmentPaths?: string[];
100100
projectMetadata?: Record<string, unknown> | null;
101+
contextBudget?: {
102+
action: string;
103+
source: string;
104+
estimatedPromptTokens: number;
105+
contextWindowTokens?: number;
106+
reservedOutputTokens?: number;
107+
inputBudgetTokens?: number;
108+
budgetRatio?: number;
109+
priorSessionInputTokens?: number;
110+
projectedInputTokens?: number;
111+
rolloverThresholdTokens?: number;
112+
compactedPromptTokens?: number;
113+
omittedTranscriptMessageBlocks?: number;
114+
};
101115
}
102116

103117
interface TraceSafeManifestResult {

‎apps/daemon/src/langfuse-trace.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,20 @@ export interface RunSummary {
134134
truncated: boolean;
135135
};
136136
diagnostics?: unknown;
137+
contextBudget?: {
138+
action: string;
139+
source: string;
140+
estimatedPromptTokens: number;
141+
contextWindowTokens?: number;
142+
reservedOutputTokens?: number;
143+
inputBudgetTokens?: number;
144+
budgetRatio?: number;
145+
priorSessionInputTokens?: number;
146+
projectedInputTokens?: number;
147+
rolloverThresholdTokens?: number;
148+
compactedPromptTokens?: number;
149+
omittedTranscriptMessageBlocks?: number;
150+
};
137151
}
138152

139153
export interface MessageSummary {
@@ -1455,6 +1469,18 @@ export function buildTracePayload(ctx: ReportContext): unknown[] {
14551469
stderr: ctx.run.stderr,
14561470
stdout: ctx.run.stdout,
14571471
diagnostics: ctx.run.diagnostics,
1472+
contextBudgetAction: ctx.run.contextBudget?.action,
1473+
contextBudgetSource: ctx.run.contextBudget?.source,
1474+
estimatedPromptTokens: ctx.run.contextBudget?.estimatedPromptTokens,
1475+
contextWindowTokens: ctx.run.contextBudget?.contextWindowTokens,
1476+
reservedOutputTokens: ctx.run.contextBudget?.reservedOutputTokens,
1477+
inputBudgetTokens: ctx.run.contextBudget?.inputBudgetTokens,
1478+
contextBudgetRatio: ctx.run.contextBudget?.budgetRatio,
1479+
priorSessionInputTokens: ctx.run.contextBudget?.priorSessionInputTokens,
1480+
projectedSessionInputTokens: ctx.run.contextBudget?.projectedInputTokens,
1481+
rolloverThresholdTokens: ctx.run.contextBudget?.rolloverThresholdTokens,
1482+
compactedPromptTokens: ctx.run.contextBudget?.compactedPromptTokens,
1483+
omittedTranscriptMessageBlocks: ctx.run.contextBudget?.omittedTranscriptMessageBlocks,
14581484
eventsSummary: ctx.eventsSummary,
14591485
tokens,
14601486
cost_usd: costBreakdown.cost_usd,

‎apps/daemon/src/routes/runs.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,21 @@ interface ChatRun {
170170
retryAttemptCount?: number;
171171
retryFinalResult?: string;
172172
retrySuppressedReason?: string;
173+
contextBudget?: {
174+
action: 'unmeasured' | 'within_budget' | 'blocked' | 'rollover';
175+
source: 'model_metadata' | 'known_model_family' | 'unknown';
176+
estimatedPromptTokens: number;
177+
contextWindowTokens?: number;
178+
reservedOutputTokens?: number;
179+
safetyMarginTokens?: number;
180+
inputBudgetTokens?: number;
181+
budgetRatio?: number;
182+
priorSessionInputTokens?: number;
183+
projectedInputTokens?: number;
184+
rolloverThresholdTokens?: number;
185+
compactedPromptTokens?: number;
186+
omittedTranscriptMessageBlocks?: number;
187+
};
173188
artifactOutcome?: {
174189
artifactCount: number;
175190
artifactsCreated?: number;
@@ -1273,6 +1288,38 @@ export function registerRunRoutes(app: Express, ctx: RegisterRunRoutesDeps) {
12731288
asked_user_question: runAskedUserQuestion(run.events),
12741289
retry_attempt_count: run.retryAttemptCount ?? 0,
12751290
retry_final_result: run.retryFinalResult ?? 'not_attempted',
1291+
context_budget_action: run.contextBudget?.action ?? 'unmeasured',
1292+
context_budget_source: run.contextBudget?.source ?? 'unknown',
1293+
...(run.contextBudget?.estimatedPromptTokens !== undefined
1294+
? { estimated_prompt_tokens: run.contextBudget.estimatedPromptTokens }
1295+
: {}),
1296+
...(run.contextBudget?.contextWindowTokens !== undefined
1297+
? { context_window_tokens: run.contextBudget.contextWindowTokens }
1298+
: {}),
1299+
...(run.contextBudget?.reservedOutputTokens !== undefined
1300+
? { reserved_output_tokens: run.contextBudget.reservedOutputTokens }
1301+
: {}),
1302+
...(run.contextBudget?.inputBudgetTokens !== undefined
1303+
? { input_budget_tokens: run.contextBudget.inputBudgetTokens }
1304+
: {}),
1305+
...(run.contextBudget?.budgetRatio !== undefined
1306+
? { context_budget_ratio: run.contextBudget.budgetRatio }
1307+
: {}),
1308+
...(run.contextBudget?.priorSessionInputTokens !== undefined
1309+
? { prior_session_input_tokens: run.contextBudget.priorSessionInputTokens }
1310+
: {}),
1311+
...(run.contextBudget?.projectedInputTokens !== undefined
1312+
? { projected_session_input_tokens: run.contextBudget.projectedInputTokens }
1313+
: {}),
1314+
...(run.contextBudget?.rolloverThresholdTokens !== undefined
1315+
? { rollover_threshold_tokens: run.contextBudget.rolloverThresholdTokens }
1316+
: {}),
1317+
...(run.contextBudget?.compactedPromptTokens !== undefined
1318+
? { compacted_prompt_tokens: run.contextBudget.compactedPromptTokens }
1319+
: {}),
1320+
...(run.contextBudget?.omittedTranscriptMessageBlocks !== undefined
1321+
? { omitted_transcript_message_blocks: run.contextBudget.omittedTranscriptMessageBlocks }
1322+
: {}),
12761323
...(run.retrySuppressedReason
12771324
? { retry_suppressed_reason: run.retrySuppressedReason }
12781325
: {}),

‎apps/daemon/src/runtimes/chat-run-messages.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type Database from 'better-sqlite3';
22
import type { PersistedAgentEvent } from '@open-design/contracts';
3+
import { scanRunEventsForUsageAnalytics } from '../run-analytics-observability.js';
34
import {
45
appendMessageAgentEvent,
56
upsertMessage,
@@ -224,9 +225,17 @@ export function daemonAgentPayloadToPersistedAgentEvent(data: unknown): Persiste
224225
}
225226
if (type === 'usage') {
226227
const usage = isRecord(data.usage) ? data.usage : {};
228+
const usageAnalytics = scanRunEventsForUsageAnalytics(
229+
[{ event: 'agent', data }],
230+
null,
231+
0,
232+
);
227233
return {
228234
kind: 'usage',
229235
...(typeof usage.input_tokens === 'number' ? { inputTokens: usage.input_tokens } : {}),
236+
...(typeof usageAnalytics.input_tokens_effective === 'number'
237+
? { inputTokensEffective: usageAnalytics.input_tokens_effective }
238+
: {}),
230239
...(typeof usage.output_tokens === 'number' ? { outputTokens: usage.output_tokens } : {}),
231240
...(typeof data.costUsd === 'number' ? { costUsd: data.costUsd } : {}),
232241
...(typeof data.durationMs === 'number' ? { durationMs: data.durationMs } : {}),

0 commit comments

Comments
 (0)