1import { assert, describe, it } from "@effect/vitest"
2import { Effect, flow, Schema, Stream } from "effect"
3import { LanguageModel, Prompt, type Response, ResponseIdTracker, Tool, Toolkit } from "effect/ai"
4
5const finishPart: Response.FinishPartEncoded = {
6 type: "finish",
7 reason: "stop",
8 usage: {
9 inputTokens: { uncached: 5, total: 5, cacheRead: undefined, cacheWrite: undefined },
10 outputTokens: { total: 5, text: undefined, reasoning: undefined }
11 },
12 response: undefined
13}
14
15const responseMetadataPart = (id: string): Response.ResponseMetadataPartEncoded => ({
16 type: "response-metadata",
17 id,
18 modelId: undefined,
19 timestamp: undefined,
20 request: undefined
21})
22
23const textMessage = (text: string) => Prompt.textPart({ text })
24
25const userMessage = (text: string) =>
26 Prompt.userMessage({
27 content: [textMessage(text)]
28 })
29
30const assistantMessage = (text: string) =>
31 Prompt.assistantMessage({
32 content: [textMessage(text)]
33 })
34
35const ApprovalTool = Tool.make("ApprovalTool", {
36 parameters: Schema.Struct({ action: Schema.String }),
37 success: Schema.Struct({ result: Schema.String }),
38 needsApproval: true
39})
40
41const ApprovalToolkit = Toolkit.make(ApprovalTool)
42
43const ApprovalToolkitLayer = ApprovalToolkit.toLayer({
44 ApprovalTool: () => Effect.succeed({ result: "approved-result" })
45})
46
47describe("LanguageModel tracker lifecycle integration", () => {
48 it.effect("tracks generateText across turns and recovers after session drop", () =>
49 Effect.gen(function*() {
50 const tracker = yield* ResponseIdTracker.make
51 const providerOptions: Array<LanguageModel.ProviderOptions> = []
52 const responseIds = ["resp_1", "resp_2", "resp_3", "resp_4"] as const
53 let responseIndex = 0
54
55 const languageModel = yield* LanguageModel.make({
56 generateText: (options) =>
57 Effect.sync(() => {
58 const responseId = responseIds[responseIndex]!
59 responseIndex += 1
60 providerOptions.push(options)
61 return [
62 responseMetadataPart(responseId),
63 { type: "text", text: `assistant-${responseId}` },
64 finishPart
65 ]
66 }),
67 streamText: () => Stream.empty
68 })
69
70 const provideModel = Effect.provideService(LanguageModel.LanguageModel, languageModel)
71
72 const system = Prompt.systemMessage({ content: "system" })
73 const user1 = userMessage("user1")
74 const assistant1 = assistantMessage("assistant1")
75 const user2 = userMessage("user2")
76 const assistant2 = assistantMessage("assistant2")
77 const user3 = userMessage("user3")
78 const assistant3 = assistantMessage("assistant3")
79 const user4 = userMessage("user4")
80
81 yield* LanguageModel.generateText({
82 prompt: Prompt.fromMessages([system, user1])
83 }).pipe(provideModel, Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker))
84
85 const firstSentPrompt = providerOptions[0]!.prompt
86
87 yield* LanguageModel.generateText({
88 prompt: Prompt.fromMessages([...firstSentPrompt.content, assistant1, user2])
89 }).pipe(provideModel, Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker))
90
91 tracker.clearUnsafe()
92
93 const secondSentPrompt = providerOptions[1]!.prompt
94
95 yield* LanguageModel.generateText({
96 prompt: Prompt.fromMessages([...secondSentPrompt.content, assistant2, user3])
97 }).pipe(provideModel, Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker))
98
99 const thirdSentPrompt = providerOptions[2]!.prompt
100 const expectedPostDropFullPrompt = Prompt.fromMessages([...secondSentPrompt.content, assistant2, user3])
101
102 yield* LanguageModel.generateText({
103 prompt: Prompt.fromMessages([...thirdSentPrompt.content, assistant3, user4])
104 }).pipe(provideModel, Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker))
105
106 assert.strictEqual(providerOptions.length, 4)
107 assert.strictEqual(providerOptions[0]!.previousResponseId, undefined)
108 assert.strictEqual(providerOptions[0]!.incrementalPrompt, undefined)
109 assert.strictEqual(providerOptions[1]!.previousResponseId, "resp_1")
110 assert.deepStrictEqual(providerOptions[1]!.incrementalPrompt, Prompt.fromMessages([user2]))
111 assert.strictEqual(providerOptions[2]!.previousResponseId, undefined)
112 assert.strictEqual(providerOptions[2]!.incrementalPrompt, undefined)
113 assert.deepStrictEqual(providerOptions[2]!.prompt, expectedPostDropFullPrompt)
114 assert.strictEqual(providerOptions[3]!.previousResponseId, "resp_3")
115 assert.deepStrictEqual(providerOptions[3]!.incrementalPrompt, Prompt.fromMessages([user4]))
116 }))
117
118 it.effect("falls back to full prompt after stale previous response id and keeps tracking", () =>
119 Effect.gen(function*() {
120 const tracker = yield* ResponseIdTracker.make
121 const providerOptions: Array<LanguageModel.ProviderOptions> = []
122 const responseIds = ["resp_object_1", "resp_object_2"] as const
123 let responseIndex = 0
124
125 const languageModel = yield* LanguageModel.make({
126 generateText: (options) =>
127 Effect.sync(() => {
128 const responseId = responseIds[responseIndex]!
129 responseIndex += 1
130 providerOptions.push(options)
131 return [
132 responseMetadataPart(responseId),
133 { type: "text", text: "{\"count\":1}" },
134 finishPart
135 ]
136 }),
137 streamText: () => Stream.empty
138 })
139
140 const provideModel = flow(
141 Effect.provideService(LanguageModel.LanguageModel, languageModel),
142 Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker)
143 )
144 const schema = Schema.Struct({ count: Schema.Number })
145
146 const system = Prompt.systemMessage({ content: "system" })
147 const user1 = userMessage("user1")
148 const assistant1 = assistantMessage("assistant1")
149 const user2 = userMessage("user2")
150
151 yield* LanguageModel.generateObject({
152 prompt: Prompt.fromMessages([system, user1]),
153 schema
154 }).pipe(provideModel)
155
156 const firstSentPrompt = providerOptions[0]!.prompt
157
158 yield* LanguageModel.generateObject({
159 prompt: Prompt.fromMessages([...firstSentPrompt.content, assistant1, user2]),
160 schema
161 }).pipe(provideModel)
162
163 assert.strictEqual(providerOptions.length, 2)
164 assert.strictEqual(providerOptions[0]!.previousResponseId, undefined)
165 assert.strictEqual(providerOptions[0]!.incrementalPrompt, undefined)
166 assert.strictEqual(providerOptions[1]!.previousResponseId, "resp_object_1")
167 assert.deepStrictEqual(providerOptions[1]!.incrementalPrompt, Prompt.fromMessages([user2]))
168 }))
169
170 it.effect("tracks streamText with metadata parts across turns", () =>
171 Effect.gen(function*() {
172 const tracker = yield* ResponseIdTracker.make
173 const providerOptions: Array<LanguageModel.ProviderOptions> = []
174 const responseIds = ["resp_stream_1", "resp_stream_2"] as const
175 let responseIndex = 0
176
177 const languageModel = yield* LanguageModel.make({
178 generateText: () => Effect.succeed([]),
179 streamText: (options) => {
180 const responseId = responseIds[responseIndex]!
181 responseIndex += 1
182 providerOptions.push(options)
183 return Stream.fromIterable([
184 responseMetadataPart(responseId),
185 finishPart
186 ])
187 }
188 })
189
190 const provideModel = flow(
191 Effect.provideService(LanguageModel.LanguageModel, languageModel),
192 Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker)
193 )
194
195 const system = Prompt.systemMessage({ content: "system" })
196 const user1 = userMessage("user1")
197 const assistant1 = assistantMessage("assistant1")
198 const user2 = userMessage("user2")
199
200 yield* LanguageModel.streamText({
201 prompt: Prompt.fromMessages([system, user1])
202 }).pipe(
203 Stream.runDrain,
204 provideModel
205 )
206
207 const firstSentPrompt = providerOptions[0]!.prompt
208
209 yield* LanguageModel.streamText({
210 prompt: Prompt.fromMessages([...firstSentPrompt.content, assistant1, user2])
211 }).pipe(
212 Stream.runDrain,
213 provideModel
214 )
215
216 assert.strictEqual(providerOptions.length, 2)
217 assert.strictEqual(providerOptions[0]!.previousResponseId, undefined)
218 assert.strictEqual(providerOptions[0]!.incrementalPrompt, undefined)
219 assert.strictEqual(providerOptions[1]!.previousResponseId, "resp_stream_1")
220 assert.deepStrictEqual(providerOptions[1]!.incrementalPrompt, Prompt.fromMessages([user2]))
221 }))
222
223 it.effect("emits finish parts in streamText when toolkit is provided and no tool calls are returned", () =>
224 Effect.gen(function*() {
225 const languageModel = yield* LanguageModel.make({
226 generateText: () => Effect.succeed([]),
227 streamText: () => Stream.fromIterable([finishPart])
228 })
229
230 const parts = yield* LanguageModel.streamText({
231 prompt: [userMessage("hello")],
232 toolkit: ApprovalToolkit
233 }).pipe(
234 Stream.runCollect,
235 Effect.map((parts) => Array.from(parts)),
236 Effect.provideService(LanguageModel.LanguageModel, languageModel),
237 Effect.provide(ApprovalToolkitLayer)
238 )
239
240 assert.strictEqual(parts.length, 1)
241 assert.strictEqual(parts[0]?.type, "finish")
242 }))
243
244 it.effect("keeps incremental provider fields undefined when tracker is omitted", () =>
245 Effect.gen(function*() {
246 const generateTextOptions: Array<LanguageModel.ProviderOptions> = []
247 const streamTextOptions: Array<LanguageModel.ProviderOptions> = []
248
249 const languageModel = yield* LanguageModel.make({
250 generateText: (options) =>
251 Effect.sync(() => {
252 generateTextOptions.push(options)
253 return [
254 responseMetadataPart("resp_none"),
255 { type: "text", text: "{\"count\":1}" },
256 finishPart
257 ]
258 }),
259 streamText: (options) => {
260 streamTextOptions.push(options)
261 return Stream.fromIterable([
262 responseMetadataPart("resp_none_stream"),
263 finishPart
264 ])
265 }
266 })
267
268 const provideModel = Effect.provideService(LanguageModel.LanguageModel, languageModel)
269 const schema = Schema.Struct({ count: Schema.Number })
270
271 const system = Prompt.systemMessage({ content: "system" })
272 const user1 = userMessage("user1")
273 const assistant1 = assistantMessage("assistant1")
274 const user2 = userMessage("user2")
275
276 yield* LanguageModel.generateText({
277 prompt: [system, user1]
278 }).pipe(provideModel)
279
280 yield* LanguageModel.generateText({
281 prompt: [system, user1, assistant1, user2]
282 }).pipe(provideModel)
283
284 yield* LanguageModel.generateObject({
285 prompt: [system, user1],
286 schema
287 }).pipe(provideModel)
288
289 yield* LanguageModel.generateObject({
290 prompt: [system, user1, assistant1, user2],
291 schema
292 }).pipe(provideModel)
293
294 yield* LanguageModel.streamText({
295 prompt: [system, user1]
296 }).pipe(
297 Stream.runDrain,
298 provideModel
299 )
300
301 yield* LanguageModel.streamText({
302 prompt: [system, user1, assistant1, user2]
303 }).pipe(
304 Stream.runDrain,
305 provideModel
306 )
307
308 for (const options of [...generateTextOptions, ...streamTextOptions]) {
309 assert.strictEqual(options.previousResponseId, undefined)
310 assert.strictEqual(options.incrementalPrompt, undefined)
311 }
312 }))
313
314 it.effect("prepares after approval pre-resolution using stripped prompt", () =>
315 Effect.gen(function*() {
316 const tracker = yield* ResponseIdTracker.make
317 const providerOptions: Array<LanguageModel.ProviderOptions> = []
318 const responseIds = ["resp_tool_1", "resp_tool_2"] as const
319 let responseIndex = 0
320
321 const languageModel = yield* LanguageModel.make({
322 generateText: (options) =>
323 Effect.sync(() => {
324 const responseId = responseIds[responseIndex]!
325 responseIndex += 1
326 providerOptions.push(options)
327 return [
328 responseMetadataPart(responseId),
329 finishPart
330 ]
331 }),
332 streamText: () => Stream.empty
333 })
334
335 const provideModel = flow(
336 Effect.provideService(LanguageModel.LanguageModel, languageModel),
337 Effect.provideService(ResponseIdTracker.ResponseIdTracker, tracker)
338 )
339
340 const system = Prompt.systemMessage({ content: "system" })
341 const user1 = userMessage("user1")
342 const toolCallId = "call-approval"
343 const approvalId = "approval-1"
344
345 yield* LanguageModel.generateText({
346 prompt: Prompt.fromMessages([system, user1])
347 }).pipe(provideModel)
348
349 const firstSentPrompt = providerOptions[0]!.prompt
350
351 const promptWithPendingApproval: Array<Prompt.Message> = [
352 ...firstSentPrompt.content,
353 Prompt.assistantMessage({
354 content: [
355 Prompt.makePart("tool-call", {
356 id: toolCallId,
357 name: "ApprovalTool",
358 params: { action: "delete" },
359 providerExecuted: false
360 }),
361 Prompt.makePart("tool-approval-request", {
362 approvalId,
363 toolCallId
364 })
365 ]
366 }),
367 Prompt.toolMessage({
368 content: [
369 Prompt.toolApprovalResponsePart({
370 approvalId,
371 approved: true
372 })
373 ]
374 })
375 ]
376
377 yield* LanguageModel.generateText({
378 prompt: Prompt.fromMessages(promptWithPendingApproval),
379 toolkit: ApprovalToolkit
380 }).pipe(
381 provideModel,
382 Effect.provide(ApprovalToolkitLayer)
383 )
384
385 assert.strictEqual(providerOptions.length, 2)
386 const secondCall = providerOptions[1]!
387 assert.strictEqual(secondCall.previousResponseId, "resp_tool_1")
388
389 const expectedIncremental = Prompt.fromMessages([
390 Prompt.toolMessage({
391 content: [
392 Prompt.toolResultPart({
393 id: toolCallId,
394 name: "ApprovalTool",
395 result: { result: "approved-result" },
396 isFailure: false,
397 providerExecuted: false
398 })
399 ]
400 })
401 ])
402
403 assert.deepStrictEqual(secondCall.incrementalPrompt, expectedIncremental)
404
405 for (const message of secondCall.prompt.content) {
406 if (message.role === "assistant") {
407 assert.strictEqual(message.content.filter((part) => part.type === "tool-approval-request").length, 0)
408 }
409 if (message.role === "tool") {
410 assert.strictEqual(message.content.filter((part) => part.type === "tool-approval-response").length, 0)
411 }
412 }
413 }))
414})