fix(pi-embed): start the extensions, and refuse a metered run that is not on the gateway model

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<CreateAgentSessionResult> =>
       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 = {

agentic/pi-embed/src/session.ts

@@ -84,6 +84,15 @@ export async function startRun(options: StartRunOptions): Promise<EmbeddedRun> {
 
   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<EmbeddedRun> {
   };
 }
 
+/**
+ * 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) => {