Skip to content

fix(runtime): enforce wall-clock budget while the event stream is silent #85

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Sthreal wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(runtime): enforce wall-clock budget while the event stream is silent #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
206 changes: 203 additions & 3 deletions packages/shared/src/kernel/run-manager.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,30 @@ import { RunManager, type RunManagerOptions } from './run-manager.ts';
let dir = '';
let clock = 1000;

interface ManualTimerEntry {
callback: () => void;
cancelled: boolean;
delayMs: number;
}

class ManualScheduler {
readonly entries: ManualTimerEntry[] = [];

setTimeout(callback: () => void, delayMs: number) {
const entry: ManualTimerEntry = { callback, cancelled: false, delayMs };
this.entries.push(entry);
return {
cancel: () => {
entry.cancelled = true;
},
};
}

fire(index: number): void {
this.entries[index]?.callback();
}
}

beforeEach(async () => {
dir = await mkdtemp(join(tmpdir(), 'finagent-kernel-'));
clock = 1000;
Expand All @@ -32,9 +56,10 @@ afterEach(async () => {

function makeKernel(
script: (input: AgentRunInput) => AsyncIterable<AgentEvent>,
extra: Partial<Omit<RunManagerOptions, 'sessions' | 'runs' | 'runtime' | 'now'>> = {}
extra: Partial<Omit<RunManagerOptions, 'sessions' | 'runs' | 'runtime' | 'now'>> = {},
ensureSessionGate?: Promise<void>
) {
const runtime = new ScriptedRuntime(script);
const runtime = new ScriptedRuntime(script, ensureSessionGate);
const store = new JsonFileStore(dir);
const sessions = new SessionManager({
sessions: new SessionRepository(store),
Expand All @@ -57,14 +82,18 @@ class ScriptedRuntime implements AgentRuntime {
ensureSessionCalls: Array<{ id: string; sessionPath?: string }> = [];
cancelCalls: Array<{ sessionId: string; runId: string }> = [];

constructor(private readonly script: (input: AgentRunInput) => AsyncIterable<AgentEvent>) {}
constructor(
private readonly script: (input: AgentRunInput) => AsyncIterable<AgentEvent>,
private readonly ensureSessionGate?: Promise<void>
) {}

async getTools(): Promise<ApiResult<ToolDefinition[]>> {
return { ok: true, data: [] };
}

async ensureSession(session: { id: string; title?: string; sessionPath?: string }): Promise<RuntimeSession> {
this.ensureSessionCalls.push(session);
await this.ensureSessionGate;
return { sessionId: session.id, status: 'active' };
}

Expand Down Expand Up @@ -445,6 +474,177 @@ describe('RunManager budgets and runaway detection (#17)', () => {
});
});

it('stops a stream that stalls before its first event', async () => {
let release!: () => void;
const stalled = new Promise<void>((resolve) => {
release = resolve;
});
const scheduler = new ManualScheduler();
const { sessions, runs, runtime } = makeKernel(
async function* () {
await stalled;
},
{ budgets: { defaults: { wallClockMs: 1_000 } }, scheduler }
);
const session = await sessions.createSession('Silent stream');

const run = await runs.startRun(session.id, 'wait forever');
expect(scheduler.entries).toHaveLength(1);
expect(scheduler.entries[0]?.delayMs).toBe(1_000);

clock += 1_000;
scheduler.fire(0);
await waitFor(async () => runtime.cancelCalls.length === 1);

release();
await waitFor(async () => !runs.isRunning());

expect(await sessions.getRun(session.id, run.id)).toMatchObject({
status: 'cancelled',
stopReason: 'budget_exhausted',
stopDetail: { key: 'wallClockMs', limit: 1_000, used: 1_000 },
});
});

it('keeps partial output when the stream stalls after a message', async () => {
let release!: () => void;
const stalled = new Promise<void>((resolve) => {
release = resolve;
});
const scheduler = new ManualScheduler();
const { sessions, runs, runtime } = makeKernel(
async function* (input) {
yield event(input.sessionId, input.runId, 'message_completed', { answer: 'partial' }, 1);
await stalled;
},
{ budgets: { defaults: { wallClockMs: 1_000 } }, scheduler }
);
const session = await sessions.createSession('Partial');
let sawPartial = false;
runs.subscribe((agentEvent) => {
if (agentEvent.type === 'message_completed') sawPartial = true;
});

const run = await runs.startRun(session.id, 'slow task');
await waitFor(async () => sawPartial);

clock += 1_000;
scheduler.fire(0);
await waitFor(async () => runtime.cancelCalls.length === 1);

release();
await waitFor(async () => !runs.isRunning());

expect(await sessions.getRun(session.id, run.id)).toMatchObject({
status: 'cancelled',
answer: 'partial',
stopReason: 'budget_exhausted',
stopDetail: { key: 'wallClockMs', limit: 1_000, used: 1_000 },
});
const messages = await sessions.listMessages(session.id);
expect(messages[1]).toMatchObject({ role: 'assistant', content: 'partial' });
});

it('does not start the runtime when the deadline expires during session setup', async () => {
let releaseEnsure!: () => void;
const ensureSessionGate = new Promise<void>((resolve) => {
releaseEnsure = resolve;
});
let runStarted = false;
const scheduler = new ManualScheduler();
const { sessions, runs, runtime } = makeKernel(
async function* () {
runStarted = true;
},
{ budgets: { defaults: { wallClockMs: 1_000 } }, scheduler },
ensureSessionGate
);
const session = await sessions.createSession('Setup race');

const run = await runs.startRun(session.id, 'slow setup');
expect(runtime.ensureSessionCalls).toHaveLength(1);

clock += 1_000;
scheduler.fire(0);
await waitFor(async () => runtime.cancelCalls.length === 1);

releaseEnsure();
await waitFor(async () => !runs.isRunning());

expect(runStarted).toBe(false);
expect(await sessions.getRun(session.id, run.id)).toMatchObject({
status: 'cancelled',
stopReason: 'budget_exhausted',
stopDetail: { key: 'wallClockMs', limit: 1_000, used: 1_000 },
});
});

it('cancels a wall-clock expiry only once even if the timer callback repeats', async () => {
let release!: () => void;
const stalled = new Promise<void>((resolve) => {
release = resolve;
});
const scheduler = new ManualScheduler();
const { sessions, runs, runtime } = makeKernel(
async function* () {
await stalled;
},
{ budgets: { defaults: { wallClockMs: 1_000 } }, scheduler }
);
const session = await sessions.createSession('Once');
const run = await runs.startRun(session.id, 'loop guard');

clock += 1_000;
scheduler.fire(0);
scheduler.fire(0);
await waitFor(async () => runtime.cancelCalls.length === 1);

release();
await waitFor(async () => !runs.isRunning());

expect(runtime.cancelCalls).toEqual([{ sessionId: session.id, runId: run.id }]);
});

it('clears a completed run deadline so a stale callback cannot cancel a later run', async () => {
let releaseLater!: () => void;
const laterStalled = new Promise<void>((resolve) => {
releaseLater = resolve;
});
const scheduler = new ManualScheduler();
const { sessions, runs, runtime } = makeKernel(
async function* (input) {
if (input.content === 'first run') {
yield event(input.sessionId, input.runId, 'run_completed', {
answer: 'done',
toolCalls: [],
});
return;
}
await laterStalled;
},
{ budgets: { defaults: { wallClockMs: 1_000 } }, scheduler }
);
const session = await sessions.createSession('Isolation');

const first = await runs.startRun(session.id, 'first run');
await waitFor(async () => !runs.isRunning());
expect(await sessions.getRun(session.id, first.id)).toMatchObject({ status: 'completed' });
expect(scheduler.entries[0]?.cancelled).toBe(true);

const later = await runs.startRun(session.id, 'second run');
expect(scheduler.entries).toHaveLength(2);

// Simulate an already-dispatched stale callback from the completed run.
scheduler.fire(0);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(runtime.cancelCalls).toEqual([]);
expect(runs.hasActiveRun(session.id)).toBe(true);

releaseLater();
await waitFor(async () => !runs.isRunning());
expect(await sessions.getRun(session.id, later.id)).toMatchObject({ status: 'completed' });
});

it('leaves a run untouched when neither budgets nor detectors are configured', async () => {
const { sessions, runs, runtime } = makeKernel(completedScript('Answer'));
const session = await sessions.createSession('Plain');
Expand Down
Loading