@alexkroman1/aai 1.0.5 → 1.1.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 (46) hide show
  1. package/.turbo/turbo-build.log +11 -11
  2. package/CHANGELOG.md +14 -0
  3. package/dist/_internal-types-CoDTiBd1.js +61 -0
  4. package/dist/host/_mock-ws.d.ts +0 -24
  5. package/dist/host/runtime-barrel.d.ts +0 -1
  6. package/dist/host/runtime-barrel.js +55 -5
  7. package/dist/host/runtime.d.ts +2 -0
  8. package/dist/host/tool-executor.d.ts +1 -0
  9. package/dist/host/ws-handler.d.ts +2 -0
  10. package/dist/sdk/manifest-barrel.d.ts +3 -5
  11. package/dist/sdk/manifest-barrel.js +2 -52
  12. package/dist/sdk/protocol.d.ts +8 -25
  13. package/dist/sdk/protocol.js +6 -3
  14. package/dist/sdk/types.d.ts +2 -0
  15. package/host/_mock-ws.ts +0 -50
  16. package/host/_test-utils.ts +1 -0
  17. package/host/runtime-barrel.ts +0 -1
  18. package/host/runtime.ts +13 -1
  19. package/host/session-ctx.test.ts +387 -0
  20. package/host/session-fixture-replay.test.ts +2 -10
  21. package/host/session.test.ts +19 -41
  22. package/host/tool-executor.test.ts +36 -0
  23. package/host/tool-executor.ts +4 -0
  24. package/host/ws-handler.ts +3 -0
  25. package/package.json +1 -1
  26. package/sdk/__snapshots__/exports.test.ts.snap +77 -0
  27. package/sdk/__snapshots__/schema-shapes.test.ts.snap +187 -0
  28. package/sdk/_test-matchers.test.ts +75 -0
  29. package/sdk/_test-matchers.ts +73 -0
  30. package/sdk/exports.test.ts +31 -0
  31. package/sdk/manifest-barrel.ts +13 -7
  32. package/sdk/manifest.test.ts +66 -2
  33. package/sdk/protocol-compat.test.ts +0 -6
  34. package/sdk/protocol-snapshot.test.ts +7 -5
  35. package/sdk/protocol.test.ts +107 -21
  36. package/sdk/protocol.ts +7 -15
  37. package/sdk/schema-alignment.test.ts +1 -27
  38. package/sdk/schema-shapes.test.ts +103 -0
  39. package/sdk/tsconfig.json +1 -1
  40. package/sdk/types.test.ts +56 -1
  41. package/sdk/types.ts +2 -0
  42. package/sdk/ws-upgrade.test.ts +8 -8
  43. package/tsconfig.build.json +8 -1
  44. package/tsconfig.json +1 -1
  45. package/vitest.config.ts +1 -0
  46. package/dist/system-prompt-nik_iavo.js +0 -92
