Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions changelog.d/boot-state-read-errors.fixed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
- Restored state is no longer taken as empty at boot. `AgentFramework.create()` now fails when a state it restores can't be read, or reads as something this version can't take for that state. The error names the state, or its recovery file, and says what to do; nothing it names is changed. Before, each was taken as empty, the framework served without it, and the next write replaced what was stored. Nothing stored still starts empty, and a store any release wrote boots as before unless one of these reads fails.
- **Host mode** (`recovery/host-mode.json`, or `framework/state` without one): a quiesced host booted serving, and landed the writes it had deferred during the maintenance.
- **Deferred writes** (`recovery/deferred-writes.json`, or `framework/deferred-writes` without one): the queue restored nothing, and the next deferral replaced it. An agent's context that can't be checked for writes that landed before their ack now fails `create()` too, where such a write landed a second time.
- **A recovery file that can't be reached** (an error other than its absence, such as `EACCES` on its directory) fails `create()` too, where it read as absent.
- **Usage totals** (`framework/state`): they started at zero, and the next inference stored zeros over them. A `framework/state` that isn't an object now fails `create()`, where every reader took it as empty.
- **An agent's persisted runtime settings** (`framework/state`): the agent booted at its recipe's settings, a `private` think-text policy as `public`, and its next settings change stored over the record. 0.22.2's notes said such a read is logged and the agent boots at the recipe's; it now fails `create()`.
- **Tool-result guards** (`framework/state`): a `toolResultGuards` that isn't an object read as no setting, and the next guard change stored over it. It's now read by each agent's own name, so an agent named `constructor` no longer fails `create()` once any guard is stored.
- **MCPL coalescing** (`mcpl/coalescing`, `mcpl/coalescing-recent`): an acknowledged deferred acceptance or a pending batch was never delivered, and the next receipt or flush replaced the stored ones. A snapshot of a version this one doesn't read was ignored and replaced; it now fails `create()` and is kept for the version that wrote it. An agent's context that can't be checked for an occurrence already delivered fails `create()` too, where the occurrence was delivered again.
- **MCPL checkpoint trees** (`mcpl/checkpoints`): a stored value without `trees` read as no checkpoints, and the next feature set registered stored over it.
- On a framework started with no MCPL servers, the coalescing state and the checkpoint trees are read at the first `connectMcplServer()`, which fails instead, with nothing of MCPL built, and the next connect reads them again.
- A host-managed MCPL feature set's state (`mcpl/state/<server>/<feature set>`) is now read after a restart: when MCPL starts, for each configured server, and before `connectMcplServer()` adds a server. It wasn't, so the state read as empty, and the feature set's next delta checkpoint applied to nothing and stored over it. A failed read of it fails `create()`, or that `connectMcplServer()`, naming it, before the server connects. A feature set with no stored tree, such as one registering for the first time or after its server was disconnected, has its state read when it registers, as before. A failed read there fails that server's connect, as before, and `create()` logs that and goes on, as for any server's failed connect. The state is now read again the next time the feature set registers; before, its tree was held after the failed read, so the state was taken as empty.
- A usage count stored as `null`, which a provider's usage without a count left as `NaN`, restores as 0, and a usage without a count no longer makes the totals `NaN`.
- A failed read of the persisted `tool_result_inline_max_chars` caps at boot is now logged. As before, they don't apply until a restart, and setting one keeps the others stored.
691 changes: 481 additions & 210 deletions src/framework.ts

Large diffs are not rendered by default.

138 changes: 114 additions & 24 deletions src/mcpl/checkpoint-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,34 @@ const CHECKPOINTS_STATE_ID = 'mcpl/checkpoints';
/** Prefix for per-feature-set host-managed state. */
const STATE_PREFIX = 'mcpl/state';

/** What the refusals below say about MCPL when stored trees or state can't be taken. */
const MCPL_STOPS = `MCPL doesn't start without them, and neither does create() with MCPL servers configured, ` +
`nor a first connectMcplServer()`;

/** A thrown value's text, bounded and never throwing, as framework.ts's thrownText. */
function thrownText(error: unknown): string {
let text: string;
try {
text = String(error instanceof Error ? error.message : error);
} catch {
text = '(a value that could not be read)';
}
return text.length <= 500 ? text : `${text.slice(0, 500)}… (${text.length} characters)`;
}

/** A stored JSON value's kind, for a refusal that names what it found without its content. */
function storedKind(value: unknown): string {
return value === undefined ? 'missing'
: value === null ? 'null'
: Array.isArray(value) ? 'an array'
: typeof value === 'object' ? 'an object'
: `a ${typeof value}`;
}

