pi-llama-cpp 0.12.0 → 0.14.0

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.
Files changed (68) hide show
  1. package/README.md +20 -9
  2. package/package.json +3 -3
  3. package/src/api/client.ts +25 -0
  4. package/src/constants.ts +8 -5
  5. package/src/enums/status.ts +0 -1
  6. package/src/interfaces/endpoints/models.ts +1 -1
  7. package/src/interfaces/settings.ts +3 -2
  8. package/src/interfaces/sortBy.ts +4 -0
  9. package/src/managers/command/models.ts +247 -0
  10. package/src/managers/command.ts +29 -459
  11. package/src/managers/events.ts +2 -2
  12. package/src/managers/server.ts +29 -5
  13. package/src/managers/settings.ts +12 -88
  14. package/src/models/baseModel.ts +27 -10
  15. package/src/models/legacyModel.ts +2 -2
  16. package/src/models/routerModel.ts +2 -2
  17. package/src/models/singleModel.ts +1 -1
  18. package/src/server.ts +35 -48
  19. package/src/sse/client.ts +113 -59
  20. package/src/sse/fetch.ts +43 -0
  21. package/src/sse/manager.ts +9 -27
  22. package/src/sse/types.ts +0 -4
  23. package/src/ui/dialog/base.ts +118 -0
  24. package/src/ui/dialog/confirm.ts +45 -0
  25. package/src/ui/dialog/factory.ts +111 -0
  26. package/src/ui/dialog/input.ts +63 -0
  27. package/src/ui/dialog/options.ts +23 -0
  28. package/src/ui/editors/editorOptions.ts +43 -0
  29. package/src/ui/editors/itemBuilder.ts +47 -0
  30. package/src/ui/editors/listEditor.ts +291 -0
  31. package/src/ui/editors/override/entry.ts +24 -0
  32. package/src/ui/editors/override/entryEditor.ts +166 -0
  33. package/src/ui/editors/override/fields.ts +294 -0
  34. package/src/ui/editors/override/handlers.ts +118 -0
  35. package/src/ui/editors/override/itemBuilder.ts +42 -0
  36. package/src/ui/editors/override/overrideList.ts +127 -0
  37. package/src/ui/editors/server/builder.ts +60 -0
  38. package/src/ui/editors/server/fields.ts +84 -0
  39. package/src/ui/editors/server/handlers.ts +32 -0
  40. package/src/ui/editors/server/itemBuilder.ts +68 -0
  41. package/src/ui/editors/server/serverEditor.ts +161 -0
  42. package/src/ui/editors/server/utils.ts +45 -0
  43. package/src/ui/editors/server/wizard.ts +110 -0
  44. package/src/ui/editors/settingField.ts +37 -0
  45. package/src/ui/editors/settingsListFactory.ts +33 -0
  46. package/src/ui/settings/index.ts +237 -0
  47. package/src/ui/strings.ts +8 -4
  48. package/src/utils/health.ts +48 -0
  49. package/src/utils/serverIds.ts +21 -0
  50. package/src/utils/settingsStore.ts +1 -1
  51. package/src/utils/urlResolver.ts +129 -0
  52. package/src/utils/urls.ts +33 -13
  53. package/tests/commandManager.test.ts +38 -7
  54. package/tests/dialog.test.ts +95 -2
  55. package/tests/health.test.ts +116 -0
  56. package/tests/legacyModel.test.ts +34 -28
  57. package/tests/overrides.test.ts +123 -68
  58. package/tests/routerModel.test.ts +67 -68
  59. package/tests/server.test.ts +47 -16
  60. package/tests/serverManager.test.ts +4 -4
  61. package/tests/settings.test.ts +10 -8
  62. package/tests/singleModel.test.ts +12 -12
  63. package/tests/sseManager.test.ts +6 -24
  64. package/src/ui/dialog.ts +0 -287
  65. package/src/ui/overrideEntryEditor.ts +0 -119
  66. package/src/ui/overrideSettingsList.ts +0 -682
  67. package/src/ui/serverListEditor.ts +0 -32
  68. package/src/ui/serverSettingsList.ts +0 -466
