fix: enforce one Pi runtime per user
This commit is contained in:
@@ -12,6 +12,7 @@ export interface SessionRuntime {
|
|||||||
rpc: RpcClient;
|
rpc: RpcClient;
|
||||||
bridge: SessionBridge;
|
bridge: SessionBridge;
|
||||||
child: ChildProcessWithoutNullStreams;
|
child: ChildProcessWithoutNullStreams;
|
||||||
|
ownerKey?: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface RuntimeOptions {
|
export interface RuntimeOptions {
|
||||||
@@ -102,6 +103,15 @@ export class PiProcessManager {
|
|||||||
|
|
||||||
get(id: string): SessionRuntime | undefined { return this.runtimes.get(id); }
|
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. */
|
/** Spawn and register a runtime synchronously, without starting a model turn. */
|
||||||
createFor(sessionId: string, o: RuntimeOptions = {}): SessionRuntime {
|
createFor(sessionId: string, o: RuntimeOptions = {}): SessionRuntime {
|
||||||
// A duplicate start must never tear down a live session: that used to send
|
// A duplicate start must never tear down a live session: that used to send
|
||||||
@@ -110,6 +120,7 @@ export class PiProcessManager {
|
|||||||
if (existing) {
|
if (existing) {
|
||||||
throw new Error(`session runtime already active: ${sessionId}`);
|
throw new Error(`session runtime already active: ${sessionId}`);
|
||||||
}
|
}
|
||||||
|
if (o.principal) this.teardownForPrincipal(o.principal);
|
||||||
if (this.runtimes.size >= this.cfg.maxPiProcesses) {
|
if (this.runtimes.size >= this.cfg.maxPiProcesses) {
|
||||||
throw new Error("max Pi processes reached");
|
throw new Error("max Pi processes reached");
|
||||||
}
|
}
|
||||||
@@ -120,7 +131,12 @@ export class PiProcessManager {
|
|||||||
try {
|
try {
|
||||||
const rpc = new RpcClient(child);
|
const rpc = new RpcClient(child);
|
||||||
const bridge = new SessionBridge(rpc);
|
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;
|
rt = runtime;
|
||||||
bridge.beginTurn();
|
bridge.beginTurn();
|
||||||
this.runtimes.set(sessionId, runtime);
|
this.runtimes.set(sessionId, runtime);
|
||||||
|
|||||||
@@ -97,14 +97,6 @@ export function sessionRoutes(
|
|||||||
return true;
|
return true;
|
||||||
};
|
};
|
||||||
|
|
||||||
const releaseFinalizedRuntimes = async (): Promise<void> => {
|
|
||||||
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 = (
|
const bindRuntime = (
|
||||||
id: string, rt: ReturnType<PiProcessManager["createFor"]>, runner: any, workspace?: string,
|
id: string, rt: ReturnType<PiProcessManager["createFor"]>, runner: any, workspace?: string,
|
||||||
) => {
|
) => {
|
||||||
@@ -208,7 +200,10 @@ export function sessionRoutes(
|
|||||||
let s: Settings;
|
let s: Settings;
|
||||||
try { s = await d.getSettings(principal); } catch { return storageFailure(reply); }
|
try { s = await d.getSettings(principal); } catch { return storageFailure(reply); }
|
||||||
const runner = runnerFor(principal);
|
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);
|
const ensure = await d.readiness.ensure(s.workspace ?? "", principal);
|
||||||
if (!ensure.ok) return reply.code(503).send({ error: READINESS_FAILURE_MESSAGE });
|
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
|
// 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<PiProcessManager["createFor"]> | undefined;
|
let rt: ReturnType<PiProcessManager["createFor"]> | undefined;
|
||||||
try {
|
try {
|
||||||
if (current) {
|
if (current) {
|
||||||
|
|||||||
@@ -95,6 +95,47 @@ test("oltre maxPiProcesses solleva errore", async () => {
|
|||||||
mgr.teardown("a");
|
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() {
|
function recordingChild() {
|
||||||
const ch: any = new EventEmitter();
|
const ch: any = new EventEmitter();
|
||||||
ch.stdout = new EventEmitter();
|
ch.stdout = new EventEmitter();
|
||||||
|
|||||||
@@ -181,15 +181,25 @@ test("POST /sessions usa i settings (workspace/provider/model/thinking) e crea+a
|
|||||||
unlinkSync(modelKey);
|
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<string, any>();
|
const runtimes = new Map<string, any>();
|
||||||
const statuses = new Map<string, string>();
|
|
||||||
const tornDown: string[] = [];
|
const tornDown: string[] = [];
|
||||||
const order: string[] = [];
|
const order: string[] = [];
|
||||||
let nextId = 0;
|
let nextId = 0;
|
||||||
const app = buildApp(loadConfig({ THT_HARNESS_DIR: "../harness" }), {
|
const app = buildApp(loadConfig({ THT_HARNESS_DIR: "../harness" }), {
|
||||||
mgr: {
|
mgr: {
|
||||||
get: (id: string) => runtimes.get(id),
|
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) => {
|
createFor: (id: string) => {
|
||||||
order.push(`create:${id}`);
|
order.push(`create:${id}`);
|
||||||
if (runtimes.size >= 1) throw new Error("max Pi processes reached");
|
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 () => {
|
sessionNew: async () => {
|
||||||
const id = `s${++nextId}`;
|
const id = `s${++nextId}`;
|
||||||
order.push(`new:${id}`);
|
order.push(`new:${id}`);
|
||||||
statuses.set(id, "open");
|
|
||||||
return { id };
|
return { id };
|
||||||
},
|
},
|
||||||
sessionShow: async (id: string) => {
|
|
||||||
order.push(`show:${id}`);
|
|
||||||
return { status: statuses.get(id) };
|
|
||||||
},
|
|
||||||
searchPack: async () => {},
|
searchPack: async () => {},
|
||||||
failSession: async () => {},
|
failSession: async () => {},
|
||||||
} as any,
|
} 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)
|
expect((await app.inject({ method: "POST", url: "/sessions", payload: { question: "one" } })).statusCode)
|
||||||
.toBe(200);
|
.toBe(200);
|
||||||
statuses.set("s1", "finalized");
|
|
||||||
order.length = 0;
|
order.length = 0;
|
||||||
|
|
||||||
const second = await app.inject({ method: "POST", url: "/sessions", payload: { question: "two" } });
|
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.statusCode).toBe(200);
|
||||||
expect(second.json()).toEqual({ id: "s2" });
|
expect(second.json()).toEqual({ id: "s2" });
|
||||||
expect(tornDown).toEqual(["s1"]);
|
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);
|
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");
|
expect(configured.thinking).toBe("low");
|
||||||
});
|
});
|
||||||
|
|
||||||
|
test("resuming a different session stops the user's previous Pi runtime", async () => {
|
||||||
|
const runtimes = new Map<string, any>();
|
||||||
|
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"])(
|
test.each(["running", "waiting"])(
|
||||||
"POST resume preserves a %s runtime",
|
"POST resume preserves a %s runtime",
|
||||||
async (state) => {
|
async (state) => {
|
||||||
|
|||||||
Reference in New Issue
Block a user