From 00761ae2cabfc7fa64614657a88f6e9448d32ae9 Mon Sep 17 00:00:00 2001 From: mptyl Date: Tue, 21 Jul 2026 16:13:17 +0200 Subject: [PATCH] fix: enforce one Pi runtime per user --- backend/src/pi/pi-process-manager.ts | 18 ++++++- backend/src/routes/sessions.ts | 17 +++---- backend/test/pi-process-manager.test.ts | 41 +++++++++++++++ backend/test/routes-sessions.test.ts | 68 +++++++++++++++++++++---- 4 files changed, 125 insertions(+), 19 deletions(-) diff --git a/backend/src/pi/pi-process-manager.ts b/backend/src/pi/pi-process-manager.ts index 9b626c30..a12750e4 100644 --- a/backend/src/pi/pi-process-manager.ts +++ b/backend/src/pi/pi-process-manager.ts @@ -12,6 +12,7 @@ export interface SessionRuntime { rpc: RpcClient; bridge: SessionBridge; child: ChildProcessWithoutNullStreams; + ownerKey?: string; } export interface RuntimeOptions { @@ -102,6 +103,15 @@ export class PiProcessManager { get(id: string): SessionRuntime | undefined { return this.runtimes.get(id); } + teardownForPrincipal(principal: PrincipalContext): string[] { + const ownerKey = `${principal.issuer}\0${principal.subject}`; + const stopped: string[] = []; + for (const [id, runtime] of [...this.runtimes.entries()]) { + if (runtime.ownerKey === ownerKey && this.teardownIfCurrent(id, runtime)) stopped.push(id); + } + return stopped; + } + /** Spawn and register a runtime synchronously, without starting a model turn. */ createFor(sessionId: string, o: RuntimeOptions = {}): SessionRuntime { // A duplicate start must never tear down a live session: that used to send @@ -110,6 +120,7 @@ export class PiProcessManager { if (existing) { throw new Error(`session runtime already active: ${sessionId}`); } + if (o.principal) this.teardownForPrincipal(o.principal); if (this.runtimes.size >= this.cfg.maxPiProcesses) { throw new Error("max Pi processes reached"); } @@ -120,7 +131,12 @@ export class PiProcessManager { try { const rpc = new RpcClient(child); const bridge = new SessionBridge(rpc); - const runtime: SessionRuntime = { rpc, bridge, child }; + const runtime: SessionRuntime = { + rpc, + bridge, + child, + ownerKey: o.principal ? `${o.principal.issuer}\0${o.principal.subject}` : undefined, + }; rt = runtime; bridge.beginTurn(); this.runtimes.set(sessionId, runtime); diff --git a/backend/src/routes/sessions.ts b/backend/src/routes/sessions.ts index 94be96e1..f7834422 100644 --- a/backend/src/routes/sessions.ts +++ b/backend/src/routes/sessions.ts @@ -97,14 +97,6 @@ export function sessionRoutes( return true; }; - const releaseFinalizedRuntimes = async (): Promise => { - await Promise.all([...boundRuntimes.entries()].map(([id, rt]) => - withSessionLifecycle(id, () => releaseIfFinalized(id, rt)).catch((error: unknown) => { - console.error(`[session:${id}] stale runtime cleanup failed:`, error); - }), - )); - }; - const bindRuntime = ( id: string, rt: ReturnType, runner: any, workspace?: string, ) => { @@ -208,7 +200,10 @@ export function sessionRoutes( let s: Settings; try { s = await d.getSettings(principal); } catch { return storageFailure(reply); } const runner = runnerFor(principal); - await releaseFinalizedRuntimes(); + // A persisted session is resumable without keeping Pi alive. New work replaces every + // runtime owned by this principal, while runtimes belonging to other users remain intact. + // Optional chaining preserves the deliberately narrow manager stubs used by route tests. + for (const id of d.mgr.teardownForPrincipal?.(principal) ?? []) boundRuntimes.delete(id); const ensure = await d.readiness.ensure(s.workspace ?? "", principal); if (!ensure.ok) return reply.code(503).send({ error: READINESS_FAILURE_MESSAGE }); // Local-only: verify the DWH is reachable BEFORE creating the session, so a dropped @@ -388,6 +383,10 @@ export function sessionRoutes( } } + for (const stoppedId of d.mgr.teardownForPrincipal?.(principal) ?? []) { + boundRuntimes.delete(stoppedId); + } + let rt: ReturnType | undefined; try { if (current) { diff --git a/backend/test/pi-process-manager.test.ts b/backend/test/pi-process-manager.test.ts index 9a552292..27c4b71a 100644 --- a/backend/test/pi-process-manager.test.ts +++ b/backend/test/pi-process-manager.test.ts @@ -95,6 +95,47 @@ test("oltre maxPiProcesses solleva errore", async () => { mgr.teardown("a"); }); +test("teardownForPrincipal stops only runtimes owned by that user", () => { + const aliceChild = recordingChild(); + const bobChild = recordingChild(); + aliceChild.kill = vi.fn(); + bobChild.kill = vi.fn(); + const children = [aliceChild, bobChild]; + const mgr = new PiProcessManager(loadConfig({}), { spawnFn: () => children.shift() as any }); + const alice = { issuer: "portal", subject: "alice", isAdmin: false }; + const bob = { issuer: "portal", subject: "bob", isAdmin: false }; + mgr.createFor("alice-session", { principal: alice }); + mgr.createFor("bob-session", { principal: bob }); + + expect(mgr.teardownForPrincipal(alice)).toEqual(["alice-session"]); + + expect(mgr.get("alice-session")).toBeUndefined(); + expect(mgr.get("bob-session")).toBeDefined(); + expect(aliceChild.kill).toHaveBeenCalledOnce(); + expect(bobChild.kill).not.toHaveBeenCalled(); + mgr.teardown("bob-session"); +}); + +test("createFor keeps at most one runtime for the same user", () => { + const firstChild = recordingChild(); + const secondChild = recordingChild(); + firstChild.kill = vi.fn(); + secondChild.kill = vi.fn(); + const children = [firstChild, secondChild]; + const mgr = new PiProcessManager(loadConfig({}), { spawnFn: () => children.shift() as any }); + const principal = { issuer: "portal", subject: "alice", isAdmin: false }; + + mgr.createFor("first", { principal }); + const second = mgr.createFor("second", { principal }); + + expect(mgr.count()).toBe(1); + expect(mgr.get("first")).toBeUndefined(); + expect(mgr.get("second")).toBe(second); + expect(firstChild.kill).toHaveBeenCalledOnce(); + expect(secondChild.kill).not.toHaveBeenCalled(); + mgr.teardown("second"); +}); + function recordingChild() { const ch: any = new EventEmitter(); ch.stdout = new EventEmitter(); diff --git a/backend/test/routes-sessions.test.ts b/backend/test/routes-sessions.test.ts index 05fc72dc..9f019541 100644 --- a/backend/test/routes-sessions.test.ts +++ b/backend/test/routes-sessions.test.ts @@ -181,15 +181,25 @@ test("POST /sessions usa i settings (workspace/provider/model/thinking) e crea+a unlinkSync(modelKey); }); -test("POST /sessions reclaims a finalized runtime before enforcing the Pi limit", async () => { +test("POST /sessions stops the user's previous open Pi runtime before creating another", async () => { const runtimes = new Map(); - const statuses = new Map(); const tornDown: string[] = []; const order: string[] = []; let nextId = 0; const app = buildApp(loadConfig({ THT_HARNESS_DIR: "../harness" }), { mgr: { get: (id: string) => runtimes.get(id), + teardownForPrincipal: () => { + const stopped: string[] = []; + for (const [id, runtime] of [...runtimes]) { + runtimes.delete(id); + stopped.push(id); + tornDown.push(id); + order.push(`teardown:${id}`); + runtime.bridge.emitClientEvent?.({ type: "system_event", event: "agent_end" }); + } + return stopped; + }, createFor: (id: string) => { order.push(`create:${id}`); if (runtimes.size >= 1) throw new Error("max Pi processes reached"); @@ -211,13 +221,8 @@ test("POST /sessions reclaims a finalized runtime before enforcing the Pi limit" sessionNew: async () => { const id = `s${++nextId}`; order.push(`new:${id}`); - statuses.set(id, "open"); return { id }; }, - sessionShow: async (id: string) => { - order.push(`show:${id}`); - return { status: statuses.get(id) }; - }, searchPack: async () => {}, failSession: async () => {}, } as any, @@ -227,7 +232,6 @@ test("POST /sessions reclaims a finalized runtime before enforcing the Pi limit" expect((await app.inject({ method: "POST", url: "/sessions", payload: { question: "one" } })).statusCode) .toBe(200); - statuses.set("s1", "finalized"); order.length = 0; const second = await app.inject({ method: "POST", url: "/sessions", payload: { question: "two" } }); @@ -235,7 +239,7 @@ test("POST /sessions reclaims a finalized runtime before enforcing the Pi limit" expect(second.statusCode).toBe(200); expect(second.json()).toEqual({ id: "s2" }); expect(tornDown).toEqual(["s1"]); - expect(order).toEqual(["show:s1", "teardown:s1", "new:s2", "create:s2"]); + expect(order).toEqual(["teardown:s1", "new:s2", "create:s2"]); expect(runtimes.has("s2")).toBe(true); }); @@ -401,6 +405,52 @@ test("POST /sessions/:id/resume usa il thinking globale se manca nel manifest", expect(configured.thinking).toBe("low"); }); +test("resuming a different session stops the user's previous Pi runtime", async () => { + const runtimes = new Map(); + const mgr = { + get: (id: string) => runtimes.get(id), + teardownForPrincipal: () => { + const stopped = [...runtimes.keys()]; + runtimes.clear(); + return stopped; + }, + createFor: (id: string) => { + if (runtimes.size >= 1) throw new Error("max Pi processes reached"); + const runtime = { bridge: { onClientEvent: () => {}, turnState: () => "running" } }; + runtimes.set(id, runtime); + return runtime; + }, + configure: async () => {}, + start: () => {}, + teardownIfCurrent: (id: string, expected: any) => { + if (runtimes.get(id) !== expected) return false; + runtimes.delete(id); + return true; + }, + } as any; + const app = buildApp(loadConfig({ THT_HARNESS_DIR: "../harness" }), { + mgr, + thtRunner: { + sessionNew: async () => ({ id: "s1" }), + sessionShow: async (id: string) => ({ id, status: "open", archived: false }), + reopenSession: async () => {}, + searchPack: async () => {}, + failSession: async () => {}, + } as any, + readiness: { ensure: async () => ({ ok: true }) } as any, + getSettings: () => ({ workspace: "local" }) as any, + }); + expect((await app.inject({ method: "POST", url: "/sessions", payload: { question: "one" } })).statusCode) + .toBe(200); + + const resumed = await app.inject({ method: "POST", url: "/sessions/s2/resume" }); + + expect(resumed.statusCode).toBe(200); + expect(resumed.json()).toEqual({ id: "s2", alreadyActive: false }); + expect(runtimes.has("s1")).toBe(false); + expect(runtimes.has("s2")).toBe(true); +}); + test.each(["running", "waiting"])( "POST resume preserves a %s runtime", async (state) => {