This commit is contained in:
Harivansh Rathi 2026-03-09 10:01:59 -07:00
parent f58ff785ce
commit 43c6b56dfa

View file

@ -39,6 +39,7 @@ import {
buildGatewaySessionStateMessages, buildGatewaySessionStateMessages,
messageContentToHistoryParts, messageContentToHistoryParts,
} from "./session-state.js"; } from "./session-state.js";
import { findMostRecentSession } from "../session-manager.js";
export { export {
createGatewaySessionManager, createGatewaySessionManager,
sanitizeSessionKey, sanitizeSessionKey,
@ -204,7 +205,7 @@ export class GatewayRuntime {
sessionKey: string, sessionKey: string,
listener: (event: GatewayEvent) => void, listener: (event: GatewayEvent) => void,
): Promise<() => void> { ): Promise<() => void> {
const managedSession = await this.ensureSession(sessionKey); const managedSession = await this.requireExistingSession(sessionKey);
managedSession.listeners.add(listener); managedSession.listeners.add(listener);
listener({ listener({
type: "hello", type: "hello",
@ -274,6 +275,32 @@ export class GatewayRuntime {
return session ? this.createSnapshot(session) : undefined; return session ? this.createSnapshot(session) : undefined;
} }
private async getOrLoadExistingSession(
sessionKey: string,
): Promise<ManagedGatewaySession | null> {
const found = this.sessions.get(sessionKey);
if (found) {
found.lastActiveAt = Date.now();
return found;
}
if (!findMostRecentSession(this.getGatewaySessionDir(sessionKey))) {
return null;
}
return this.ensureSession(sessionKey);
}
private async requireExistingSession(
sessionKey: string,
): Promise<ManagedGatewaySession> {
const managedSession = await this.getOrLoadExistingSession(sessionKey);
if (!managedSession) {
throw new HttpError(404, `Session not found: ${sessionKey}`);
}
return managedSession;
}
private async ensureSession( private async ensureSession(
sessionKey: string, sessionKey: string,
existingSession?: AgentSession, existingSession?: AgentSession,
@ -377,14 +404,6 @@ export class GatewayRuntime {
} }
} }
private getManagedSessionOrThrow(sessionKey: string): ManagedGatewaySession {
const managedSession = this.sessions.get(sessionKey);
if (!managedSession) {
throw new HttpError(404, `Session not found: ${sessionKey}`);
}
return managedSession;
}
private rejectQueuedMessages( private rejectQueuedMessages(
managedSession: ManagedGatewaySession, managedSession: ManagedGatewaySession,
error: string, error: string,
@ -807,7 +826,7 @@ export class GatewayRuntime {
const action = sessionMatch[2]; const action = sessionMatch[2];
if (!action && method === "GET") { if (!action && method === "GET") {
const session = await this.ensureSession(sessionKey); const session = await this.requireExistingSession(sessionKey);
this.writeJson(response, 200, { session: this.createSnapshot(session) }); this.writeJson(response, 200, { session: this.createSnapshot(session) });
return; return;
} }
@ -852,13 +871,13 @@ export class GatewayRuntime {
} }
if (action === "abort" && method === "POST") { if (action === "abort" && method === "POST") {
this.getManagedSessionOrThrow(sessionKey); await this.requireExistingSession(sessionKey);
this.writeJson(response, 200, { ok: this.abortSession(sessionKey) }); this.writeJson(response, 200, { ok: this.abortSession(sessionKey) });
return; return;
} }
if (action === "reset" && method === "POST") { if (action === "reset" && method === "POST") {
this.getManagedSessionOrThrow(sessionKey); await this.requireExistingSession(sessionKey);
await this.resetSession(sessionKey); await this.resetSession(sessionKey);
this.writeJson(response, 200, { ok: true }); this.writeJson(response, 200, { ok: true });
return; return;
@ -875,7 +894,7 @@ export class GatewayRuntime {
} }
if (action === "state" && method === "GET") { if (action === "state" && method === "GET") {
const session = await this.ensureSession(sessionKey); const session = await this.requireExistingSession(sessionKey);
this.writeJson(response, 200, this.createSessionState(session)); this.writeJson(response, 200, this.createSessionState(session));
return; return;
} }
@ -1121,7 +1140,7 @@ export class GatewayRuntime {
provider: string, provider: string,
modelId: string, modelId: string,
): Promise<{ ok: true; model: { provider: string; modelId: string } }> { ): Promise<{ ok: true; model: { provider: string; modelId: string } }> {
const managed = await this.ensureSession(sessionKey); const managed = await this.requireExistingSession(sessionKey);
const found = managed.session.modelRegistry.find(provider, modelId); const found = managed.session.modelRegistry.find(provider, modelId);
if (!found) { if (!found) {
throw new HttpError(404, `Model not found: ${provider}/${modelId}`); throw new HttpError(404, `Model not found: ${provider}/${modelId}`);
@ -1137,7 +1156,7 @@ export class GatewayRuntime {
if (limit !== undefined && (!Number.isFinite(limit) || limit < 1)) { if (limit !== undefined && (!Number.isFinite(limit) || limit < 1)) {
throw new HttpError(400, "History limit must be a positive integer"); throw new HttpError(400, "History limit must be a positive integer");
} }
const managed = await this.ensureSession(sessionKey); const managed = await this.requireExistingSession(sessionKey);
const rawMessages = managed.session.messages; const rawMessages = managed.session.messages;
const messages: HistoryMessage[] = []; const messages: HistoryMessage[] = [];
for (const [index, msg] of rawMessages.entries()) { for (const [index, msg] of rawMessages.entries()) {
@ -1162,7 +1181,7 @@ export class GatewayRuntime {
sessionKey: string, sessionKey: string,
patch: { name?: string }, patch: { name?: string },
): Promise<void> { ): Promise<void> {
const managed = await this.ensureSession(sessionKey); const managed = await this.requireExistingSession(sessionKey);
if (patch.name !== undefined) { if (patch.name !== undefined) {
// Labels in pi-mono are per-entry; we label the current leaf entry // Labels in pi-mono are per-entry; we label the current leaf entry
const leafId = managed.session.sessionManager.getLeafId(); const leafId = managed.session.sessionManager.getLeafId();
@ -1180,7 +1199,7 @@ export class GatewayRuntime {
if (sessionKey === this.primarySessionKey) { if (sessionKey === this.primarySessionKey) {
throw new HttpError(400, "Cannot delete primary session"); throw new HttpError(400, "Cannot delete primary session");
} }
const managed = this.sessions.get(sessionKey); const managed = await this.requireExistingSession(sessionKey);
if (managed) { if (managed) {
if (managed.processing) { if (managed.processing) {
await managed.session.abort(); await managed.session.abort();
@ -1239,7 +1258,7 @@ export class GatewayRuntime {
} }
private async handleReloadSession(sessionKey: string): Promise<void> { private async handleReloadSession(sessionKey: string): Promise<void> {
const managed = await this.ensureSession(sessionKey); const managed = await this.requireExistingSession(sessionKey);
// Reloading config by calling settingsManager.reload() on the session // Reloading config by calling settingsManager.reload() on the session
managed.session.settingsManager.reload(); managed.session.settingsManager.reload();
} }