function isObject(value: unknown): value is Record<string, unknown> {
return !!value && typeof value === 'object' && !Array.isArray(value);
}

// ============================================================================
// Internal types
// ============================================================================
Expand Down Expand Up @@ -214,31 +242,21 @@ export class CheckpointManager {
opts: { hostState: boolean; rollback: boolean },
): void {
const key = this.key(serverId, featureSet);
if (this.trees.has(key)) return; // Already registered (e.g. from loaded state)
// Already registered, or restored from the store by loadFromStore, whose
// host-managed state loadHostStates reads before its server connects.
if (this.trees.has(key)) return;

this.trees.set(key, {
const tree: FeatureSetState = {
hostState: opts.hostState,
rollback: opts.rollback,
currentCheckpoint: null,
currentState: undefined,
nodes: new Map(),
});

// Register per-feature-set state slot for host-managed state
if (opts.hostState) {
const stateId = this.stateId(serverId, featureSet);
try {
this.store.registerState({ id: stateId, strategy: 'snapshot' });
} catch {
// Already registered
}
// Load existing state if any
const existing = this.store.getStateJson(stateId);
if (existing !== null && existing !== undefined) {
const tree = this.trees.get(key)!;
tree.currentState = existing;
}
}
};
// Read before the tree is held, so a registration that fails here is
// made whole by the next one.
if (opts.hostState) this.loadHostState(serverId, featureSet, tree, 'registration');
this.trees.set(key, tree);

this.emitTraceFn({
type: 'mcpl:state_registered',
Expand Down Expand Up @@ -540,12 +558,84 @@ export class CheckpointManager {
this.store.setStateJson(CHECKPOINTS_STATE_ID, serialized);
}

/** Load checkpoint tree from Chronicle on startup. */
/** A host-managed feature set's stored state, read into its tree; nothing
* stored leaves it unset. A read that fails throws, naming the state:
* taken as nothing stored, the next patch would apply to nothing and
* store over it. `when` says what the refusal stops: a restored tree's
* read comes before its server connects, and a feature set with no tree
* is read when it registers, once its server has connected. */
private loadHostState(
serverId: string,
featureSet: string,
tree: FeatureSetState,
when: 'before-connect' | 'registration',
): void {
const stateId = this.stateId(serverId, featureSet);
try {
this.store.registerState({ id: stateId, strategy: 'snapshot' });
} catch {
// Already registered
}
let existing: unknown;
try {
existing = this.store.getStateJson(stateId);
} catch (error) {
throw new Error(
`the host-managed state of feature set "${featureSet}" on MCPL server "${serverId}" (${stateId}) could not ` +
`be read: ${thrownText(error)}. ` +
(when === 'before-connect'
? `Neither create() with the server configured nor a connectMcplServer() of it goes on without it, since ` +
`the feature set's next checkpoint would store over it. It's unchanged, and they go on once it can be read.`
: `The feature set isn't registered without it, since its next checkpoint would store over it, so the ` +
`server's connect fails. It's unchanged, and it's read again when the feature set next registers.`),
{ cause: error },
);
}
if (existing !== null && existing !== undefined) tree.currentState = existing;
}

/** Read the stored state of every host-managed tree held for `serverId`
* with none in memory, whatever its feature sets declare now. The
* framework calls it at MCPL's start for each configured server and
* before a runtime connect adds its server, so no server connects with a
* restored tree whose state its next checkpoint would store over unread.
* A tree whose state is in memory, as a connected server's is, keeps it.
* A server's trees are the keys `${serverId}:`, as removeServer finds
* them. */
loadHostStates(serverId: string): void {
const prefix = `${serverId}:`;
for (const [key, tree] of this.trees) {
if (!tree.hostState || tree.currentState !== undefined || !key.startsWith(prefix)) continue;
this.loadHostState(serverId, key.slice(prefix.length), tree, 'before-connect');
}
}

/** Load checkpoint trees from Chronicle on startup. Nothing stored is no
* trees. Trees that can't be read, or a stored value that isn't an object
* whose trees is an object, throw: taken as no trees, the next feature set
* registered would store over them. */
private loadFromStore(): void {
const data = this.store.getStateJson(CHECKPOINTS_STATE_ID) as SerializedTrees | null;
if (!data?.trees) return;
const what = `the MCPL checkpoint trees (${CHECKPOINTS_STATE_ID})`;
const stops = `${MCPL_STOPS}, since the next feature set registered would store over them.`;
let data: unknown;
try {
data = this.store.getStateJson(CHECKPOINTS_STATE_ID);
} catch (error) {
throw new Error(
`${what} could not be read: ${thrownText(error)}. ${stops} They're unchanged, and MCPL starts once they can be read.`,
{ cause: error },
);
}
if (data === null || data === undefined) return;
const trees = isObject(data) ? data.trees : undefined;
if (!isObject(trees)) {
throw new Error(
`${what} are ${isObject(data) ? `an object whose trees is ${storedKind(trees)}` : storedKind(data)}, not ` +
`checkpoint trees. ${stops} They're unchanged, and starting again won't help until they're repaired.`,
);
}

for (const [key, entry] of Object.entries(data.trees)) {
for (const [key, entry] of Object.entries(trees as SerializedTrees['trees'])) {
const nodes = new Map<string, CheckpointNode>();
for (const [id, nodeData] of Object.entries(entry.nodes)) {
nodes.set(id, {
Expand All @@ -561,7 +651,7 @@ export class CheckpointManager {
hostState: entry.hostState,
rollback: entry.rollback,
currentCheckpoint: entry.current,
currentState: undefined, // Will be loaded per-feature-set during registerFeatureSet
currentState: undefined, // a host-managed tree's state is read by loadHostStates, before its server connects
nodes,
});
}
Expand Down
40 changes: 32 additions & 8 deletions src/usage/usage-tracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,14 @@ function accumulateUsage(
tokens: { inputTokens: number; outputTokens: number; cacheCreationTokens: number; cacheReadTokens: number },
estimatedCost?: { total: number; currency: string }
): void {
target.inputTokens += tokens.inputTokens;
target.outputTokens += tokens.outputTokens;
target.cacheCreationTokens += tokens.cacheCreationTokens;
target.cacheReadTokens += tokens.cacheReadTokens;
// A provider's usage can lack a count (membrane 0.5.87's Bedrock stream
// takes a message_delta's output count as reported); adding it would make
// the total NaN, which is stored as null.
const count = (n: number): number => (Number.isFinite(n) ? n : 0);
target.inputTokens += count(tokens.inputTokens);
target.outputTokens += count(tokens.outputTokens);
target.cacheCreationTokens += count(tokens.cacheCreationTokens);
target.cacheReadTokens += count(tokens.cacheReadTokens);
if (estimatedCost) {
if (!target.estimatedCost) {
target.estimatedCost = { total: 0, currency: estimatedCost.currency };
Expand All @@ -31,10 +35,16 @@ function accumulateUsage(
function parseUsage(raw: unknown): SessionUsage | null {
if (!raw || typeof raw !== 'object') return null;
const o = raw as Record<string, unknown>;
if (typeof o.inputTokens !== 'number' || typeof o.outputTokens !== 'number') return null;
// A count stored as null is a NaN total, as JSON stores one: an earlier
// release could write it when a provider's usage lacked a count. The sum
// before it was lost when it became NaN, so it reads as 0.
const count = (v: unknown): number | undefined => (v === null ? 0 : typeof v === 'number' ? v : undefined);
const inputTokens = count(o.inputTokens);
const outputTokens = count(o.outputTokens);
if (inputTokens === undefined || outputTokens === undefined) return null;
return {
inputTokens: o.inputTokens,
outputTokens: o.outputTokens as number,
inputTokens,
outputTokens,
cacheCreationTokens: typeof o.cacheCreationTokens === 'number' ? o.cacheCreationTokens : 0,
cacheReadTokens: typeof o.cacheReadTokens === 'number' ? o.cacheReadTokens : 0,
estimatedCost: o.estimatedCost && typeof o.estimatedCost === 'object'
Expand All @@ -49,6 +59,18 @@ export interface PersistedUsageState {
inferenceCount: number;
}

/** What keeps a stored usage record from restoring, or undefined when it
* restores: it must be an object whose totals are token counts and whose
* byAgent, if it has one, is a list. Its per-agent entries are checked one
* by one as they restore, and one that isn't an agent's usage is skipped. */
export function persistedUsageProblem(value: unknown): 'not a record' | 'totals' | 'byAgent' | undefined {
if (!value || typeof value !== 'object' || Array.isArray(value)) return 'not a record';
const { totals, byAgent } = value as { totals?: unknown; byAgent?: unknown };
if (parseUsage(totals) === null) return 'totals';
if (byAgent !== undefined && !Array.isArray(byAgent)) return 'byAgent';
return undefined;
}

export class UsageTracker {
private totals: SessionUsage = { inputTokens: 0, outputTokens: 0, cacheCreationTokens: 0, cacheReadTokens: 0 };
private byAgent: Map<string, { usage: SessionUsage; inferenceCount: number }> = new Map();
Expand All @@ -71,7 +93,9 @@ export class UsageTracker {
this.totals = totals;
this.inferenceCount = typeof state.inferenceCount === 'number' ? state.inferenceCount : 0;
if (Array.isArray(state.byAgent)) {
for (const entry of state.byAgent) {
for (const item of state.byAgent as unknown[]) {
if (!item || typeof item !== 'object') continue;
const entry = item as { agentName?: unknown; usage?: unknown; inferenceCount?: unknown };
const usage = parseUsage(entry.usage);
if (usage && typeof entry.agentName === 'string') {
this.byAgent.set(entry.agentName, {
Expand Down
39 changes: 29 additions & 10 deletions test/agent-runtime-settings.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import { JsStore } from '@animalabs/chronicle';

import { AutobiographicalStrategy, ContextManager } from '@animalabs/context-manager';
import { Agent, AgentFramework, BudgetPreflightError, type AgentRuntimeSettingsPatch } from '../src/index.js';
import { rejectsBuild } from './helpers/failed-create.js';

const membrane = {} as any;

Expand Down Expand Up @@ -853,7 +854,7 @@ describe('a reached budget across a restart', () => {
});
});

it("a record the boot can't read is logged, and the agent boots at the recipe's", async () => {
it("a record the boot can't read fails create(), naming it, and the next boot restores it", async () => {
await run(async ({ framework, restart }) => {
framework().updateAgentRuntimeSettings('agent', { contextBudgetTokens: 60_000, immediate: true });
// Only the boot's read of the record fails.
Expand All @@ -863,19 +864,37 @@ describe('a reached budget across a restart', () => {
if (new Error().stack?.includes('readAgentRuntimeSettings')) throw new Error('store unreadable');
return original.apply(this, args);
};
const log = quietErrors();
let restarted: AgentFramework;
try {
restarted = await restart();
// Booted at the recipe's, its next settings change would store over the record.
await rejectsBuild(restart, /^the persisted settings of agent "agent" \(framework\/state, agentRuntimeSettings\) can't be restored: store unreadable\. The framework doesn't start without them, since the agent would boot at its recipe's settings and the next settings change would store over them\. They're unchanged, and the framework starts once they can be read\.$/);
} finally {
proto.getStateJson = original;
log.restore();
}
assert.deepEqual(budgets(restarted), { target: 100_000, effective: 100_000, transition: 'stable' });
assert.ok(
log.lines.some((l) => l.includes("agent: its persisted settings can't be read (store unreadable); it boots at the recipe's")),
log.lines.join('\n'),
);
const restarted = await restart();
assert.equal(budgets(restarted).target, 60_000);
assert.deepEqual(persisted(restarted), { contextBudgetTokens: 60_000 });
});
});

it("what's stored for an agent that isn't a settings record fails create(), naming what it is", async () => {
await run(async ({ framework, restart, config }) => {
writeRecord(framework(), ['contextBudgetTokens', 60_000] as never);
await rejectsBuild(restart, /^the persisted settings of agent "agent" \(framework\/state, agentRuntimeSettings\) can't be restored: what's stored for the agent is an array, not a settings record\. The framework doesn't start without them, since the agent would boot at its recipe's settings and the next settings change would store over them\. They're unchanged, and starting again won't help until they're repaired\.$/);
// The refused boots closed their stores, so write the next form through one of the test's own.
const store = JsStore.openOrCreate({ path: config().storePath });
try {
const state = store.getStateJson('framework/state') as { agentRuntimeSettings?: object };
store.setStateJson('framework/state', { ...state, agentRuntimeSettings: { ...state.agentRuntimeSettings, agent: null } });
} finally { store.close(); }
await rejectsBuild(restart, /can't be restored: what's stored for the agent is null, not a settings record\./);
});
});

it("a map of agents' settings that isn't an object fails create(), naming what it is", async () => {
await run(async ({ framework, restart }) => {
const store = framework().getStore();
store.setStateJson('framework/state', { ...(store.getStateJson('framework/state') as object), agentRuntimeSettings: 'agent' });
await rejectsBuild(restart, /^the persisted settings of agent "agent" \(framework\/state, agentRuntimeSettings\) can't be restored: the map that holds every agent's settings is a string, not an object\. /);
});
});

Expand Down
Loading
Loading