1import { assert, describe, it } from "@effect/vitest"
2import { Effect, Option } from "effect"
3import { Prompt, ResponseIdTracker } from "effect/ai"
4
5const systemMessage = (text: string) => Prompt.systemMessage({ content: text })
6
7const userMessage = (text: string) => Prompt.userMessage({ content: [Prompt.textPart({ text })] })
8
9const assistantMessage = (text: string) => Prompt.assistantMessage({ content: [Prompt.textPart({ text })] })
10
11const assistantToolCallMessage = (id: string) =>
12 Prompt.assistantMessage({
13 content: [
14 Prompt.makePart("tool-call", {
15 id,
16 name: "test_tool",
17 params: { query: "query" },
18 providerExecuted: false
19 })
20 ]
21 })
22
23const toolResultMessage = (id: string) =>
24 Prompt.toolMessage({
25 content: [
26 Prompt.toolResultPart({
27 id,
28 name: "test_tool",
29 result: { ok: true },
30 isFailure: false,
31 providerExecuted: false
32 })
33 ]
34 })
35
36const assertPreparedSome = (
37 prepared: Option.Option<ResponseIdTracker.PrepareResult>,
38 previousResponseId: string,
39 prompt: Prompt.Prompt
40) => {
41 assert.isTrue(Option.isSome(prepared))
42 if (Option.isSome(prepared)) {
43 assert.strictEqual(prepared.value.previousResponseId, previousResponseId)
44 assert.deepStrictEqual(prepared.value.prompt, prompt)
45 }
46}
47
48describe("ResponseIdTracker", () => {
49 it.effect("returns None for a fresh tracker", () =>
50 Effect.gen(function*() {
51 const tracker = yield* ResponseIdTracker.make
52 const msg1 = userMessage("msg1")
53
54 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1]))
55
56 assert.isTrue(Option.isNone(prepared))
57 }))
58
59 it.effect("prepares only the new user message after a tracked first turn", () =>
60 Effect.gen(function*() {
61 const tracker = yield* ResponseIdTracker.make
62 const msg1 = userMessage("msg1")
63 const asst = assistantMessage("done")
64 const msg2 = userMessage("msg2")
65
66 tracker.markParts([msg1], "resp_123")
67 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1, asst, msg2]))
68
69 assertPreparedSome(prepared, "resp_123", Prompt.fromMessages([msg2]))
70 }))
71
72 it.effect("shares tracker state through its Context service layer", () =>
73 Effect.gen(function*() {
74 const tracker = yield* ResponseIdTracker.make
75 const msg1 = userMessage("msg1")
76 const msg1New = userMessage("msg1-new")
77 const asst = assistantMessage("done")
78 const msg2 = userMessage("msg2")
79
80 tracker.markParts([msg1], "resp_123")
81 tracker.clearUnsafe()
82 tracker.markParts([msg1New], "resp_456")
83 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1New, asst, msg2]))
84
85 assertPreparedSome(prepared, "resp_456", Prompt.fromMessages([msg2]))
86 }))
87
88 it.effect("keeps latest response id when markParts is called twice", () =>
89 Effect.gen(function*() {
90 const tracker = yield* ResponseIdTracker.make
91 const msg1 = userMessage("msg1")
92 const asst = assistantMessage("done")
93 const msg2 = userMessage("msg2")
94
95 tracker.markParts([msg1], "resp_1")
96 tracker.markParts([msg1], "resp_2")
97 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1, asst, msg2]))
98
99 assertPreparedSome(prepared, "resp_2", Prompt.fromMessages([msg2]))
100 }))
101
102 it.effect("returns None after clear", () =>
103 Effect.gen(function*() {
104 const tracker = yield* ResponseIdTracker.make
105 const msg1 = userMessage("msg1")
106 const asst = assistantMessage("done")
107 const msg2 = userMessage("msg2")
108
109 tracker.markParts([msg1], "resp_123")
110 tracker.clearUnsafe()
111 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1, asst, msg2]))
112
113 assert.isTrue(Option.isNone(prepared))
114 }))
115
116 it.effect("recovers after clear with new marks", () =>
117 Effect.gen(function*() {
118 const tracker = yield* ResponseIdTracker.make
119 const msg1 = userMessage("msg1")
120 const msg1New = userMessage("msg1-new")
121 const asst = assistantMessage("done")
122 const msg2 = userMessage("msg2")
123
124 tracker.markParts([msg1], "resp_123")
125 tracker.clearUnsafe()
126 tracker.markParts([msg1New], "resp_456")
127 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1New, asst, msg2]))
128
129 assertPreparedSome(prepared, "resp_456", Prompt.fromMessages([msg2]))
130 }))
131
132 it.effect("remains usable after concurrent markParts and clear", () =>
133 Effect.gen(function*() {
134 const tracker = yield* ResponseIdTracker.make
135 const raceMessage = userMessage("race")
136
137 const markPartsRace = Effect.gen(function*() {
138 for (let i = 0; i < 200; i++) {
139 tracker.markParts([raceMessage], "resp_race")
140 yield* Effect.yieldNow
141 }
142 })
143
144 const clearRace = Effect.gen(function*() {
145 for (let i = 0; i < 200; i++) {
146 tracker.clearUnsafe()
147 yield* Effect.yieldNow
148 }
149 })
150
151 yield* Effect.all([
152 markPartsRace,
153 clearRace
154 ], {
155 concurrency: "unbounded",
156 discard: true
157 })
158
159 const msg1 = userMessage("msg1")
160 const asst = assistantMessage("done")
161 const msg2 = userMessage("msg2")
162 tracker.markParts([msg1], "resp_final")
163 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1, asst, msg2]))
164
165 assertPreparedSome(prepared, "resp_final", Prompt.fromMessages([msg2]))
166 }))
167
168 it.effect("returns None for first turn with system and user messages", () =>
169 Effect.gen(function*() {
170 const tracker = yield* ResponseIdTracker.make
171 const sys = systemMessage("sys")
172 const user1 = userMessage("user1")
173
174 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1]))
175
176 assert.isTrue(Option.isNone(prepared))
177 }))
178
179 it.effect("returns None when tracked prompt has no assistant message", () =>
180 Effect.gen(function*() {
181 const tracker = yield* ResponseIdTracker.make
182 const sys = systemMessage("sys")
183 const user1 = userMessage("user1")
184
185 tracker.markParts([sys, user1], "resp_1")
186 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1]))
187
188 assert.isTrue(Option.isNone(prepared))
189 }))
190
191 it.effect("returns incremental user follow-up after assistant", () =>
192 Effect.gen(function*() {
193 const tracker = yield* ResponseIdTracker.make
194 const sys = systemMessage("sys")
195 const user1 = userMessage("user1")
196 const asst1 = assistantMessage("assistant1")
197 const user2 = userMessage("user2")
198
199 tracker.markParts([sys, user1], "resp_1")
200 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, user2]))
201
202 assertPreparedSome(prepared, "resp_1", Prompt.fromMessages([user2]))
203 }))
204
205 it.effect("returns incremental tool results after assistant tool call", () =>
206 Effect.gen(function*() {
207 const tracker = yield* ResponseIdTracker.make
208 const sys = systemMessage("sys")
209 const user1 = userMessage("user1")
210 const asst1 = assistantToolCallMessage("call_1")
211 const tool1 = toolResultMessage("call_1")
212
213 tracker.markParts([sys, user1], "resp_1")
214 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, tool1]))
215
216 assertPreparedSome(prepared, "resp_1", Prompt.fromMessages([tool1]))
217 }))
218
219 it.effect("returns incremental tool results plus user message in multi-step flow", () =>
220 Effect.gen(function*() {
221 const tracker = yield* ResponseIdTracker.make
222 const sys = systemMessage("sys")
223 const user1 = userMessage("user1")
224 const asst1 = assistantToolCallMessage("call_1")
225 const tool1 = toolResultMessage("call_1")
226 const user2 = userMessage("user2")
227
228 tracker.markParts([sys, user1], "resp_1")
229 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, tool1, user2]))
230
231 assertPreparedSome(prepared, "resp_1", Prompt.fromMessages([tool1, user2]))
232 }))
233
234 it.effect("returns None when no new messages follow assistant", () =>
235 Effect.gen(function*() {
236 const tracker = yield* ResponseIdTracker.make
237 const sys = systemMessage("sys")
238 const user1 = userMessage("user1")
239 const asst1 = assistantMessage("assistant1")
240
241 tracker.markParts([sys, user1], "resp_1")
242 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1]))
243
244 assert.isTrue(Option.isNone(prepared))
245 }))
246
247 it.effect("returns None in multi-turn prompt with no messages after last assistant", () =>
248 Effect.gen(function*() {
249 const tracker = yield* ResponseIdTracker.make
250 const sys = systemMessage("sys")
251 const user1 = userMessage("user1")
252 const asst1 = assistantMessage("assistant1")
253 const user2 = userMessage("user2")
254 const asst2 = assistantMessage("assistant2")
255
256 tracker.markParts([sys, user1, asst1, user2], "resp_2")
257 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, user2, asst2]))
258
259 assert.isTrue(Option.isNone(prepared))
260 }))
261
262 it.effect("returns None for empty prompt", () =>
263 Effect.gen(function*() {
264 const tracker = yield* ResponseIdTracker.make
265
266 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([]))
267
268 assert.isTrue(Option.isNone(prepared))
269 }))
270
271 it.effect("returns None when system prompt object changes", () =>
272 Effect.gen(function*() {
273 const tracker = yield* ResponseIdTracker.make
274 const sys = systemMessage("sys")
275 const sysNew = systemMessage("sys-new")
276 const user1 = userMessage("user1")
277 const asst1 = assistantMessage("assistant1")
278 const user2 = userMessage("user2")
279
280 tracker.markParts([sys, user1], "resp_1")
281 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sysNew, user1, asst1, user2]))
282
283 assert.isTrue(Option.isNone(prepared))
284 }))
285
286 it.effect("returns None when user message object changes", () =>
287 Effect.gen(function*() {
288 const tracker = yield* ResponseIdTracker.make
289 const sys = systemMessage("sys")
290 const user1 = userMessage("user1")
291 const user1Edited = userMessage("user1-edited")
292 const asst1 = assistantMessage("assistant1")
293 const user2 = userMessage("user2")
294
295 tracker.markParts([sys, user1], "resp_1")
296 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1Edited, asst1, user2]))
297
298 assert.isTrue(Option.isNone(prepared))
299 }))
300
301 it.effect("returns None when multiple prefix messages diverge", () =>
302 Effect.gen(function*() {
303 const tracker = yield* ResponseIdTracker.make
304 const sys = systemMessage("sys")
305 const user1 = userMessage("user1")
306 const sysNew = systemMessage("sys-new")
307 const user1Edited = userMessage("user1-edited")
308 const asst1 = assistantMessage("assistant1")
309 const user2 = userMessage("user2")
310
311 tracker.markParts([sys, user1], "resp_1")
312 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sysNew, user1Edited, asst1, user2]))
313
314 assert.isTrue(Option.isNone(prepared))
315 }))
316
317 it.effect("clears stale state after divergence and recovers when the new prefix is re-marked", () =>
318 Effect.gen(function*() {
319 const tracker = yield* ResponseIdTracker.make
320 const sys = systemMessage("sys")
321 const user1 = userMessage("user1")
322 const sysNew = systemMessage("sys-new")
323 const asst1 = assistantMessage("assistant1")
324 const user2 = userMessage("user2")
325
326 tracker.markParts([sys, user1], "resp_1")
327 const diverged = tracker.prepareUnsafe(Prompt.fromMessages([sysNew, user1, asst1, user2]))
328 assert.isTrue(Option.isNone(diverged))
329
330 tracker.markParts([sysNew, user1], "resp_2")
331 const recovered = tracker.prepareUnsafe(Prompt.fromMessages([sysNew, user1, asst1, user2]))
332
333 assertPreparedSome(recovered, "resp_2", Prompt.fromMessages([user2]))
334 }))
335
336 it.effect("supports full lifecycle across two turns", () =>
337 Effect.gen(function*() {
338 const tracker = yield* ResponseIdTracker.make
339 const sys = systemMessage("sys")
340 const user1 = userMessage("user1")
341 const asst1 = assistantMessage("assistant1")
342 const user2 = userMessage("user2")
343 const asst2 = assistantMessage("assistant2")
344 const user3 = userMessage("user3")
345
346 tracker.markParts([sys, user1], "resp_1")
347 const firstPrepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, user2]))
348 assertPreparedSome(firstPrepared, "resp_1", Prompt.fromMessages([user2]))
349
350 tracker.markParts([sys, user1, asst1, user2], "resp_2")
351 const secondPrepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, user2, asst2, user3]))
352 assertPreparedSome(secondPrepared, "resp_2", Prompt.fromMessages([user3]))
353 }))
354
355 it.effect("returns None for structurally-equal but identity-different prefix message", () =>
356 Effect.gen(function*() {
357 const tracker = yield* ResponseIdTracker.make
358 const msg1 = userMessage("msg1")
359 const msg1Copy = userMessage("msg1")
360 const asst1 = assistantMessage("assistant1")
361 const msg2 = userMessage("msg2")
362
363 tracker.markParts([msg1], "resp_1")
364 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([msg1Copy, asst1, msg2]))
365
366 assert.isTrue(Option.isNone(prepared))
367 }))
368
369 it.effect("returns None after explicit clear", () =>
370 Effect.gen(function*() {
371 const tracker = yield* ResponseIdTracker.make
372 const sys = systemMessage("sys")
373 const user1 = userMessage("user1")
374 const asst1 = assistantMessage("assistant1")
375 const user2 = userMessage("user2")
376
377 tracker.markParts([sys, user1], "resp_1")
378 tracker.clearUnsafe()
379 const prepared = tracker.prepareUnsafe(Prompt.fromMessages([sys, user1, asst1, user2]))
380
381 assert.isTrue(Option.isNone(prepared))
382 }))
383
384 it.effect("isolates two conversations on the same tracker", () =>
385 Effect.gen(function*() {
386 const tracker = yield* ResponseIdTracker.make
387 const sysA = systemMessage("sysA")
388 const userA1 = userMessage("userA1")
389 const asstA1 = assistantMessage("asstA1")
390 const userA2 = userMessage("userA2")
391 const sysB = systemMessage("sysB")
392 const userB1 = userMessage("userB1")
393 const asstB1 = assistantMessage("asstB1")
394 const userB2 = userMessage("userB2")
395
396 tracker.markParts([sysA, userA1], "resp_A1")
397 tracker.markParts([sysB, userB1], "resp_B1")
398
399 const preparedA = tracker.prepareUnsafe(Prompt.fromMessages([sysA, userA1, asstA1, userA2]))
400 assertPreparedSome(preparedA, "resp_A1", Prompt.fromMessages([userA2]))
401
402 const preparedB = tracker.prepareUnsafe(Prompt.fromMessages([sysB, userB1, asstB1, userB2]))
403 assertPreparedSome(preparedB, "resp_B1", Prompt.fromMessages([userB2]))
404 }))
405
406 it.effect("does not leak tracked state across conversations without assistant turns", () =>
407 Effect.gen(function*() {
408 const tracker = yield* ResponseIdTracker.make
409 const sysA = systemMessage("sysA")
410 const userA1 = userMessage("userA1")
411 const asstA1 = assistantMessage("asstA1")
412 const userA2 = userMessage("userA2")
413 const sysB = systemMessage("sysB")
414 const userB1 = userMessage("userB1")
415
416 tracker.markParts([sysA, userA1], "resp_A1")
417 tracker.markParts([sysB, userB1], "resp_B1")
418
419 const preparedA = tracker.prepareUnsafe(Prompt.fromMessages([sysA, userA1, asstA1, userA2]))
420 assertPreparedSome(preparedA, "resp_A1", Prompt.fromMessages([userA2]))
421
422 const preparedB = tracker.prepareUnsafe(Prompt.fromMessages([sysB, userB1]))
423 assert.isTrue(Option.isNone(preparedB))
424 }))
425})