From a87479815a3f4074e50372214e72f4304885eb7c Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Sun, 28 Jun 2026 16:46:37 +0200 Subject: [PATCH] fix: inherit dispatcher model for spawned sessions --- ...rit-dispatch-model-for-spawned-sessions.md | 5 ++ src/server/sessions/piSessionService.test.ts | 57 +++++++++++++ src/server/sessions/piSessionService.ts | 84 ++++++++++++++----- src/server/sessions/spawnSessionTool.test.ts | 21 +++-- src/server/sessions/spawnSessionTool.ts | 15 +++- .../sessions/spawnSubsessionTool.test.ts | 26 +++++- src/server/sessions/spawnSubsessionTool.ts | 15 +++- 7 files changed, 189 insertions(+), 34 deletions(-) create mode 100644 .changeset/inherit-dispatch-model-for-spawned-sessions.md diff --git a/.changeset/inherit-dispatch-model-for-spawned-sessions.md b/.changeset/inherit-dispatch-model-for-spawned-sessions.md new file mode 100644 index 0000000..c2d9a4b --- /dev/null +++ b/.changeset/inherit-dispatch-model-for-spawned-sessions.md @@ -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. diff --git a/src/server/sessions/piSessionService.test.ts b/src/server/sessions/piSessionService.test.ts index d2bcdbf..fe63392 100644 --- a/src/server/sessions/piSessionService.test.ts +++ b/src/server/sessions/piSessionService.test.ts @@ -52,6 +52,12 @@ function sessionRef(id: string, cwd = "/workspace") { return { id, cwd }; } +function testModel(): NonNullable { + 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 = {}) { 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 }[] = []; diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index 7a3d8dd..df36016 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -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; +type AgentModel = NonNullable; type ModelRegistryInstance = ReturnType; export interface PiSessionManager { @@ -223,29 +228,56 @@ interface CreateAgentRuntimeOptions { cwd: string; agentDir: string; sessionManager: PiSessionManager; + initialModel?: AgentModel; } -type CreateAgentRuntime = (createRuntime: CreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions) => Promise; +type PiWebCreateAgentSessionRuntimeFactory = ( + options: Parameters[0] & { initialModel?: AgentModel } +) => ReturnType; -function defaultCreateAgentRuntime(createRuntime: CreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions): Promise { +type CreateAgentRuntime = (createRuntime: PiWebCreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions) => Promise; + +function defaultCreateAgentRuntime(createRuntime: PiWebCreateAgentSessionRuntimeFactory, options: CreateAgentRuntimeOptions): Promise { 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; -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 | undefined; @@ -464,8 +496,12 @@ export class PiSessionService { return [...unarchivedSessions, ...archivedSessions]; } - async start(cwd: string, parentSession?: string): Promise { - const active = await this.create(this.sessionManager.create(cwd, parentSession === undefined ? undefined : { parentSession }), cwd); + async start(cwd: string, options: StartSessionOptions = {}): Promise { + 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> { - const runtime = await this.createAgentRuntime(this.createRuntime, { cwd, agentDir: this.agentDir, sessionManager }); + private async create(sessionManager: PiSessionManager, cwd: string, options: Pick = {}): Promise> { + 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 = { runtime, unsubscribe: noop }; this.bindRuntime(active); diff --git a/src/server/sessions/spawnSessionTool.test.ts b/src/server/sessions/spawnSessionTool.test.ts index 078e634..fbc8c95 100644 --- a/src/server/sessions/spawnSessionTool.test.ts +++ b/src/server/sessions/spawnSessionTool.test.ts @@ -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"); }); }); diff --git a/src/server/sessions/spawnSessionTool.ts b/src/server/sessions/spawnSessionTool.ts index dba84e1..9696520 100644 --- a/src/server/sessions/spawnSessionTool.ts +++ b/src/server/sessions/spawnSessionTool.ts @@ -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; + 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, diff --git a/src/server/sessions/spawnSubsessionTool.test.ts b/src/server/sessions/spawnSubsessionTool.test.ts index 3bfa036..3201653 100644 --- a/src/server/sessions/spawnSubsessionTool.test.ts +++ b/src/server/sessions/spawnSubsessionTool.test.ts @@ -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) { @@ -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 }, diff --git a/src/server/sessions/spawnSubsessionTool.ts b/src/server/sessions/spawnSubsessionTool.ts index 9297396..68792e6 100644 --- a/src/server/sessions/spawnSubsessionTool.ts +++ b/src/server/sessions/spawnSubsessionTool.ts @@ -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; + 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,