@@ -9,23 +9,22 @@ import { join } from "node:path";
9
9
  import {
10
10
  API_KEY_PLACEHOLDER,
11
11
  AUTOLOAD_ON_MESSAGE,
12
- LLAMA_SERVER_URL,
13
12
  POLLING_TIMEOUT,
14
13
  REACT_TO_MODEL_SELECT,
15
14
  SERVER_TIMEOUT,
16
15
  SETTINGS_KEY,
17
16
  SORT_BY,
18
17
  THINKING_BUDGETS,
19
- type SortBy,
20
18
  } from "../constants";
21
19
  import {
22
20
  LlamaServer,
23
21
  LlamaSettings,
24
22
  ModelOverride,
25
23
  } from "../interfaces/settings";
24
+ import type { SortBy } from "../interfaces/sortBy";
26
25
  import { Server } from "../server";
27
26
  import { SettingsStore } from "../utils/settingsStore";
28
- import { isValidServerUrl, normalizeUrl } from "../utils/urls";
27
+ import { UrlResolver } from "../utils/urlResolver";
29
28
 
30
29
  export class LlamaSettingsManager {
31
30
  private settingsManager = SettingsManager.create(process.cwd());
@@ -47,16 +46,17 @@ export class LlamaSettingsManager {
47
46
  }
48
47
  }
49
48
 
50
- /** Warnings collected during URL resolution (dropped invalid entries). */
51
- private warnings: string[] = [];
49
+ /** Delegated multi-source URL resolution chain (see `utils/urlResolver`). */
50
+ private urlResolver = new UrlResolver({
51
+ getLlamaSettings: () => this.getLlamaSettings(),
52
+ getMergedSettings: () => this.getMergedSettings(),
53
+ });
52
54
 
53
55
  /**
54
56
  * Returns and clears warnings collected during URL resolution.
55
57
  */
56
58
  takeWarnings(): string[] {
57
- const warnings = [...this.warnings];
58
- this.warnings.length = 0;
59
- return warnings;
59
+ return this.urlResolver.takeWarnings();
60
60
  }
61
61
 
62
62
  /**
@@ -75,7 +75,7 @@ export class LlamaSettingsManager {
75
75
  * Convenience method for the `llamaSettings` key.
76
76
  * Reloads settings from disk before reading.
77
77
  */
