session-entry.test.ts 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681
  1. import { describe, expect, test } from "bun:test"
  2. import * as DateTime from "effect/DateTime"
  3. import * as FastCheck from "effect/testing/FastCheck"
  4. import { SessionEntry } from "../../src/v2/session-entry"
  5. import { SessionEvent } from "../../src/v2/session-event"
  6. const time = (n: number) => DateTime.makeUnsafe(n)
  7. const word = FastCheck.string({ minLength: 1, maxLength: 8 })
  8. const text = FastCheck.string({ maxLength: 16 })
  9. const texts = FastCheck.array(text, { maxLength: 8 })
  10. const val = FastCheck.oneof(FastCheck.boolean(), FastCheck.integer(), FastCheck.string({ maxLength: 12 }))
  11. const dict = FastCheck.dictionary(word, val, { maxKeys: 4 })
  12. const files = FastCheck.array(
  13. word.map((x) => SessionEvent.FileAttachment.create({ uri: `file://${encodeURIComponent(x)}`, mime: "text/plain" })),
  14. { maxLength: 2 },
  15. )
  16. function maybe<A>(arb: FastCheck.Arbitrary<A>) {
  17. return FastCheck.oneof(FastCheck.constant(undefined), arb)
  18. }
  19. function assistant() {
  20. return new SessionEntry.Assistant({
  21. id: SessionEvent.ID.create(),
  22. type: "assistant",
  23. time: { created: time(0) },
  24. content: [],
  25. })
  26. }
  27. function history() {
  28. const state: SessionEntry.History = {
  29. entries: [],
  30. pending: [],
  31. }
  32. return state
  33. }
  34. function active() {
  35. const state: SessionEntry.History = {
  36. entries: [assistant()],
  37. pending: [],
  38. }
  39. return state
  40. }
  41. function run(events: SessionEvent.Event[], state = history()) {
  42. return events.reduce<SessionEntry.History>((state, event) => SessionEntry.step(state, event), state)
  43. }
  44. function last(state: SessionEntry.History) {
  45. const entry = [...state.pending, ...state.entries].reverse().find((x) => x.type === "assistant")
  46. expect(entry?.type).toBe("assistant")
  47. return entry?.type === "assistant" ? entry : undefined
  48. }
  49. function texts_of(state: SessionEntry.History) {
  50. const entry = last(state)
  51. if (!entry) return []
  52. return entry.content.filter((x): x is SessionEntry.AssistantText => x.type === "text")
  53. }
  54. function reasons(state: SessionEntry.History) {
  55. const entry = last(state)
  56. if (!entry) return []
  57. return entry.content.filter((x): x is SessionEntry.AssistantReasoning => x.type === "reasoning")
  58. }
  59. function tools(state: SessionEntry.History) {
  60. const entry = last(state)
  61. if (!entry) return []
  62. return entry.content.filter((x): x is SessionEntry.AssistantTool => x.type === "tool")
  63. }
  64. function tool(state: SessionEntry.History, callID: string) {
  65. return tools(state).find((x) => x.callID === callID)
  66. }
  67. describe("session-entry step", () => {
  68. describe("seeded pending assistant", () => {
  69. test("stores prompts in entries when no assistant is pending", () => {
  70. FastCheck.assert(
  71. FastCheck.property(word, (body) => {
  72. const next = SessionEntry.step(history(), SessionEvent.Prompt.create({ text: body, timestamp: time(1) }))
  73. expect(next.entries).toHaveLength(1)
  74. expect(next.entries[0]?.type).toBe("user")
  75. if (next.entries[0]?.type !== "user") return
  76. expect(next.entries[0].text).toBe(body)
  77. }),
  78. { numRuns: 50 },
  79. )
  80. })
  81. test("stores prompts in pending when an assistant is pending", () => {
  82. FastCheck.assert(
  83. FastCheck.property(word, (body) => {
  84. const next = SessionEntry.step(active(), SessionEvent.Prompt.create({ text: body, timestamp: time(1) }))
  85. expect(next.pending).toHaveLength(1)
  86. expect(next.pending[0]?.type).toBe("user")
  87. if (next.pending[0]?.type !== "user") return
  88. expect(next.pending[0].text).toBe(body)
  89. }),
  90. { numRuns: 50 },
  91. )
  92. })
  93. test("accumulates text deltas on the latest text part", () => {
  94. FastCheck.assert(
  95. FastCheck.property(texts, (parts) => {
  96. const next = parts.reduce(
  97. (state, part, i) =>
  98. SessionEntry.step(state, SessionEvent.Text.Delta.create({ delta: part, timestamp: time(i + 2) })),
  99. SessionEntry.step(active(), SessionEvent.Text.Started.create({ timestamp: time(1) })),
  100. )
  101. expect(texts_of(next)).toEqual([
  102. {
  103. type: "text",
  104. text: parts.join(""),
  105. },
  106. ])
  107. }),
  108. { numRuns: 100 },
  109. )
  110. })
  111. test("routes later text deltas to the latest text segment", () => {
  112. FastCheck.assert(
  113. FastCheck.property(texts, texts, (a, b) => {
  114. const next = run(
  115. [
  116. SessionEvent.Text.Started.create({ timestamp: time(1) }),
  117. ...a.map((x, i) => SessionEvent.Text.Delta.create({ delta: x, timestamp: time(i + 2) })),
  118. SessionEvent.Text.Started.create({ timestamp: time(a.length + 2) }),
  119. ...b.map((x, i) => SessionEvent.Text.Delta.create({ delta: x, timestamp: time(i + a.length + 3) })),
  120. ],
  121. active(),
  122. )
  123. expect(texts_of(next)).toEqual([
  124. { type: "text", text: a.join("") },
  125. { type: "text", text: b.join("") },
  126. ])
  127. }),
  128. { numRuns: 50 },
  129. )
  130. })
  131. test("reasoning.ended replaces buffered reasoning text", () => {
  132. FastCheck.assert(
  133. FastCheck.property(texts, text, (parts, end) => {
  134. const next = run(
  135. [
  136. SessionEvent.Reasoning.Started.create({ timestamp: time(1) }),
  137. ...parts.map((x, i) => SessionEvent.Reasoning.Delta.create({ delta: x, timestamp: time(i + 2) })),
  138. SessionEvent.Reasoning.Ended.create({ text: end, timestamp: time(parts.length + 2) }),
  139. ],
  140. active(),
  141. )
  142. expect(reasons(next)).toEqual([
  143. {
  144. type: "reasoning",
  145. text: end,
  146. },
  147. ])
  148. }),
  149. { numRuns: 100 },
  150. )
  151. })
  152. test("tool.success completes the latest running tool", () => {
  153. FastCheck.assert(
  154. FastCheck.property(
  155. word,
  156. word,
  157. dict,
  158. maybe(text),
  159. maybe(dict),
  160. maybe(files),
  161. texts,
  162. (callID, title, input, output, metadata, attachments, parts) => {
  163. const next = run(
  164. [
  165. SessionEvent.Tool.Input.Started.create({ callID, name: "bash", timestamp: time(1) }),
  166. ...parts.map((x, i) =>
  167. SessionEvent.Tool.Input.Delta.create({ callID, delta: x, timestamp: time(i + 2) }),
  168. ),
  169. SessionEvent.Tool.Called.create({
  170. callID,
  171. tool: "bash",
  172. input,
  173. provider: { executed: true },
  174. timestamp: time(parts.length + 2),
  175. }),
  176. SessionEvent.Tool.Success.create({
  177. callID,
  178. title,
  179. output,
  180. metadata,
  181. attachments,
  182. provider: { executed: true },
  183. timestamp: time(parts.length + 3),
  184. }),
  185. ],
  186. active(),
  187. )
  188. const match = tool(next, callID)
  189. expect(match?.state.status).toBe("completed")
  190. if (match?.state.status !== "completed") return
  191. expect(match.time.ran).toEqual(time(parts.length + 2))
  192. expect(match.state.input).toEqual(input)
  193. expect(match.state.output).toBe(output ?? "")
  194. expect(match.state.title).toBe(title)
  195. expect(match.state.metadata).toEqual(metadata ?? {})
  196. expect(match.state.attachments).toEqual(attachments ?? [])
  197. },
  198. ),
  199. { numRuns: 50 },
  200. )
  201. })
  202. test("tool.error completes the latest running tool with an error", () => {
  203. FastCheck.assert(
  204. FastCheck.property(word, dict, word, maybe(dict), (callID, input, error, metadata) => {
  205. const next = run(
  206. [
  207. SessionEvent.Tool.Input.Started.create({ callID, name: "bash", timestamp: time(1) }),
  208. SessionEvent.Tool.Called.create({
  209. callID,
  210. tool: "bash",
  211. input,
  212. provider: { executed: true },
  213. timestamp: time(2),
  214. }),
  215. SessionEvent.Tool.Error.create({
  216. callID,
  217. error,
  218. metadata,
  219. provider: { executed: true },
  220. timestamp: time(3),
  221. }),
  222. ],
  223. active(),
  224. )
  225. const match = tool(next, callID)
  226. expect(match?.state.status).toBe("error")
  227. if (match?.state.status !== "error") return
  228. expect(match.time.ran).toEqual(time(2))
  229. expect(match.state.input).toEqual(input)
  230. expect(match.state.error).toBe(error)
  231. expect(match.state.metadata).toEqual(metadata ?? {})
  232. }),
  233. { numRuns: 50 },
  234. )
  235. })
  236. test("tool.success is ignored before tool.called promotes the tool to running", () => {
  237. FastCheck.assert(
  238. FastCheck.property(word, word, (callID, title) => {
  239. const next = run(
  240. [
  241. SessionEvent.Tool.Input.Started.create({ callID, name: "bash", timestamp: time(1) }),
  242. SessionEvent.Tool.Success.create({
  243. callID,
  244. title,
  245. provider: { executed: true },
  246. timestamp: time(2),
  247. }),
  248. ],
  249. active(),
  250. )
  251. const match = tool(next, callID)
  252. expect(match?.state).toEqual({
  253. status: "pending",
  254. input: "",
  255. })
  256. }),
  257. { numRuns: 50 },
  258. )
  259. })
  260. test("step.ended copies completion fields onto the pending assistant", () => {
  261. FastCheck.assert(
  262. FastCheck.property(FastCheck.integer({ min: 1, max: 1000 }), (n) => {
  263. const event = SessionEvent.Step.Ended.create({
  264. reason: "stop",
  265. cost: 1,
  266. tokens: {
  267. input: 1,
  268. output: 2,
  269. reasoning: 3,
  270. cache: {
  271. read: 4,
  272. write: 5,
  273. },
  274. },
  275. timestamp: time(n),
  276. })
  277. const next = SessionEntry.step(active(), event)
  278. const entry = last(next)
  279. expect(entry).toBeDefined()
  280. if (!entry) return
  281. expect(entry.time.completed).toEqual(event.timestamp)
  282. expect(entry.cost).toBe(event.cost)
  283. expect(entry.tokens).toEqual(event.tokens)
  284. }),
  285. { numRuns: 50 },
  286. )
  287. })
  288. })
  289. describe("known reducer gaps", () => {
  290. test("prompt appends immutably when no assistant is pending", () => {
  291. FastCheck.assert(
  292. FastCheck.property(word, (body) => {
  293. const old = history()
  294. const next = SessionEntry.step(old, SessionEvent.Prompt.create({ text: body, timestamp: time(1) }))
  295. expect(old).not.toBe(next)
  296. expect(old.entries).toHaveLength(0)
  297. expect(next.entries).toHaveLength(1)
  298. }),
  299. { numRuns: 50 },
  300. )
  301. })
  302. test("prompt appends immutably when an assistant is pending", () => {
  303. FastCheck.assert(
  304. FastCheck.property(word, (body) => {
  305. const old = active()
  306. const next = SessionEntry.step(old, SessionEvent.Prompt.create({ text: body, timestamp: time(1) }))
  307. expect(old).not.toBe(next)
  308. expect(old.pending).toHaveLength(0)
  309. expect(next.pending).toHaveLength(1)
  310. }),
  311. { numRuns: 50 },
  312. )
  313. })
  314. test("step.started creates an assistant consumed by follow-up events", () => {
  315. FastCheck.assert(
  316. FastCheck.property(texts, (parts) => {
  317. const next = run([
  318. SessionEvent.Step.Started.create({
  319. model: {
  320. id: "model",
  321. providerID: "provider",
  322. },
  323. timestamp: time(1),
  324. }),
  325. SessionEvent.Text.Started.create({ timestamp: time(2) }),
  326. ...parts.map((x, i) => SessionEvent.Text.Delta.create({ delta: x, timestamp: time(i + 3) })),
  327. SessionEvent.Step.Ended.create({
  328. reason: "stop",
  329. cost: 1,
  330. tokens: {
  331. input: 1,
  332. output: 2,
  333. reasoning: 3,
  334. cache: {
  335. read: 4,
  336. write: 5,
  337. },
  338. },
  339. timestamp: time(parts.length + 3),
  340. }),
  341. ])
  342. const entry = last(next)
  343. expect(entry).toBeDefined()
  344. if (!entry) return
  345. expect(entry.content).toEqual([
  346. {
  347. type: "text",
  348. text: parts.join(""),
  349. },
  350. ])
  351. expect(entry.time.completed).toEqual(time(parts.length + 3))
  352. }),
  353. { numRuns: 100 },
  354. )
  355. })
  356. test("replays prompt -> step -> text -> step.ended", () => {
  357. FastCheck.assert(
  358. FastCheck.property(word, texts, (body, parts) => {
  359. const next = run([
  360. SessionEvent.Prompt.create({ text: body, timestamp: time(0) }),
  361. SessionEvent.Step.Started.create({
  362. model: {
  363. id: "model",
  364. providerID: "provider",
  365. },
  366. timestamp: time(1),
  367. }),
  368. SessionEvent.Text.Started.create({ timestamp: time(2) }),
  369. ...parts.map((x, i) => SessionEvent.Text.Delta.create({ delta: x, timestamp: time(i + 3) })),
  370. SessionEvent.Step.Ended.create({
  371. reason: "stop",
  372. cost: 1,
  373. tokens: {
  374. input: 1,
  375. output: 2,
  376. reasoning: 3,
  377. cache: {
  378. read: 4,
  379. write: 5,
  380. },
  381. },
  382. timestamp: time(parts.length + 3),
  383. }),
  384. ])
  385. expect(next.entries).toHaveLength(2)
  386. expect(next.entries[0]?.type).toBe("user")
  387. expect(next.entries[1]?.type).toBe("assistant")
  388. if (next.entries[1]?.type !== "assistant") return
  389. expect(next.entries[1].content).toEqual([
  390. {
  391. type: "text",
  392. text: parts.join(""),
  393. },
  394. ])
  395. expect(next.entries[1].time.completed).toEqual(time(parts.length + 3))
  396. }),
  397. { numRuns: 50 },
  398. )
  399. })
  400. test("replays prompt -> step -> reasoning -> tool -> success -> step.ended", () => {
  401. FastCheck.assert(
  402. FastCheck.property(
  403. word,
  404. texts,
  405. text,
  406. dict,
  407. word,
  408. maybe(text),
  409. maybe(dict),
  410. maybe(files),
  411. (body, reason, end, input, title, output, metadata, attachments) => {
  412. const callID = "call"
  413. const next = run([
  414. SessionEvent.Prompt.create({ text: body, timestamp: time(0) }),
  415. SessionEvent.Step.Started.create({
  416. model: {
  417. id: "model",
  418. providerID: "provider",
  419. },
  420. timestamp: time(1),
  421. }),
  422. SessionEvent.Reasoning.Started.create({ timestamp: time(2) }),
  423. ...reason.map((x, i) => SessionEvent.Reasoning.Delta.create({ delta: x, timestamp: time(i + 3) })),
  424. SessionEvent.Reasoning.Ended.create({ text: end, timestamp: time(reason.length + 3) }),
  425. SessionEvent.Tool.Input.Started.create({ callID, name: "bash", timestamp: time(reason.length + 4) }),
  426. SessionEvent.Tool.Called.create({
  427. callID,
  428. tool: "bash",
  429. input,
  430. provider: { executed: true },
  431. timestamp: time(reason.length + 5),
  432. }),
  433. SessionEvent.Tool.Success.create({
  434. callID,
  435. title,
  436. output,
  437. metadata,
  438. attachments,
  439. provider: { executed: true },
  440. timestamp: time(reason.length + 6),
  441. }),
  442. SessionEvent.Step.Ended.create({
  443. reason: "stop",
  444. cost: 1,
  445. tokens: {
  446. input: 1,
  447. output: 2,
  448. reasoning: 3,
  449. cache: {
  450. read: 4,
  451. write: 5,
  452. },
  453. },
  454. timestamp: time(reason.length + 7),
  455. }),
  456. ])
  457. expect(next.entries.at(-1)?.type).toBe("assistant")
  458. const entry = next.entries.at(-1)
  459. if (entry?.type !== "assistant") return
  460. expect(entry.content).toHaveLength(2)
  461. expect(entry.content[0]).toEqual({
  462. type: "reasoning",
  463. text: end,
  464. })
  465. expect(entry.content[1]?.type).toBe("tool")
  466. if (entry.content[1]?.type !== "tool") return
  467. expect(entry.content[1].state.status).toBe("completed")
  468. expect(entry.time.completed).toEqual(time(reason.length + 7))
  469. },
  470. ),
  471. { numRuns: 50 },
  472. )
  473. })
  474. test("starting a new step completes the old assistant and appends a new active assistant", () => {
  475. const next = run(
  476. [
  477. SessionEvent.Step.Started.create({
  478. model: {
  479. id: "model",
  480. providerID: "provider",
  481. },
  482. timestamp: time(1),
  483. }),
  484. ],
  485. active(),
  486. )
  487. expect(next.entries).toHaveLength(2)
  488. expect(next.entries[0]?.type).toBe("assistant")
  489. expect(next.entries[1]?.type).toBe("assistant")
  490. if (next.entries[0]?.type !== "assistant" || next.entries[1]?.type !== "assistant") return
  491. expect(next.entries[0].time.completed).toEqual(time(1))
  492. expect(next.entries[1].time.created).toEqual(time(1))
  493. expect(next.entries[1].time.completed).toBeUndefined()
  494. })
  495. test("handles sequential tools independently", () => {
  496. FastCheck.assert(
  497. FastCheck.property(dict, dict, word, word, (a, b, title, error) => {
  498. const next = run(
  499. [
  500. SessionEvent.Tool.Input.Started.create({ callID: "a", name: "bash", timestamp: time(1) }),
  501. SessionEvent.Tool.Called.create({
  502. callID: "a",
  503. tool: "bash",
  504. input: a,
  505. provider: { executed: true },
  506. timestamp: time(2),
  507. }),
  508. SessionEvent.Tool.Success.create({
  509. callID: "a",
  510. title,
  511. output: "done",
  512. provider: { executed: true },
  513. timestamp: time(3),
  514. }),
  515. SessionEvent.Tool.Input.Started.create({ callID: "b", name: "grep", timestamp: time(4) }),
  516. SessionEvent.Tool.Called.create({
  517. callID: "b",
  518. tool: "bash",
  519. input: b,
  520. provider: { executed: true },
  521. timestamp: time(5),
  522. }),
  523. SessionEvent.Tool.Error.create({
  524. callID: "b",
  525. error,
  526. provider: { executed: true },
  527. timestamp: time(6),
  528. }),
  529. ],
  530. active(),
  531. )
  532. const first = tool(next, "a")
  533. const second = tool(next, "b")
  534. expect(first?.state.status).toBe("completed")
  535. if (first?.state.status !== "completed") return
  536. expect(first.state.input).toEqual(a)
  537. expect(first.state.output).toBe("done")
  538. expect(first.state.title).toBe(title)
  539. expect(second?.state.status).toBe("error")
  540. if (second?.state.status !== "error") return
  541. expect(second.state.input).toEqual(b)
  542. expect(second.state.error).toBe(error)
  543. }),
  544. { numRuns: 50 },
  545. )
  546. })
  547. test("routes tool events by callID when tool streams interleave", () => {
  548. FastCheck.assert(
  549. FastCheck.property(dict, dict, word, word, text, text, (a, b, titleA, titleB, deltaA, deltaB) => {
  550. const next = run(
  551. [
  552. SessionEvent.Tool.Input.Started.create({ callID: "a", name: "bash", timestamp: time(1) }),
  553. SessionEvent.Tool.Input.Started.create({ callID: "b", name: "grep", timestamp: time(2) }),
  554. SessionEvent.Tool.Input.Delta.create({ callID: "a", delta: deltaA, timestamp: time(3) }),
  555. SessionEvent.Tool.Input.Delta.create({ callID: "b", delta: deltaB, timestamp: time(4) }),
  556. SessionEvent.Tool.Called.create({
  557. callID: "a",
  558. tool: "bash",
  559. input: a,
  560. provider: { executed: true },
  561. timestamp: time(5),
  562. }),
  563. SessionEvent.Tool.Called.create({
  564. callID: "b",
  565. tool: "grep",
  566. input: b,
  567. provider: { executed: true },
  568. timestamp: time(6),
  569. }),
  570. SessionEvent.Tool.Success.create({
  571. callID: "a",
  572. title: titleA,
  573. output: "done-a",
  574. provider: { executed: true },
  575. timestamp: time(7),
  576. }),
  577. SessionEvent.Tool.Success.create({
  578. callID: "b",
  579. title: titleB,
  580. output: "done-b",
  581. provider: { executed: true },
  582. timestamp: time(8),
  583. }),
  584. ],
  585. active(),
  586. )
  587. const first = tool(next, "a")
  588. const second = tool(next, "b")
  589. expect(first?.state.status).toBe("completed")
  590. expect(second?.state.status).toBe("completed")
  591. if (first?.state.status !== "completed" || second?.state.status !== "completed") return
  592. expect(first.state.input).toEqual(a)
  593. expect(second.state.input).toEqual(b)
  594. expect(first.state.title).toBe(titleA)
  595. expect(second.state.title).toBe(titleB)
  596. }),
  597. { numRuns: 50 },
  598. )
  599. })
  600. test("records synthetic events", () => {
  601. FastCheck.assert(
  602. FastCheck.property(word, (body) => {
  603. const next = SessionEntry.step(history(), SessionEvent.Synthetic.create({ text: body, timestamp: time(1) }))
  604. expect(next.entries).toHaveLength(1)
  605. expect(next.entries[0]?.type).toBe("synthetic")
  606. if (next.entries[0]?.type !== "synthetic") return
  607. expect(next.entries[0].text).toBe(body)
  608. }),
  609. { numRuns: 50 },
  610. )
  611. })
  612. test("records compaction events", () => {
  613. FastCheck.assert(
  614. FastCheck.property(FastCheck.boolean(), maybe(FastCheck.boolean()), (auto, overflow) => {
  615. const next = SessionEntry.step(
  616. history(),
  617. SessionEvent.Compacted.create({ auto, overflow, timestamp: time(1) }),
  618. )
  619. expect(next.entries).toHaveLength(1)
  620. expect(next.entries[0]?.type).toBe("compaction")
  621. if (next.entries[0]?.type !== "compaction") return
  622. expect(next.entries[0].auto).toBe(auto)
  623. expect(next.entries[0].overflow).toBe(overflow)
  624. }),
  625. { numRuns: 50 },
  626. )
  627. })
  628. })
  629. })