fix(core): Observe agent execution tail when observation cursor moved in the middle of the run (#38907)

This commit is contained in:
Anne Aguirre 2026-09-17 12:18:00 +00:00 committed by GitHub
parent 74334d2d85
commit d1c805a595
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 159 additions and 8 deletions

View File

@ -445,7 +445,7 @@ describe('MemoryOrchestrator.saveToMemory observer gating', () => {
expect(await store.getCursor(THREAD_ID)).not.toBeNull();
});
it('deduplicates an in-flight mid-run task without blocking later post-turn observation', async () => {
it('queues the tail observer behind an in-flight mid-run task at end of turn', async () => {
const store = new InMemoryMemory();
const pendingObservation = deferred<string>();
const observe = vi.fn(async () => await pendingObservation.promise);
@ -457,19 +457,155 @@ describe('MemoryOrchestrator.saveToMemory observer gating', () => {
const list = new AgentMessageList();
list.addInput([userMsg('x'.repeat(750))]);
// Soft crossing schedules the mid-run observer in the background.
await orchestrator.maybeObserveMidRun(list, runOptions());
await vi.waitFor(() => expect(observe).toHaveBeenCalledTimes(1));
// The post-boundary tail is persisted at end of turn. The mid-run task
// only covers up to its boundary, so the tail observer must queue
// behind it on the per-scope runner instead of being joined and dropped.
list.addResponse([assistantMsg('y'.repeat(750))]);
await orchestrator.saveToMemory(list, runOptions());
pendingObservation.resolve('* CRITICAL (14:30) Observed.');
await tracker.flush();
// Mid-run observer plus the queued tail observer both ran, and the
// cursor now points at the turn's last message.
expect(observe).toHaveBeenCalledTimes(2);
const cursor = await store.getCursor(THREAD_ID);
expect(cursor?.lastObservedMessageId).toBe(list.messages().at(-1)?.id);
// The follow-up turn starts clean: loadInto resets the per-run flag and
// history loads since the cursor (now at the previous turn's last
// message), so no catch-up observer is needed and the gate applies.
const nextList = new AgentMessageList();
await orchestrator.loadInto(nextList, runOptions());
nextList.addInput([userMsg('z'.repeat(250))]);
await orchestrator.saveToMemory(nextList, runOptions());
await tracker.flush();
expect(observe).toHaveBeenCalledTimes(2);
});
it('schedules the tail observer after a settled, activated mid-run observation even when the tail is below threshold', async () => {
const store = new InMemoryMemory();
const observe = vi.fn(async () => await Promise.resolve('* CRITICAL (14:30) Observed.'));
const { orchestrator, tracker } = buildOrchestrator(store, {
observerThresholdTokens: 1000,
observe,
observationLogTailLimit: 20,
});
const list = new AgentMessageList();
list.addInput([userMsg('x'.repeat(750))]);
list.addResponse([assistantMsg('y'.repeat(750))]);
// Hard crossing runs the observer synchronously, advances the cursor,
// and masks the window — no in-flight mid-run task remains.
await orchestrator.maybeObserveMidRun(list, runOptions());
expect(observe).toHaveBeenCalledTimes(1);
expect(list.forLlm('base').messages).toEqual([
{ role: 'user', content: OBSERVATION_CONTINUATION_REMINDER },
]);
// A small post-boundary tail is added before the turn ends. Its budget
// is below the threshold, so the gate alone would skip it.
list.addResponse([assistantMsg('z'.repeat(200))]);
await orchestrator.saveToMemory(list, runOptions());
await tracker.flush();
// The per-run flag forces the tail observer past the budget gate.
expect(observe).toHaveBeenCalledTimes(2);
const cursor = await store.getCursor(THREAD_ID);
expect(cursor?.lastObservedMessageId).toBe(list.messages().at(-1)?.id);
});
it('schedules the tail observer after a soft mid-run task settled without a later boundary activating it', async () => {
const store = new InMemoryMemory();
const observe = vi.fn(async () => await Promise.resolve('* CRITICAL (14:30) Observed.'));
const { orchestrator, tracker } = buildOrchestrator(store, {
observerThresholdTokens: 1000,
observe,
observationLogTailLimit: 20,
});
const list = new AgentMessageList();
list.addInput([userMsg('x'.repeat(750))]);
// Soft crossing schedules in the background; the task settles and
// advances the cursor, but no later boundary activates it, so the
// window is never masked.
await orchestrator.maybeObserveMidRun(list, runOptions());
await tracker.flush();
expect(await store.getCursor(THREAD_ID)).not.toBeNull();
expect(list.forLlm('base').messages).toHaveLength(1);
// The whole unmasked turn (750 + 200) is still below the threshold, so
// the budget gate alone would skip the tail.
list.addResponse([assistantMsg('z'.repeat(200))]);
await orchestrator.saveToMemory(list, runOptions());
await tracker.flush();
expect(observe).toHaveBeenCalledTimes(2);
const cursor = await store.getCursor(THREAD_ID);
expect(cursor?.lastObservedMessageId).toBe(list.messages().at(-1)?.id);
});
it('detaches a settled mid-run task before gating a new turn', async () => {
const store = new InMemoryMemory();
const observe = vi.fn(async () => await Promise.resolve('* CRITICAL (14:30) Observed.'));
const { orchestrator, tracker } = buildOrchestrator(store, {
observerThresholdTokens: 1000,
observe,
observationLogTailLimit: 20,
});
const list = new AgentMessageList();
list.addInput([userMsg('x'.repeat(750))]);
await orchestrator.maybeObserveMidRun(list, runOptions());
await tracker.flush();
expect(observe).toHaveBeenCalledTimes(1);
const nextList = new AgentMessageList();
await orchestrator.loadInto(nextList, runOptions());
nextList.addInput([userMsg('z'.repeat(250))]);
nextList.addInput([userMsg('z'.repeat(100))]);
await orchestrator.saveToMemory(nextList, runOptions());
await tracker.flush();
expect(observe).toHaveBeenCalledTimes(1);
});
it('keeps the gate on the next turn when the previous mid-run task settles after it started', async () => {
const store = new InMemoryMemory();
const pendingObservation = deferred<string>();
const observe = vi
.fn(async () => await Promise.resolve('* CRITICAL (14:30) Tail observed.'))
.mockImplementationOnce(async () => await pendingObservation.promise);
const { orchestrator, tracker } = buildOrchestrator(store, {
observerThresholdTokens: 1000,
observe,
observationLogTailLimit: 20,
});
const list = new AgentMessageList();
list.addInput([userMsg('x'.repeat(750))]);
await orchestrator.maybeObserveMidRun(list, runOptions());
await vi.waitFor(() => expect(observe).toHaveBeenCalledTimes(1));
list.addResponse([assistantMsg('y'.repeat(50))]);
await orchestrator.saveToMemory(list, runOptions());
// The follow-up turn starts while the previous mid-run task is still in
// flight, so it loads the full previous turn (no cursor yet). The task
// settles only afterwards; that settlement must not mark the new run as
// having advanced the cursor.
const nextList = new AgentMessageList();
await orchestrator.loadInto(nextList, runOptions());
pendingObservation.resolve('* CRITICAL (14:30) Observed.');
await tracker.flush();
expect(observe).toHaveBeenCalledTimes(2);
// 800 chars of history plus 100 of input stay below the threshold: with
// no mid-run observation this turn, the gate applies and skips.
nextList.addInput([userMsg('z'.repeat(100))]);
await orchestrator.saveToMemory(nextList, runOptions());
await tracker.flush();

View File

@ -151,6 +151,12 @@ export class MemoryOrchestrator {
/** In-flight background mid-run observer task; `result` set on settlement. */
private midRunObserverTask: MidRunObserverTask | undefined;
/**
* Set when a mid-run observer advanced the cursor this run. The post-turn
* gate reads the masked window, so the tail must be observed unconditionally.
*/
private cursorAdvancedThisRun = false;
/** Consecutive mid-run observer attempts that did not advance the cursor. */
private midRunNonAdvancingAttempts = 0;
@ -475,10 +481,12 @@ export class MemoryOrchestrator {
/** Reset per-run observation state on generate and resume entry. */
private resetRunState(): void {
this.midRunObserverTask = undefined;
this.midRunNonAdvancingAttempts = 0;
this.lastPersistedTurnKeyset = undefined;
this.visibleTokenEstimates.clear();
this.visibleTokenEstimateTotal = 0;
this.cursorAdvancedThisRun = false;
}
/**
@ -492,6 +500,7 @@ export class MemoryOrchestrator {
): boolean {
if (didAdvanceCursor(result)) {
this.midRunNonAdvancingAttempts = 0;
this.cursorAdvancedThisRun = true;
return true;
}
this.midRunNonAdvancingAttempts += 1;
@ -695,12 +704,6 @@ export class MemoryOrchestrator {
const scope = this.getObservationLogScope(persistence);
const runner = this.getMemoryTaskRunner(memory, observationalMemory.lockTtlMs);
// A mid-run task still in flight for this scope already covers the
// messages persisted at its boundary: join it instead of queueing a
// second observer behind it — the post-boundary tail waits for the next
// turn's gate. A task that settled after the last boundary was never
// activated; the run is over, so drop it and let the gauge decide.
const midRunTask = this.midRunObserverTask;
if (midRunTask?.result) this.midRunObserverTask = undefined;
if (
@ -708,9 +711,21 @@ export class MemoryOrchestrator {
!midRunTask.result &&
midRunTask.handle.observationScopeId === scope.observationScopeId
) {
// A mid-run task still in flight only covers up to its boundary.
// Queue the tail observer behind it on the per-scope runner instead
// of joining and dropping the post-boundary tail.
void midRunTask.handle.done.then(() => {
if (this.midRunObserverTask === midRunTask) this.midRunObserverTask = undefined;
});
this.scheduleObserverTask(persistence, executionCounter, telemetry);
} else if (
this.cursorAdvancedThisRun ||
(midRunTask?.result !== undefined && didAdvanceCursor(midRunTask.result))
) {
// Mid-run advanced the cursor this run — either activated at a boundary
// or settled after the last one: observe the tail regardless of the
// visible-window budget, which the mask shrank below threshold.
this.scheduleObserverTask(persistence, executionCounter, telemetry);
} else if (await this.shouldScheduleObserver(list, persistence.threadId)) {
this.scheduleObserverTask(persistence, executionCounter, telemetry);
}