78
- async getLlamaSettings(): Promise<LlamaSettings> {
78
+ private async getLlamaSettings(): Promise<LlamaSettings> {
79
79
  return (await this.getMergedSettings())[SETTINGS_KEY] ?? {};
80
80
  }
81
81
 
@@ -88,89 +88,13 @@ export class LlamaSettingsManager {
88
88
  }
89
89
 
90
90
  /**
91
- * Resolves the server URLs to use in the following order:
92
- *
93
- * - `LLAMA_SERVER_URL` env variable
94
- * - `llamaSettings` key (current - project, then global)
95
- * - `llamaServerUrl` key (legacy - project, then global)
96
- * - Default URL
91
+ * Resolves the server URLs to use. Delegates to the URL resolver chain
92
+ * (env → settings → legacy → default, see `utils/urlResolver`).
97
93
  *
98
94
  * @returns The list of URLs to use
99
95
  */
100
96
  async resolveUrls(): Promise<string[]> {
101
- let response = this.resolveEnvUrls();
102
- if (response.length > 0) return response;
103
-
104
- response = await this.resolveServerUrls();
105
- if (response.length > 0) return response;
106
-
107
- response = await this.resolveLegacyUrls();
108
- if (response.length > 0) return response;
109
-
110
- return [LLAMA_SERVER_URL];
111
- }
112
-
113
- /**
114
- * Resolves the llama-server URLs from the environment variable.
115
- *
116
- * @returns A list of detected URLs
117
- */
118
- private resolveEnvUrls(): string[] {
119
- const raw = process.env.LLAMA_SERVER_URL;
120
- if (!raw) return [];
121
-
122
- return this.parseUrls(raw);
123
- }
124
-
125
- /**
126
- * Resolves the llama-server URLs from `llamaSettings.servers`.
127
- * Settings are merged, prioritizing project over global settings.
128
- * Reloads settings from disk before reading.
129
- *
130
- * @returns A list of detected URLs
131
- */
132
- private async resolveServerUrls(): Promise<string[]> {
133
- const { servers = [] } = await this.getLlamaSettings();
134
- return servers.map((s) => this.parseUrls(s.url)).flat();
135
- }
136
-
137
- /**
138
- * Resolves the llama-server URLs from `llamaServerUrl` legacy key.
139
- * Settings are merged, prioritizing project over global settings.
140
- * Reloads settings from disk before reading.
141
- *
142
- * @returns A list of detected URLs
143
- */
144
- private async resolveLegacyUrls(): Promise<string[]> {
145
- const { llamaServerUrl = null } = await this.getMergedSettings();
146
- if (!llamaServerUrl) return [];
147
-
148
- return this.parseUrls(llamaServerUrl);
149
- }
150
-
151
- /**
152
- * Parses a raw URL string into an array of cleaned URLs.
153
- * Splits on semicolons, trims whitespace, filters empty strings, strips
154
- * trailing slashes, and drops entries without an http(s) scheme —
155
- * collecting a warning for each dropped entry (same validation the
156
- * `/models servers` editor applies).
157
- *
158
- * @returns A sanitized URL
159
- */
160
- private parseUrls(raw: string): string[] {
161
- return raw
162
- .split(";")
163
- .map(normalizeUrl)
164
- .filter((u) => {
165
- if (u.length === 0) return false;
166
- if (!isValidServerUrl(u)) {
167
- this.warnings.push(
168
- `Ignoring invalid server URL '${u}' (needs http(s)://)`,
169
- );
170
- return false;
171
- }
172
- return true;
173
- });
97
+ return this.urlResolver.resolveUrls();
174
98
  }
175
99
 
176
100
  /**
@@ -23,7 +23,6 @@ export abstract class BaseModel {
23
23
  [Status.FAILED]: "🔴",
24
24
  [Status.SLEEPING]: "🔵",
25
25
  [Status.UNLOADED]: "⚪",
26
- [Status.UNAUTHORIZED]: "⛔",
27
26
  };
28
27
 
29
28
  abstract get mode(): Mode;
@@ -72,7 +71,7 @@ export abstract class BaseModel {
72
71
  *
73
72
  * @returns An array of capabilities, as expected by Pi
74
73
  */
75
- async getCapabilities(): Promise<("text" | "image")[]> {
74
+ protected async getCapabilities(): Promise<("text" | "image")[]> {
76
75
  const overridden = this.server.findOverrideForModel(this.id)?.capabilities;
77
76
  if (overridden) return overridden;
78
77
 
@@ -108,7 +107,6 @@ export abstract class BaseModel {
108
107
 
109
108
  if (is_sleeping) return Status.SLEEPING;
110
109
  if (!error) return Status.LOADED;
111
- if (error.code === 401) return Status.UNAUTHORIZED;
112
110
  if (error.code === 503) return Status.LOADING;
113
111
  if (error.code === 400 && error.message === "model is not loaded")
114
112
  return Status.UNLOADED;
@@ -128,14 +126,13 @@ export abstract class BaseModel {
128
126
  *
129
127
  * @returns The context size in tokens
130
128
  */
131
- async getContextSize(): Promise<number> {
129
+ protected async getContextSize(): Promise<number> {
132
130
  const overridden = this.server.findOverrideForModel(this.id)?.contextSize;
133
131
  if (overridden && overridden > 0) return overridden;
134
132
 
135
133
  try {
136
134
  const { data } = await this.server.fetchModels();
137
- const { n_ctx } = data.find((m) => m.id === this.id)?.meta!;
138
-
135
+ const n_ctx = data.find((m) => m.id === this.id)?.meta?.n_ctx;
139
136
  return n_ctx ?? FALLBACK_CTX;
140
137
  } catch {
141
138
  return FALLBACK_CTX;
@@ -187,6 +184,10 @@ export abstract class BaseModel {
187
184
  cacheWrite: userCost.cacheWrite ?? 0,
188
185
  };
189
186
 
187
+ const input = await this.getCapabilities();
188
+ const contextWindow = await this.getContextSize();
189
+ const maxTokens = this.getMaxTokens(contextWindow);
190
+
190
191
  const response: ProviderModelConfig = {
191
192
  id: this.id,
192
193
  name: this.name,
@@ -199,10 +200,10 @@ export abstract class BaseModel {
199
200
  xhigh: "xhigh",
200
201
  max: "max",
201
202
  },
202
- input: await this.getCapabilities(),
203
- contextWindow: await this.getContextSize(),
203
+ input,
204
+ contextWindow,
204
205
  cost,
205
- maxTokens: override.maxTokens ?? (await this.getContextSize()),
206
+ maxTokens,
206
207
  };
207
208
 
208
209
  // Add compat if the override specifies it
@@ -213,6 +214,22 @@ export abstract class BaseModel {
213
214
  return response;
214
215
  }
215
216
 
217
+ /**
218
+ * Gets the maximum number of tokens the model can generate.
219
+ *
220
+ * An override's `maxTokens` (when set and `> 0`) replaces the value;
221
+ * otherwise it falls back to the context size — a model cannot generate
222
+ * more tokens than its context holds. A stored `0` behaves as if the
223
+ * key were absent.
224
+ *
225
+ * @param contextSize - The already-resolved context size, used as fallback
226
+ * @returns The maximum number of tokens
227
+ */
228
+ protected getMaxTokens(contextSize: number): number {
229
+ const overridden = this.server.findOverrideForModel(this.id)?.maxTokens;
230
+ return overridden && overridden > 0 ? overridden : contextSize;
231
+ }
232
+
216
233
  /**
217
234
  * Loads the model in llama-server.
218
235
  * Uses SSE status events when available, falling back to polling.
@@ -269,7 +286,7 @@ export abstract class BaseModel {
269
286
  * @param timeout The maximum amount of ms before timeout. Defaults to server's pollingTimeout
270
287
  * @param interval The polling interval. Defaults to POLLING_INTERVAL
271
288
  */
272
- async pollStatus(
289
+ protected async pollStatus(
273
290
  startTime: number = Date.now(),
274
291
  timeout?: number,
275
292
  interval: number = POLLING_INTERVAL,
@@ -13,7 +13,7 @@ export class LegacyModel extends SingleModel {
13
13
  *
14
14
  * @returns The context size
15
15
  */
16
- async getContextSize(): Promise<number> {
16
+ protected async getContextSize(): Promise<number> {
17
17
  const props = await this.server.fetchModelProps(this.id);
18
18
  const models = await this.server.fetchModels();
19
19
 
@@ -34,7 +34,7 @@ export class LegacyModel extends SingleModel {
34
34
  *
35
35
  * @returns An array of capabilities, as expected by Pi
36
36
  */
37
- async getCapabilities(): Promise<("text" | "image")[]> {
37
+ protected async getCapabilities(): Promise<("text" | "image")[]> {
38
38
  try {
39
39
  return await super.getCapabilities();
40
40
  } catch {
@@ -26,7 +26,7 @@ export class RouterModel extends BaseModel {
26
26
  *
27
27
  * In exchange, it will allow unloaded models to be correctly shown as "unloaded".
28
28
  */
29
- async pollStatus(startTime = Date.now()): Promise<void> {
29
+ protected async pollStatus(startTime = Date.now()): Promise<void> {
30
30
  let elapsed = 0;
31
31
  const limit = 5000;
32
32
 
@@ -52,7 +52,7 @@ export class RouterModel extends BaseModel {
52
52
  *
53
53
  * @returns The context size in tokens
54
54
  */
55
- async getContextSize(): Promise<number> {
55
+ protected async getContextSize(): Promise<number> {
56
56
  // We can get a more accurate context size if the model is already loaded
57
57
  if ((await this.getStatus()) === Status.LOADED) {
58
58
  return super.getContextSize();
@@ -6,7 +6,7 @@ export class SingleModel extends BaseModel {
6
6
  return Mode.SINGLE;
7
7
  }
8
8
 
9
- async getCapabilities(): Promise<("text" | "image")[]> {
9
+ protected async getCapabilities(): Promise<("text" | "image")[]> {
10
10
  try {
11
11
  return await super.getCapabilities();
12
12
  } catch {
package/src/server.ts CHANGED
@@ -1,12 +1,11 @@
1
1
  import { ApiClient } from "./api/client";
2
2
  import {
3
3
  API_KEY_PLACEHOLDER,
4
+ ENDPOINT_PREFIX,
4
5
  PROVIDER_NAME,
5
- PROVIDER_PREFIX,
6
6
  } from "./constants";
7
7
  import { Mode } from "./enums/mode";
8
8
  import { ServerStatus } from "./enums/serverStatus";
9
- import { HealthEndpoint } from "./interfaces/endpoints/health";
10
9
  import { ModelsEndpoint } from "./interfaces/endpoints/models";
11
10
  import {
12
11
  PropsEndpoint,
@@ -20,20 +19,23 @@ import { LegacyModel } from "./models/legacyModel";
20
19
  import { RouterModel } from "./models/routerModel";
21
20
  import { SingleModel } from "./models/singleModel";
22
21
  import { SSEManager } from "./sse/manager";
22
+ import { checkServerHealth } from "./utils/health";
23
+ import { ServerIds } from "./utils/serverIds";
23
24
 
24
25
  /**
25
26
  * Optional constructor collaborators for {@link Server} — the seam tests use
26
27
  * to run the real Server against fake clients.
27
28
  *
28
29
  * Both are factories because their arguments only exist around construction:
29
- * the API key is (re-)resolved by the Server, and SSEManager needs its owner.
30
- * Factories must stay pure functions of their arguments — `initialize()`
31
- * re-invokes both on every scan (the ApiClient rebuild picks up a fresh key,
32
- * by design), so captured per-server state would leak across re-scans.
30
+ * the API key is (re-)resolved by the Server, and SSEManager needs its owner
31
+ * (it reads the key and timeouts live through it). Factories must stay pure
32
+ * functions of their arguments — `initialize()` re-invokes both on every scan
33
+ * (the ApiClient rebuild picks up a fresh key, by design), so captured
34
+ * per-server state would leak across re-scans.
33
35
  */
34
36
  export type ServerDeps = {
35
37
  createApiClient?: (apiKey: string) => ApiClient;
36
- createSSEManager?: (server: Server, apiKey: string) => SSEManager;
38
+ createSSEManager?: (server: Server) => SSEManager;
37
39
  };
38
40
 
39
41
  export class Server {
@@ -59,6 +61,19 @@ export class Server {
59
61
  return this.options.baseUrl;
60
62
  }
61
63
 
64
+ /**
65
+ * Base URL of the OpenAI-compatible API: {@link baseUrl} joined with
66
+ * {@link ENDPOINT_PREFIX}. Idempotent — a baseUrl that already ends with
67
+ * the prefix is returned untouched, so proxied setups exposing the API
68
+ * under a /v1 path don't double up.
69
+ */
70
+ get apiBaseUrl(): string {
71
+ const { baseUrl } = this.options;
72
+ return baseUrl.endsWith(ENDPOINT_PREFIX)
73
+ ? baseUrl
74
+ : `${baseUrl}${ENDPOINT_PREFIX}`;
75
+ }
76
+
62
77
  /**
63
78
  * Maximum time (ms) for server verification and SSE support probe.
64
79
  * Resolved live from the injected settings manager.
@@ -87,7 +102,7 @@ export class Server {
87
102
  * Uses custom ID if provided, otherwise falls back to URL-based ID.
88
103
  */
89
104
  get providerId(): string {
90
- return this.options.customId ?? `${PROVIDER_PREFIX}=${this.baseUrl}`;
105
+ return ServerIds.resolve(this.baseUrl, this.options.customId);
91
106
  }
92
107
 
93
108
  /**
@@ -114,7 +129,7 @@ export class Server {
114
129
  if (key !== API_KEY_PLACEHOLDER) return key;
115
130
  }
116
131
  // Fall back to URL-based ID
117
- return this.settings.resolveApiKey(`${PROVIDER_PREFIX}=${this.baseUrl}`);
132
+ return this.settings.resolveApiKey(ServerIds.fromUrl(this.baseUrl));
118
133
  }
119
134
 
120
135
  /**
@@ -126,9 +141,7 @@ export class Server {
126
141
  this.apiClient =
127
142
  this.deps.createApiClient?.(apiKey) ??
128
143
  new ApiClient(this.baseUrl, apiKey);
129
- this.sse =
130
- this.deps.createSSEManager?.(this, apiKey) ??
131
- new SSEManager(this, apiKey);
144
+ this.sse = this.deps.createSSEManager?.(this) ?? new SSEManager(this);
132
145
  const { data } = await this.fetchModels();
133
146
  const mode = await this.detectServerMode(data);
134
147
 
@@ -163,37 +176,16 @@ export class Server {
163
176
  /**
164
177
  * Checks if the server is ready, with a timeout.
165
178
  *
179
+ * Delegates to the shared health probe (`utils/health`) — bypasses the
180
+ * `ApiClient` on purpose: each caller probes once per scan, so the
181
+ * client's cache/dedup would be dead weight, and a plain `fetch` with
182
+ * `AbortSignal.timeout` actually cancels the request on timeout.
183
+ *
166
184
  * @param timeout Maximum time to wait for the health check
167
185
  * @returns The server status
168
186
  */
169
187
  async isReady(timeout: number): Promise<ServerStatus> {
170
- try {
171
- const timeoutPromise = new Promise<never>((_, reject) =>
172
- setTimeout(() => reject(new Error("timeout")), timeout),
173
- );
174
- const health = await Promise.race([
175
- this.fetchServerHealth(),
176
- timeoutPromise,
177
- ]);
178
- if (health.status === "ok") {
179
- return ServerStatus.READY;
180
- }
181
- return ServerStatus.UNREACHABLE;
182
- } catch (error) {
183
- if (error instanceof Error && error.message === "timeout") {
184
- return ServerStatus.TIMEOUT;
185
- }
186
- return ServerStatus.UNREACHABLE;
187
- }
188
- }
189
-
190
- /**
191
- * Retrieves the health status of the server
192
- *
193
- * @returns The health status
194
- */
195
- async fetchServerHealth(): Promise<HealthEndpoint> {
196
- return await this.apiClient.get<HealthEndpoint>("/health");
188
+ return checkServerHealth(this.baseUrl, timeout, this.getApiKey());
197
189
  }
198
190
 
199
191
  /**
@@ -202,7 +194,9 @@ export class Server {
202
194
  * @return The models from the server
203
195
  */
204
196
  async fetchModels(): Promise<ModelsEndpoint> {
205
- return await this.apiClient.get<ModelsEndpoint>("/v1/models");
197
+ return await this.apiClient.get<ModelsEndpoint>(
198
+ `${ENDPOINT_PREFIX}/models`,
199
+ );
206
200
  }
207
201
 
208
202
  /**
@@ -226,13 +220,6 @@ export class Server {
226
220
  );
227
221
  }
228
222
 
229
- /**
230
- * Returns the per-model override configuration for this server.
231
- */
232
- getOverrides(): Record<string, ModelOverride> {
233
- return this.options.overrides ?? {};
234
- }
235
-
236
223
  /**
237
224
  * Resolves the override for a given model ID using prefix matching.
238
225
  *
@@ -244,7 +231,7 @@ export class Server {
244
231
  * @returns The matching override, or `undefined` if no key matches.
245
232
  */
246
233
  findOverrideForModel(modelId: string): ModelOverride | undefined {
247
- const overrides = this.getOverrides();
234
+ const overrides = this.options.overrides ?? {};
248
235
  let best: ModelOverride | undefined;
249
236
  let bestLen = 0;
250
237
 
package/src/sse/client.ts CHANGED
@@ -1,30 +1,22 @@
1
1
  import { POLLING_INTERVAL } from "../constants";
2
+ import { openSSEStream } from "./fetch";
2
3
  import type { SSECallback, SSECleanup, SSEEvent } from "./types";
3
4
 
4
- /**
5
- * Builds the full SSE endpoint URL, appending the API key as a query
6
- * parameter when one is set. Shared by {@link SSEClient} and
7
- * {@link SSEManager.probeSSE} so the two can't drift.
8
- */
9
- export const buildSSEUrl = (endpoint: string, apiKey?: string): string => {
10
- if (apiKey) {
11
- return `${endpoint}?api_key=${encodeURIComponent(apiKey)}`;
12
- }
13
- return endpoint;
14
- };
15
-
16
5
  /**
17
6
  * SSE client for llama-server's /models/sse endpoint.
18
7
  *
19
- * Uses a single shared EventSource per server instance.
8
+ * Uses a single shared stream per server instance, consumed with `fetch`
9
+ * (EventSource cannot send the API key as a header, and llama-server no
10
+ * longer accepts it as a query parameter — see `fetch.ts`).
20
11
  * Supports multiple model subscriptions with automatic event routing.
21
12
  * Handles reconnection by re-subscribing all callbacks.
22
13
  */
23
14
  export class SSEClient {
24
- private eventSource: EventSource | null = null;
15
+ private abortController: AbortController | null = null;
16
+ private disposed: boolean = false;
25
17
  private subscribers: Map<string, SSECallback> = new Map();
26
18
  private connected: boolean = false;
27
- private reconnecting: boolean = false; // tracks if EventSource auto-reconnect is in progress
19
+ private reconnecting: boolean = false; // tracks if reconnect is in progress
28
20
  /**
29
21
  * Single shared slot — each setOnConnectFailed call overwrites the
30
22
  * previous callback (see there for the constraint this imposes).
@@ -42,7 +34,15 @@ export class SSEClient {
42
34
  ) {}
43
35
 
44
36
  /**
45
- * Connects to the SSE endpoint.
37
+ * Waits the given amount of time (used between reconnection attempts).
38
+ */
39
+ private delay(ms: number): Promise<void> {
40
+ return new Promise<void>((resolve) => setTimeout(resolve, ms));
41
+ }
42
+
43
+ /**
44
+ * Connects to the SSE endpoint and keeps it open, reconnecting with a
45
+ * fixed delay until `disconnect()` is called.
46
46
  *
47
47
  * No current caller consumes the result: `subscribe()` triggers the
48
48
  * connection without awaiting it, and connection failures before the
@@ -50,56 +50,110 @@ export class SSEClient {
50
50
  *
51
51
  * @returns true if the connection was established successfully
52
52
  */
53
- async connect(): Promise<boolean> {
53
+ private async connect(): Promise<boolean> {
54
54
  if (this.connected) return true;
55
+ this.disposed = false;
55
56
 
56
- const url = buildSSEUrl(this.sseEndpoint, this.apiKey);
57
+ // Loop-local so a stale, already-aborted controller from a previous
58
+ // connection can't be confused with the current one after a revival.
59
+ let abortController: AbortController | null = this.abortController;
57
60
 
58
- try {
59
- this.eventSource = new EventSource(url);
60
- } catch {
61
- this.connected = false;
62
- return false;
63
- }
61
+ while (!this.disposed) {
62
+ if (abortController?.signal.aborted) return false;
63
+ abortController = new AbortController();
64
+ this.abortController = abortController;
64
65
 
65
- this.eventSource.onerror = () => {
66
- // EventSource will auto-reconnect; we just track state
67
- this.connected = false;
66
+ let body: ReadableStream<Uint8Array>;
67
+ try {
68
+ body = await openSSEStream(
69
+ this.sseEndpoint,
70
+ this.apiKey,
71
+ abortController.signal,
72
+ );
73
+ } catch {
74
+ if (this.disposed || abortController.signal.aborted) return false;
75
+ this.notifyConnectFailed();
76
+ this.reconnecting = true;
77
+ await this.delay(POLLING_INTERVAL);
78
+ continue;
79
+ }
80
+
81
+ this.connected = true;
82
+ this.reconnecting = false;
83
+
84
+ await this.consume(body);
85
+
86
+ if (this.disposed || abortController.signal.aborted) return false;
87
+ // Stream ended (server closed or network error): reconnect
68
88
  this.reconnecting = true;
89
+ await this.delay(POLLING_INTERVAL);
90
+ }
69
91
 
70
- // Notify subscriber if connection fails before any event is received
71
- if (!this._hasReceivedEvents && this._onConnectFailed) {
72
- this._onConnectFailed();
73
- }
74
- };
92
+ return false;
93
+ }
75
94
 
76
- this.eventSource.onmessage = (event: MessageEvent) => {
77
- try {
78
- const data = JSON.parse(event.data);
79
- const sseEvent: SSEEvent = {
80
- event: data.event ?? "unknown",
81
- model: data.model ?? "*",
82
- data: data.data,
83
- };
84
- this._hasReceivedEvents = true;
85
- this.dispatch(sseEvent);
86
- } catch {
87
- // Invalid JSON, ignore
95
+ /**
96
+ * Reads the raw byte stream, parses the SSE framing and dispatches
97
+ * `data:` payloads. Resolves when the stream ends or errors.
98
+ */
99
+ private async consume(body: ReadableStream<Uint8Array>): Promise<void> {
100
+ const reader = body.getReader();
101
+ const decoder = new TextDecoder();
102
+ let buffer = "";
103
+
104
+ try {
105
+ for (;;) {
106
+ const { done, value } = await reader.read();
107
+ if (done) break;
108
+
109
+ buffer += decoder.decode(value, { stream: true });
110
+ let newlineIndex: number;
111
+ while ((newlineIndex = buffer.indexOf("\n")) !== -1) {
112
+ const line = buffer.slice(0, newlineIndex);
113
+ buffer = buffer.slice(newlineIndex + 1);
114
+ this.handleLine(line);
115
+ }
88
116
  }
89
- };
117
+ // Flush a trailing data line that lacked its event-terminating newline
118
+ if (buffer) this.handleLine(buffer);
119
+ } catch {
120
+ // Stream aborted or network error — handled by the reconnect logic
121
+ } finally {
122
+ reader.releaseLock();
123
+ }
124
+ }
90
125
 
91
- // Wait a bit for the connection to establish
92
- await new Promise<void>((resolve) => {
93
- const timeout = setTimeout(() => resolve(), POLLING_INTERVAL);
94
- this.eventSource!.onopen = () => {
95
- clearTimeout(timeout);
96
- this.connected = true;
97
- this.reconnecting = false;
98
- resolve();
126
+ /**
127
+ * Handles a single SSE line, dispatching `data:` payloads as events.
128
+ */
129
+ private handleLine(line: string): void {
130
+ if (!line.startsWith("data:")) return;
131
+
132
+ const payload = line.slice(5).trim();
133
+ if (!payload || payload === "[DONE]") return;
134
+
135
+ try {
136
+ const data = JSON.parse(payload);
137
+ const sseEvent: SSEEvent = {
138
+ event: data.event ?? "unknown",
139
+ model: data.model ?? "*",
140
+ data: data.data,
99
141
  };
100
- });
142
+ this._hasReceivedEvents = true;
143
+ this.dispatch(sseEvent);
144
+ } catch {
145
+ // Invalid JSON, ignore
146
+ }
147
+ }
101
148
 
102
- return this.connected;
149
+ /**
150
+ * Notifies the connect-failed callback if the connection failed before
151
+ * any event was received.
152
+ */
153
+ private notifyConnectFailed(): void {
154
+ if (!this._hasReceivedEvents && this._onConnectFailed) {
155
+ this._onConnectFailed();
156
+ }
103
157
  }
104
158
 
105
159
  /**
@@ -143,11 +197,11 @@ export class SSEClient {
143
197
  * Disconnects from the SSE endpoint and clears all subscriptions.
144
198
  */
145
199
  disconnect(): void {
146
- if (this.eventSource) {
147
- this.eventSource.close();
148
- this.eventSource = null;
149
- }
200
+ this.disposed = true;
201
+ this.abortController?.abort();
202
+ this.abortController = null;
150
203
  this.connected = false;
204
+ this.reconnecting = false;
151
205
  this.subscribers.clear();
152
206
  }
153
207