fix: inherit dispatcher model for spawned sessions

2026-06-28 16:46:37 +02:00
parent 7063c2c3b1
commit a87479815a
7 changed files with 189 additions and 34 deletions
@@ -0,0 +1,5 @@
+---
+"@jmfederico/pi-web": patch
+---
+
+Make spawned and tracked subsessions inherit the dispatching session's current model instead of falling back to the last globally selected model.
@@ -52,6 +52,12 @@ function sessionRef(id: string, cwd = "/workspace") {
  return { id, cwd };
 }

+function testModel(): NonNullable<PiAgentSession["model"]> {
+  const model = ModelRegistry.inMemory(AuthStorage.inMemory()).find("anthropic", "claude-3-5-sonnet-20241022");
+  if (model === undefined) throw new Error("test model not found");
+  return model;
+}
+
 function fakeRuntime(sessionId = "session-1", patch: Partial<TestSession> = {}) {
  const promptCalls: { text: string; options: unknown }[] = [];
  const customMessageCalls: { message: { customType: string; content: string; display: boolean; details?: unknown }; options: unknown }[] = [];
@@ -908,6 +914,28 @@ describe("PiSessionService", () => {
      await service.dispose();
    });

+    it("uses the dispatching session's model as the spawned session's initial model", async () => {
+      const fake = fakeRuntime("spawned-1", { sessionFile: "/tmp/spawned-1.jsonl" });
+      const model = testModel();
+      let initialModel: PiAgentSession["model"];
+      const createAgentRuntime: RuntimeCreator = async (_createRuntime, options) => {
+        await Promise.resolve();
+        initialModel = options.initialModel;
+        return fake.runtime;
+      };
+      const service = new PiSessionService(new CapturingSessionEventHub(), {
+        createAgentRuntime,
+        sessionManager: sessionGateway([]),
+        spawnTargets: { resolveSpawnTarget: () => Promise.resolve({ allowed: true, cwd: "/workspace-feature" }) },
+        heartbeatIntervalMs: 60_000,
+      });
+
+      await service.spawnSession({ spawningCwd: "/workspace", prompt: "continue", cwd: "/workspace-feature", model });
+
+      expect(initialModel).toBe(model);
+      await service.dispose();
+    });
+
    it("rejects an out-of-project target without starting a session", async () => {
      const { fake, service } = spawnService({ allowed: false, reason: "out-of-project", allowedCwds: ["/workspace"] });

@@ -989,6 +1017,35 @@ describe("PiSessionService", () => {
      await service.dispose();
    });

+    it("uses the parent session's model as the tracked child's initial model", async () => {
+      const parent = fakeRuntime("parent-1", { sessionFile: "/tmp/parent-1.jsonl" });
+      const child = fakeRuntime("child-1", { sessionFile: "/tmp/child-1.jsonl", sessionManager: fakeSessionManager("/workspace-feature") });
+      const model = testModel();
+      const initialModels: PiAgentSession["model"][] = [];
+      const runtimes = [parent.runtime, child.runtime];
+      let index = 0;
+      const createAgentRuntime: RuntimeCreator = async (_createRuntime, options) => {
+        await Promise.resolve();
+        initialModels.push(options.initialModel);
+        const runtime = runtimes[index] ?? child.runtime;
+        index += 1;
+        return runtime;
+      };
+      const service = new PiSessionService(new CapturingSessionEventHub(), {
+        createAgentRuntime,
+        sessionManager: sessionGateway([]),
+        archiveStore: emptyArchiveStore(),
+        spawnTargets: { resolveSpawnTarget: () => Promise.resolve({ allowed: true, cwd: "/workspace-feature" }) },
+        heartbeatIntervalMs: 60_000,
+      });
+
+      await service.start("/workspace");
+      await service.spawnSubsession({ spawningCwd: "/workspace", parentSessionId: "parent-1", parentSessionFile: "/tmp/parent-1.jsonl", prompt: "do the slice", cwd: "/workspace-feature", model });
+
+      expect(initialModels).toEqual([undefined, model]);
+      await service.dispose();
+    });
+
    it("persists tracked child links in the parent and child sessions", async () => {
      const parentPersisted: { customType: string; data?: unknown }[] = [];
      const childPersisted: { customType: string; data?: unknown }[] = [];
@@ -1,5 +1,5 @@
 import { open, readFile, writeFile } from "node:fs/promises";
-import type { Api, ImageContent, Model } from "@earendil-works/pi-ai";
+import type { ImageContent } from "@earendil-works/pi-ai";
 import {
  AuthStorage,
  createAgentSessionFromServices,
@@ -102,6 +102,11 @@ interface PersistedChildSubsessionLink {
  spawnedSessionId: string;
 }

+interface StartSessionOptions {
+  parentSession?: string;
+  initialModel?: AgentModel;
+}
+
 function requirePromptText(value: unknown): string {
  if (typeof value !== "string") throw new Error("Prompt text is required");
  return value;
@@ -138,7 +143,7 @@ interface WorkspaceArchiveCandidate extends SessionArchiveTreeCandidate {
  activeSession?: PiAgentSession;
 }

-type AgentModel = Model<Api>;
+type AgentModel = NonNullable<SpawnSessionInvocation["model"]>;
 type ModelRegistryInstance = ReturnType<typeof ModelRegistry.create>;

 export interface PiSessionManager {
@@ -223,29 +228,56 @@ interface CreateAgentRuntimeOptions {
  cwd: string;
  agentDir: string;
  sessionManager: PiSessionManager;
+  initialModel?: AgentModel;
 }

-type CreateAgentRuntime = (createRuntime: CreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions) => Promise<PiSessionRuntime>;
+type PiWebCreateAgentSessionRuntimeFactory = (
+  options: Parameters<CreateAgentSessionRuntimeFactory>[0] & { initialModel?: AgentModel }
+) => ReturnType<CreateAgentSessionRuntimeFactory>;

-function defaultCreateAgentRuntime(createRuntime: CreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions): Promise<PiSessionRuntime> {
+type CreateAgentRuntime = (createRuntime: PiWebCreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions) => Promise<PiSessionRuntime>;
+
+function defaultCreateAgentRuntime(createRuntime: PiWebCreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions): Promise<PiSessionRuntime> {
  if (!(options.sessionManager instanceof SessionManager)) throw new Error("Default runtime creation requires an SDK SessionManager");
-  return createAgentSessionRuntime(createRuntime, { ...options, sessionManager: options.sessionManager });
+  const runtimeFactory = createRuntimeWithOneShotInitialModel(createRuntime, options.initialModel);
+  return createAgentSessionRuntime(runtimeFactory, {
+    cwd: options.cwd,
+    agentDir: options.agentDir,
+    sessionManager: options.sessionManager,
+  });
+}
+
+function createRuntimeWithOneShotInitialModel(createRuntime: PiWebCreateAgentSessionRuntimeFactory, initialModel: AgentModel | undefined): CreateAgentSessionRuntimeFactory {
+  // The inherited model belongs only to the session being spawned. Do not keep
+  // reapplying it if that runtime later creates/forks/switches sessions itself.
+  let pendingInitialModel = initialModel;
+  return async (options) => {
+    const model = pendingInitialModel;
+    pendingInitialModel = undefined;
+    return createRuntime({
+      ...options,
+      ...(model === undefined ? {} : { initialModel: model }),
+    });
+  };
 }

 type SpawnSessionFn = (input: SpawnSessionInvocation) => Promise<SpawnSessionResult>;

-function createDefaultRuntimeFactory(authStorage: AuthStorage, modelRegistry: ModelRegistryInstance, spawn?: SpawnSessionFn, subsessions?: SubsessionToolDeps): CreateAgentSessionRuntimeFactory {
-  return async ({ cwd, agentDir, sessionManager, sessionStartEvent }) => {
+function createDefaultRuntimeFactory(authStorage: AuthStorage, modelRegistry: ModelRegistryInstance, spawn?: SpawnSessionFn, subsessions?: SubsessionToolDeps): PiWebCreateAgentSessionRuntimeFactory {
+  return async ({ cwd, agentDir, sessionManager, sessionStartEvent, initialModel }) => {
    const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry });
    const customTools = [
      createPiWebEditToolDefinition(cwd),
      ...(spawn === undefined ? [] : [createSpawnSessionToolDefinition(cwd, { spawn })]),
      ...(subsessions === undefined ? [] : createSubsessionToolDefinitions(cwd, subsessions)),
    ];
-    const options = sessionStartEvent === undefined
-      ? { services, sessionManager, customTools }
-      : { services, sessionManager, sessionStartEvent, customTools };
-    const result = await createAgentSessionFromServices(options);
+    const result = await createAgentSessionFromServices({
+      services,
+      sessionManager,
+      customTools,
+      ...(sessionStartEvent === undefined ? {} : { sessionStartEvent }),
+      ...(initialModel === undefined ? {} : { model: initialModel }),
+    });
    return { ...result, services, diagnostics: services.diagnostics };
  };
 }
@@ -278,7 +310,7 @@ export interface PiSessionServiceDependencies {
  archiveStore?: SessionArchiveRepository;
  agentDir?: string;
  sessionManager?: PiSessionManagerGateway;
-  createRuntime?: CreateAgentSessionRuntimeFactory;
+  createRuntime?: PiWebCreateAgentSessionRuntimeFactory;
  createAgentRuntime?: CreateAgentRuntime;
  modelRegistry?: ModelRegistryInstance;
  heartbeatIntervalMs?: number;
@@ -328,7 +360,7 @@ export class PiSessionService {
  private readonly archiveStore: SessionArchiveRepository;
  private readonly agentDir: string;
  private readonly sessionManager: PiSessionManagerGateway;
-  private readonly createRuntime: CreateAgentSessionRuntimeFactory;
+  private readonly createRuntime: PiWebCreateAgentSessionRuntimeFactory;
  private readonly createAgentRuntime: CreateAgentRuntime;
  private readonly modelRegistry: ModelRegistryInstance;
  private readonly workspaceActivity: Pick<WorkspaceActivityService, "applySessionStatus" | "applySessionActivity" | "removeSession" | "reconcileSessionActivity"> | undefined;
@@ -464,8 +496,12 @@ export class PiSessionService {
    return [...unarchivedSessions, ...archivedSessions];
  }

-  async start(cwd: string, parentSession?: string): Promise<ClientSession> {
-    const active = await this.create(this.sessionManager.create(cwd, parentSession === undefined ? undefined : { parentSession }), cwd);
+  async start(cwd: string, options: StartSessionOptions = {}): Promise<ClientSession> {
+    const active = await this.create(
+      this.sessionManager.create(cwd, options.parentSession === undefined ? undefined : { parentSession: options.parentSession }),
+      cwd,
+      options.initialModel === undefined ? {} : { initialModel: options.initialModel },
+    );
    const { session } = active.runtime;
    const created: ClientSession = {
      id: session.sessionId,
@@ -477,7 +513,7 @@ export class PiSessionService {
      firstMessage: "",
      // Include the parent so listeners can nest the new session in the tree
      // immediately, instead of showing it flat until the next reload.
-      ...(parentSession === undefined ? {} : { parentSessionPath: parentSession }),
+      ...(options.parentSession === undefined ? {} : { parentSessionPath: options.parentSession }),
    };
    // Broadcast so other clients (and the spawning agent's UI) can add the new
    // session to their list without a manual reload.
@@ -494,7 +530,7 @@ export class PiSessionService {
    if (this.spawnTargets === undefined) throw new Error("Spawning sessions is disabled");
    const decision = await this.spawnTargets.resolveSpawnTarget(input.spawningCwd, input.cwd);
    if (!decision.allowed) throw spawnTargetError(decision);
-    const created = await this.start(decision.cwd);
+    const created = await this.start(decision.cwd, input.model === undefined ? {} : { initialModel: input.model });
    await this.prompt(created.id, input.prompt);
    this.logger.info(
      { spawningCwd: input.spawningCwd, sessionId: created.id, cwd: decision.cwd, promptLength: input.prompt.length },
@@ -513,7 +549,10 @@ export class PiSessionService {
    if (this.spawnTargets === undefined) throw new Error("Spawning sessions is disabled");
    const decision = await this.spawnTargets.resolveSpawnTarget(input.spawningCwd, input.cwd);
    if (!decision.allowed) throw spawnTargetError(decision);
-    const created = await this.start(decision.cwd, input.parentSessionFile);
+    const created = await this.start(decision.cwd, {
+      ...(input.parentSessionFile === undefined ? {} : { parentSession: input.parentSessionFile }),
+      ...(input.model === undefined ? {} : { initialModel: input.model }),
+    });
    const parentSessionFile = nonEmptyString(input.parentSessionFile);
    const link: TrackedSubsessionLink = {
      parentSessionId: input.parentSessionId,
@@ -1305,8 +1344,13 @@ export class PiSessionService {
    return undefined;
  }

-  private async create(sessionManager: PiSessionManager, cwd: string): Promise<ActiveSession<PiSessionRuntime>> {
-    const runtime = await this.createAgentRuntime(this.createRuntime, { cwd, agentDir: this.agentDir, sessionManager });
+  private async create(sessionManager: PiSessionManager, cwd: string, options: Pick<StartSessionOptions, "initialModel"> = {}): Promise<ActiveSession<PiSessionRuntime>> {
+    const runtime = await this.createAgentRuntime(this.createRuntime, {
+      cwd,
+      agentDir: this.agentDir,
+      sessionManager,
+      ...(options.initialModel === undefined ? {} : { initialModel: options.initialModel }),
+    });
    await this.bindSessionExtensions(runtime.session);
    const active: ActiveSession<PiSessionRuntime> = { runtime, unsubscribe: noop };
    this.bindRuntime(active);
@@ -2,18 +2,20 @@ import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { describe, expect, it, vi } from "vitest";
 import { createSpawnSessionToolDefinition } from "./spawnSessionTool.js";

-// The spawn tool's execute() never reads ctx, so an empty stub is sufficient.
-// eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- test stub; execute() does not use ctx.
+// eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- test stub with the minimal surface the tool reads.
 const ctx = {} as ExtensionContext;
+const dispatchModel = { provider: "anthropic", id: "claude-sonnet" };
+// eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- test stub with the minimal surface the tool reads.
+const ctxWithModel = { model: dispatchModel } as ExtensionContext;

 describe("createSpawnSessionToolDefinition", () => {
  it("passes the spawning cwd and params to the spawn callback and reports success", async () => {
    const spawn = vi.fn(() => Promise.resolve({ sessionId: "new-1", cwd: "/repos/a-feature" }));
    const tool = createSpawnSessionToolDefinition("/repos/a", { spawn });

-    const result = await tool.execute("call-1", { prompt: "do the thing", cwd: "/repos/a-feature" }, undefined, undefined, ctx);
+    const result = await tool.execute("call-1", { prompt: "do the thing", cwd: "/repos/a-feature" }, undefined, undefined, ctxWithModel);

-    expect(spawn).toHaveBeenCalledWith({ spawningCwd: "/repos/a", prompt: "do the thing", cwd: "/repos/a-feature" });
+    expect(spawn).toHaveBeenCalledWith({ spawningCwd: "/repos/a", prompt: "do the thing", cwd: "/repos/a-feature", model: dispatchModel });
    expect(result.details).toEqual({ sessionId: "new-1", cwd: "/repos/a-feature" });
    expect(result.content[0]).toMatchObject({ type: "text", text: "Started session new-1 in /repos/a-feature." });
  });
@@ -27,11 +29,20 @@ describe("createSpawnSessionToolDefinition", () => {
    expect(spawn).toHaveBeenCalledWith({ spawningCwd: "/repos/a", prompt: "continue", cwd: undefined });
  });

+  it("omits the inherited model when the dispatching session has no current model", async () => {
+    const spawn = vi.fn(() => Promise.resolve({ sessionId: "new-3", cwd: "/repos/a" }));
+    const tool = createSpawnSessionToolDefinition("/repos/a", { spawn });
+
+    await tool.execute("call-3", { prompt: "continue" }, undefined, undefined, ctx);
+
+    expect(spawn).toHaveBeenCalledWith({ spawningCwd: "/repos/a", prompt: "continue", cwd: undefined });
+  });
+
  it("propagates the spawn callback error so the agent loop reports it", async () => {
    const spawn = vi.fn(() => Promise.reject(new Error("cwd must be a workspace of this project. Allowed: /repos/a")));
    const tool = createSpawnSessionToolDefinition("/repos/a", { spawn });

-    await expect(tool.execute("call-3", { prompt: "x", cwd: "/elsewhere" }, undefined, undefined, ctx))
+    await expect(tool.execute("call-4", { prompt: "x", cwd: "/elsewhere" }, undefined, undefined, ctx))
      .rejects.toThrow("cwd must be a workspace of this project. Allowed: /repos/a");
  });
 });
@@ -1,15 +1,19 @@
 import { Type } from "typebox";
-import { defineTool } from "@earendil-works/pi-coding-agent";
+import { defineTool, type ExtensionContext } from "@earendil-works/pi-coding-agent";

 export interface SpawnSessionResult {
  sessionId: string;
  cwd: string;
 }

+export type SpawnSessionModel = NonNullable<ExtensionContext["model"]>;
+
 export interface SpawnSessionInvocation {
  spawningCwd: string;
  prompt: string;
  cwd: string | undefined;
+  /** Current model from the dispatching session, used as the spawned session's default. */
+  model?: SpawnSessionModel;
 }

 export interface SpawnSessionToolDeps {
@@ -40,11 +44,16 @@ export function createSpawnSessionToolDefinition(spawningCwd: string, deps: Spaw
    description: "Start a new, independent pi-web session and send it an initial prompt. Use this to dispatch a fresh agent to continue work or follow a plan. The new session runs on its own and a human can interact with it; you do not receive its output.",
    promptSnippet: "spawn_session: start a new independent session with a first prompt",
    parameters: SpawnSessionParams,
-    async execute(_toolCallId, params) {
+    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      // Failures throw: the agent loop turns the thrown message into an error
      // tool result the model sees, so the spawning agent can adapt (e.g. pick a
      // valid workspace) rather than crash.
-      const result = await deps.spawn({ spawningCwd, prompt: params.prompt, cwd: params.cwd });
+      const result = await deps.spawn({
+        spawningCwd,
+        prompt: params.prompt,
+        cwd: params.cwd,
+        ...(ctx.model === undefined ? {} : { model: ctx.model }),
+      });
      return {
        content: [{ type: "text", text: `Started session ${result.sessionId} in ${result.cwd}.` }],
        details: result,
@@ -3,11 +3,13 @@ import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { describe, expect, it, vi } from "vitest";
 import { createSubsessionToolDefinitions, type SubsessionToolDeps } from "./spawnSubsessionTool.js";

-function ctxFor(sessionId: string, sessionFile: string | undefined): ExtensionContext {
+const dispatchModel = { provider: "anthropic", id: "claude-sonnet" };
+
+function ctxFor(sessionId: string, sessionFile: string | undefined, model?: unknown): ExtensionContext {
  const sessionManager = { getSessionId: () => sessionId, getSessionFile: () => sessionFile };
-  // The subsession tools only read sessionManager.getSessionId/getSessionFile.
+  // The subsession tools only read sessionManager.getSessionId/getSessionFile and model.
  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions -- test stub with the minimal surface the tools use.
-  return { sessionManager } as unknown as ExtensionContext;
+  return { sessionManager, ...(model === undefined ? {} : { model }) } as unknown as ExtensionContext;
 }

 function tools(deps: Partial<SubsessionToolDeps>) {
@@ -36,7 +38,7 @@ describe("createSubsessionToolDefinitions", () => {
    const spawn = vi.fn(() => Promise.resolve({ sessionId: "child-1", cwd: "/repos/a-feature" }));
    const { spawn: spawnTool } = tools({ spawn });

-    const result = await spawnTool.execute("call-1", { prompt: "do it", cwd: "/repos/a-feature" }, undefined, undefined, ctxFor("parent-1", "/sessions/parent-1.jsonl"));
+    const result = await spawnTool.execute("call-1", { prompt: "do it", cwd: "/repos/a-feature" }, undefined, undefined, ctxFor("parent-1", "/sessions/parent-1.jsonl", dispatchModel));

    expect(spawn).toHaveBeenCalledWith({
      spawningCwd: "/repos/a",
@@ -44,11 +46,27 @@ describe("createSubsessionToolDefinitions", () => {
      parentSessionFile: "/sessions/parent-1.jsonl",
      prompt: "do it",
      cwd: "/repos/a-feature",
+      model: dispatchModel,
    });
    expect(result.details).toEqual({ sessionId: "child-1", cwd: "/repos/a-feature" });
    expect(firstText(result.content)).toContain("Started subsession child-1");
  });

+  it("spawn_subsession omits the inherited model when the dispatching session has no current model", async () => {
+    const spawn = vi.fn(() => Promise.resolve({ sessionId: "child-2", cwd: "/repos/a" }));
+    const { spawn: spawnTool } = tools({ spawn });
+
+    await spawnTool.execute("call-modeless", { prompt: "do it" }, undefined, undefined, ctxFor("parent-1", undefined));
+
+    expect(spawn).toHaveBeenCalledWith({
+      spawningCwd: "/repos/a",
+      parentSessionId: "parent-1",
+      parentSessionFile: undefined,
+      prompt: "do it",
+      cwd: undefined,
+    });
+  });
+
  it("list_subsessions reports the caller's subsessions and their status", async () => {
    const list = vi.fn(() => Promise.resolve([
      { sessionId: "child-1", cwd: "/repos/a", status: "working" as const },
@@ -1,5 +1,5 @@
 import { Type } from "typebox";
-import { defineTool } from "@earendil-works/pi-coding-agent";
+import { defineTool, type ExtensionContext } from "@earendil-works/pi-coding-agent";
 import type { TranscriptContentKind, TranscriptEntry, TranscriptRole, TranscriptView } from "./subsessionTranscript.js";

 /** Lifecycle phase of a tracked subsession as seen by its parent. */
@@ -10,6 +10,8 @@ export interface SpawnSubsessionResult {
  cwd: string;
 }

+export type SpawnSubsessionModel = NonNullable<ExtensionContext["model"]>;
+
 export interface SpawnSubsessionInvocation {
  /** cwd of the session that invoked the tool (used for project-scope checks). */
  spawningCwd: string;
@@ -19,6 +21,8 @@ export interface SpawnSubsessionInvocation {
  parentSessionFile: string | undefined;
  prompt: string;
  cwd: string | undefined;
+  /** Current model from the dispatching session, used as the spawned session's default. */
+  model?: SpawnSubsessionModel;
 }

 export interface SubsessionSummary {
@@ -180,7 +184,14 @@ export function createSubsessionToolDefinitions(spawningCwd: string, deps: Subse
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const parentSessionId = ctx.sessionManager.getSessionId();
      const parentSessionFile = ctx.sessionManager.getSessionFile() ?? undefined;
-      const result = await deps.spawn({ spawningCwd, parentSessionId, parentSessionFile, prompt: params.prompt, cwd: params.cwd });
+      const result = await deps.spawn({
+        spawningCwd,
+        parentSessionId,
+        parentSessionFile,
+        prompt: params.prompt,
+        cwd: params.cwd,
+        ...(ctx.model === undefined ? {} : { model: ctx.model }),
+      });
      return {
        content: [{ type: "text", text: `Started subsession ${result.sessionId} in ${result.cwd}. You will be notified when it stops working.` }],
        details: result,