@ekodb/ekodb-client 0.18.2 → 0.20.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.
package/dist/client.js CHANGED
@@ -69,6 +69,19 @@ var MergeStrategy;
69
69
  MergeStrategy["LatestOnly"] = "LatestOnly";
70
70
  MergeStrategy["Interleaved"] = "Interleaved";
71
71
  })(MergeStrategy || (exports.MergeStrategy = MergeStrategy = {}));
72
+ /**
73
+ * Strip trailing slashes from a base URL so path concatenation
74
+ * (`${base}/api/...`) never yields a double-slash path. Uses a linear scan
75
+ * rather than a regex like `/\/+$/`, which CodeQL flags as polynomial-time
76
+ * backtracking on caller-supplied input.
77
+ */
78
+ function stripTrailingSlashes(url) {
79
+ let end = url.length;
80
+ while (end > 0 && url.charCodeAt(end - 1) === 47 /* "/" */) {
81
+ end--;
82
+ }
83
+ return end === url.length ? url : url.slice(0, end);
84
+ }
72
85
  class EkoDBClient {
73
86
  constructor(config, apiKey) {
74
87
  this.token = null;
@@ -76,14 +89,16 @@ class EkoDBClient {
76
89
  this.rateLimitInfo = null;
77
90
  // Support both old (baseURL, apiKey) and new (config object) signatures
78
91
  if (typeof config === "string") {
79
- this.baseURL = config;
92
+ // Strip trailing slashes so `${baseURL}/api/...` never produces a
93
+ // double-slash path (some servers/proxies reject `//api/...`).
94
+ this.baseURL = stripTrailingSlashes(config);
80
95
  this.apiKey = apiKey;
81
96
  this.shouldRetry = true;
82
97
  this.maxRetries = 3;
83
98
  this.format = SerializationFormat.MessagePack; // Default to MessagePack for 2-3x performance
84
99
  }
85
100
  else {
86
- this.baseURL = config.baseURL;
101
+ this.baseURL = stripTrailingSlashes(config.baseURL);
87
102
  this.apiKey = config.apiKey;
88
103
  this.shouldRetry = config.shouldRetry ?? true;
89
104
  this.maxRetries = config.maxRetries ?? 3;
@@ -223,6 +238,38 @@ class EkoDBClient {
223
238
  sleep(seconds) {
224
239
  return new Promise((resolve) => setTimeout(resolve, seconds * 1000));
225
240
  }
241
+ /**
242
+ * Parse a `Retry-After` header into a non-negative delay in seconds.
243
+ *
244
+ * Per RFC 9110 the value is either delay-seconds (an integer) or an
245
+ * HTTP-date. Anything that doesn't resolve to a finite, non-negative number
246
+ * (missing header, garbage, a past date) falls back to `defaultSecs`.
247
+ */
248
+ parseRetryAfter(header, defaultSecs = 60) {
249
+ if (!header)
250
+ return defaultSecs;
251
+ // delay-seconds form: a bare integer.
252
+ const secs = Number(header.trim());
253
+ if (Number.isFinite(secs))
254
+ return Math.max(0, secs);
255
+ // HTTP-date form: compute the delay from now.
256
+ const dateMs = Date.parse(header);
257
+ if (Number.isFinite(dateMs)) {
258
+ return Math.max(0, (dateMs - Date.now()) / 1000);
259
+ }
260
+ return defaultSecs;
261
+ }
262
+ /**
263
+ * Backoff delay (in seconds) for a 0-indexed retry attempt: a capped
264
+ * exponential schedule (0.2s → 5s) with full jitter, so concurrent clients
265
+ * don't retry in lockstep. Returns a value in [d/2, d].
266
+ */
267
+ backoffSeconds(attempt) {
268
+ const base = 0.2;
269
+ const max = 5;
270
+ const d = Math.min(base * Math.pow(2, Math.max(0, attempt)), max);
271
+ return d / 2 + Math.random() * (d / 2);
272
+ }
226
273
  /**
227
274
  * Helper to determine if a path should use JSON
228
275
  * Only CRUD operations (insert/update/delete/batch) use MessagePack
@@ -292,10 +339,13 @@ class EkoDBClient {
292
339
  }
293
340
  // Handle rate limiting (429)
294
341
  if (response.status === 429) {
295
- const retryAfter = parseInt(response.headers.get("retry-after") || "60", 10);
342
+ const retryAfter = this.parseRetryAfter(response.headers.get("retry-after"));
296
343
  if (this.shouldRetry && attempt < this.maxRetries) {
297
- console.log(`Rate limited. Retrying after ${retryAfter} seconds...`);
298
- await this.sleep(retryAfter);
344
+ // Honor the server's Retry-After, but cap it so a hostile/large value
345
+ // can't pin the client for minutes.
346
+ const wait = Math.min(retryAfter, 60);
347
+ console.log(`Rate limited. Retrying after ${wait} seconds...`);
348
+ await this.sleep(wait);
299
349
  return this.makeRequest(method, path, data, attempt + 1, forceJson);
300
350
  }
301
351
  throw new RateLimitError(retryAfter);
@@ -311,8 +361,8 @@ class EkoDBClient {
311
361
  if (response.status === 503 &&
312
362
  this.shouldRetry &&
313
363
  attempt < this.maxRetries) {
314
- const retryDelay = 10;
315
- console.log(`Service unavailable. Retrying after ${retryDelay} seconds...`);
364
+ const retryDelay = this.backoffSeconds(attempt);
365
+ console.log(`Service unavailable. Retrying after ${retryDelay.toFixed(2)}s...`);
316
366
  await this.sleep(retryDelay);
317
367
  return this.makeRequest(method, path, data, attempt + 1, forceJson);
318
368
  }
@@ -325,8 +375,8 @@ class EkoDBClient {
325
375
  if (error instanceof TypeError &&
326
376
  this.shouldRetry &&
327
377
  attempt < this.maxRetries) {
328
- const retryDelay = 3;
329
- console.log(`Network error. Retrying after ${retryDelay} seconds...`);
378
+ const retryDelay = this.backoffSeconds(attempt);
379
+ console.log(`Network error. Retrying after ${retryDelay.toFixed(2)}s...`);
330
380
  await this.sleep(retryDelay);
331
381
  return this.makeRequest(method, path, data, attempt + 1, forceJson);
332
382
  }
@@ -1125,10 +1175,10 @@ class EkoDBClient {
1125
1175
  const stream = new EventStream();
1126
1176
  (async () => {
1127
1177
  try {
1128
- let token = this.getToken();
1178
+ let token = await this.getToken();
1129
1179
  if (!token) {
1130
1180
  await this.refreshToken();
1131
- token = this.getToken();
1181
+ token = await this.getToken();
1132
1182
  }
1133
1183
  const url = `${this.baseURL}/api/chat/${chatId}/messages/stream`;
1134
1184
  const response = await fetch(url, {
@@ -1149,13 +1199,12 @@ class EkoDBClient {
1149
1199
  stream.close();
1150
1200
  return;
1151
1201
  }
1152
- const body = await response.text();
1153
- for (const line of body.split("\n")) {
1202
+ const emitLine = (line) => {
1154
1203
  if (!line.startsWith("data:"))
1155
- continue;
1204
+ return;
1156
1205
  const dataStr = line.slice(5).trim();
1157
1206
  if (!dataStr)
1158
- continue;
1207
+ return;
1159
1208
  try {
1160
1209
  const eventData = JSON.parse(dataStr);
1161
1210
  if (eventData.error) {
@@ -1184,6 +1233,33 @@ class EkoDBClient {
1184
1233
  catch {
1185
1234
  // skip malformed SSE data
1186
1235
  }
1236
+ };
1237
+ const reader = response.body?.getReader?.();
1238
+ if (reader) {
1239
+ // True incremental streaming: decode and emit each SSE line as soon as
1240
+ // it arrives, rather than buffering the entire response body first.
1241
+ const decoder = new TextDecoder();
1242
+ let buffer = "";
1243
+ for (;;) {
1244
+ const { done, value } = await reader.read();
1245
+ if (done)
1246
+ break;
1247
+ buffer += decoder.decode(value, { stream: true });
1248
+ let nl;
1249
+ while ((nl = buffer.indexOf("\n")) >= 0) {
1250
+ emitLine(buffer.slice(0, nl));
1251
+ buffer = buffer.slice(nl + 1);
1252
+ }
1253
+ }
1254
+ buffer += decoder.decode();
1255
+ if (buffer)
1256
+ emitLine(buffer);
1257
+ }
1258
+ else {
1259
+ // Fallback for environments/tests without a readable body stream.
1260
+ const body = await response.text();
1261
+ for (const line of body.split("\n"))
1262
+ emitLine(line);
1187
1263
  }
1188
1264
  stream.close();
1189
1265
  }
@@ -1277,6 +1353,23 @@ class EkoDBClient {
1277
1353
  async toggleForgottenMessage(sessionId, messageId, forgotten) {
1278
1354
  await this.makeRequest("PATCH", `/api/chat/${sessionId}/messages/${messageId}/forgotten`, { forgotten }, 0, true);
1279
1355
  }
1356
+ /**
1357
+ * Compact a chat session's history on demand.
1358
+ *
1359
+ * Folds older messages into a summary while preserving the most recent
1360
+ * messages verbatim, reducing context size for long-running sessions.
1361
+ *
1362
+ * @param chatId - Chat session ID
1363
+ * @param keepRecent - Number of recent messages to preserve verbatim (optional)
1364
+ * @returns Compaction result with counts and the summary message ID
1365
+ */
1366
+ async compactChat(chatId, keepRecent) {
1367
+ const body = {};
1368
+ if (keepRecent !== undefined) {
1369
+ body.keep_recent = keepRecent;
1370
+ }
1371
+ return this.makeRequest("POST", `/api/chat/${chatId}/compact`, body, 0, true);
1372
+ }
1280
1373
  /**
1281
1374
  * Merge multiple chat sessions into one
1282
1375
  */
@@ -1719,10 +1812,18 @@ class EkoDBClient {
1719
1812
  return stream;
1720
1813
  }
1721
1814
  /**
1722
- * Create a WebSocket client
1815
+ * Create a WebSocket client.
1816
+ *
1817
+ * The token is supplied as a provider bound to this client's
1818
+ * {@link getToken}, so every (re)connect re-evaluates (and proactively
1819
+ * refreshes) the auth token instead of snapshotting it once. This means a
1820
+ * reconnect after a token rotation uses the current token.
1821
+ *
1822
+ * @param wsURL - The WebSocket URL (e.g. `wss://host`); `/api/ws` is appended if absent.
1823
+ * @param options - Optional reconnect/timeout tunables.
1723
1824
  */
1724
- websocket(wsURL) {
1725
- return new WebSocketClient(wsURL, this.token);
1825
+ websocket(wsURL, options) {
1826
+ return new WebSocketClient(wsURL, () => this.getToken(), options);
1726
1827
  }
1727
1828
  // ========== RAG Helper Methods ==========
1728
1829
  /**
@@ -1971,50 +2072,108 @@ function extractRecordId(record, extraCandidates = []) {
1971
2072
  const val = record[key];
1972
2073
  if (typeof val === "string")
1973
2074
  return val;
1974
- if (val && typeof val === "object" && "value" in val)
2075
+ // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
2076
+ // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
2077
+ if (val && typeof val === "object" && "type" in val && "value" in val)
1975
2078
  return String(val.value);
1976
2079
  }
1977
2080
  for (const key of ["id", "_id"]) {
1978
2081
  const val = record[key];
1979
2082
  if (typeof val === "string")
1980
2083
  return val;
1981
- if (val && typeof val === "object" && "value" in val)
2084
+ // Unwrap only a genuine typed wrapper (both "type" and "value"), matching
2085
+ // getValue's rule so a user object like { value: 1 } isn't mistaken for one.
2086
+ if (val && typeof val === "object" && "type" in val && "value" in val)
1982
2087
  return String(val.value);
1983
2088
  }
1984
2089
  return undefined;
1985
2090
  }
1986
2091
  class WebSocketClient {
1987
- constructor(wsURL, token) {
2092
+ /**
2093
+ * @param wsURL - WebSocket URL; `/api/ws` is appended if absent.
2094
+ * @param token - A static token string OR a {@link TokenProvider} function
2095
+ * re-evaluated on every (re)connect (so a refreshed token is used after a drop).
2096
+ * @param options - Optional reconnect/timeout tunables.
2097
+ */
2098
+ constructor(wsURL, token, options = {}) {
1988
2099
  this.ws = null;
1989
2100
  this.dispatcherRunning = false;
1990
2101
  this.schemaCache = null;
2102
+ // Reconnect state
2103
+ /** Set while close() is in progress so the close handler doesn't reconnect. */
2104
+ this.closed = false;
2105
+ this.reconnectAttempts = 0;
2106
+ this.reconnecting = false;
2107
+ this.connectPromise = null;
1991
2108
  // Dispatcher state
1992
2109
  this.pendingRequests = new Map();
1993
2110
  this.subscriptions = new Map();
2111
+ /** Bookkeeping so subscriptions can be replayed on reconnect. */
2112
+ this.subscriptionParams = new Map();
1994
2113
  this.chatStreams = new Map();
1995
2114
  this.registerToolsAck = null;
1996
2115
  this.messageCounter = 0;
1997
- this.wsURL = wsURL;
1998
- this.token = token;
2116
+ // Strip trailing slashes so appending `/api/ws` can't yield `//api/ws`,
2117
+ // which warp's exact path match (`api / ws`) would reject.
2118
+ this.wsURL = stripTrailingSlashes(wsURL);
2119
+ this.tokenProvider = typeof token === "function" ? token : () => token;
2120
+ this.autoReconnect = options.autoReconnect ?? true;
2121
+ this.reconnectInitialDelayMs = options.reconnectInitialDelayMs ?? 200;
2122
+ this.reconnectMaxDelayMs = options.reconnectMaxDelayMs ?? 5000;
2123
+ this.reconnectMaxAttempts = options.reconnectMaxAttempts ?? 0;
2124
+ this.requestTimeoutMs = options.requestTimeoutMs ?? 30000;
1999
2125
  }
2000
2126
  genMessageId() {
2001
2127
  const counter = this.messageCounter++;
2002
2128
  return `${Date.now()}-${counter}-${Math.random().toString(36).slice(2, 8)}`;
2003
2129
  }
2004
2130
  /**
2005
- * Connect and start the dispatcher.
2131
+ * Compute the capped exponential backoff (with jitter) for a reconnect
2132
+ * attempt. attempt 0 -> ~initial, growing x2 each time up to the max cap.
2133
+ * Jitter is +/-25% to avoid thundering-herd reconnect storms.
2134
+ * @internal exposed for testing
2135
+ */
2136
+ computeBackoff(attempt) {
2137
+ const base = Math.min(this.reconnectInitialDelayMs * 2 ** attempt, this.reconnectMaxDelayMs);
2138
+ const jitter = base * 0.25 * (Math.random() * 2 - 1);
2139
+ return Math.max(0, Math.round(base + jitter));
2140
+ }
2141
+ /**
2142
+ * Connect and start the dispatcher. Re-evaluates the token provider so the
2143
+ * current/refreshed token is used for this socket.
2006
2144
  */
2007
2145
  async ensureConnected() {
2008
2146
  if (this.ws && this.dispatcherRunning)
2009
2147
  return;
2148
+ // Coalesce concurrent connect attempts onto a single in-flight promise.
2149
+ if (this.connectPromise)
2150
+ return this.connectPromise;
2151
+ // Clear the intentional-close flag only for user-initiated connects. During
2152
+ // a reconnect cycle this stays untouched so a concurrent close() can't be
2153
+ // undone and have the reconnect proceed against the user's intent.
2154
+ if (!this.reconnecting)
2155
+ this.closed = false;
2156
+ this.connectPromise = this.openSocket().finally(() => {
2157
+ this.connectPromise = null;
2158
+ });
2159
+ return this.connectPromise;
2160
+ }
2161
+ async openSocket() {
2010
2162
  const WebSocket = (await Promise.resolve().then(() => __importStar(require("ws")))).default;
2011
2163
  let url = this.wsURL;
2012
2164
  if (!url.endsWith("/api/ws")) {
2013
2165
  url += "/api/ws";
2014
2166
  }
2167
+ // Re-evaluate the token on every (re)connect — never a stale snapshot.
2168
+ const token = await this.tokenProvider();
2169
+ if (!token) {
2170
+ // Fail fast with a clear error instead of sending `Bearer null`, which
2171
+ // would surface as a confusing 401 from the server.
2172
+ throw new Error("WebSocket auth token is unavailable (the token provider returned null/empty)");
2173
+ }
2015
2174
  this.ws = new WebSocket(url, {
2016
2175
  headers: {
2017
- Authorization: `Bearer ${this.token}`,
2176
+ Authorization: `Bearer ${token}`,
2018
2177
  },
2019
2178
  });
2020
2179
  await new Promise((resolve, reject) => {
@@ -2027,7 +2186,13 @@ class WebSocketClient {
2027
2186
  if (this.dispatcherRunning)
2028
2187
  return;
2029
2188
  this.dispatcherRunning = true;
2030
- this.ws.on("message", (data) => {
2189
+ // Capture the socket this dispatcher is bound to. After a reconnect, the old
2190
+ // socket may still emit late close/error events; ignore them so they don't
2191
+ // tear down the replacement connection.
2192
+ const socket = this.ws;
2193
+ socket.on("message", (data) => {
2194
+ if (this.ws !== socket)
2195
+ return;
2031
2196
  try {
2032
2197
  const msg = JSON.parse(data.toString());
2033
2198
  this.routeMessage(msg);
@@ -2036,26 +2201,149 @@ class WebSocketClient {
2036
2201
  // Ignore malformed messages
2037
2202
  }
2038
2203
  });
2039
- this.ws.on("close", () => {
2040
- this.dispatcherRunning = false;
2041
- // Notify all pending requests
2042
- for (const [, pending] of this.pendingRequests) {
2043
- pending.reject(new Error("WebSocket connection closed"));
2044
- }
2045
- this.pendingRequests.clear();
2046
- // Close all chat streams
2047
- for (const [, stream] of this.chatStreams) {
2048
- stream.emit("event", { type: "error", error: "Connection closed" });
2049
- stream.close();
2050
- }
2051
- this.chatStreams.clear();
2052
- // Close all subscriptions
2204
+ // Both "close" and "error" mean this socket is dead. ws typically emits
2205
+ // "error" followed by "close", so route both through one handler and let the
2206
+ // identity check dedupe: the first to fire nulls this.ws, the second no-ops.
2207
+ const onDown = () => {
2208
+ if (this.ws !== socket)
2209
+ return;
2210
+ this.handleDisconnect();
2211
+ };
2212
+ socket.on("close", onDown);
2213
+ socket.on("error", onDown);
2214
+ }
2215
+ /**
2216
+ * Reject in-flight requests and tear down the dead socket. If the close was
2217
+ * unexpected (not an explicit `close()`) and auto-reconnect is enabled,
2218
+ * schedule a reconnect that re-sends the active subscriptions.
2219
+ */
2220
+ handleDisconnect() {
2221
+ this.dispatcherRunning = false;
2222
+ this.ws = null;
2223
+ // Reject all in-flight pending requests so callers don't hang forever.
2224
+ for (const [, pending] of this.pendingRequests) {
2225
+ if (pending.timer)
2226
+ clearTimeout(pending.timer);
2227
+ pending.reject(new Error("WebSocket connection closed"));
2228
+ }
2229
+ this.pendingRequests.clear();
2230
+ if (this.registerToolsAck) {
2231
+ this.registerToolsAck.reject(new Error("WebSocket connection closed"));
2232
+ this.registerToolsAck = null;
2233
+ }
2234
+ // Close all chat streams (they are one-shot; not replayed on reconnect).
2235
+ for (const [, stream] of this.chatStreams) {
2236
+ stream.emit("event", { type: "error", error: "Connection closed" });
2237
+ stream.close();
2238
+ }
2239
+ this.chatStreams.clear();
2240
+ const shouldReconnect = this.autoReconnect && !this.closed && this.subscriptionParams.size > 0;
2241
+ if (shouldReconnect) {
2242
+ this.scheduleReconnect();
2243
+ }
2244
+ else {
2245
+ // No reconnect: tear down subscriptions too.
2053
2246
  for (const [, stream] of this.subscriptions) {
2054
2247
  stream.close();
2055
2248
  }
2056
2249
  this.subscriptions.clear();
2057
- this.ws = null;
2058
- });
2250
+ this.subscriptionParams.clear();
2251
+ }
2252
+ }
2253
+ /**
2254
+ * Reconnect with capped exponential backoff + jitter, then re-send the
2255
+ * subscribe messages for every active subscription so the SAME EventStream
2256
+ * keeps delivering mutations after a transient drop.
2257
+ */
2258
+ scheduleReconnect() {
2259
+ if (this.reconnecting)
2260
+ return;
2261
+ this.reconnecting = true;
2262
+ const attempt = async () => {
2263
+ // Bail if the client was closed, or if every subscription was torn down
2264
+ // (e.g. unsubscribed) while a reconnect was in-flight — reconnect was only
2265
+ // opted into because subscriptions existed, so there's nothing to restore.
2266
+ if (this.closed || this.subscriptionParams.size === 0) {
2267
+ this.reconnecting = false;
2268
+ return;
2269
+ }
2270
+ if (this.reconnectMaxAttempts > 0 &&
2271
+ this.reconnectAttempts >= this.reconnectMaxAttempts) {
2272
+ // Give up: tear down subscriptions and notify consumers.
2273
+ this.reconnecting = false;
2274
+ for (const [, stream] of this.subscriptions) {
2275
+ stream.emit("error", "WebSocket reconnect failed");
2276
+ stream.close();
2277
+ }
2278
+ this.subscriptions.clear();
2279
+ this.subscriptionParams.clear();
2280
+ return;
2281
+ }
2282
+ const delay = this.computeBackoff(this.reconnectAttempts);
2283
+ this.reconnectAttempts++;
2284
+ await new Promise((r) => setTimeout(r, delay));
2285
+ // Re-check after the backoff delay: close() or a full unsubscribe may have
2286
+ // happened while we were waiting, in which case skip reopening the socket.
2287
+ if (this.closed || this.subscriptionParams.size === 0) {
2288
+ this.reconnecting = false;
2289
+ return;
2290
+ }
2291
+ try {
2292
+ // Route through ensureConnected() so a request-driven connect and this
2293
+ // reconnect share one in-flight connectPromise/socket — opening two live
2294
+ // sockets would misroute responses.
2295
+ await this.ensureConnected();
2296
+ // close() may have been called while the connect was in-flight; if so,
2297
+ // tear down the freshly-opened socket instead of leaving it orphaned.
2298
+ if (this.closed) {
2299
+ try {
2300
+ this.ws?.close?.();
2301
+ }
2302
+ catch {
2303
+ /* already closing */
2304
+ }
2305
+ this.ws = null;
2306
+ this.dispatcherRunning = false;
2307
+ this.reconnecting = false;
2308
+ return;
2309
+ }
2310
+ // Success — reset backoff and replay subscriptions.
2311
+ this.reconnectAttempts = 0;
2312
+ this.reconnecting = false;
2313
+ await this.resubscribeAll();
2314
+ }
2315
+ catch {
2316
+ // Connect failed — schedule the next attempt WITHOUT recursive await so
2317
+ // a prolonged outage can't build an unbounded promise chain.
2318
+ setTimeout(() => void attempt(), 0);
2319
+ }
2320
+ };
2321
+ void attempt();
2322
+ }
2323
+ /** Re-send Subscribe frames for every tracked subscription after a reconnect. */
2324
+ async resubscribeAll() {
2325
+ for (const [collection, options] of this.subscriptionParams) {
2326
+ const stream = this.subscriptions.get(collection);
2327
+ if (!stream || stream.closed)
2328
+ continue;
2329
+ const messageId = this.genMessageId();
2330
+ const request = {
2331
+ type: "Subscribe",
2332
+ messageId,
2333
+ payload: {
2334
+ collection,
2335
+ ...(options?.filterField && { filter_field: options.filterField }),
2336
+ ...(options?.filterValue && { filter_value: options.filterValue }),
2337
+ },
2338
+ };
2339
+ try {
2340
+ await this.sendRequest(request);
2341
+ }
2342
+ catch {
2343
+ // If the re-subscribe ack fails, leave it tracked; the next
2344
+ // disconnect/reconnect cycle will attempt it again.
2345
+ }
2346
+ }
2059
2347
  }
2060
2348
  routeMessage(msg) {
2061
2349
  switch (msg.type) {
@@ -2068,15 +2356,7 @@ class WebSocketClient {
2068
2356
  msg.payload?.messageId;
2069
2357
  let matched = false;
2070
2358
  if (messageId && this.pendingRequests.has(messageId)) {
2071
- const pending = this.pendingRequests.get(messageId);
2072
- this.pendingRequests.delete(messageId);
2073
- if (msg.type === "Error") {
2074
- pending.reject(new Error(msg.message || "Unknown error"));
2075
- }
2076
- else {
2077
- pending.resolve(msg.payload);
2078
- }
2079
- matched = true;
2359
+ matched = this.settlePending(messageId, msg.type === "Error", msg);
2080
2360
  }
2081
2361
  if (!matched && this.registerToolsAck) {
2082
2362
  const ack = this.registerToolsAck;
@@ -2089,19 +2369,14 @@ class WebSocketClient {
2089
2369
  }
2090
2370
  matched = true;
2091
2371
  }
2092
- // Server doesn't echo messageId — if there's exactly one pending
2372
+ // Server doesn't echo messageId at all — if there's exactly one pending
2093
2373
  // request, deliver the response to it (sequential request/response).
2094
- if (!matched && this.pendingRequests.size === 1) {
2095
- const entry = this.pendingRequests.entries().next().value;
2096
- const key = entry[0];
2097
- const pending = entry[1];
2098
- this.pendingRequests.delete(key);
2099
- if (msg.type === "Error") {
2100
- pending.reject(new Error(msg.message || "Unknown error"));
2101
- }
2102
- else {
2103
- pending.resolve(msg.payload);
2104
- }
2374
+ // Only when messageId is absent: a present-but-unmatched id means a late
2375
+ // response for an already-settled/timed-out request, which must NOT be
2376
+ // misrouted to whatever request happens to still be pending.
2377
+ if (!matched && !messageId && this.pendingRequests.size === 1) {
2378
+ const key = this.pendingRequests.keys().next().value;
2379
+ this.settlePending(key, msg.type === "Error", msg);
2105
2380
  }
2106
2381
  break;
2107
2382
  }
@@ -2188,16 +2463,46 @@ class WebSocketClient {
2188
2463
  await this.ensureConnected();
2189
2464
  const messageId = request.messageId || request.message_id;
2190
2465
  return new Promise((resolve, reject) => {
2191
- this.pendingRequests.set(messageId, { resolve, reject });
2466
+ // Per-request timeout: reject if no response arrives in the window so a
2467
+ // dropped/never-answered response can't leave the promise pending forever.
2468
+ let timer;
2469
+ if (this.requestTimeoutMs > 0) {
2470
+ timer = setTimeout(() => {
2471
+ if (this.pendingRequests.delete(messageId)) {
2472
+ reject(new Error(`WebSocket request "${request.type}" timed out after ${this.requestTimeoutMs}ms`));
2473
+ }
2474
+ }, this.requestTimeoutMs);
2475
+ // Don't keep the process alive just for this timer.
2476
+ timer?.unref?.();
2477
+ }
2478
+ this.pendingRequests.set(messageId, { resolve, reject, timer });
2192
2479
  try {
2193
2480
  this.ws.send(JSON.stringify(request));
2194
2481
  }
2195
2482
  catch (err) {
2196
2483
  this.pendingRequests.delete(messageId);
2484
+ if (timer)
2485
+ clearTimeout(timer);
2197
2486
  reject(err);
2198
2487
  }
2199
2488
  });
2200
2489
  }
2490
+ /** Resolve/reject a pending request, clearing its timeout timer. */
2491
+ settlePending(messageId, isError, msg) {
2492
+ const pending = this.pendingRequests.get(messageId);
2493
+ if (!pending)
2494
+ return false;
2495
+ this.pendingRequests.delete(messageId);
2496
+ if (pending.timer)
2497
+ clearTimeout(pending.timer);
2498
+ if (isError) {
2499
+ pending.reject(new Error(msg.message || "Unknown error"));
2500
+ }
2501
+ else {
2502
+ pending.resolve(msg.payload);
2503
+ }
2504
+ return true;
2505
+ }
2201
2506
  /**
2202
2507
  * Find all records in a collection via WebSocket.
2203
2508
  */
@@ -2222,6 +2527,8 @@ class WebSocketClient {
2222
2527
  const messageId = this.genMessageId();
2223
2528
  const stream = new EventStream();
2224
2529
  this.subscriptions.set(collection, stream);
2530
+ // Track params so the subscription can be replayed on reconnect.
2531
+ this.subscriptionParams.set(collection, options);
2225
2532
  const request = {
2226
2533
  type: "Subscribe",
2227
2534
  messageId,
@@ -2237,10 +2544,23 @@ class WebSocketClient {
2237
2544
  }
2238
2545
  catch (err) {
2239
2546
  this.subscriptions.delete(collection);
2547
+ this.subscriptionParams.delete(collection);
2240
2548
  throw err;
2241
2549
  }
2242
2550
  return stream;
2243
2551
  }
2552
+ /**
2553
+ * Unsubscribe from a collection's mutation notifications. This is an
2554
+ * intentional teardown, so the subscription is NOT replayed on reconnect.
2555
+ */
2556
+ unsubscribe(collection) {
2557
+ const stream = this.subscriptions.get(collection);
2558
+ this.subscriptions.delete(collection);
2559
+ this.subscriptionParams.delete(collection);
2560
+ if (stream && !stream.closed) {
2561
+ stream.close();
2562
+ }
2563
+ }
2244
2564
  /**
2245
2565
  * Send a chat message and receive a streaming response.
2246
2566
  * Returns an EventStream that emits "event" with ChatStreamEvent objects.
@@ -2464,8 +2784,42 @@ class WebSocketClient {
2464
2784
  }
2465
2785
  /**
2466
2786
  * Close the WebSocket connection.
2787
+ *
2788
+ * This is an INTENTIONAL close: it disables auto-reconnect, rejects any
2789
+ * in-flight requests, and tears down all subscriptions/chat streams so
2790
+ * nothing is replayed afterward.
2467
2791
  */
2468
2792
  close() {
2793
+ // Mark intentional so the close handler doesn't trigger a reconnect.
2794
+ this.closed = true;
2795
+ this.reconnecting = false;
2796
+ // Reject any in-flight requests and clear their timers.
2797
+ for (const [, pending] of this.pendingRequests) {
2798
+ if (pending.timer)
2799
+ clearTimeout(pending.timer);
2800
+ pending.reject(new Error("WebSocket connection closed"));
2801
+ }
2802
+ this.pendingRequests.clear();
2803
+ // Tear down subscriptions + their replay bookkeeping.
2804
+ for (const [, stream] of this.subscriptions) {
2805
+ if (!stream.closed)
2806
+ stream.close();
2807
+ }
2808
+ this.subscriptions.clear();
2809
+ this.subscriptionParams.clear();
2810
+ // Reject any in-flight tool registration ack. Done here (not just in the
2811
+ // ws "close" handler) so it's cleaned up even when this.ws is already null.
2812
+ if (this.registerToolsAck) {
2813
+ this.registerToolsAck.reject(new Error("WebSocket connection closed"));
2814
+ this.registerToolsAck = null;
2815
+ }
2816
+ // Tear down chat streams immediately; they are one-shot and not replayed,
2817
+ // and we can't rely on the underlying ws "close" event having fired.
2818
+ for (const [, stream] of this.chatStreams) {
2819
+ stream.emit("event", { type: "error", error: "Connection closed" });
2820
+ stream.close();
2821
+ }
2822
+ this.chatStreams.clear();
2469
2823
  if (this.ws) {
2470
2824
  this.ws.close();
2471
2825
  this.ws = null;