diff --git a/packages/runtime-host/src/__tests__/context-protocol.test.ts b/packages/runtime-host/src/__tests__/context-protocol.test.ts index 2e555ff508..d557db1917 100644 --- a/packages/runtime-host/src/__tests__/context-protocol.test.ts +++ b/packages/runtime-host/src/__tests__/context-protocol.test.ts @@ -71,6 +71,12 @@ test('context operations preserve bounded exact wire values', () => { turnCount: 2, estimatedTokens: 12, }, + requestPreservation: { + status: 'diverged', + previousSegmentCount: 8, + preservedSegmentCount: 2, + firstDivergentSegment: { kind: 'message', index: 2, role: '' }, + }, }, }), { @@ -96,6 +102,12 @@ test('context operations preserve bounded exact wire values', () => { turnCount: 2, estimatedTokens: 12, }, + requestPreservation: { + status: 'diverged', + previousSegmentCount: 8, + preservedSegmentCount: 2, + firstDivergentSegment: { kind: 'message', index: 2, role: '' }, + }, }, }, ); diff --git a/packages/runtime-host/src/__tests__/execution-model-composition.test.ts b/packages/runtime-host/src/__tests__/execution-model-composition.test.ts index a4fe2d3c7a..1a8e7578c2 100644 --- a/packages/runtime-host/src/__tests__/execution-model-composition.test.ts +++ b/packages/runtime-host/src/__tests__/execution-model-composition.test.ts @@ -45,7 +45,7 @@ import type { BackendCompactHistoryInput } from '@maka/core/backend-types'; import { decodeCanonicalToolResultContent } from '@maka/core/tool-result-record-schema'; import { type ModelCallAttempt, type ModelCallKind } from '@maka/core/model-call-attempt'; import { type RuntimeEvent } from '@maka/core/runtime-event'; -import { createDefaultRuntimePolicy } from '@maka/core/runtime-policy'; +import { createDefaultRuntimePolicy, type RuntimePolicy } from '@maka/core/runtime-policy'; import type { PlanSessionState, PlanStore } from '@maka/core/plan'; import type { SessionTodoToolStore } from '@maka/runtime/session-todo-tools'; import { @@ -184,6 +184,31 @@ test('prepared backend activation builds from its admitted provider snapshot', a await backend.dispose(); }); +test('connection and proxy changes qualify the admitted provider identity', async () => { + const prepare = async (revision: number, networkProxy?: RuntimePolicy['networkProxy']) => { + const input = backendCreationFixture({ + abortSignal: new AbortController().signal, + resolveExecutionConnection: async () => + readyExecutionConnection(undefined, { revision, networkProxy }), + readPricing: async () => ({ revision: 0, overrides: [] }), + }); + const { context, ...dependencies } = input; + return prepareHostAiSdkBackend({ context, ...dependencies }); + }; + + const original = (await prepare(1)).providerStateIdentity; + assert.notEqual(original, (await prepare(2)).providerStateIdentity); + assert.notEqual( + original, + ( + await prepare(1, { + ...createDefaultRuntimePolicy().networkProxy, + enabled: true, + }) + ).providerStateIdentity, + ); +}); + test('backend creation aborts a stalled canonical connection read', async () => { const abort = new AbortController(); const creating = createHostAiSdkBackend( @@ -4047,12 +4072,15 @@ function readyExecutionConnection( customization: { readonly requestHeaders?: Readonly>; readonly requestBodyOverlay?: Readonly>; + readonly networkProxy?: RuntimePolicy['networkProxy']; + readonly revision?: number; readonly vision?: boolean; } = {}, ) { return { kind: 'ready', connection: { + revision: customization.revision ?? 1, slug: 'backend-creation-connection', providerType: 'moonshot', ...(baseUrl ? { baseUrl } : {}), @@ -4073,7 +4101,7 @@ function readyExecutionConnection( }, ], }, - networkProxy: { enabled: false }, + networkProxy: customization.networkProxy ?? { enabled: false }, secretMaterial: { connection: { secret: API_KEY }, ...(customization.requestHeaders diff --git a/packages/runtime-host/src/protocol/context.ts b/packages/runtime-host/src/protocol/context.ts index 6d0e18c56c..456d2cb2b0 100644 --- a/packages/runtime-host/src/protocol/context.ts +++ b/packages/runtime-host/src/protocol/context.ts @@ -28,6 +28,7 @@ import { invalidProtocolFrame } from './errors.js'; import { defineOperation } from './operation-spec.js'; import { decodeContextCompactionOutcome, decodeTurnSnapshot, type TurnSnapshot } from './turn.js'; import type { ContextCompactionOutcome } from '@maka/core/events'; +import { isRequestPreservation, type RequestPreservation } from '@maka/runtime/context-diagnostics'; export interface ContextDiagnosticsQueryInput { readonly sessionId: string; @@ -71,6 +72,8 @@ export interface ContextDiagnosticsComposition { readonly unlabelledToolBytes?: number; } +export type ContextDiagnosticsRequestPreservation = RequestPreservation; + export type ContextDiagnosticsResult = | { readonly status: 'unavailable'; @@ -98,6 +101,8 @@ export type ContextDiagnosticsResult = readonly turnCount: number; readonly estimatedTokens: number; }; + /** Runtime-owned verdict; Host and Desktop must not recompute it. */ + readonly requestPreservation?: ContextDiagnosticsRequestPreservation; }; const QUERY_ERRORS = [ @@ -174,6 +179,7 @@ function decodeContextDiagnosticsResult(value: unknown): ContextDiagnosticsResul 'contextWindow', 'composition', 'compaction', + 'requestPreservation', ], ); if (record.status === 'unavailable') { @@ -196,7 +202,14 @@ function decodeContextDiagnosticsResult(value: unknown): ContextDiagnosticsResul record, 'Available context diagnostics', ['status', 'providerId', 'modelId', 'completedAt'], - ['inputTokens', 'cacheReadInputTokens', 'contextWindow', 'composition', 'compaction'], + [ + 'inputTokens', + 'cacheReadInputTokens', + 'contextWindow', + 'composition', + 'compaction', + 'requestPreservation', + ], ); return { status: 'available', @@ -223,9 +236,25 @@ function decodeContextDiagnosticsResult(value: unknown): ContextDiagnosticsResul ...(available.compaction === undefined ? {} : { compaction: decodeContextDiagnosticsCompaction(available.compaction) }), + ...(available.requestPreservation === undefined + ? {} + : { + requestPreservation: decodeContextDiagnosticsRequestPreservation( + available.requestPreservation, + ), + }), }; } +function decodeContextDiagnosticsRequestPreservation( + value: unknown, +): ContextDiagnosticsRequestPreservation { + if (!isRequestPreservation(value)) { + throw invalidProtocolFrame('Invalid context diagnostics request preservation'); + } + return value; +} + /** * The tool list is bounded on the wire for the same reason the evidence reads * are: a Host answer is built in memory, and a registry that grew without limit diff --git a/packages/runtime-host/src/protocol/index.ts b/packages/runtime-host/src/protocol/index.ts index 8aff63a16a..4ffe415b9a 100644 --- a/packages/runtime-host/src/protocol/index.ts +++ b/packages/runtime-host/src/protocol/index.ts @@ -95,7 +95,9 @@ export const RUNTIME_HOST_REGISTRATION_SCHEMA_VERSION = 1 as const; export const RUNTIME_HOST_PROTOCOL_VERSION = 0 as const; // Increment when the same protocol version no longer guarantees safe Client-Host // interoperability. Mismatches are rejected before domain commands are admitted. -export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 87 as const; +export const RUNTIME_HOST_COMPATIBILITY_EPOCH = 88 as const; +// 88: Context diagnostics may carry Runtime-owned request-preservation results. +// Older peers reject this optional field on the closed diagnostics shape. // 87: The connection catalog projects each model as the Host resolved it — // a `catalog_entry` item per model, counted by the connection header. Clients // render those entries instead of merging the stored row against their own diff --git a/packages/runtime-host/src/server/execution-model-authority.ts b/packages/runtime-host/src/server/execution-model-authority.ts index 72e71a03e1..214c30563d 100644 --- a/packages/runtime-host/src/server/execution-model-authority.ts +++ b/packages/runtime-host/src/server/execution-model-authority.ts @@ -763,15 +763,21 @@ function providerStateIdentityForResolvedExecution( { kind: 'ready' } >, ): `sha256:${string}` { - const credentialBasis = (material: typeof resolved.secretMaterial.connection) => - material ? { credentialId: material.credentialId, revision: material.revision } : null; + const credentialBasis = ( + material: { readonly credentialId?: string; readonly revision?: number } | null | undefined, + ) => (material ? { credentialId: material.credentialId, revision: material.revision } : null); + const proxy = resolved.networkProxy.enabled ? resolved.networkProxy : null; return stableHash({ protocol: 'provider_state_identity_v1', connectionId: resolved.connection.connectionId, + connectionRevision: resolved.connection.revision, providerType: resolved.connection.providerType, endpoint: new URL(effectiveBaseUrl(resolved.connection)).toString(), credential: credentialBasis(resolved.secretMaterial.connection), requestHeaders: credentialBasis(resolved.secretMaterial.requestHeaders), + proxy, + proxyCredential: + proxy?.authEnabled === true ? credentialBasis(resolved.secretMaterial.networkProxy) : null, }); } diff --git a/packages/runtime/src/__tests__/context-diagnostics.test.ts b/packages/runtime/src/__tests__/context-diagnostics.test.ts index 9354d762a3..86d9744f11 100644 --- a/packages/runtime/src/__tests__/context-diagnostics.test.ts +++ b/packages/runtime/src/__tests__/context-diagnostics.test.ts @@ -68,6 +68,24 @@ test('rejects v2 snapshots that the canonical writer cannot produce', () => { estimatedTokens: 10, }, }, + { + ...base, + requestPreservation: { + status: 'diverged', + previousSegmentCount: 1, + preservedSegmentCount: 0, + firstDivergentSegment: { kind: 'message', index: 0, role: {} }, + }, + }, + { + ...base, + requestPreservation: { + status: 'preserved', + previousSegmentCount: 1, + preservedSegmentCount: 1, + unexpected: true, + }, + }, ]; for (const snapshot of impossible) { @@ -104,6 +122,139 @@ test('serves the sealed snapshot without reading a single run', async () => { } }); +test('repairs a pending preservation result after its admitted predecessor becomes durable', async () => { + const root = await mkdtemp(join(tmpdir(), 'maka-context-diagnostics-')); + try { + const store = createSqliteAgentRunStore(root); + const previousRun = { + ...runHeader('run-1', 1), + status: 'running' as const, + providerStateIdentity: `sha256:${'1'.repeat(64)}` as const, + llmConnectionSlug: 'connection', + }; + const currentRun = { + ...runHeader('run-2', 2), + providerStateIdentity: `sha256:${'1'.repeat(64)}` as const, + llmConnectionSlug: 'connection', + }; + await store.createRun(previousRun); + await store.createRun(currentRun); + await admitTurn(store, 'turn-run-1', 'run-1', null, 1); + await admitTurn(store, 'turn-run-2', 'run-2', 'turn-run-1', 2); + + const observation = requestObservation([ + { + kind: 'message', + index: 0, + cacheable: true, + comparison: 'exact', + digest: `sha256:${'a'.repeat(64)}`, + bytes: 10, + }, + ]); + const currentEvent = meteringEvent('run-2', 'current', 20, 'model', 40, 200, { + logicalCallId: 'current-call', + connectionSlug: 'connection', + requestObservation: observation, + }); + const currentProjection = latestContext('current', 20); + Object.assign(currentProjection.snapshot, { requestPreservationPending: true }); + await store.appendEvent('session-1', 'run-2', currentEvent, { + durable: true, + latestContext: currentProjection, + }); + + const before = await readLatestContextDiagnostics(store, 'session-1'); + assert.equal(before.status, 'available'); + if (before.status !== 'available') return; + assert.deepEqual(before.requestPreservation, { + status: 'unavailable', + previousSegmentCount: 0, + preservedSegmentCount: 0, + }); + + const previousEvent = meteringEvent('run-1', 'previous', 10, 'model', 40, 200, { + logicalCallId: 'previous-call', + connectionSlug: 'connection', + requestObservation: observation, + }); + await store.appendEvent('session-1', 'run-1', previousEvent, { + durable: true, + latestContext: latestContext('previous', 10), + }); + + const repaired = await readLatestContextDiagnostics(store, 'session-1'); + assert.equal(repaired.status, 'available'); + if (repaired.status !== 'available') return; + assert.equal(repaired.modelId, 'model'); + assert.deepEqual(repaired.requestPreservation, { + status: 'preserved', + previousSegmentCount: 1, + preservedSegmentCount: 1, + }); + } finally { + await rm(root, { recursive: true, force: true }); + } +}); + +test('settles unavailable when a terminal predecessor has no canonical attempt', async () => { + const root = await mkdtemp(join(tmpdir(), 'maka-context-diagnostics-')); + try { + const store = createSqliteAgentRunStore(root); + const providerStateIdentity = `sha256:${'1'.repeat(64)}` as const; + await store.createRun({ + ...runHeader('run-1', 1), + providerStateIdentity, + llmConnectionSlug: 'connection', + }); + await store.createRun({ + ...runHeader('run-2', 2), + providerStateIdentity, + llmConnectionSlug: 'connection', + }); + await admitTurn(store, 'turn-run-1', 'run-1', null, 1); + await admitTurn(store, 'turn-run-2', 'run-2', 'turn-run-1', 2); + + const currentProjection = latestContext('current', 20); + Object.assign(currentProjection.snapshot, { requestPreservationPending: true }); + await store.appendEvent( + 'session-1', + 'run-2', + meteringEvent('run-2', 'current', 20, 'model', 40, 200, { + connectionSlug: 'connection', + requestObservation: requestObservation([ + { + kind: 'message', + index: 0, + cacheable: true, + comparison: 'exact', + digest: `sha256:${'a'.repeat(64)}`, + bytes: 10, + }, + ]), + }), + { durable: true, latestContext: currentProjection }, + ); + + let scanned = 0; + const counted = countingStore(store, () => { + scanned += 1; + }); + const first = await readLatestContextDiagnostics(counted, 'session-1'); + assert.equal(first.status, 'available'); + if (first.status !== 'available') return; + assert.equal(first.requestPreservation?.status, 'unavailable'); + assert.ok(scanned > 0, 'the pending projection is rebuilt once'); + + scanned = 0; + const second = await readLatestContextDiagnostics(counted, 'session-1'); + assert.equal(second.status, 'available'); + assert.equal(scanned, 0, 'terminal unavailability becomes a settled warm projection'); + } finally { + await rm(root, { recursive: true, force: true }); + } +}); + test('does not trust a pre-observation projection over its canonical attempt', async () => { const root = await mkdtemp(join(tmpdir(), 'maka-context-diagnostics-')); try { @@ -229,6 +380,34 @@ test('a failed call does not replace the last good snapshot', async () => { } }); +test('an identity-mismatched canonical event cannot become the latest request', async () => { + const root = await mkdtemp(join(tmpdir(), 'maka-context-diagnostics-')); + try { + const writer = createSqliteAgentRunStore(root); + await writer.createRun(runHeader('run-1', 1)); + await writer.appendEvent( + 'session-1', + 'run-1', + meteringEvent('run-1', 'attempt-1', 10, 'model-valid', 40, 200), + ); + await writer.appendEvent( + 'session-1', + 'run-1', + meteringEvent('run-1', 'attempt-2', 20, 'model-mismatched', 40, 200, { + runId: 'another-run', + }), + ); + + const diagnostics = await readLatestContextDiagnostics(writer, 'session-1'); + + assert.equal(diagnostics.status, 'available'); + if (diagnostics.status !== 'available') return; + assert.equal(diagnostics.modelId, 'model-valid'); + } finally { + await rm(root, { recursive: true, force: true }); + } +}); + test("a subagent's run never becomes the session's context", async () => { const root = await mkdtemp(join(tmpdir(), 'maka-context-diagnostics-')); try { @@ -1010,9 +1189,30 @@ function countingStore( readEventLedgerRevision: (sessionId) => reader.readEventLedgerRevision(sessionId), repairEventProjection: (sessionId, type, event, options) => reader.repairEventProjection(sessionId, type, event, options), + readRootTurnAdmission: (sessionId, turnId) => reader.readRootTurnAdmission(sessionId, turnId), }; } +async function admitTurn( + store: ReturnType, + turnId: string, + runId: string, + previousRootTurnId: string | null, + admittedAt: number, +): Promise { + await store.admitRootTurn({ + sessionId: 'session-1', + turnId, + proposedRunId: runId, + proposedUserMessageId: `message-${turnId}`, + execution: { kind: 'external_message' }, + previousRootTurnId, + normalizedInput: { text: turnId }, + sourceMessages: [], + admittedAt, + }); +} + /** * The derived row as the canonical append commits it — the same shape the * store writes inside that transaction, never a separate ledger record. @@ -1120,7 +1320,7 @@ function meteringEvent( const turnId = `turn-${runId}`; return { type: 'model_call_attempt_recorded', - id: `metering-${attemptId}`, + id: attemptId, runId, sessionId: 'session-1', turnId, diff --git a/packages/runtime/src/__tests__/latest-context-commit.test.ts b/packages/runtime/src/__tests__/latest-context-commit.test.ts index 5075d6cbb9..5666a5b875 100644 --- a/packages/runtime/src/__tests__/latest-context-commit.test.ts +++ b/packages/runtime/src/__tests__/latest-context-commit.test.ts @@ -48,7 +48,7 @@ import { BackendRegistry, SessionManager } from '../session-manager.js'; import { readLatestContextDiagnostics } from '../context-diagnostics.js'; import { createTestAiSdkBackend } from './execution-boundary-test-helpers.js'; -test('a real send seals its observation into SQLite and reconstructs it after restart', async () => { +test('real consecutive sends preserve the previous request across restart', async () => { // Tracker → backend → the kernel seam a backend is actually built with → // AgentRun → the storage transaction. Every layer in that list once had a // signature that compiled while dropping the row, and no test crossed all of @@ -64,30 +64,32 @@ test('a real send seals its observation into SQLite and reconstructs it after re let clock = 1_000; const now = () => (clock += 1); - backends.register('ai-sdk', (ctx) => - createTestAiSdkBackend({ - sessionId: ctx.sessionId, - header: ctx.header, - appendMessage: async () => {}, - connection: { - slug: 'mock-main', - providerType: 'anthropic', - defaultModel: 'mock-model-id', - models: [{ id: 'mock-model-id', contextWindow: 200_000 }], - }, - apiKey: 'sk-test', - modelId: 'mock-model-id', - modelFactory: () => answeringModel(), - tools: [], - // The seams the kernel hands a real backend, forwarded exactly as the - // production composition forwards them — this is the hop that broke. - ...(ctx.recordModelCallAttempt - ? { recordModelCallAttempt: ctx.recordModelCallAttempt } - : {}), - newId, - now, + backends.register('ai-sdk', { + prepare: async () => ({ + providerStateIdentity: `sha256:${'1'.repeat(64)}`, + build: (ctx) => + createTestAiSdkBackend({ + sessionId: ctx.sessionId, + header: ctx.header, + appendMessage: ctx.appendMessage ?? (async () => {}), + connection: { + slug: 'mock-main', + providerType: 'anthropic', + defaultModel: 'mock-model-id', + models: [{ id: 'mock-model-id', contextWindow: 200_000 }], + }, + apiKey: 'sk-test', + modelId: 'mock-model-id', + modelFactory: () => answeringModel(), + tools: [], + ...(ctx.recordModelCallAttempt + ? { recordModelCallAttempt: ctx.recordModelCallAttempt } + : {}), + newId, + now, + }), }), - ); + }); const manager = new SessionManager({ store: sessionStore, @@ -102,10 +104,12 @@ test('a real send seals its observation into SQLite and reconstructs it after re llmConnectionSlug: 'mock-main', permissionMode: 'bypass', }); - for await (const _event of manager.sendMessage(session.id, { - turnId: 'turn-1', - text: 'what is my context made of?', - })) { + await admitTurn(runStore, session.id, 'turn-1', 'run-1', null, now()); + for await (const _event of manager.sendMessage( + session.id, + { turnId: 'turn-1', text: 'what is my context made of?' }, + { runId: 'run-1', userMessageId: 'message-turn-1', durability: 'required' }, + )) { // Drain the turn so its run reaches the durable ledger. } @@ -133,8 +137,31 @@ test('a real send seals its observation into SQLite and reconstructs it after re diagnostics.composition?.segments.some((segment) => segment.kind === 'messages'), 'and the request describes what it was made of', ); + assert.deepEqual(diagnostics.requestPreservation, { + status: 'no_predecessor', + previousSegmentCount: 0, + preservedSegmentCount: 0, + }); assert.equal(scanned, 0, 'the row was committed by the send, not rebuilt by the read'); + await admitTurn(runStore, session.id, 'turn-2', 'run-2', 'turn-1', now()); + for await (const _event of manager.sendMessage( + session.id, + { turnId: 'turn-2', text: 'what changed?' }, + { runId: 'run-2', userMessageId: 'message-turn-2', durability: 'required' }, + )) { + // Drain the successor through the same public Runtime boundary. + } + const successor = await readLatestContextDiagnostics(runStore, session.id); + assert.equal(successor.status, 'available'); + if (successor.status !== 'available') return; + assert.ok(successor.requestPreservation); + assert.equal(successor.requestPreservation.status, 'preserved'); + assert.equal( + successor.requestPreservation.preservedSegmentCount, + successor.requestPreservation.previousSegmentCount, + ); + await manager.stopSession(session.id, { source: 'stop_button' }); runStore.close?.(); @@ -151,13 +178,18 @@ test('a real send seals its observation into SQLite and reconstructs it after re }), ) ).flat(); - assert.equal(canonicalAttempts.length, 1); - const observation = canonicalAttempts[0]?.requestObservation; + assert.equal(canonicalAttempts.length, 2); + const observation = canonicalAttempts[1]?.requestObservation; assert.ok(observation); assert.ok(observation.segments.length <= PREPARED_REQUEST_OBSERVATION_MAX_SEGMENTS); assert.ok(observation.segments.length > 0); assert.ok(observation.segments.every((segment) => segment.comparison === 'exact')); + const restarted = await readLatestContextDiagnostics(reopened, session.id); + assert.equal(restarted.status, 'available'); + if (restarted.status !== 'available') return; + assert.deepEqual(restarted.requestPreservation, successor.requestPreservation); + let coldScans = 0; const cold = await readLatestContextDiagnostics( { @@ -168,6 +200,8 @@ test('a real send seals its observation into SQLite and reconstructs it after re }, repairEventProjection: (sessionId, type, event, options) => reopened.repairEventProjection(sessionId, type, event, options), + readRootTurnAdmission: (sessionId, turnId) => + reopened.readRootTurnAdmission(sessionId, turnId), }, session.id, ); @@ -175,7 +209,8 @@ test('a real send seals its observation into SQLite and reconstructs it after re assert.ok(coldScans > 0, 'omitting the projection reader forces a restart-safe ledger fold'); assert.equal(cold.status, 'available'); if (cold.status !== 'available') return; - assert.deepEqual(cold.composition, diagnostics.composition); + assert.deepEqual(cold.composition, successor.composition); + assert.deepEqual(cold.requestPreservation, successor.requestPreservation); } finally { reopened.close?.(); } @@ -293,3 +328,24 @@ function answeringModel(): MockLanguageModelV4 { }), }); } + +async function admitTurn( + runStore: ReturnType, + sessionId: string, + turnId: string, + runId: string, + previousRootTurnId: string | null, + admittedAt: number, +): Promise { + await runStore.admitRootTurn({ + sessionId, + turnId, + proposedRunId: runId, + proposedUserMessageId: `message-${turnId}`, + execution: { kind: 'external_message' }, + previousRootTurnId, + normalizedInput: { text: turnId }, + sourceMessages: [], + admittedAt, + }); +} diff --git a/packages/runtime/src/__tests__/request-preservation.test.ts b/packages/runtime/src/__tests__/request-preservation.test.ts new file mode 100644 index 0000000000..3df09e56ad --- /dev/null +++ b/packages/runtime/src/__tests__/request-preservation.test.ts @@ -0,0 +1,442 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import { test } from 'node:test'; +import type { + ModelCallAttempt, + PreparedRequestObservation, + PreparedRequestObservationSegment, +} from '@maka/core/model-call-attempt'; +import type { AgentRunEvent, AgentRunHeader } from '@maka/core/agent-run'; +import { + deriveAttemptRequestPreservation, + deriveRequestPreservation, +} from '../request-preservation.js'; + +test('keeps every earlier cacheable segment when the current request only appends', () => { + const previous = observation([message(0, 'system'), message(1, 'user-1')]); + const current = observation([ + message(0, 'system'), + message(1, 'user-1'), + message(2, 'assistant-1'), + ]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'preserved', + previousSegmentCount: 2, + preservedSegmentCount: 2, + }); +}); + +test('does not treat opaque digests as evidence of divergence', () => { + const previous = observation([{ ...message(0, 'redacted-a'), comparison: 'opaque' }]); + const current = observation([{ ...message(0, 'redacted-b'), comparison: 'opaque' }]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'unknown', + previousSegmentCount: 1, + preservedSegmentCount: 0, + }); +}); + +test('reports the first changed earlier segment', () => { + const previous = observation([message(0, 'system'), message(1, 'user-1')]); + const current = observation([message(0, 'system'), message(1, 'edited-user-1')]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'diverged', + previousSegmentCount: 2, + preservedSegmentCount: 1, + firstDivergentSegment: { kind: 'message', index: 1 }, + }); +}); + +test('reports the first removed earlier segment', () => { + const previous = observation([message(0, 'system'), message(1, 'user-1')]); + const current = observation([message(0, 'system')]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'diverged', + previousSegmentCount: 2, + preservedSegmentCount: 1, + firstDivergentSegment: { kind: 'message', index: 1 }, + }); +}); + +test('reports the current segment at the first middle deletion', () => { + const previous = observation([ + message(0, 'system'), + message(1, 'user-1'), + message(2, 'assistant-1'), + ]); + const current = observation([message(0, 'system'), message(2, 'assistant-1')]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'diverged', + previousSegmentCount: 3, + preservedSegmentCount: 1, + firstDivergentSegment: { kind: 'message', index: 2 }, + }); +}); + +test('reports an inserted segment as the first divergence', () => { + const previous = observation([message(0, 'system'), message(2, 'assistant-1')]); + const current = observation([ + message(0, 'system'), + message(1, 'inserted-user'), + message(2, 'assistant-1'), + ]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'diverged', + previousSegmentCount: 2, + preservedSegmentCount: 1, + firstDivergentSegment: { kind: 'message', index: 1 }, + }); +}); + +test('reports the first moved segment after a reorder', () => { + const previous = observation([ + message(0, 'system'), + message(1, 'user-1'), + message(2, 'assistant-1'), + ]); + const current = observation([ + message(0, 'system'), + message(2, 'assistant-1'), + message(1, 'user-1'), + ]); + + assert.deepEqual(deriveRequestPreservation(current, previous), { + status: 'diverged', + previousSegmentCount: 3, + preservedSegmentCount: 1, + firstDivergentSegment: { kind: 'message', index: 2 }, + }); +}); + +test('uses the preceding physical retry as its durable baseline', async () => { + const previous = attempt({ attemptId: 'attempt-0', attempt: 0 }); + const current = attempt({ attemptId: 'attempt-1', attempt: 1 }); + assert.equal( + await statusOf(current, [run('run-1', 'turn-1', PROVIDER_STATE)], [previous]), + 'preserved', + ); +}); + +test('does not compare across provider execution identities', async () => { + const previous = attempt({ attemptId: 'previous', runId: 'run-1', turnId: 'turn-1' }); + const current = attempt({ attemptId: 'current', runId: 'run-2', turnId: 'turn-2' }); + + assert.equal( + await statusOf( + current, + [run('run-1', 'turn-1', OTHER_PROVIDER_STATE), run('run-2', 'turn-2', PROVIDER_STATE)], + [previous], + { runId: 'run-2', previousRootTurnId: 'turn-1' }, + ), + 'unavailable', + ); +}); + +test('does not compare attempts whose connection identity is missing', async () => { + const previous = attempt({ attemptId: 'previous', runId: 'run-1', turnId: 'turn-1' }); + const current = attempt({ + attemptId: 'current', + runId: 'run-2', + turnId: 'turn-2', + connectionSlug: undefined, + }); + previous.connectionSlug = undefined; + + assert.equal( + await statusOf( + current, + [run('run-1', 'turn-1', PROVIDER_STATE), run('run-2', 'turn-2', PROVIDER_STATE)], + [previous], + { runId: 'run-2', previousRootTurnId: 'turn-1' }, + ), + 'unavailable', + ); +}); + +test('does not use an attempt whose durable turn disagrees with its run', async () => { + const previous = attempt({ attemptId: 'previous', runId: 'run-1', turnId: 'wrong-turn' }); + const current = attempt({ attemptId: 'current', runId: 'run-2', turnId: 'turn-2' }); + + assert.equal( + await statusOf( + current, + [run('run-1', 'turn-1', PROVIDER_STATE), run('run-2', 'turn-2', PROVIDER_STATE)], + [previous], + { runId: 'run-2', previousRootTurnId: 'turn-1' }, + ), + 'unavailable', + ); +}); + +test('does not choose between overlapping durable predecessor runs', async () => { + const current = attempt({ attemptId: 'current', runId: 'run-3', turnId: 'turn-2' }); + const headers = [ + run('run-1', 'turn-1', PROVIDER_STATE), + run('run-2', 'turn-1', PROVIDER_STATE), + run('run-3', 'turn-2', PROVIDER_STATE), + ]; + + assert.equal( + await statusOf( + current, + headers, + [ + attempt({ attemptId: 'previous-1', runId: 'run-1', turnId: 'turn-1' }), + attempt({ attemptId: 'previous-2', runId: 'run-2', turnId: 'turn-1' }), + ], + { runId: 'run-3', previousRootTurnId: 'turn-1' }, + ), + 'unavailable', + ); +}); + +test('uses the unique continuation tip for the previous root turn', async () => { + const previous = attempt({ attemptId: 'continued', runId: 'run-2', turnId: 'turn-1' }); + const current = attempt({ attemptId: 'current', runId: 'run-3', turnId: 'turn-2' }); + + assert.equal( + await statusOf( + current, + [ + run('run-1', 'turn-1', PROVIDER_STATE), + run('run-2', 'turn-1', PROVIDER_STATE, { + parentRunId: 'run-1', + continuationSource: continuationSource('run-1'), + }), + run('run-3', 'turn-2', PROVIDER_STATE), + ], + [previous], + { runId: 'run-3', previousRootTurnId: 'turn-1' }, + ), + 'preserved', + ); +}); + +test('leaves an admitted predecessor unresolved until its canonical attempt is durable', async () => { + const current = attempt({ attemptId: 'current', runId: 'run-2', turnId: 'turn-2' }); + + assert.equal( + await deriveAttemptRequestPreservation({ + current, + store: preservationStore( + [ + run('run-1', 'turn-1', PROVIDER_STATE, { status: 'running' }), + run('run-2', 'turn-2', PROVIDER_STATE), + ], + [], + { runId: 'run-2', previousRootTurnId: 'turn-1' }, + ), + }), + undefined, + ); +}); + +test('leaves durable lineage unresolved while an expected run header is not readable', async () => { + const current = attempt({ attemptId: 'current', runId: 'run-2', turnId: 'turn-2' }); + const admission = { runId: 'run-2', previousRootTurnId: 'turn-1' }; + + assert.equal( + await deriveAttemptRequestPreservation({ + current, + store: preservationStore([run('run-2', 'turn-2', PROVIDER_STATE)], [], admission), + }), + undefined, + ); + assert.equal( + await deriveAttemptRequestPreservation({ + current, + store: preservationStore([run('run-1', 'turn-1', PROVIDER_STATE)], [], admission), + }), + undefined, + ); +}); + +test('compares later local turns without inheriting a copied session baseline', async () => { + const previous = attempt({ attemptId: 'local-1', runId: 'run-1', turnId: 'turn-1' }); + const current = attempt({ attemptId: 'local-2', runId: 'run-2', turnId: 'turn-2' }); + const runs = [ + run('copied-run', 'copied-turn', PROVIDER_STATE), + run('run-1', 'turn-1', PROVIDER_STATE), + run('run-2', 'turn-2', PROVIDER_STATE), + ]; + + assert.equal( + await statusOf(previous, runs, [], { runId: 'run-1', previousRootTurnId: null }), + 'unavailable', + ); + + assert.equal( + await statusOf(current, runs, [previous], { + runId: 'run-2', + previousRootTurnId: 'turn-1', + }), + 'preserved', + ); +}); + +test('does not compare a child run against its parent session run', async () => { + const previous = attempt({ attemptId: 'parent-attempt', runId: 'run-parent' }); + const current = attempt({ attemptId: 'child-attempt', runId: 'run-child' }); + assert.equal( + await statusOf( + current, + [ + run('run-parent', 'turn-1', PROVIDER_STATE), + run('run-child', 'turn-1', PROVIDER_STATE, { parentRunId: 'run-parent' }), + ], + [previous], + ), + 'unavailable', + ); +}); + +async function statusOf( + current: ModelCallAttempt, + runs: AgentRunHeader[], + attempts: ModelCallAttempt[], + admission?: { runId: string; previousRootTurnId: string | null }, +) { + return ( + await deriveAttemptRequestPreservation({ + current, + store: preservationStore(runs, attempts, admission), + }) + )?.status; +} + +function observation(segments: PreparedRequestObservationSegment[]): PreparedRequestObservation { + return { + schemaVersion: 1, + digest: digestFor(segments.map((segment) => segment.digest).join(':')), + bytes: 1, + segments, + }; +} + +function message(index: number, digest: string): PreparedRequestObservationSegment { + return { + kind: 'message', + index, + cacheable: true, + comparison: 'exact', + digest: digestFor(digest), + bytes: 1, + }; +} + +function digestFor(value: string): string { + return `sha256:${createHash('sha256').update(value).digest('hex')}`; +} + +const PROVIDER_STATE = `sha256:${'1'.repeat(64)}` as const; +const OTHER_PROVIDER_STATE = `sha256:${'2'.repeat(64)}` as const; + +function attempt(overrides: Partial): ModelCallAttempt { + return { + schemaVersion: 1, + logicalCallId: 'logical-1', + attemptId: 'attempt-0', + traceId: 'trace-1', + sessionId: 'session-1', + runId: 'run-1', + turnId: 'turn-1', + step: 0, + attempt: 0, + callKind: 'main', + connectionSlug: 'connection', + providerId: 'anthropic', + modelId: 'model', + requestObservation: observation([message(0, 'same')]), + startedAt: 1, + completedAt: 2, + latencyMs: 1, + status: 'completed', + usageBasis: 'missing', + costBasis: 'unpriced', + ...overrides, + }; +} + +function run( + runId: string, + turnId: string, + providerStateIdentity: `sha256:${string}`, + overrides: Partial = {}, +): AgentRunHeader { + return { + runId, + sessionId: 'session-1', + turnId, + status: 'completed', + backendKind: 'ai-sdk', + providerStateIdentity, + llmConnectionSlug: 'connection', + modelId: 'model', + cwd: '/repo', + permissionMode: 'ask', + createdAt: 1, + updatedAt: 2, + ...overrides, + }; +} + +function continuationSource( + sourceRunId: string, +): NonNullable { + return { + sourceInvocationId: sourceRunId, + sourceRunId, + sourceTurnId: 'turn-1', + sourceRuntimeEventHighWater: 1, + }; +} + +function preservationStore( + runs: AgentRunHeader[], + attempts: ModelCallAttempt[], + admission?: { runId: string; previousRootTurnId: string | null }, +) { + return { + listSessionRuns: async () => runs, + readEvents: async (_sessionId: string, runId: string) => + attempts.filter((item) => item.runId === runId).map(attemptEvent), + readRootTurnAdmission: async () => admission, + }; +} + +function attemptEvent(attempt: ModelCallAttempt): AgentRunEvent { + return { + type: 'model_call_attempt_recorded', + id: attempt.attemptId, + runId: attempt.runId, + sessionId: attempt.sessionId, + turnId: attempt.turnId, + ts: attempt.completedAt, + data: { ...attempt }, + }; +} diff --git a/packages/runtime/src/agent-run.ts b/packages/runtime/src/agent-run.ts index cff8aaed16..d40602f177 100644 --- a/packages/runtime/src/agent-run.ts +++ b/packages/runtime/src/agent-run.ts @@ -91,6 +91,7 @@ import { import { DEFAULT_TOOL_MODE, isToolMode, type ToolMode } from '@maka/core/tool-mode'; import { materializeRuntimeEventTranscriptProjection } from './runtime-ledger-repair.js'; import { cloneAndFreezeRuntimeSnapshot } from './runtime-snapshot.js'; +import { deriveAttemptRequestPreservation } from './request-preservation.js'; export interface AgentRunActiveSession { sessionId: string; @@ -474,7 +475,25 @@ export class AgentRun { const { attempt, latestContext } = commit; if (!this.input.runStore) return Promise.resolve(); return this.enqueueRequiredRunStoreWrite('append model call attempt', async () => { - await this.input.runStore?.appendEvent( + const runStore = this.input.runStore; + if (!runStore) return; + let projected = latestContext; + if (projected && attempt.callKind === 'main') { + const requestPreservation = await deriveAttemptRequestPreservation({ + current: attempt, + store: runStore, + }).catch(() => undefined); + projected = { + ...projected, + snapshot: { + ...projected.snapshot, + ...(requestPreservation + ? { requestPreservation } + : { requestPreservationPending: true }), + }, + }; + } + await runStore.appendEvent( this.sessionId, this.runId, { @@ -489,7 +508,7 @@ export class AgentRun { // The latest-context projection rides this durable append rather than // racing it: one commit for the request, and derived state that cannot // survive a metering write that failed (#2323). - { durable: true, ...(latestContext ? { latestContext } : {}) }, + { durable: true, ...(projected ? { latestContext: projected } : {}) }, ); }); } diff --git a/packages/runtime/src/context-diagnostics.ts b/packages/runtime/src/context-diagnostics.ts index bc45466ab0..756f107332 100644 --- a/packages/runtime/src/context-diagnostics.ts +++ b/packages/runtime/src/context-diagnostics.ts @@ -21,6 +21,7 @@ import { isSessionInlineRun, supersedesLatestContext, type AgentRunEvent, + type AgentRunHeader, type AgentRunStore, } from '@maka/core/agent-run'; import { decodeModelCallAttempt, type ModelCallAttempt } from '@maka/core/model-call-attempt'; @@ -39,6 +40,13 @@ import { validateHistoryCompactCheckpointShape, type HistoryCompactCheckpoint, } from './history-compact-checkpoint.js'; +import { + deriveAttemptRequestPreservation, + isCanonicalAttemptForRun, + type RequestPreservation, +} from './request-preservation.js'; + +export { isRequestPreservation, type RequestPreservation } from './request-preservation.js'; export type ContextDiagnosticsUnavailableReason = 'no_completed_request' | 'trace_unavailable'; @@ -98,6 +106,8 @@ export type ContextDiagnostics = */ composition?: ContextDiagnosticsComposition; compaction?: ContextDiagnosticsCompaction; + /** Absent only for historical projections written before this diagnostic existed. */ + requestPreservation?: RequestPreservation; }; export interface ContextDiagnosticsComposition { @@ -117,7 +127,12 @@ type ContextRunStore = Pick< | 'readEventProjection' | 'readEventLedgerRevision' | 'repairEventProjection' ->; +> & { + readRootTurnAdmission?( + sessionId: string, + turnId: string, + ): Promise<{ runId: string; previousRootTurnId?: string | null } | undefined>; +}; /** * What the session's context is made of right now (#1580, reshaped for #2323). @@ -158,7 +173,7 @@ export async function readLatestContextDiagnostics( .catch(() => undefined); if (projected === null) return { status: 'unavailable', reason: 'no_completed_request' }; const snapshot = readLatestContextSnapshot(projected ?? undefined); - if (snapshot) return availableFrom(snapshot); + if (snapshot && !snapshot.requestPreservationPending) return availableFrom(snapshot); if (projected) replaceProjectionId = projected.id; } const ledgerRevision = @@ -211,7 +226,7 @@ async function rebuildContextFromLedger( for (const event of await runStore.readEvents(sessionId, run.runId)) { if (event.type === METERING_EVENT_TYPE) { sawCanonicalRecord = true; - const candidate = meteringAnchor(event); + const candidate = meteringAnchor(event, run); if (candidate && supersedesLatestContext(candidate, anchor)) anchor = candidate; continue; } @@ -247,6 +262,13 @@ async function rebuildContextFromLedger( (anchor && !anchor.hasRequestObservation ? exactHistoricalComposition(anchor, historicalAttempts) : undefined); + const requestPreservation = + anchor?.attempt && anchor.run + ? await deriveAttemptRequestPreservation({ + current: anchor.attempt, + store: runStore, + }) + : undefined; const snapshot: LatestContextSnapshot = { schemaVersion: LATEST_CONTEXT_SNAPSHOT_SCHEMA_VERSION, attemptId: resolved.attemptId, @@ -260,6 +282,11 @@ async function rebuildContextFromLedger( ...(resolved.contextWindow !== undefined ? { contextWindow: resolved.contextWindow } : {}), ...(composition ? { composition } : {}), ...(boundary ? { compaction: contextDiagnosticsCompactionOf(boundary.checkpoint) } : {}), + ...(requestPreservation + ? { requestPreservation } + : anchor?.attempt + ? { requestPreservationPending: true as const } + : {}), }; // Repair on the way out, so this scan happens once per session rather than // on every panel refresh. Best-effort: the caller already has its answer, @@ -345,6 +372,15 @@ function legacyProviderAnchor(event: AgentRunEvent): MeteringAnchor | undefined type LegacyProviderAnchor = MeteringAnchor; function availableFrom(snapshot: LatestContextSnapshot): ContextDiagnostics { + const requestPreservation = + snapshot.requestPreservation ?? + (snapshot.requestPreservationPending + ? { + status: 'unavailable' as const, + previousSegmentCount: 0 as const, + preservedSegmentCount: 0 as const, + } + : undefined); return { status: 'available', providerId: snapshot.providerId, @@ -357,6 +393,7 @@ function availableFrom(snapshot: LatestContextSnapshot): ContextDiagnostics { ...(snapshot.contextWindow !== undefined ? { contextWindow: snapshot.contextWindow } : {}), ...(snapshot.composition ? { composition: snapshot.composition } : {}), ...(snapshot.compaction ? { compaction: snapshot.compaction } : {}), + ...(requestPreservation ? { requestPreservation } : {}), }; } @@ -375,6 +412,8 @@ interface MeteringAnchor { contextWindow?: number; composition?: ContextDiagnosticsComposition; hasRequestObservation: boolean; + attempt?: ModelCallAttempt; + run?: AgentRunHeader; } interface CheckpointCandidate { @@ -384,14 +423,20 @@ interface CheckpointCandidate { } /** Only a completed MAIN call describes the conversation's own context. */ -function meteringAnchor(event: AgentRunEvent): MeteringAnchor | undefined { +function meteringAnchor(event: AgentRunEvent, run: AgentRunHeader): MeteringAnchor | undefined { let attempt: ModelCallAttempt; try { attempt = decodeModelCallAttempt(event.data); } catch { return undefined; } - if (attempt.callKind !== 'main' || attempt.status !== 'completed') return undefined; + if ( + attempt.callKind !== 'main' || + attempt.status !== 'completed' || + !isCanonicalAttemptForRun(event, attempt, run) + ) { + return undefined; + } const composition = attempt.requestObservation ? foldPromptComposition(attempt.requestObservation.segments) : undefined; @@ -403,6 +448,8 @@ function meteringAnchor(event: AgentRunEvent): MeteringAnchor | undefined { startedAt: attempt.startedAt, completedAt: attempt.completedAt, hasRequestObservation: attempt.requestObservation !== undefined, + attempt, + run, ...(attempt.inputTokens !== undefined ? { inputTokens: attempt.inputTokens } : {}), ...(attempt.cacheReadInputTokens !== undefined ? { cacheReadInputTokens: attempt.cacheReadInputTokens } diff --git a/packages/runtime/src/latest-context-snapshot.ts b/packages/runtime/src/latest-context-snapshot.ts index d2f53652d6..080d0a8e8c 100644 --- a/packages/runtime/src/latest-context-snapshot.ts +++ b/packages/runtime/src/latest-context-snapshot.ts @@ -29,6 +29,7 @@ import type { ContextDiagnosticsComposition, } from './context-diagnostics.js'; import { foldPromptComposition, type SizedRequestSegment } from './prompt-composition.js'; +import { isRequestPreservation, type RequestPreservation } from './request-preservation.js'; /** * One request's context, frozen by the transaction that committed it (#2323). @@ -67,6 +68,10 @@ export interface LatestContextSnapshot { composition?: ContextDiagnosticsComposition; /** The boundary that applied when this request was built, if any. */ compaction?: ContextDiagnosticsCompaction; + /** Runtime-owned conclusion; consumers must not select or compare observations. */ + requestPreservation?: RequestPreservation; + /** Internal retry marker; never projected through Runtime Host. */ + requestPreservationPending?: true; } /** @@ -129,7 +134,15 @@ export function readLatestContextSnapshot( const record = shapedRecord( event.data, ['schemaVersion', 'attemptId', 'providerId', 'modelId', 'completedAt'], - ['inputTokens', 'cacheReadInputTokens', 'contextWindow', 'composition', 'compaction'], + [ + 'inputTokens', + 'cacheReadInputTokens', + 'contextWindow', + 'composition', + 'compaction', + 'requestPreservation', + 'requestPreservationPending', + ], ); if (!record) return undefined; if ( @@ -143,7 +156,12 @@ export function readLatestContextSnapshot( (record.contextWindow !== undefined && (!isCount(record.contextWindow) || record.contextWindow === 0)) || (record.composition !== undefined && !isContextDiagnosticsComposition(record.composition)) || - (record.compaction !== undefined && !isContextDiagnosticsCompaction(record.compaction)) + (record.compaction !== undefined && !isContextDiagnosticsCompaction(record.compaction)) || + (record.requestPreservation !== undefined && + !isRequestPreservation(record.requestPreservation)) || + (record.requestPreservationPending !== undefined && + record.requestPreservationPending !== true) || + (record.requestPreservation !== undefined && record.requestPreservationPending === true) ) { return undefined; } diff --git a/packages/runtime/src/request-preservation.ts b/packages/runtime/src/request-preservation.ts new file mode 100644 index 0000000000..07daa2f96b --- /dev/null +++ b/packages/runtime/src/request-preservation.ts @@ -0,0 +1,451 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import type { + ModelCallAttempt, + PreparedRequestObservation, + PreparedRequestObservationSegment, +} from '@maka/core/model-call-attempt'; +import { decodeModelCallAttempt } from '@maka/core/model-call-attempt'; +import { + isSessionInlineRun, + type AgentRunEvent, + type AgentRunHeader, + type AgentRunStore, +} from '@maka/core/agent-run'; + +export type RequestPreservation = + | { + status: 'no_predecessor' | 'unavailable'; + previousSegmentCount: 0; + preservedSegmentCount: 0; + } + | { + status: 'preserved' | 'unknown'; + previousSegmentCount: number; + preservedSegmentCount: number; + } + | { + status: 'diverged'; + previousSegmentCount: number; + preservedSegmentCount: number; + firstDivergentSegment: RequestPreservationSegmentRef; + }; + +export type RequestPreservationSegmentRef = Pick< + PreparedRequestObservationSegment, + 'kind' | 'index' | 'role' | 'label' +>; + +type PreservationStore = Pick; + +interface RootAdmissionReader { + readRootTurnAdmission?( + sessionId: string, + turnId: string, + ): Promise<{ runId: string; previousRootTurnId?: string | null } | undefined>; +} + +interface AttemptPreservationInput { + current: ModelCallAttempt; + store: PreservationStore; +} + +const PREDECESSOR_PENDING = Symbol('predecessor_pending'); +type Predecessor = { attempt: ModelCallAttempt; run: AgentRunHeader }; +type PredecessorResolution = Predecessor | null | undefined | typeof PREDECESSOR_PENDING; +interface AttemptCandidates { + attempts: ModelCallAttempt[]; + sawInvalidAttemptEvent: boolean; +} + +export async function deriveAttemptRequestPreservation( + input: AttemptPreservationInput, +): Promise { + const { current, store } = input; + if (current.callKind !== 'main' || !current.requestObservation) { + return unavailable(); + } + const runs = await store.listSessionRuns(current.sessionId); + const currentRun = runs.find((run) => run.runId === current.runId); + if (!currentRun) return undefined; + if (!isSessionInlineRun(currentRun)) return unavailable(); + + const predecessor = await predecessorAttempt(current, currentRun, runs, store); + if (predecessor === PREDECESSOR_PENDING) return undefined; + if (predecessor === null) { + return { status: 'no_predecessor', previousSegmentCount: 0, preservedSegmentCount: 0 }; + } + if (!predecessor || !sameDomain(current, currentRun, predecessor)) return unavailable(); + return deriveRequestPreservation( + current.requestObservation, + predecessor.attempt.requestObservation!, + ); +} + +export function deriveRequestPreservation( + currentObservation: PreparedRequestObservation, + previousObservation: PreparedRequestObservation, +): RequestPreservation { + const current = currentObservation.segments.filter((segment) => segment.cacheable); + const previous = previousObservation.segments.filter((segment) => segment.cacheable); + const previousSegmentCount = representedCount(previous); + + for (let index = 0; index < previous.length; index += 1) { + const before = previous[index]!; + const after = current[index]; + if (!after || !sameIdentity(before, after)) { + return { + status: 'diverged', + previousSegmentCount, + preservedSegmentCount: representedCount(previous.slice(0, index)), + firstDivergentSegment: segmentRef(after ?? before), + }; + } + if (before.comparison === 'opaque' || after.comparison === 'opaque') { + return { + status: 'unknown', + previousSegmentCount, + preservedSegmentCount: representedCount(previous.slice(0, index)), + }; + } + if (before.digest !== after.digest) { + return { + status: 'diverged', + previousSegmentCount, + preservedSegmentCount: representedCount(previous.slice(0, index)), + firstDivergentSegment: segmentRef(after), + }; + } + } + + return { + status: 'preserved', + previousSegmentCount, + preservedSegmentCount: previousSegmentCount, + }; +} + +function sameIdentity( + left: PreparedRequestObservationSegment, + right: PreparedRequestObservationSegment, +): boolean { + return ( + left.kind === right.kind && + left.index === right.index && + left.role === right.role && + left.label === right.label + ); +} + +function representedCount(segments: readonly PreparedRequestObservationSegment[]): number { + return segments.reduce((count, segment) => count + (segment.representedSegments ?? 1), 0); +} + +function segmentRef(segment: PreparedRequestObservationSegment): RequestPreservationSegmentRef { + return { + kind: segment.kind, + index: segment.index, + ...(segment.role === undefined ? {} : { role: segment.role }), + ...(segment.label === undefined ? {} : { label: segment.label }), + }; +} + +async function predecessorAttempt( + current: ModelCallAttempt, + currentRun: AgentRunHeader, + runs: readonly AgentRunHeader[], + store: PreservationStore, +): Promise { + if (current.attempt > 0) { + return expectedAttempt( + currentRun, + await attemptsFor(store, currentRun), + (candidate) => + candidate.logicalCallId === current.logicalCallId && + candidate.attempt === current.attempt - 1, + ); + } + if (current.step > 0) { + return expectedLatestAttempt( + currentRun, + filterAttempts( + await attemptsFor(store, currentRun), + (candidate) => candidate.turnId === current.turnId && candidate.step === current.step - 1, + ), + ); + } + if (currentRun.parentRunId) { + const parent = runs.find((run) => run.runId === currentRun.parentRunId); + return parent + ? expectedLatestAttempt(parent, await attemptsFor(store, parent)) + : PREDECESSOR_PENDING; + } + if ( + currentRun.parentSessionId || + currentRun.parentTurnId || + currentRun.retriedFromTurnId || + currentRun.regeneratedFromTurnId || + currentRun.branchOfTurnId + ) { + return undefined; + } + + const admission = await ( + store as PreservationStore & RootAdmissionReader + ).readRootTurnAdmission?.(current.sessionId, current.turnId); + if ( + !admission || + admission.runId !== current.runId || + admission.previousRootTurnId === undefined + ) { + return undefined; + } + if (admission.previousRootTurnId === null) { + return runs.some((run) => run.runId !== current.runId && isSessionInlineRun(run)) + ? undefined + : null; + } + const previousRuns = runs.filter( + (run) => run.turnId === admission.previousRootTurnId && isSessionInlineRun(run), + ); + if (previousRuns.length === 0) return PREDECESSOR_PENDING; + const previous = uniqueDurableRunTip(previousRuns); + return previous ? expectedLatestAttempt(previous, await attemptsFor(store, previous)) : undefined; +} + +function uniqueDurableRunTip(runs: readonly AgentRunHeader[]): AgentRunHeader | undefined { + const byId = new Map(runs.map((run) => [run.runId, run])); + if (byId.size !== runs.length) return undefined; + const childByParent = new Map(); + for (const run of runs) { + if (!run.parentRunId || !byId.has(run.parentRunId)) continue; + if (childByParent.has(run.parentRunId)) return undefined; + childByParent.set(run.parentRunId, run.runId); + } + const tips = runs.filter((run) => !childByParent.has(run.runId)); + if (tips.length !== 1) return undefined; + + const visited = new Set(); + let cursor: AgentRunHeader | undefined = tips[0]; + while (cursor) { + if (visited.has(cursor.runId)) return undefined; + visited.add(cursor.runId); + cursor = cursor.parentRunId ? byId.get(cursor.parentRunId) : undefined; + } + return visited.size === runs.length ? tips[0] : undefined; +} + +async function attemptsFor( + store: PreservationStore, + run: AgentRunHeader, +): Promise { + const attempts: ModelCallAttempt[] = []; + let sawInvalidAttemptEvent = false; + for (const event of await store.readEvents(run.sessionId, run.runId)) { + if (event.type !== 'model_call_attempt_recorded') continue; + try { + const attempt = decodeModelCallAttempt(event.data); + if (attempt.callKind === 'main' && isCanonicalAttemptForRun(event, attempt, run)) { + attempts.push(attempt); + } else { + sawInvalidAttemptEvent = true; + } + } catch { + // An unreadable attempt cannot become a guessed baseline. + sawInvalidAttemptEvent = true; + } + } + return { attempts, sawInvalidAttemptEvent }; +} + +function filterAttempts( + candidates: AttemptCandidates, + predicate: (attempt: ModelCallAttempt) => boolean, +): AttemptCandidates { + return { ...candidates, attempts: candidates.attempts.filter(predicate) }; +} + +export function isCanonicalAttemptForRun( + event: AgentRunEvent, + attempt: ModelCallAttempt, + run: AgentRunHeader, +): boolean { + return ( + event.sessionId === run.sessionId && + event.runId === run.runId && + event.turnId === run.turnId && + attempt.sessionId === run.sessionId && + attempt.runId === run.runId && + attempt.turnId === run.turnId && + attempt.attemptId === event.id + ); +} + +function latestAttempt( + run: AgentRunHeader, + attempts: readonly ModelCallAttempt[], +): { attempt: ModelCallAttempt; run: AgentRunHeader } | undefined { + if (attempts.length === 0) return undefined; + const step = Math.max(...attempts.map((attempt) => attempt.step)); + const onStep = attempts.filter((attempt) => attempt.step === step); + const physical = Math.max(...onStep.map((attempt) => attempt.attempt)); + return uniqueAttempt(run, onStep, (attempt) => attempt.attempt === physical); +} + +function expectedLatestAttempt( + run: AgentRunHeader, + candidates: AttemptCandidates, +): Predecessor | undefined | typeof PREDECESSOR_PENDING { + if (candidates.attempts.length === 0) { + return candidates.sawInvalidAttemptEvent || isTerminalRun(run) + ? undefined + : PREDECESSOR_PENDING; + } + return latestAttempt(run, candidates.attempts); +} + +function expectedAttempt( + run: AgentRunHeader, + candidates: AttemptCandidates, + predicate: (attempt: ModelCallAttempt) => boolean, +): Predecessor | undefined | typeof PREDECESSOR_PENDING { + const matches = candidates.attempts.filter(predicate); + if (matches.length === 0) { + return candidates.sawInvalidAttemptEvent || isTerminalRun(run) + ? undefined + : PREDECESSOR_PENDING; + } + return matches.length === 1 ? { attempt: matches[0]!, run } : undefined; +} + +function isTerminalRun(run: AgentRunHeader): boolean { + return run.status === 'completed' || run.status === 'failed' || run.status === 'cancelled'; +} + +function uniqueAttempt( + run: AgentRunHeader, + attempts: readonly ModelCallAttempt[], + predicate: (attempt: ModelCallAttempt) => boolean, +): { attempt: ModelCallAttempt; run: AgentRunHeader } | undefined { + const matches = attempts.filter(predicate); + return matches.length === 1 ? { attempt: matches[0]!, run } : undefined; +} + +function sameDomain( + current: ModelCallAttempt, + currentRun: AgentRunHeader, + previous: { attempt: ModelCallAttempt; run: AgentRunHeader }, +): boolean { + const before = previous.attempt; + if (!before.requestObservation) return false; + const currentPartition = exactProviderPartition(current.requestObservation!); + const previousPartition = exactProviderPartition(before.requestObservation); + return ( + currentRun.providerStateIdentity !== undefined && + currentRun.providerStateIdentity === previous.run.providerStateIdentity && + current.sessionId === before.sessionId && + current.connectionSlug !== undefined && + current.connectionSlug === before.connectionSlug && + current.providerId === before.providerId && + current.modelId === before.modelId && + currentPartition !== undefined && + currentPartition === previousPartition + ); +} + +function exactProviderPartition(observation: PreparedRequestObservation): string | undefined { + const segments = observation.segments.filter((segment) => segment.kind === 'provider_options'); + if (segments.some((segment) => segment.comparison === 'opaque')) return undefined; + return segments.map((segment) => `${segment.index}:${segment.digest}`).join('|'); +} + +function unavailable(): RequestPreservation { + return { status: 'unavailable', previousSegmentCount: 0, preservedSegmentCount: 0 }; +} + +export function isRequestPreservation(value: unknown): value is RequestPreservation { + if (!value || typeof value !== 'object') return false; + const candidate = value as { + status?: unknown; + previousSegmentCount?: unknown; + preservedSegmentCount?: unknown; + firstDivergentSegment?: unknown; + }; + if ( + !hasOnlyKeys(candidate, [ + 'status', + 'previousSegmentCount', + 'preservedSegmentCount', + 'firstDivergentSegment', + ]) + ) { + return false; + } + if ( + typeof candidate.previousSegmentCount !== 'number' || + typeof candidate.preservedSegmentCount !== 'number' || + !Number.isSafeInteger(candidate.previousSegmentCount) || + !Number.isSafeInteger(candidate.preservedSegmentCount) || + candidate.previousSegmentCount < 0 || + candidate.preservedSegmentCount < 0 || + candidate.preservedSegmentCount > candidate.previousSegmentCount + ) { + return false; + } + if (candidate.status === 'no_predecessor' || candidate.status === 'unavailable') { + return ( + candidate.firstDivergentSegment === undefined && + candidate.previousSegmentCount === 0 && + candidate.preservedSegmentCount === 0 + ); + } + if (candidate.status === 'preserved') { + return ( + candidate.firstDivergentSegment === undefined && + candidate.preservedSegmentCount === candidate.previousSegmentCount + ); + } + if (candidate.status === 'unknown') return candidate.firstDivergentSegment === undefined; + if (candidate.status !== 'diverged') return false; + const segment = candidate.firstDivergentSegment as + | { kind?: unknown; index?: unknown; role?: unknown; label?: unknown } + | undefined; + return Boolean( + segment && + hasOnlyKeys(segment, ['kind', 'index', 'role', 'label']) && + (segment.kind === 'tool_schema' || + segment.kind === 'system_prompt' || + segment.kind === 'message' || + segment.kind === 'provider_options') && + typeof segment.index === 'number' && + Number.isSafeInteger(segment.index) && + segment.index >= 0 && + (segment.role === undefined || isBoundedText(segment.role)) && + (segment.label === undefined || isBoundedText(segment.label)), + ); +} + +function hasOnlyKeys(value: object, allowed: readonly string[]): boolean { + return Object.keys(value).every((key) => allowed.includes(key)); +} + +function isBoundedText(value: unknown): value is string { + return typeof value === 'string' && value.length <= 256; +}