53 lines
3.5 KiB
JavaScript
53 lines
3.5 KiB
JavaScript
import { calculateContextTokens, estimateTokens, getLastAssistantUsage } from '@earendil-works/pi-coding-agent';
|
|
|
|
/** Usage is scoped to one model and one project. Never turn a local safety
|
|
* budget, a billing tier or a missing receipt into a claimed model capacity. */
|
|
export const count = value => Number.isSafeInteger(value) && value >= 0 ? value : null;
|
|
export const contextTextLength = context => JSON.stringify(context, (_key,value)=>value?.type==='image'&&typeof value.data==='string'?{...value,data:''}:value).length;
|
|
export function modelCapacity(model) {
|
|
const limit = count(model.contextWindow);
|
|
return limit > 0 && typeof model.contextWindowSource === 'string' && model.contextWindowSource.length > 0
|
|
? { limit, source: model.contextWindowSource } : { limit: null, source: null };
|
|
}
|
|
export function projectedUsage(context) {
|
|
// Preflight only. Use Pi's own estimator, never show this safety budget as
|
|
// the session meter or replace an available provider receipt with it.
|
|
return { budgetTokens: context.messages.reduce((sum, message) => sum + estimateTokens(message), 0), inputChars: contextTextLength(context) };
|
|
}
|
|
export function piContextUsage(session, model) {
|
|
const native = session.getContextUsage();
|
|
const manager = session.sessionManager;
|
|
const branch = manager.getBranch();
|
|
const boundary = branch.findLastIndex(entry => entry.type === 'compaction' || entry.type === 'context_edit');
|
|
const eligible = new Set(branch.slice(boundary + 1).map(entry => entry.id));
|
|
const entries = manager.buildSessionProjection().entries
|
|
.filter(entry => eligible.has(entry.sourceEntry.id))
|
|
.flatMap(entry => entry.messages.map(message => ({ type: 'message', message })));
|
|
const last = entries.findLast(entry => getLastAssistantUsage([entry]));
|
|
// A prior model's receipt is not the current model's context measurement.
|
|
if (last && last.message.model !== model.id) return { contextTokens: null, contextSource: 'pi-model-changed' };
|
|
if (native) return { contextTokens: count(native.tokens), contextSource: 'pi-session', piContextUsage: native };
|
|
// Pi 1.3.0 deliberately returns undefined when contextWindow <= 0. Preserve
|
|
// its real latest assistant usage without fabricating a limit or percentage.
|
|
// Exclude receipts invalidated by compaction/context edits, just like Pi.
|
|
const usage = getLastAssistantUsage(entries);
|
|
return { contextTokens: usage ? count(calculateContextTokens(usage)) : null,
|
|
contextSource: usage ? 'pi-last-response' : 'pi-awaiting-usage' };
|
|
}
|
|
export function contextSnapshot(identity, model, session, previous = null, receipt = null, phase = 'ready') {
|
|
const capacity = modelCapacity(model);
|
|
const same = previous?.scope === identity.scope && previous?.project === identity.project && previous?.model === model.id;
|
|
return { version: 1, ...identity, model: model.id, modelKey: model.key ?? null, protocol: model.protocol,
|
|
...piContextUsage(session, model), contextWindow: capacity.limit, capacitySource: capacity.source,
|
|
receipt: receipt ?? (same ? previous.receipt : null), phase, at: Date.now() };
|
|
}
|
|
export function inputBudget(model) {
|
|
const capacity = modelCapacity(model).limit;
|
|
const reserve = Math.min(8192, capacity ? Math.max(1, Math.floor(capacity / 4)) : 8192);
|
|
return { reserve, inputLimit: capacity ? capacity - reserve : null };
|
|
}
|
|
export function exceedsModelBudget(context, model, fraction = 1) {
|
|
const { inputLimit } = inputBudget(model);
|
|
return inputLimit !== null && projectedUsage(context).budgetTokens > inputLimit * fraction;
|
|
}
|