Skip to content

Commit 9b0ef73

Browse files
PerishCodePerishCode
authored
Merge pull request #5918 from nexu-io/backport-5816-to-release/v0.16.0
[backport release/v0.16.0] fix(daemon): roll over near-limit agent sessions
2 parents 158b32d + a1c64cf commit 9b0ef73

25 files changed

Lines changed: 918 additions & 15 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/integrations/provider-models.ts‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ function openAiModelOption(item: unknown): ProviderModelOption | null {
141141
const obj = item as { id?: unknown; metadata?: unknown };
142142
const id = typeof obj?.id === 'string' ? obj.id : '';
143143
if (!id || !isOpenAiChatModelId(id)) return null;
144-
const metadata = extractModelMetadata(obj.metadata);
144+
const metadata = extractModelMetadata(obj);
145145
return {
146146
id,
147147
label: id,
@@ -151,16 +151,38 @@ function openAiModelOption(item: unknown): ProviderModelOption | null {
151151

152152
function extractModelMetadata(value: unknown): ModelMetadata | null {
153153
if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
154-
const metadata = value as { cost?: unknown; capability?: unknown };
154+
const record = value as Record<string, unknown>;
155+
const metadata = (
156+
record.metadata && typeof record.metadata === 'object' && !Array.isArray(record.metadata)
157+
? record.metadata
158+
: record
159+
) as Record<string, unknown>;
155160
const cost = parseModelCost(metadata.cost);
156161
const capability = parseModelCapability(metadata.capability);
157-
if (!cost && !capability) return null;
162+
const contextWindowTokens = positiveInteger(
163+
metadata.contextWindowTokens ??
164+
metadata.context_window_tokens ??
165+
record.context_length ??
166+
record.contextWindow,
167+
);
168+
const maxOutputTokens = positiveInteger(
169+
metadata.maxOutputTokens ?? metadata.max_output_tokens ?? record.max_output_tokens,
170+
);
171+
if (!cost && !capability && !contextWindowTokens && !maxOutputTokens) return null;
158172
return {
159173
...(cost ? { cost } : {}),
160174
...(capability ? { capability } : {}),
175+
...(contextWindowTokens ? { contextWindowTokens } : {}),
176+
...(maxOutputTokens ? { maxOutputTokens } : {}),
161177
};
162178
}
163179

180+
function positiveInteger(value: unknown): number | null {
181+
return typeof value === 'number' && Number.isInteger(value) && value > 0
182+
? value
183+
: null;
184+
}
185+
164186
function parseModelCost(value: unknown): ModelCost | null {
165187
return value === 'low' ||
166188
value === 'medium' ||

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

Lines changed: 15 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 {
@@ -1086,6 +1100,7 @@ export async function reportRunCompletedFromDaemon(
10861100
...(stderr ? { stderr } : {}),
10871101
...(stdout ? { stdout } : {}),
10881102
diagnostics,
1103+
...(run.contextBudget ? { contextBudget: run.contextBudget } : {}),
10891104
},
10901105
message: {
10911106
messageId: run.assistantMessageId ?? '',

‎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 } : {}),

‎apps/daemon/src/runtimes/defs/amr.ts‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -198,13 +198,45 @@ function extractModelMetadata(item: unknown): ModelMetadata | null {
198198
const metadata = isRecord(item.metadata) ? item.metadata : item;
199199
const cost = parseModelCost(metadata.cost);
200200
const capability = parseModelCapability(metadata.capability);
201-
if (!cost && !capability) return null;
201+
const contextWindowTokens = extractPositiveModelLimit(item, [
202+
'contextWindowTokens',
203+
'context_window_tokens',
204+
'context_length',
205+
'contextLength',
206+
'context',
207+
]);
208+
const maxOutputTokens = extractPositiveModelLimit(item, [
209+
'maxOutputTokens',
210+
'max_output_tokens',
211+
'output',
212+
]);
213+
if (!cost && !capability && !contextWindowTokens && !maxOutputTokens) return null;
202214
return {
203215
...(cost ? { cost } : {}),
204216
...(capability ? { capability } : {}),
217+
...(contextWindowTokens ? { contextWindowTokens } : {}),
218+
...(maxOutputTokens ? { maxOutputTokens } : {}),
205219
};
206220
}
207221

222+
function extractPositiveModelLimit(
223+
item: Record<string, unknown>,
224+
keys: string[],
225+
): number | null {
226+
const metadata = isRecord(item.metadata) ? item.metadata : null;
227+
const limit = isRecord(item.limit) ? item.limit : null;
228+
for (const source of [metadata, limit, item]) {
229+
if (!source) continue;
230+
for (const key of keys) {
231+
const value = source[key];
232+
if (typeof value === 'number' && Number.isInteger(value) && value > 0) {
233+
return value;
234+
}
235+
}
236+
}
237+
return null;
238+
}
239+
208240
function withPriceDerivedCostMetadata(
209241
metadata: ModelMetadata | null,
210242
inputPriceUsdPerMillion: number | undefined,
@@ -462,14 +494,20 @@ function openCodeModelPrice(
462494
> | null {
463495
if (!isRecord(model)) return null;
464496
const inputPriceUsdPerMillion = extractInputPriceUsdPerMillion(model);
465-
if (inputPriceUsdPerMillion === undefined) return null;
466497
const outputPriceUsdPerMillion = extractOutputPriceUsdPerMillion(model);
467498
const metadata = withPriceDerivedCostMetadata(
468499
extractModelMetadata(model),
469500
inputPriceUsdPerMillion,
470501
);
502+
if (
503+
inputPriceUsdPerMillion === undefined &&
504+
outputPriceUsdPerMillion === undefined &&
505+
metadata === null
506+
) {
507+
return null;
508+
}
471509
return {
472-
inputPriceUsdPerMillion,
510+
...(inputPriceUsdPerMillion === undefined ? {} : { inputPriceUsdPerMillion }),
473511
...(outputPriceUsdPerMillion === undefined ? {} : { outputPriceUsdPerMillion }),
474512
...(metadata === null ? {} : { metadata }),
475513
};

0 commit comments

Comments
 (0)