From dfb459ad436a2a686e35987e61beb57cd575bc9c Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Sat, 15 Aug 2026 00:39:06 +0000 Subject: [PATCH] fix(pi-embed): start the extensions, and refuse a metered run that is not on the gateway model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pi emits `session_start` from `bindExtensions`, which its CLI modes call once their UI exists — `createAgentSession` never does. An embedder that skips it loads the extensions but never starts them, and the metered lane selects its model from exactly that event: so a run stayed on whatever provider key the process happened to hold, outside the gateway and unmetered, while looking healthy. `startRun` now fires it, then asserts the session really is on `/` and fails otherwise — under-reported usage is the one failure this lane exists to prevent. --- agentic/pi-embed/__tests__/session.test.ts | 54 +++++++++++++++++++++- agentic/pi-embed/src/session.ts | 28 +++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) diff --git a/agentic/pi-embed/__tests__/session.test.ts b/agentic/pi-embed/__tests__/session.test.ts index 3acd2dfc7..fb0b5bf95 100644 --- a/agentic/pi-embed/__tests__/session.test.ts +++ b/agentic/pi-embed/__tests__/session.test.ts @@ -9,8 +9,19 @@ import { type PiModule, startRun } from '../src'; const fakeLoader = {} as ResourceLoader; -const fakePi = () => { - const session = { dispose: jest.fn() }; +const gatewayMetering = { + mode: 'gateway' as const, + gatewayUrl: 'https://agentic.example.com', + identity: { databaseId: 'db-1' }, + models: [{ id: 'deepseek/deepseek-chat', contextWindow: 128_000, maxTokens: 8_192 }] +}; + +const fakePi = (model?: { provider: string; id: string }) => { + const session = { + dispose: jest.fn(), + bindExtensions: jest.fn(() => Promise.resolve()), + model + }; const createAgentSession = jest.fn( (_options: CreateAgentSessionOptions): Promise => Promise.resolve({ @@ -36,6 +47,8 @@ const fakePi = () => { return { pi, session, createAgentSession, loaderOptions, reload }; }; +const meteredSession = () => fakePi({ provider: 'constructive-gateway', id: 'deepseek/deepseek-chat' }); + describe('startRun', () => { it('hands the composed lanes to the resource loader, since that is how pi takes extensions', async () => { const { pi, createAgentSession } = fakePi(); @@ -72,6 +85,43 @@ describe('startRun', () => { expect(reload).toHaveBeenCalled(); }); + it('binds the extensions, since pi emits session_start from there and not from createAgentSession', async () => { + const { pi, session } = fakePi(); + + await startRun({ runId: 'run-1', pi, log: { store: new MemoryRunLogStore() }, createResourceLoader: () => fakeLoader }); + + expect(session.bindExtensions).toHaveBeenCalledWith({}); + }); + + it('refuses a metered run whose session did not end up on the gateway model', async () => { + const { pi } = fakePi({ provider: 'deepseek', id: 'deepseek-chat' }); + + await expect( + startRun({ runId: 'run-1', pi, metering: gatewayMetering, createResourceLoader: () => fakeLoader }) + ).rejects.toThrow(/would leave outside the gateway and go unmetered/); + }); + + it('refuses a metered run that selected no model at all', async () => { + const { pi } = fakePi(); + + await expect( + startRun({ runId: 'run-1', pi, metering: gatewayMetering, createResourceLoader: () => fakeLoader }) + ).rejects.toThrow(/no model/); + }); + + it('accepts a metered run once the session is on the gateway model', async () => { + const { pi } = meteredSession(); + + const embedded = await startRun({ + runId: 'run-1', + pi, + metering: gatewayMetering, + createResourceLoader: () => fakeLoader + }); + + expect(embedded.run.lanes.meteredModel?.selectedModel).toBe('deepseek/deepseek-chat'); + }); + it('demands an agentDir when the injected pi module cannot supply one', async () => { const { pi } = fakePi(); const withoutAgentDir: PiModule = { diff --git a/agentic/pi-embed/src/session.ts b/agentic/pi-embed/src/session.ts index 737cb28f5..31af491e1 100644 --- a/agentic/pi-embed/src/session.ts +++ b/agentic/pi-embed/src/session.ts @@ -84,6 +84,15 @@ export async function startRun(options: StartRunOptions): Promise { const result = await options.pi.createAgentSession({ ...options.session, cwd, agentDir, resourceLoader }); + // pi emits `session_start` from `bindExtensions`, which its own CLI modes call + // once their UI exists — `createAgentSession` never does. An embedder that + // skips it loads the extensions but never starts them: the metered lane never + // selects the gateway model, so a cloud run's calls leave on whatever provider + // key the process happens to hold, unmetered, and the log lane never binds on + // resume. So the embedding, not the host, fires it. + await result.session.bindExtensions({}); + assertMeteredModelSelected(run, result.session); + return { run, session: result.session, @@ -103,6 +112,25 @@ export async function startRun(options: StartRunOptions): Promise { }; } +/** + * The metered lane's whole purpose is that usage cannot be under-reported, so a + * session that ended up on some other provider must fail the run rather than + * quietly bill nothing. + */ +function assertMeteredModelSelected(run: ComposedRun, session: CreateAgentSessionResult['session']): void { + const lane = run.lanes.meteredModel; + if (!lane?.selectedModel) return; + + const model = session.model; + if (model?.provider === lane.providerName && model.id === lane.selectedModel) return; + + throw new Error( + `pi-embed: the metered lane selected "${lane.providerName}/${lane.selectedModel}" but the ` + + `session is on "${model ? `${model.provider}/${model.id}` : 'no model'}" — model calls would ` + + 'leave outside the gateway and go unmetered' + ); +} + const defaultResourceLoader = (pi: PiModule): CreateResourceLoader => async (request) => {