@@ -0,0 +1,387 @@
1
+ // Copyright 2025 the AAI authors. MIT license.
2
+
3
+ import { describe, expect, it, vi } from "vitest";
4
+ import { DEFAULT_MAX_HISTORY } from "../sdk/constants.ts";
5
+ import type { Message } from "../sdk/types.ts";
6
+ import { toolError } from "../sdk/utils.ts";
7
+ import { flush, makeClient, makeConfig, silentLogger } from "./_test-utils.ts";
8
+ import { buildCtx } from "./session-ctx.ts";
9
+
10
+ function makeBuildCtxOpts(overrides?: Record<string, unknown>) {
11
+ return {
12
+ id: "session-1",
13
+ agent: "test-agent",
14
+ client: makeClient(),
15
+ agentConfig: makeConfig({ maxSteps: 3 }),
16
+ executeTool: vi.fn(async () => "ok"),
17
+ log: silentLogger,
18
+ ...overrides,
19
+ };
20
+ }
21
+
22
+ describe("buildCtx", () => {
23
+ it("returns ctx with the correct session id", () => {
24
+ const ctx = buildCtx(makeBuildCtxOpts({ id: "my-session" }));
25
+ expect(ctx.id).toBe("my-session");
26
+ });
27
+
28
+ it("returns ctx with the correct agent name", () => {
29
+ const ctx = buildCtx(makeBuildCtxOpts({ agent: "my-agent" }));
30
+ expect(ctx.agent).toBe("my-agent");
31
+ });
32
+
33
+ it("initializes with empty conversation messages", () => {
34
+ const ctx = buildCtx(makeBuildCtxOpts());
35
+ expect(ctx.conversationMessages).toEqual([]);
36
+ });
37
+
38
+ it("initializes with null s2s handle", () => {
39
+ const ctx = buildCtx(makeBuildCtxOpts());
40
+ expect(ctx.s2s).toBeNull();
41
+ });
42
+
43
+ it("initializes with null turnPromise", () => {
44
+ const ctx = buildCtx(makeBuildCtxOpts());
45
+ expect(ctx.turnPromise).toBeNull();
46
+ });
47
+
48
+ it("initializes reply state with empty pendingTools, zero toolCallCount, and null replyId", () => {
49
+ const ctx = buildCtx(makeBuildCtxOpts());
50
+ expect(ctx.reply).toEqual({
51
+ pendingTools: [],
52
+ toolCallCount: 0,
53
+ currentReplyId: null,
54
+ });
55
+ });
56
+
57
+ it("defaults maxHistory to DEFAULT_MAX_HISTORY when not provided", () => {
58
+ const ctx = buildCtx(makeBuildCtxOpts());
59
+ expect(ctx.maxHistory).toBe(DEFAULT_MAX_HISTORY);
60
+ });
61
+
62
+ it("uses custom maxHistory when provided", () => {
63
+ const ctx = buildCtx(makeBuildCtxOpts({ maxHistory: 50 }));
64
+ expect(ctx.maxHistory).toBe(50);
65
+ });
66
+
67
+ it("passes through the agentConfig, executeTool, and log dependencies", () => {
68
+ const config = makeConfig({ maxSteps: 7 });
69
+ const executeTool = vi.fn(async () => "done");
70
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: config, executeTool }));
71
+ expect(ctx.agentConfig).toBe(config);
72
+ expect(ctx.executeTool).toBe(executeTool);
73
+ expect(ctx.log).toBe(silentLogger);
74
+ });
75
+ });
76
+
77
+ describe("consumeToolCallStep", () => {
78
+ it("returns null (success) when tool call is within maxSteps", () => {
79
+ const ctx = buildCtx(makeBuildCtxOpts());
80
+ ctx.beginReply("reply-1");
81
+ const result = ctx.consumeToolCallStep("my-tool", "reply-1");
82
+ expect(result).toBeNull();
83
+ });
84
+
85
+ it("increments toolCallCount on each call", () => {
86
+ const ctx = buildCtx(makeBuildCtxOpts());
87
+ ctx.beginReply("reply-1");
88
+
89
+ ctx.consumeToolCallStep("tool-a", "reply-1");
90
+ expect(ctx.reply.toolCallCount).toBe(1);
91
+
92
+ ctx.consumeToolCallStep("tool-b", "reply-1");
93
+ expect(ctx.reply.toolCallCount).toBe(2);
94
+ });
95
+
96
+ it("allows exactly maxSteps tool calls", () => {
97
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig({ maxSteps: 2 }) }));
98
+ ctx.beginReply("reply-1");
99
+
100
+ expect(ctx.consumeToolCallStep("tool-1", "reply-1")).toBeNull();
101
+ expect(ctx.consumeToolCallStep("tool-2", "reply-1")).toBeNull();
102
+ });
103
+
104
+ it("rejects when tool call count exceeds maxSteps", () => {
105
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig({ maxSteps: 2 }) }));
106
+ ctx.beginReply("reply-1");
107
+
108
+ ctx.consumeToolCallStep("tool-1", "reply-1");
109
+ ctx.consumeToolCallStep("tool-2", "reply-1");
110
+ const result = ctx.consumeToolCallStep("tool-3", "reply-1");
111
+ expect(result).toBe(toolError("Maximum tool steps reached. Please respond to the user now."));
112
+ });
113
+
114
+ it("logs when maxSteps is exceeded", () => {
115
+ const log = { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() };
116
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig({ maxSteps: 1 }), log }));
117
+ ctx.beginReply("reply-1");
118
+
119
+ ctx.consumeToolCallStep("tool-1", "reply-1"); // ok
120
+ ctx.consumeToolCallStep("tool-2", "reply-1"); // exceeds
121
+
122
+ expect(log.info).toHaveBeenCalledWith("maxSteps exceeded, refusing tool call", {
123
+ toolCallCount: 2,
124
+ maxSteps: 1,
125
+ });
126
+ });
127
+
128
+ it("rejects with stale replyId (mismatched)", () => {
129
+ const ctx = buildCtx(makeBuildCtxOpts());
130
+ ctx.beginReply("reply-1");
131
+
132
+ const result = ctx.consumeToolCallStep("my-tool", "stale-reply");
133
+ expect(result).toBe(toolError("Reply was interrupted. Discarding stale tool call."));
134
+ });
135
+
136
+ it("rejects when replyId is null", () => {
137
+ const ctx = buildCtx(makeBuildCtxOpts());
138
+ ctx.beginReply("reply-1");
139
+
140
+ const result = ctx.consumeToolCallStep("my-tool", null);
141
+ expect(result).toBe(toolError("Reply was interrupted. Discarding stale tool call."));
142
+ });
143
+
144
+ it("rejects when no reply has been started (currentReplyId is null)", () => {
145
+ const ctx = buildCtx(makeBuildCtxOpts());
146
+ // No beginReply — currentReplyId stays null
147
+ const result = ctx.consumeToolCallStep("my-tool", "some-reply");
148
+ expect(result).toBe(toolError("Reply was interrupted. Discarding stale tool call."));
149
+ });
150
+
151
+ it("allows unlimited tool calls when maxSteps is undefined", () => {
152
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig() }));
153
+ ctx.beginReply("reply-1");
154
+
155
+ // makeConfig() without maxSteps leaves it undefined
156
+ for (let i = 0; i < 100; i++) {
157
+ expect(ctx.consumeToolCallStep(`tool-${i}`, "reply-1")).toBeNull();
158
+ }
159
+ });
160
+ });
161
+
162
+ describe("pushMessages", () => {
163
+ it("appends messages to conversationMessages", () => {
164
+ const ctx = buildCtx(makeBuildCtxOpts());
165
+ const msg1: Message = { role: "user", content: "hello" };
166
+ const msg2: Message = { role: "assistant", content: "hi" };
167
+
168
+ ctx.pushMessages(msg1);
169
+ expect(ctx.conversationMessages).toEqual([msg1]);
170
+
171
+ ctx.pushMessages(msg2);
172
+ expect(ctx.conversationMessages).toEqual([msg1, msg2]);
173
+ });
174
+
175
+ it("accepts multiple messages at once", () => {
176
+ const ctx = buildCtx(makeBuildCtxOpts());
177
+ const msg1: Message = { role: "user", content: "a" };
178
+ const msg2: Message = { role: "assistant", content: "b" };
179
+ const msg3: Message = { role: "tool", content: "c" };
180
+
181
+ ctx.pushMessages(msg1, msg2, msg3);
182
+ expect(ctx.conversationMessages).toEqual([msg1, msg2, msg3]);
183
+ });
184
+
185
+ it("trims to maxHistory keeping the most recent messages", () => {
186
+ const ctx = buildCtx(makeBuildCtxOpts({ maxHistory: 3 }));
187
+
188
+ ctx.pushMessages(
189
+ { role: "user", content: "1" },
190
+ { role: "assistant", content: "2" },
191
+ { role: "user", content: "3" },
192
+ );
193
+ expect(ctx.conversationMessages).toHaveLength(3);
194
+
195
+ ctx.pushMessages({ role: "assistant", content: "4" });
196
+ expect(ctx.conversationMessages).toHaveLength(3);
197
+ expect(ctx.conversationMessages.map((m) => m.content)).toEqual(["2", "3", "4"]);
198
+ });
199
+
200
+ it("trims correctly when pushing multiple messages that exceed maxHistory", () => {
201
+ const ctx = buildCtx(makeBuildCtxOpts({ maxHistory: 2 }));
202
+
203
+ ctx.pushMessages(
204
+ { role: "user", content: "a" },
205
+ { role: "assistant", content: "b" },
206
+ { role: "user", content: "c" },
207
+ { role: "assistant", content: "d" },
208
+ );
209
+
210
+ expect(ctx.conversationMessages).toHaveLength(2);
211
+ expect(ctx.conversationMessages.map((m) => m.content)).toEqual(["c", "d"]);
212
+ });
213
+
214
+ it("does not trim when maxHistory is 0 (disabled)", () => {
215
+ const ctx = buildCtx(makeBuildCtxOpts({ maxHistory: 0 }));
216
+
217
+ for (let i = 0; i < 300; i++) {
218
+ ctx.pushMessages({ role: "user", content: `msg-${i}` });
219
+ }
220
+ expect(ctx.conversationMessages).toHaveLength(300);
221
+ });
222
+ });
223
+
224
+ describe("cancelReply", () => {
225
+ it("resets pendingTools and toolCallCount", () => {
226
+ const ctx = buildCtx(makeBuildCtxOpts());
227
+ ctx.beginReply("reply-1");
228
+ ctx.consumeToolCallStep("tool-1", "reply-1");
229
+ ctx.reply.pendingTools.push({ callId: "c1", result: "r1" });
230
+
231
+ expect(ctx.reply.toolCallCount).toBe(1);
232
+ expect(ctx.reply.pendingTools).toHaveLength(1);
233
+
234
+ ctx.cancelReply();
235
+
236
+ expect(ctx.reply.toolCallCount).toBe(0);
237
+ expect(ctx.reply.pendingTools).toEqual([]);
238
+ expect(ctx.reply.currentReplyId).toBeNull();
239
+ });
240
+
241
+ it("allows a new reply to start fresh after cancel", () => {
242
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig({ maxSteps: 1 }) }));
243
+ ctx.beginReply("reply-1");
244
+ ctx.consumeToolCallStep("tool-1", "reply-1"); // uses the single step
245
+
246
+ ctx.cancelReply();
247
+ ctx.beginReply("reply-2");
248
+
249
+ // Should succeed because toolCallCount was reset
250
+ const result = ctx.consumeToolCallStep("tool-1", "reply-2");
251
+ expect(result).toBeNull();
252
+ });
253
+ });
254
+
255
+ describe("beginReply", () => {
256
+ it("resets reply state with the given replyId", () => {
257
+ const ctx = buildCtx(makeBuildCtxOpts());
258
+ ctx.beginReply("reply-1");
259
+
260
+ expect(ctx.reply.currentReplyId).toBe("reply-1");
261
+ expect(ctx.reply.pendingTools).toEqual([]);
262
+ expect(ctx.reply.toolCallCount).toBe(0);
263
+ });
264
+
265
+ it("clears turnPromise", () => {
266
+ const ctx = buildCtx(makeBuildCtxOpts());
267
+ ctx.chainTurn(Promise.resolve());
268
+ expect(ctx.turnPromise).not.toBeNull();
269
+
270
+ ctx.beginReply("reply-1");
271
+ expect(ctx.turnPromise).toBeNull();
272
+ });
273
+
274
+ it("resets toolCallCount from a previous reply", () => {
275
+ const ctx = buildCtx(makeBuildCtxOpts({ agentConfig: makeConfig({ maxSteps: 2 }) }));
276
+ ctx.beginReply("reply-1");
277
+ ctx.consumeToolCallStep("tool-a", "reply-1");
278
+ ctx.consumeToolCallStep("tool-b", "reply-1");
279
+ expect(ctx.reply.toolCallCount).toBe(2);
280
+
281
+ ctx.beginReply("reply-2");
282
+ expect(ctx.reply.toolCallCount).toBe(0);
283
+
284
+ // Can now use maxSteps again
285
+ expect(ctx.consumeToolCallStep("tool-a", "reply-2")).toBeNull();
286
+ expect(ctx.consumeToolCallStep("tool-b", "reply-2")).toBeNull();
287
+ });
288
+
289
+ it("invalidates tool calls from the previous reply", () => {
290
+ const ctx = buildCtx(makeBuildCtxOpts());
291
+ ctx.beginReply("reply-1");
292
+ ctx.beginReply("reply-2");
293
+
294
+ // Tool call using old replyId should be rejected
295
+ const result = ctx.consumeToolCallStep("my-tool", "reply-1");
296
+ expect(result).toBe(toolError("Reply was interrupted. Discarding stale tool call."));
297
+ });
298
+ });
299
+
300
+ describe("chainTurn", () => {
301
+ it("sets turnPromise on first call", () => {
302
+ const ctx = buildCtx(makeBuildCtxOpts());
303
+ expect(ctx.turnPromise).toBeNull();
304
+
305
+ ctx.chainTurn(Promise.resolve());
306
+ expect(ctx.turnPromise).not.toBeNull();
307
+ });
308
+
309
+ it("chains promises sequentially", async () => {
310
+ const ctx = buildCtx(makeBuildCtxOpts());
311
+ const order: number[] = [];
312
+
313
+ ctx.chainTurn(
314
+ new Promise<void>((resolve) => {
315
+ queueMicrotask(() => {
316
+ order.push(1);
317
+ resolve();
318
+ });
319
+ }),
320
+ );
321
+
322
+ ctx.chainTurn(
323
+ new Promise<void>((resolve) => {
324
+ queueMicrotask(() => {
325
+ order.push(2);
326
+ resolve();
327
+ });
328
+ }),
329
+ );
330
+
331
+ await ctx.turnPromise;
332
+ await flush();
333
+ expect(order).toEqual([1, 2]);
334
+ });
335
+
336
+ it("continues the chain even if a prior turn rejects", async () => {
337
+ const ctx = buildCtx(makeBuildCtxOpts());
338
+ const order: string[] = [];
339
+
340
+ ctx.chainTurn(
341
+ new Promise<void>((_, reject) => {
342
+ queueMicrotask(() => {
343
+ order.push("fail");
344
+ reject(new Error("boom"));
345
+ });
346
+ }),
347
+ );
348
+
349
+ ctx.chainTurn(
350
+ new Promise<void>((resolve) => {
351
+ queueMicrotask(() => {
352
+ order.push("success");
353
+ resolve();
354
+ });
355
+ }),
356
+ );
357
+
358
+ // The chain uses .then() which means rejection propagates.
359
+ // We need to catch the final promise to avoid unhandled rejection.
360
+ try {
361
+ await ctx.turnPromise;
362
+ } catch {
363
+ // expected
364
+ }
365
+ await flush();
366
+
367
+ expect(order).toContain("fail");
368
+ });
369
+
370
+ it("allows awaiting turnPromise to wait for all chained turns", async () => {
371
+ const ctx = buildCtx(makeBuildCtxOpts());
372
+ let completed = false;
373
+
374
+ ctx.chainTurn(
375
+ new Promise<void>((resolve) => {
376
+ setTimeout(() => {
377
+ completed = true;
378
+ resolve();
379
+ }, 10);
380
+ }),
381
+ );
382
+
383
+ expect(completed).toBe(false);
384
+ await ctx.turnPromise;
385
+ expect(completed).toBe(true);
386
+ });
387
+ });
@@ -83,16 +83,8 @@ describe("fixture replay through session", () => {
83
83
  );
84
84
 
85
85
  // Client received tool_call and tool_call_done events
86
- const toolStart = client.events.find((e) => (e as { type: string }).type === "tool_call") as
87
- | { toolName: string; args: Record<string, unknown> }
88
- | undefined;
89
- expect(toolStart).toBeDefined();
90
- expect(toolStart?.toolName).toBe("get_weather");
91
-
92
- const toolDone = client.events.find((e) => (e as { type: string }).type === "tool_call_done") as
93
- | { result: string }
94
- | undefined;
95
- expect(toolDone).toBeDefined();
86
+ expect(client.events).toContainEvent("tool_call", { toolName: "get_weather" });
87
+ expect(client.events).toContainEvent("tool_call_done");
96
88
 
97
89
  // Tool result was sent back to S2S after replyDone
98
90
  await vi.waitFor(() => expect(mockHandle.sendToolResult).toHaveBeenCalled());
@@ -83,14 +83,14 @@ describe("createS2sSession", () => {
83
83
  const { session, client } = setup();
84
84
  await session.start();
85
85
  session.onCancel();
86
- expect(client.events).toContainEqual(expect.objectContaining({ type: "cancelled" }));
86
+ expect(client.events).toContainEvent("cancelled");
87
87
  });
88
88
 
89
89
  test("onReset clears state and emits reset event", async () => {
90
90
  const { session, client, mockHandle } = setup();
91
91
  await session.start();
92
92
  session.onReset();
93
- expect(client.events).toContainEqual(expect.objectContaining({ type: "reset" }));
93
+ expect(client.events).toContainEvent("reset");
94
94
  expect(mockHandle.close).toHaveBeenCalled();
95
95
  });
96
96
 
@@ -119,10 +119,7 @@ describe("createS2sSession", () => {
119
119
  mockHandle._fire("event", { type: "user_transcript", text: "Hello there" });
120
120
  await flush();
121
121
 
122
- expect(client.events).toContainEqual({
123
- type: "user_transcript",
124
- text: "Hello there",
125
- });
122
+ expect(client.events).toContainEvent("user_transcript", { text: "Hello there" });
126
123
  });
127
124
 
128
125
  test("audio event forwards audio to client", async () => {
@@ -145,10 +142,7 @@ describe("createS2sSession", () => {
145
142
  _interrupted: false,
146
143
  });
147
144
 
148
- expect(client.events).toContainEqual({
149
- type: "agent_transcript",
150
- text: "Full response",
151
- });
145
+ expect(client.events).toContainEvent("agent_transcript", { text: "Full response" });
152
146
  });
153
147
 
154
148
  test("speech_started and speech_stopped events are forwarded", async () => {
@@ -158,8 +152,8 @@ describe("createS2sSession", () => {
158
152
  mockHandle._fire("event", { type: "speech_started" });
159
153
  mockHandle._fire("event", { type: "speech_stopped" });
160
154
 
161
- expect(client.events).toContainEqual({ type: "speech_started" });
162
- expect(client.events).toContainEqual({ type: "speech_stopped" });
155
+ expect(client.events).toContainEvent("speech_started");
156
+ expect(client.events).toContainEvent("speech_stopped");
163
157
  });
164
158
 
165
159
  test("reply_started resets tool call count", async () => {
@@ -177,7 +171,7 @@ describe("createS2sSession", () => {
177
171
  mockHandle._fire("event", { type: "reply_done" });
178
172
 
179
173
  expect(client.audioDoneCount).toBe(1);
180
- expect(client.events).toContainEqual({ type: "reply_done" });
174
+ expect(client.events).toContainEvent("reply_done");
181
175
  });
182
176
 
183
177
  test("cancelled event emits cancelled", async () => {
@@ -186,7 +180,7 @@ describe("createS2sSession", () => {
186
180
 
187
181
  mockHandle._fire("event", { type: "cancelled" });
188
182
 
189
- expect(client.events).toContainEqual({ type: "cancelled" });
183
+ expect(client.events).toContainEvent("cancelled");
190
184
  });
191
185
 
192
186
  test("error event emits error to client and closes handle", async () => {
@@ -195,11 +189,7 @@ describe("createS2sSession", () => {
195
189
 
196
190
  mockHandle._fire("error", new Error("Something broke"));
197
191
 
198
- expect(client.events).toContainEqual({
199
- type: "error",
200
- code: "internal",
201
- message: "Something broke",
202
- });
192
+ expect(client.events).toContainEvent("error", { code: "internal", message: "Something broke" });
203
193
  expect(mockHandle.close).toHaveBeenCalled();
204
194
  });
205
195
 
@@ -227,14 +217,12 @@ describe("createS2sSession", () => {
227
217
  "session-1",
228
218
  expect.any(Array),
229
219
  );
230
- expect(client.events).toContainEqual({
231
- type: "tool_call",
220
+ expect(client.events).toContainEvent("tool_call", {
232
221
  toolCallId: "call-1",
233
222
  toolName: "my_tool",
234
223
  args: { key: "value" },
235
224
  });
236
- expect(client.events).toContainEqual({
237
- type: "tool_call_done",
225
+ expect(client.events).toContainEvent("tool_call_done", {
238
226
  toolCallId: "call-1",
239
227
  result: "tool-output",
240
228
  });
@@ -318,13 +306,7 @@ describe("createS2sSession", () => {
318
306
 
319
307
  await session.start();
320
308
 
321
- expect(client.events).toContainEqual(
322
- expect.objectContaining({
323
- type: "error",
324
- code: "internal",
325
- message: "connect failed",
326
- }),
327
- );
309
+ expect(client.events).toContainEvent("error", { code: "internal", message: "connect failed" });
328
310
 
329
311
  spy.mockRestore();
330
312
  });
@@ -503,10 +485,6 @@ describe("createS2sSession", () => {
503
485
 
504
486
  // ─── Idle timeout tests ──────────────────────────────────────────────
505
487
 
506
- function hasIdleTimeout(events: unknown[]): boolean {
507
- return events.some((e) => (e as Record<string, unknown>).type === "idle_timeout");
508
- }
509
-
510
488
  test("idle timeout fires after configured period of inactivity", async () => {
511
489
  vi.useFakeTimers();
512
490
  const { session, client, mockHandle } = setup({
@@ -519,7 +497,7 @@ describe("createS2sSession", () => {
519
497
  });
520
498
  await session.start();
521
499
  vi.advanceTimersByTime(10_000);
522
- expect(hasIdleTimeout(client.events)).toBe(true);
500
+ expect(client.events).toContainEvent("idle_timeout");
523
501
  expect(mockHandle.close).toHaveBeenCalled();
524
502
  vi.useRealTimers();
525
503
  });
@@ -538,9 +516,9 @@ describe("createS2sSession", () => {
538
516
  vi.advanceTimersByTime(8000);
539
517
  session.onAudio(new Uint8Array([1, 2, 3]));
540
518
  vi.advanceTimersByTime(8000);
541
- expect(hasIdleTimeout(client.events)).toBe(false);
519
+ expect(client.events).not.toContainEvent("idle_timeout");
542
520
  vi.advanceTimersByTime(2000);
543
- expect(hasIdleTimeout(client.events)).toBe(true);
521
+ expect(client.events).toContainEvent("idle_timeout");
544
522
  vi.useRealTimers();
545
523
  });
546
524
 
@@ -556,7 +534,7 @@ describe("createS2sSession", () => {
556
534
  });
557
535
  await session.start();
558
536
  vi.advanceTimersByTime(600_000);
559
- expect(hasIdleTimeout(client.events)).toBe(false);
537
+ expect(client.events).not.toContainEvent("idle_timeout");
560
538
  vi.useRealTimers();
561
539
  });
562
540
 
@@ -573,7 +551,7 @@ describe("createS2sSession", () => {
573
551
  await session.start();
574
552
  await session.stop();
575
553
  vi.advanceTimersByTime(20_000);
576
- expect(hasIdleTimeout(client.events)).toBe(false);
554
+ expect(client.events).not.toContainEvent("idle_timeout");
577
555
  vi.useRealTimers();
578
556
  });
579
557
 
@@ -582,9 +560,9 @@ describe("createS2sSession", () => {
582
560
  const { session, client } = setup();
583
561
  await session.start();
584
562
  vi.advanceTimersByTime(240_000);
585
- expect(hasIdleTimeout(client.events)).toBe(false);
563
+ expect(client.events).not.toContainEvent("idle_timeout");
586
564
  vi.advanceTimersByTime(60_000);
587
- expect(hasIdleTimeout(client.events)).toBe(true);
565
+ expect(client.events).toContainEvent("idle_timeout");
588
566
  vi.useRealTimers();
589
567
  });
590
568
  });
@@ -121,4 +121,40 @@ describe("executeToolCall", () => {
121
121
  expect(result).toBe(JSON.stringify({ error: 'Tool "slow" timed out after 30000ms' }));
122
122
  vi.useRealTimers();
123
123
  });
124
+
125
+ test("ctx.send calls the send callback", async () => {
126
+ const sends: Array<{ event: string; data: unknown }> = [];
127
+ const tool: ToolDef = {
128
+ description: "sends an event",
129
+ parameters: z.object({}),
130
+ execute: (_args, ctx) => {
131
+ ctx.send("game_state", { hp: 10 });
132
+ return "ok";
133
+ },
134
+ };
135
+ const result = await executeToolCall(
136
+ "sender",
137
+ {},
138
+ {
139
+ tool,
140
+ env: {},
141
+ send: (event, data) => sends.push({ event, data }),
142
+ },
143
+ );
144
+ expect(result).toBe("ok");
145
+ expect(sends).toEqual([{ event: "game_state", data: { hp: 10 } }]);
146
+ });
147
+
148
+ test("ctx.send is a no-op when no send callback provided", async () => {
149
+ const tool: ToolDef = {
150
+ description: "sends an event",
151
+ parameters: z.object({}),
152
+ execute: (_args, ctx) => {
153
+ ctx.send("test", {});
154
+ return "ok";
155
+ },
156
+ };
157
+ const result = await executeToolCall("sender", {}, { tool, env: {} });
158
+ expect(result).toBe("ok");
159
+ });
124
160
  });
@@ -27,6 +27,7 @@ export type ExecuteToolCallOptions = {
27
27
  kv?: Kv | undefined;
28
28
  messages?: readonly Message[] | undefined;
29
29
  logger?: Logger | undefined;
30
+ send?: ((event: string, data: unknown) => void) | undefined;
30
31
  };
31
32
 
32
33
  function buildToolContext(opts: ExecuteToolCallOptions): ToolContext {
@@ -40,6 +41,9 @@ function buildToolContext(opts: ExecuteToolCallOptions): ToolContext {
40
41
  },
41
42
  messages: messages ?? [],
42
43
  sessionId: sessionId ?? "",
44
+ send(event: string, data: unknown): void {
45
+ opts.send?.(event, data);
46
+ },
43
47
  };
44
48
  }
45
49
 
@@ -47,6 +47,8 @@ export type WsSessionOptions = {
47
47
  onClose?: () => void;
48
48
  /** Callback invoked with the session ID after session cleanup. */
49
49
  onSessionEnd?: (sessionId: string) => void;
50
+ /** Callback invoked with the session ID and client sink after session setup. */
51
+ onSinkCreated?: (sessionId: string, sink: ClientSink) => void;
50
52
  /** Logger instance. Defaults to console. */
51
53
  logger?: Logger;
52
54
  /** Timeout in ms for session.start(). Defaults to 10 000 (10s). */
@@ -184,6 +186,7 @@ export function wireSessionSocket(ws: SessionWebSocket, opts: WsSessionOptions):
184
186
  const client = createClientSink(ws, log);
185
187
  session = opts.createSession(sessionId, client);
186
188
  sessions.set(sessionId, session);
189
+ opts.onSinkCreated?.(sessionId, client);
187
190
 
188
191
  // Send config immediately — zero RTT. Include sessionId so the client
189
192
  // can reconnect with ?sessionId=<id> to resume a persisted session.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@alexkroman1/aai",
3
- "version": "1.0.5",
3
+ "version": "1.1.0",
4
4
  "type": "module",
5
5
  "exports": {
6
6
  ".": {