@plannotator/pi-extension 0.15.0 → 0.15.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +1 -1
- package/generated/ai/base-session.ts +95 -0
- package/generated/ai/context.ts +212 -0
- package/generated/ai/endpoints.ts +309 -0
- package/generated/ai/index.ts +106 -0
- package/generated/ai/provider.ts +104 -0
- package/generated/ai/providers/claude-agent-sdk.ts +441 -0
- package/generated/ai/providers/codex-sdk.ts +430 -0
- package/generated/ai/providers/opencode-sdk.ts +491 -0
- package/generated/ai/providers/pi-events.ts +111 -0
- package/generated/ai/providers/pi-sdk-node.ts +377 -0
- package/generated/ai/providers/pi-sdk.ts +442 -0
- package/generated/ai/session-manager.ts +196 -0
- package/generated/ai/types.ts +370 -0
- package/generated/resolve-file.ts +28 -0
- package/index.ts +74 -45
- package/package.json +2 -2
- package/plannotator.html +70 -70
- package/review-editor.html +2 -2
- package/server/serverAnnotate.ts +2 -1
- package/server/serverReview.ts +5 -5
|
@@ -0,0 +1,442 @@
|
|
|
1
|
+
// @generated — DO NOT EDIT. Source: packages/ai/providers/pi-sdk.ts
|
|
2
|
+
/**
|
|
3
|
+
* Pi SDK provider — bridges Plannotator's AI layer with Pi's coding agent.
|
|
4
|
+
*
|
|
5
|
+
* Spawns `pi --mode rpc` as a subprocess and communicates via JSONL over
|
|
6
|
+
* stdio. No Pi SDK is imported — this is a thin protocol adapter.
|
|
7
|
+
*
|
|
8
|
+
* One subprocess per session. The user must have the `pi` CLI installed.
|
|
9
|
+
*/
|
|
10
|
+
|
|
11
|
+
import { BaseSession } from "../base-session.ts";
|
|
12
|
+
import { buildEffectivePrompt, buildSystemPrompt } from "../context.ts";
|
|
13
|
+
import type {
|
|
14
|
+
AIMessage,
|
|
15
|
+
AIProvider,
|
|
16
|
+
AIProviderCapabilities,
|
|
17
|
+
CreateSessionOptions,
|
|
18
|
+
PiSDKConfig,
|
|
19
|
+
} from "../types.ts";
|
|
20
|
+
|
|
21
|
+
// ---------------------------------------------------------------------------
|
|
22
|
+
// Constants
|
|
23
|
+
// ---------------------------------------------------------------------------
|
|
24
|
+
|
|
25
|
+
const PROVIDER_NAME = "pi-sdk";
|
|
26
|
+
|
|
27
|
+
// ---------------------------------------------------------------------------
|
|
28
|
+
// JSONL subprocess wrapper
|
|
29
|
+
// ---------------------------------------------------------------------------
|
|
30
|
+
|
|
31
|
+
type EventListener = (event: Record<string, unknown>) => void;
|
|
32
|
+
|
|
33
|
+
class PiProcess {
|
|
34
|
+
private proc: ReturnType<typeof Bun.spawn> | null = null;
|
|
35
|
+
private listeners: EventListener[] = [];
|
|
36
|
+
private pendingRequests = new Map<
|
|
37
|
+
string,
|
|
38
|
+
{
|
|
39
|
+
resolve: (data: Record<string, unknown>) => void;
|
|
40
|
+
reject: (err: Error) => void;
|
|
41
|
+
}
|
|
42
|
+
>();
|
|
43
|
+
private nextId = 0;
|
|
44
|
+
private buffer = "";
|
|
45
|
+
private _alive = false;
|
|
46
|
+
|
|
47
|
+
async spawn(piPath: string, cwd: string): Promise<void> {
|
|
48
|
+
this.proc = Bun.spawn([piPath, "--mode", "rpc"], {
|
|
49
|
+
cwd,
|
|
50
|
+
stdin: "pipe",
|
|
51
|
+
stdout: "pipe",
|
|
52
|
+
stderr: "pipe",
|
|
53
|
+
});
|
|
54
|
+
this._alive = true;
|
|
55
|
+
|
|
56
|
+
this.readStream();
|
|
57
|
+
|
|
58
|
+
this.proc.exited.then(() => {
|
|
59
|
+
this._alive = false;
|
|
60
|
+
for (const [, pending] of this.pendingRequests) {
|
|
61
|
+
pending.reject(new Error("Pi process exited unexpectedly"));
|
|
62
|
+
}
|
|
63
|
+
this.pendingRequests.clear();
|
|
64
|
+
// Signal active query listeners so the drain loop exits with an error
|
|
65
|
+
for (const listener of this.listeners) {
|
|
66
|
+
listener({ type: "process_exited" });
|
|
67
|
+
}
|
|
68
|
+
});
|
|
69
|
+
}
|
|
70
|
+
|
|
71
|
+
private async readStream(): Promise<void> {
|
|
72
|
+
if (!this.proc?.stdout || typeof this.proc.stdout === "number") return;
|
|
73
|
+
const reader = (this.proc.stdout as ReadableStream<Uint8Array>).getReader();
|
|
74
|
+
const decoder = new TextDecoder();
|
|
75
|
+
|
|
76
|
+
try {
|
|
77
|
+
while (true) {
|
|
78
|
+
const { done, value } = await reader.read();
|
|
79
|
+
if (done) break;
|
|
80
|
+
|
|
81
|
+
this.buffer += decoder.decode(value, { stream: true });
|
|
82
|
+
const lines = this.buffer.split("\n");
|
|
83
|
+
this.buffer = lines.pop() ?? "";
|
|
84
|
+
|
|
85
|
+
for (const line of lines) {
|
|
86
|
+
const trimmed = line.replace(/\r$/, "");
|
|
87
|
+
if (!trimmed) continue;
|
|
88
|
+
try {
|
|
89
|
+
const parsed = JSON.parse(trimmed);
|
|
90
|
+
this.routeMessage(parsed);
|
|
91
|
+
} catch {
|
|
92
|
+
// Ignore malformed lines
|
|
93
|
+
}
|
|
94
|
+
}
|
|
95
|
+
}
|
|
96
|
+
} catch {
|
|
97
|
+
// Stream closed
|
|
98
|
+
}
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
private routeMessage(msg: Record<string, unknown>): void {
|
|
102
|
+
// Response to a command we sent
|
|
103
|
+
if (msg.type === "response" && typeof msg.id === "string") {
|
|
104
|
+
const pending = this.pendingRequests.get(msg.id);
|
|
105
|
+
if (pending) {
|
|
106
|
+
this.pendingRequests.delete(msg.id);
|
|
107
|
+
if (msg.success === false) {
|
|
108
|
+
pending.reject(new Error((msg.error as string) ?? "RPC error"));
|
|
109
|
+
} else {
|
|
110
|
+
pending.resolve((msg.data as Record<string, unknown>) ?? {});
|
|
111
|
+
}
|
|
112
|
+
return;
|
|
113
|
+
}
|
|
114
|
+
}
|
|
115
|
+
|
|
116
|
+
// Agent event — forward to listeners
|
|
117
|
+
for (const listener of this.listeners) {
|
|
118
|
+
listener(msg);
|
|
119
|
+
}
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
/** Send a command without waiting for a response. */
|
|
123
|
+
send(command: Record<string, unknown>): void {
|
|
124
|
+
if (!this.proc?.stdin || typeof this.proc.stdin === "number") return;
|
|
125
|
+
// Bun.spawn stdin is a FileSink with .write(), not a WritableStream
|
|
126
|
+
const sink = this.proc.stdin as { write(data: string): void; flush(): void };
|
|
127
|
+
sink.write(`${JSON.stringify(command)}\n`);
|
|
128
|
+
sink.flush();
|
|
129
|
+
}
|
|
130
|
+
|
|
131
|
+
/** Send a command and wait for the correlated response. */
|
|
132
|
+
sendAndWait(
|
|
133
|
+
command: Record<string, unknown>,
|
|
134
|
+
): Promise<Record<string, unknown>> {
|
|
135
|
+
const id = `req_${++this.nextId}`;
|
|
136
|
+
return new Promise((resolve, reject) => {
|
|
137
|
+
this.pendingRequests.set(id, { resolve, reject });
|
|
138
|
+
this.send({ ...command, id });
|
|
139
|
+
});
|
|
140
|
+
}
|
|
141
|
+
|
|
142
|
+
/** Register a listener for agent events (non-response messages). */
|
|
143
|
+
onEvent(listener: EventListener): () => void {
|
|
144
|
+
this.listeners.push(listener);
|
|
145
|
+
return () => {
|
|
146
|
+
const idx = this.listeners.indexOf(listener);
|
|
147
|
+
if (idx >= 0) this.listeners.splice(idx, 1);
|
|
148
|
+
};
|
|
149
|
+
}
|
|
150
|
+
|
|
151
|
+
get alive(): boolean {
|
|
152
|
+
return this._alive;
|
|
153
|
+
}
|
|
154
|
+
|
|
155
|
+
kill(): void {
|
|
156
|
+
this._alive = false;
|
|
157
|
+
if (this.proc) {
|
|
158
|
+
this.proc.kill();
|
|
159
|
+
this.proc = null;
|
|
160
|
+
}
|
|
161
|
+
this.listeners.length = 0;
|
|
162
|
+
for (const [, pending] of this.pendingRequests) {
|
|
163
|
+
pending.reject(new Error("Process killed"));
|
|
164
|
+
}
|
|
165
|
+
this.pendingRequests.clear();
|
|
166
|
+
}
|
|
167
|
+
}
|
|
168
|
+
|
|
169
|
+
// ---------------------------------------------------------------------------
|
|
170
|
+
// Provider
|
|
171
|
+
// ---------------------------------------------------------------------------
|
|
172
|
+
|
|
173
|
+
export class PiSDKProvider implements AIProvider {
|
|
174
|
+
readonly name = PROVIDER_NAME;
|
|
175
|
+
readonly capabilities: AIProviderCapabilities = {
|
|
176
|
+
fork: false,
|
|
177
|
+
resume: false,
|
|
178
|
+
streaming: true,
|
|
179
|
+
tools: true,
|
|
180
|
+
};
|
|
181
|
+
models?: Array<{ id: string; label: string; default?: boolean }>;
|
|
182
|
+
|
|
183
|
+
private config: PiSDKConfig;
|
|
184
|
+
private sessions = new Map<string, PiSDKSession>();
|
|
185
|
+
|
|
186
|
+
constructor(config: PiSDKConfig) {
|
|
187
|
+
this.config = config;
|
|
188
|
+
}
|
|
189
|
+
|
|
190
|
+
async createSession(options: CreateSessionOptions): Promise<PiSDKSession> {
|
|
191
|
+
const session = new PiSDKSession({
|
|
192
|
+
systemPrompt: buildSystemPrompt(options.context),
|
|
193
|
+
cwd: options.cwd ?? this.config.cwd ?? process.cwd(),
|
|
194
|
+
parentSessionId: null,
|
|
195
|
+
piExecutablePath: this.config.piExecutablePath ?? "pi",
|
|
196
|
+
model: options.model ?? this.config.model,
|
|
197
|
+
});
|
|
198
|
+
this.sessions.set(session.id, session);
|
|
199
|
+
return session;
|
|
200
|
+
}
|
|
201
|
+
|
|
202
|
+
async forkSession(): Promise<never> {
|
|
203
|
+
throw new Error(
|
|
204
|
+
"Pi does not support session forking. " +
|
|
205
|
+
"The endpoint layer should fall back to createSession().",
|
|
206
|
+
);
|
|
207
|
+
}
|
|
208
|
+
|
|
209
|
+
async resumeSession(): Promise<never> {
|
|
210
|
+
throw new Error("Pi does not support session resuming.");
|
|
211
|
+
}
|
|
212
|
+
|
|
213
|
+
dispose(): void {
|
|
214
|
+
for (const session of this.sessions.values()) {
|
|
215
|
+
session.killProcess();
|
|
216
|
+
}
|
|
217
|
+
this.sessions.clear();
|
|
218
|
+
}
|
|
219
|
+
|
|
220
|
+
/** Fetch available models from Pi. Call before registering the provider. */
|
|
221
|
+
async fetchModels(): Promise<void> {
|
|
222
|
+
const piPath = this.config.piExecutablePath ?? "pi";
|
|
223
|
+
|
|
224
|
+
let proc: PiProcess | undefined;
|
|
225
|
+
|
|
226
|
+
try {
|
|
227
|
+
proc = new PiProcess();
|
|
228
|
+
await proc.spawn(piPath, this.config.cwd ?? process.cwd());
|
|
229
|
+
|
|
230
|
+
const data = await Promise.race([
|
|
231
|
+
proc.sendAndWait({ type: "get_available_models" }),
|
|
232
|
+
new Promise<never>((_, reject) =>
|
|
233
|
+
setTimeout(() => reject(new Error("Timeout")), 10_000),
|
|
234
|
+
),
|
|
235
|
+
]);
|
|
236
|
+
|
|
237
|
+
const rawModels = (
|
|
238
|
+
data as {
|
|
239
|
+
models?: Array<{ provider: string; id: string; name?: string }>;
|
|
240
|
+
}
|
|
241
|
+
).models;
|
|
242
|
+
if (rawModels && rawModels.length > 0) {
|
|
243
|
+
this.models = rawModels.map((m, i) => ({
|
|
244
|
+
id: `${m.provider}/${m.id}`,
|
|
245
|
+
label: m.name ?? m.id,
|
|
246
|
+
...(i === 0 && { default: true }),
|
|
247
|
+
}));
|
|
248
|
+
}
|
|
249
|
+
} catch {
|
|
250
|
+
// Pi not configured or no models available
|
|
251
|
+
} finally {
|
|
252
|
+
proc?.kill();
|
|
253
|
+
}
|
|
254
|
+
}
|
|
255
|
+
}
|
|
256
|
+
|
|
257
|
+
// ---------------------------------------------------------------------------
|
|
258
|
+
// Session
|
|
259
|
+
// ---------------------------------------------------------------------------
|
|
260
|
+
|
|
261
|
+
interface SessionConfig {
|
|
262
|
+
systemPrompt: string;
|
|
263
|
+
cwd: string;
|
|
264
|
+
parentSessionId: string | null;
|
|
265
|
+
piExecutablePath: string;
|
|
266
|
+
/** Model in "provider/modelId" format, e.g. "anthropic/claude-haiku-4-5". */
|
|
267
|
+
model?: string;
|
|
268
|
+
}
|
|
269
|
+
|
|
270
|
+
class PiSDKSession extends BaseSession {
|
|
271
|
+
private config: SessionConfig;
|
|
272
|
+
private process: PiProcess | null = null;
|
|
273
|
+
|
|
274
|
+
constructor(config: SessionConfig) {
|
|
275
|
+
super({ parentSessionId: config.parentSessionId });
|
|
276
|
+
this.config = config;
|
|
277
|
+
}
|
|
278
|
+
|
|
279
|
+
async *query(prompt: string): AsyncIterable<AIMessage> {
|
|
280
|
+
const started = this.startQuery();
|
|
281
|
+
if (!started) {
|
|
282
|
+
yield BaseSession.BUSY_ERROR;
|
|
283
|
+
return;
|
|
284
|
+
}
|
|
285
|
+
const { gen } = started;
|
|
286
|
+
|
|
287
|
+
try {
|
|
288
|
+
// Lazy-spawn subprocess
|
|
289
|
+
if (!this.process || !this.process.alive) {
|
|
290
|
+
this.process = new PiProcess();
|
|
291
|
+
await this.process.spawn(this.config.piExecutablePath, this.config.cwd);
|
|
292
|
+
|
|
293
|
+
// Set model if specified (format: "provider/modelId")
|
|
294
|
+
if (this.config.model) {
|
|
295
|
+
const [provider, ...rest] = this.config.model.split("/");
|
|
296
|
+
const modelId = rest.join("/");
|
|
297
|
+
if (provider && modelId) {
|
|
298
|
+
try {
|
|
299
|
+
await this.process.sendAndWait({
|
|
300
|
+
type: "set_model",
|
|
301
|
+
provider,
|
|
302
|
+
modelId,
|
|
303
|
+
});
|
|
304
|
+
} catch {
|
|
305
|
+
// Continue with Pi's default model
|
|
306
|
+
}
|
|
307
|
+
}
|
|
308
|
+
}
|
|
309
|
+
|
|
310
|
+
// Get session ID
|
|
311
|
+
try {
|
|
312
|
+
const state = await this.process.sendAndWait({ type: "get_state" });
|
|
313
|
+
if (typeof state.sessionId === "string") {
|
|
314
|
+
this.resolveId(state.sessionId);
|
|
315
|
+
}
|
|
316
|
+
} catch {
|
|
317
|
+
// Continue with placeholder ID
|
|
318
|
+
}
|
|
319
|
+
|
|
320
|
+
// If subprocess died during startup, surface the error immediately
|
|
321
|
+
if (!this.process.alive) {
|
|
322
|
+
yield {
|
|
323
|
+
type: "error",
|
|
324
|
+
error:
|
|
325
|
+
"Pi process exited during startup. Check that Pi is configured correctly (API keys, models).",
|
|
326
|
+
code: "pi_startup_error",
|
|
327
|
+
};
|
|
328
|
+
return;
|
|
329
|
+
}
|
|
330
|
+
}
|
|
331
|
+
|
|
332
|
+
// Build effective prompt (prepend system prompt on first query)
|
|
333
|
+
const effectivePrompt = buildEffectivePrompt(
|
|
334
|
+
prompt,
|
|
335
|
+
this.config.systemPrompt,
|
|
336
|
+
this._firstQuerySent,
|
|
337
|
+
);
|
|
338
|
+
|
|
339
|
+
// Set up async queue to bridge callback events → async iterable
|
|
340
|
+
const queue: AIMessage[] = [];
|
|
341
|
+
let resolve: (() => void) | null = null;
|
|
342
|
+
let done = false;
|
|
343
|
+
|
|
344
|
+
const push = (msg: AIMessage) => {
|
|
345
|
+
queue.push(msg);
|
|
346
|
+
resolve?.();
|
|
347
|
+
};
|
|
348
|
+
|
|
349
|
+
const finish = () => {
|
|
350
|
+
done = true;
|
|
351
|
+
resolve?.();
|
|
352
|
+
};
|
|
353
|
+
|
|
354
|
+
const unsubscribe = this.process.onEvent((event) => {
|
|
355
|
+
const mapped = mapPiEvent(event, this.id);
|
|
356
|
+
for (const msg of mapped) {
|
|
357
|
+
push(msg);
|
|
358
|
+
if (
|
|
359
|
+
msg.type === "result" ||
|
|
360
|
+
(msg.type === "error" &&
|
|
361
|
+
(event.type === "agent_end" || event.type === "process_exited"))
|
|
362
|
+
) {
|
|
363
|
+
finish();
|
|
364
|
+
}
|
|
365
|
+
}
|
|
366
|
+
});
|
|
367
|
+
|
|
368
|
+
// Send prompt — use sendAndWait to catch RPC-level rejections
|
|
369
|
+
// (e.g. expired credentials, invalid session)
|
|
370
|
+
try {
|
|
371
|
+
await this.process.sendAndWait({
|
|
372
|
+
type: "prompt",
|
|
373
|
+
message: effectivePrompt,
|
|
374
|
+
});
|
|
375
|
+
} catch (err) {
|
|
376
|
+
unsubscribe();
|
|
377
|
+
yield {
|
|
378
|
+
type: "error",
|
|
379
|
+
error: `Pi rejected prompt: ${err instanceof Error ? err.message : String(err)}`,
|
|
380
|
+
code: "pi_prompt_rejected",
|
|
381
|
+
};
|
|
382
|
+
return;
|
|
383
|
+
}
|
|
384
|
+
this._firstQuerySent = true;
|
|
385
|
+
|
|
386
|
+
// Drain queue
|
|
387
|
+
try {
|
|
388
|
+
while (!done || queue.length > 0) {
|
|
389
|
+
if (queue.length > 0) {
|
|
390
|
+
yield queue.shift()!;
|
|
391
|
+
} else {
|
|
392
|
+
await new Promise<void>((r) => {
|
|
393
|
+
resolve = r;
|
|
394
|
+
});
|
|
395
|
+
resolve = null;
|
|
396
|
+
}
|
|
397
|
+
}
|
|
398
|
+
} finally {
|
|
399
|
+
unsubscribe();
|
|
400
|
+
}
|
|
401
|
+
} catch (err) {
|
|
402
|
+
yield {
|
|
403
|
+
type: "error",
|
|
404
|
+
error: err instanceof Error ? err.message : String(err),
|
|
405
|
+
code: "provider_error",
|
|
406
|
+
};
|
|
407
|
+
} finally {
|
|
408
|
+
this.endQuery(gen);
|
|
409
|
+
}
|
|
410
|
+
}
|
|
411
|
+
|
|
412
|
+
abort(): void {
|
|
413
|
+
if (this.process?.alive) {
|
|
414
|
+
this.process.send({ type: "abort" });
|
|
415
|
+
}
|
|
416
|
+
super.abort();
|
|
417
|
+
}
|
|
418
|
+
|
|
419
|
+
/** Kill the subprocess. Called by the provider on dispose. */
|
|
420
|
+
killProcess(): void {
|
|
421
|
+
this.process?.kill();
|
|
422
|
+
this.process = null;
|
|
423
|
+
}
|
|
424
|
+
}
|
|
425
|
+
|
|
426
|
+
// ---------------------------------------------------------------------------
|
|
427
|
+
// Event mapping — shared with pi-sdk-node.ts
|
|
428
|
+
// ---------------------------------------------------------------------------
|
|
429
|
+
|
|
430
|
+
import { mapPiEvent } from "./pi-events.ts";
|
|
431
|
+
export { mapPiEvent } from "./pi-events.ts";
|
|
432
|
+
|
|
433
|
+
// ---------------------------------------------------------------------------
|
|
434
|
+
// Factory registration
|
|
435
|
+
// ---------------------------------------------------------------------------
|
|
436
|
+
|
|
437
|
+
import { registerProviderFactory } from "../provider.ts";
|
|
438
|
+
|
|
439
|
+
registerProviderFactory(
|
|
440
|
+
PROVIDER_NAME,
|
|
441
|
+
async (config) => new PiSDKProvider(config as PiSDKConfig),
|
|
442
|
+
);
|
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
// @generated — DO NOT EDIT. Source: packages/ai/session-manager.ts
|
|
2
|
+
/**
|
|
3
|
+
* Session manager — tracks active and historical AI sessions.
|
|
4
|
+
*
|
|
5
|
+
* Each Plannotator server instance (plan review, code review, annotate)
|
|
6
|
+
* gets its own SessionManager. It tracks:
|
|
7
|
+
*
|
|
8
|
+
* - Active sessions (currently streaming or idle but resumable)
|
|
9
|
+
* - The lineage from forked sessions back to their parent
|
|
10
|
+
* - Metadata for UI display (timestamps, mode, status)
|
|
11
|
+
*
|
|
12
|
+
* This is an in-memory store scoped to the server's lifetime. Sessions
|
|
13
|
+
* are not persisted to disk by the manager (the underlying provider
|
|
14
|
+
* handles its own persistence via the agent SDK).
|
|
15
|
+
*/
|
|
16
|
+
|
|
17
|
+
import type { AISession, AIContextMode } from "./types.ts";
|
|
18
|
+
|
|
19
|
+
// ---------------------------------------------------------------------------
|
|
20
|
+
// Types
|
|
21
|
+
// ---------------------------------------------------------------------------
|
|
22
|
+
|
|
23
|
+
export interface SessionEntry {
|
|
24
|
+
/** The live session handle (if still active). */
|
|
25
|
+
session: AISession;
|
|
26
|
+
/** What mode this session was created for. */
|
|
27
|
+
mode: AIContextMode;
|
|
28
|
+
/** The parent session ID this was forked from (null if standalone). */
|
|
29
|
+
parentSessionId: string | null;
|
|
30
|
+
/** When this session was created. */
|
|
31
|
+
createdAt: number;
|
|
32
|
+
/** When the last query was sent. */
|
|
33
|
+
lastActiveAt: number;
|
|
34
|
+
/** Short description for UI display (e.g., the user's first question). */
|
|
35
|
+
label?: string;
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
export interface SessionManagerOptions {
|
|
39
|
+
/**
|
|
40
|
+
* Maximum number of sessions to keep in the manager.
|
|
41
|
+
* Oldest idle sessions are evicted when the limit is reached.
|
|
42
|
+
* Default: 20.
|
|
43
|
+
*/
|
|
44
|
+
maxSessions?: number;
|
|
45
|
+
}
|
|
46
|
+
|
|
47
|
+
// ---------------------------------------------------------------------------
|
|
48
|
+
// Implementation
|
|
49
|
+
// ---------------------------------------------------------------------------
|
|
50
|
+
|
|
51
|
+
export class SessionManager {
|
|
52
|
+
private sessions = new Map<string, SessionEntry>();
|
|
53
|
+
private aliases = new Map<string, string>();
|
|
54
|
+
private maxSessions: number;
|
|
55
|
+
|
|
56
|
+
constructor(options: SessionManagerOptions = {}) {
|
|
57
|
+
this.maxSessions = options.maxSessions ?? 20;
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
/**
|
|
61
|
+
* Track a newly created session.
|
|
62
|
+
*
|
|
63
|
+
* If the session supports ID resolution (e.g., the real SDK session ID
|
|
64
|
+
* arrives after the first query), call `remapId()` to update the key.
|
|
65
|
+
*/
|
|
66
|
+
track(session: AISession, mode: AIContextMode, label?: string): SessionEntry {
|
|
67
|
+
this.evictIfNeeded();
|
|
68
|
+
|
|
69
|
+
const entry: SessionEntry = {
|
|
70
|
+
session,
|
|
71
|
+
mode,
|
|
72
|
+
parentSessionId: session.parentSessionId,
|
|
73
|
+
createdAt: Date.now(),
|
|
74
|
+
lastActiveAt: Date.now(),
|
|
75
|
+
label,
|
|
76
|
+
};
|
|
77
|
+
this.sessions.set(session.id, entry);
|
|
78
|
+
|
|
79
|
+
// Wire up ID remapping so providers can resolve the real session ID later
|
|
80
|
+
session.onIdResolved = (oldId, newId) => this.remapId(oldId, newId);
|
|
81
|
+
|
|
82
|
+
return entry;
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
/**
|
|
86
|
+
* Remap a session from one ID to another.
|
|
87
|
+
* Used when the real session ID is resolved after initial tracking.
|
|
88
|
+
*/
|
|
89
|
+
remapId(oldId: string, newId: string): void {
|
|
90
|
+
const entry = this.sessions.get(oldId);
|
|
91
|
+
if (entry) {
|
|
92
|
+
this.sessions.delete(oldId);
|
|
93
|
+
this.sessions.set(newId, entry);
|
|
94
|
+
// Keep the old ID as an alias so clients using the original ID still work
|
|
95
|
+
this.aliases.set(oldId, newId);
|
|
96
|
+
}
|
|
97
|
+
}
|
|
98
|
+
|
|
99
|
+
/** Resolve an alias to the canonical ID, or return the ID as-is. */
|
|
100
|
+
private resolve(sessionId: string): string {
|
|
101
|
+
return this.aliases.get(sessionId) ?? sessionId;
|
|
102
|
+
}
|
|
103
|
+
|
|
104
|
+
/**
|
|
105
|
+
* Get a tracked session by ID (or alias).
|
|
106
|
+
*/
|
|
107
|
+
get(sessionId: string): SessionEntry | undefined {
|
|
108
|
+
return this.sessions.get(this.resolve(sessionId));
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
/**
|
|
112
|
+
* Mark a session as recently active (updates lastActiveAt).
|
|
113
|
+
*/
|
|
114
|
+
touch(sessionId: string): void {
|
|
115
|
+
const entry = this.sessions.get(this.resolve(sessionId));
|
|
116
|
+
if (entry) {
|
|
117
|
+
entry.lastActiveAt = Date.now();
|
|
118
|
+
}
|
|
119
|
+
}
|
|
120
|
+
|
|
121
|
+
/**
|
|
122
|
+
* Remove a session from tracking.
|
|
123
|
+
* Does NOT abort the session — call session.abort() first if needed.
|
|
124
|
+
*/
|
|
125
|
+
remove(sessionId: string): void {
|
|
126
|
+
const canonical = this.resolve(sessionId);
|
|
127
|
+
this.sessions.delete(canonical);
|
|
128
|
+
// Clean up any aliases pointing to this session
|
|
129
|
+
for (const [alias, target] of this.aliases) {
|
|
130
|
+
if (target === canonical) this.aliases.delete(alias);
|
|
131
|
+
}
|
|
132
|
+
}
|
|
133
|
+
|
|
134
|
+
/**
|
|
135
|
+
* List all tracked sessions, newest first.
|
|
136
|
+
*/
|
|
137
|
+
list(): SessionEntry[] {
|
|
138
|
+
return [...this.sessions.values()].sort(
|
|
139
|
+
(a, b) => b.lastActiveAt - a.lastActiveAt
|
|
140
|
+
);
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
/**
|
|
144
|
+
* List sessions forked from a specific parent.
|
|
145
|
+
*/
|
|
146
|
+
forksOf(parentSessionId: string): SessionEntry[] {
|
|
147
|
+
return this.list().filter(
|
|
148
|
+
(e) => e.parentSessionId === parentSessionId
|
|
149
|
+
);
|
|
150
|
+
}
|
|
151
|
+
|
|
152
|
+
/**
|
|
153
|
+
* Get the number of tracked sessions.
|
|
154
|
+
*/
|
|
155
|
+
get size(): number {
|
|
156
|
+
return this.sessions.size;
|
|
157
|
+
}
|
|
158
|
+
|
|
159
|
+
/**
|
|
160
|
+
* Abort all active sessions and clear tracking.
|
|
161
|
+
*/
|
|
162
|
+
disposeAll(): void {
|
|
163
|
+
for (const entry of this.sessions.values()) {
|
|
164
|
+
if (entry.session.isActive) {
|
|
165
|
+
entry.session.abort();
|
|
166
|
+
}
|
|
167
|
+
}
|
|
168
|
+
this.sessions.clear();
|
|
169
|
+
this.aliases.clear();
|
|
170
|
+
}
|
|
171
|
+
|
|
172
|
+
// -------------------------------------------------------------------------
|
|
173
|
+
// Internal
|
|
174
|
+
// -------------------------------------------------------------------------
|
|
175
|
+
|
|
176
|
+
private evictIfNeeded(): void {
|
|
177
|
+
if (this.sessions.size < this.maxSessions) return;
|
|
178
|
+
|
|
179
|
+
// Find the oldest idle session to evict
|
|
180
|
+
let oldest: { id: string; at: number } | null = null;
|
|
181
|
+
for (const [id, entry] of this.sessions) {
|
|
182
|
+
if (entry.session.isActive) continue; // don't evict active sessions
|
|
183
|
+
if (!oldest || entry.lastActiveAt < oldest.at) {
|
|
184
|
+
oldest = { id, at: entry.lastActiveAt };
|
|
185
|
+
}
|
|
186
|
+
}
|
|
187
|
+
|
|
188
|
+
if (oldest) {
|
|
189
|
+
this.sessions.delete(oldest.id);
|
|
190
|
+
// Clean up aliases pointing to the evicted session
|
|
191
|
+
for (const [alias, target] of this.aliases) {
|
|
192
|
+
if (target === oldest.id) this.aliases.delete(alias);
|
|
193
|
+
}
|
|
194
|
+
}
|
|
195
|
+
}
|
|
196
|
+
}
|