message-v2.test.ts 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686
  1. import { describe, expect, test } from "bun:test"
  2. import { MessageV2 } from "../../src/session/message-v2"
  3. import type { ToolSet } from "ai"
  4. const sessionID = "session"
  5. // Mock tool that transforms output to content format with media support
  6. function createMockTools(): ToolSet {
  7. return {
  8. bash: {
  9. description: "mock bash tool",
  10. inputSchema: { type: "object", properties: {} } as any,
  11. toModelOutput(result: { output: string; attachments?: MessageV2.FilePart[] }) {
  12. return {
  13. type: "content" as const,
  14. value: [
  15. { type: "text" as const, text: result.output },
  16. ...(result.attachments?.map((attachment) => {
  17. const base64 = attachment.url.startsWith("data:") ? attachment.url.split(",", 2)[1] : attachment.url
  18. return {
  19. type: "media" as const,
  20. data: base64,
  21. mediaType: attachment.mime,
  22. }
  23. }) ?? []),
  24. ],
  25. }
  26. },
  27. },
  28. } as ToolSet
  29. }
  30. function userInfo(id: string): MessageV2.User {
  31. return {
  32. id,
  33. sessionID,
  34. role: "user",
  35. time: { created: 0 },
  36. agent: "user",
  37. model: { providerID: "test", modelID: "test" },
  38. tools: {},
  39. mode: "",
  40. } as unknown as MessageV2.User
  41. }
  42. function assistantInfo(id: string, parentID: string, error?: MessageV2.Assistant["error"]): MessageV2.Assistant {
  43. return {
  44. id,
  45. sessionID,
  46. role: "assistant",
  47. time: { created: 0 },
  48. error,
  49. parentID,
  50. modelID: "model",
  51. providerID: "provider",
  52. mode: "",
  53. agent: "agent",
  54. path: { cwd: "/", root: "/" },
  55. cost: 0,
  56. tokens: {
  57. input: 0,
  58. output: 0,
  59. reasoning: 0,
  60. cache: { read: 0, write: 0 },
  61. },
  62. } as unknown as MessageV2.Assistant
  63. }
  64. function basePart(messageID: string, id: string) {
  65. return {
  66. id,
  67. sessionID,
  68. messageID,
  69. }
  70. }
  71. describe("session.message-v2.toModelMessage", () => {
  72. test("filters out messages with no parts", () => {
  73. const input: MessageV2.WithParts[] = [
  74. {
  75. info: userInfo("m-empty"),
  76. parts: [],
  77. },
  78. {
  79. info: userInfo("m-user"),
  80. parts: [
  81. {
  82. ...basePart("m-user", "p1"),
  83. type: "text",
  84. text: "hello",
  85. },
  86. ] as MessageV2.Part[],
  87. },
  88. ]
  89. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  90. {
  91. role: "user",
  92. content: [{ type: "text", text: "hello" }],
  93. },
  94. ])
  95. })
  96. test("filters out messages with only ignored parts", () => {
  97. const messageID = "m-user"
  98. const input: MessageV2.WithParts[] = [
  99. {
  100. info: userInfo(messageID),
  101. parts: [
  102. {
  103. ...basePart(messageID, "p1"),
  104. type: "text",
  105. text: "ignored",
  106. ignored: true,
  107. },
  108. ] as MessageV2.Part[],
  109. },
  110. ]
  111. expect(MessageV2.toModelMessage(input)).toStrictEqual([])
  112. })
  113. test("includes synthetic text parts", () => {
  114. const messageID = "m-user"
  115. const input: MessageV2.WithParts[] = [
  116. {
  117. info: userInfo(messageID),
  118. parts: [
  119. {
  120. ...basePart(messageID, "p1"),
  121. type: "text",
  122. text: "hello",
  123. synthetic: true,
  124. },
  125. ] as MessageV2.Part[],
  126. },
  127. {
  128. info: assistantInfo("m-assistant", messageID),
  129. parts: [
  130. {
  131. ...basePart("m-assistant", "a1"),
  132. type: "text",
  133. text: "assistant",
  134. synthetic: true,
  135. },
  136. ] as MessageV2.Part[],
  137. },
  138. ]
  139. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  140. {
  141. role: "user",
  142. content: [{ type: "text", text: "hello" }],
  143. },
  144. {
  145. role: "assistant",
  146. content: [{ type: "text", text: "assistant" }],
  147. },
  148. ])
  149. })
  150. test("converts user text/file parts and injects compaction/subtask prompts", () => {
  151. const messageID = "m-user"
  152. const input: MessageV2.WithParts[] = [
  153. {
  154. info: userInfo(messageID),
  155. parts: [
  156. {
  157. ...basePart(messageID, "p1"),
  158. type: "text",
  159. text: "hello",
  160. },
  161. {
  162. ...basePart(messageID, "p2"),
  163. type: "text",
  164. text: "ignored",
  165. ignored: true,
  166. },
  167. {
  168. ...basePart(messageID, "p3"),
  169. type: "file",
  170. mime: "image/png",
  171. filename: "img.png",
  172. url: "https://example.com/img.png",
  173. },
  174. {
  175. ...basePart(messageID, "p4"),
  176. type: "file",
  177. mime: "text/plain",
  178. filename: "note.txt",
  179. url: "https://example.com/note.txt",
  180. },
  181. {
  182. ...basePart(messageID, "p5"),
  183. type: "file",
  184. mime: "application/x-directory",
  185. filename: "dir",
  186. url: "https://example.com/dir",
  187. },
  188. {
  189. ...basePart(messageID, "p6"),
  190. type: "compaction",
  191. auto: true,
  192. },
  193. {
  194. ...basePart(messageID, "p7"),
  195. type: "subtask",
  196. prompt: "prompt",
  197. description: "desc",
  198. agent: "agent",
  199. },
  200. ] as MessageV2.Part[],
  201. },
  202. ]
  203. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  204. {
  205. role: "user",
  206. content: [
  207. { type: "text", text: "hello" },
  208. {
  209. type: "file",
  210. mediaType: "image/png",
  211. filename: "img.png",
  212. data: "https://example.com/img.png",
  213. },
  214. { type: "text", text: "What did we do so far?" },
  215. { type: "text", text: "The following tool was executed by the user" },
  216. ],
  217. },
  218. ])
  219. })
  220. test("converts assistant tool completion into tool-call + tool-result messages and emits attachment message", () => {
  221. const userID = "m-user"
  222. const assistantID = "m-assistant"
  223. const input: MessageV2.WithParts[] = [
  224. {
  225. info: userInfo(userID),
  226. parts: [
  227. {
  228. ...basePart(userID, "u1"),
  229. type: "text",
  230. text: "run tool",
  231. },
  232. ] as MessageV2.Part[],
  233. },
  234. {
  235. info: assistantInfo(assistantID, userID),
  236. parts: [
  237. {
  238. ...basePart(assistantID, "a1"),
  239. type: "text",
  240. text: "done",
  241. metadata: { openai: { assistant: "meta" } },
  242. },
  243. {
  244. ...basePart(assistantID, "a2"),
  245. type: "tool",
  246. callID: "call-1",
  247. tool: "bash",
  248. state: {
  249. status: "completed",
  250. input: { cmd: "ls" },
  251. output: "ok",
  252. title: "Bash",
  253. metadata: {},
  254. time: { start: 0, end: 1 },
  255. attachments: [
  256. {
  257. ...basePart(assistantID, "file-1"),
  258. type: "file",
  259. mime: "image/png",
  260. filename: "attachment.png",
  261. url: "https://example.com/attachment.png",
  262. },
  263. ],
  264. },
  265. metadata: { openai: { tool: "meta" } },
  266. },
  267. ] as MessageV2.Part[],
  268. },
  269. ]
  270. expect(MessageV2.toModelMessage(input, { tools: createMockTools() })).toStrictEqual([
  271. {
  272. role: "user",
  273. content: [{ type: "text", text: "run tool" }],
  274. },
  275. {
  276. role: "assistant",
  277. content: [
  278. { type: "text", text: "done", providerOptions: { openai: { assistant: "meta" } } },
  279. {
  280. type: "tool-call",
  281. toolCallId: "call-1",
  282. toolName: "bash",
  283. input: { cmd: "ls" },
  284. providerExecuted: undefined,
  285. providerOptions: { openai: { tool: "meta" } },
  286. },
  287. ],
  288. },
  289. {
  290. role: "tool",
  291. content: [
  292. {
  293. type: "tool-result",
  294. toolCallId: "call-1",
  295. toolName: "bash",
  296. output: {
  297. type: "content",
  298. value: [
  299. { type: "text", text: "ok" },
  300. { type: "media", data: "https://example.com/attachment.png", mediaType: "image/png" },
  301. ],
  302. },
  303. providerOptions: { openai: { tool: "meta" } },
  304. },
  305. ],
  306. },
  307. ])
  308. })
  309. test("replaces compacted tool output with placeholder", () => {
  310. const userID = "m-user"
  311. const assistantID = "m-assistant"
  312. const input: MessageV2.WithParts[] = [
  313. {
  314. info: userInfo(userID),
  315. parts: [
  316. {
  317. ...basePart(userID, "u1"),
  318. type: "text",
  319. text: "run tool",
  320. },
  321. ] as MessageV2.Part[],
  322. },
  323. {
  324. info: assistantInfo(assistantID, userID),
  325. parts: [
  326. {
  327. ...basePart(assistantID, "a1"),
  328. type: "tool",
  329. callID: "call-1",
  330. tool: "bash",
  331. state: {
  332. status: "completed",
  333. input: { cmd: "ls" },
  334. output: "this should be cleared",
  335. title: "Bash",
  336. metadata: {},
  337. time: { start: 0, end: 1, compacted: 1 },
  338. },
  339. },
  340. ] as MessageV2.Part[],
  341. },
  342. ]
  343. expect(MessageV2.toModelMessage(input, { tools: createMockTools() })).toStrictEqual([
  344. {
  345. role: "user",
  346. content: [{ type: "text", text: "run tool" }],
  347. },
  348. {
  349. role: "assistant",
  350. content: [
  351. {
  352. type: "tool-call",
  353. toolCallId: "call-1",
  354. toolName: "bash",
  355. input: { cmd: "ls" },
  356. providerExecuted: undefined,
  357. },
  358. ],
  359. },
  360. {
  361. role: "tool",
  362. content: [
  363. {
  364. type: "tool-result",
  365. toolCallId: "call-1",
  366. toolName: "bash",
  367. output: {
  368. type: "content",
  369. value: [{ type: "text", text: "[Old tool result content cleared]" }],
  370. },
  371. },
  372. ],
  373. },
  374. ])
  375. })
  376. test("converts assistant tool error into error-text tool result", () => {
  377. const userID = "m-user"
  378. const assistantID = "m-assistant"
  379. const input: MessageV2.WithParts[] = [
  380. {
  381. info: userInfo(userID),
  382. parts: [
  383. {
  384. ...basePart(userID, "u1"),
  385. type: "text",
  386. text: "run tool",
  387. },
  388. ] as MessageV2.Part[],
  389. },
  390. {
  391. info: assistantInfo(assistantID, userID),
  392. parts: [
  393. {
  394. ...basePart(assistantID, "a1"),
  395. type: "tool",
  396. callID: "call-1",
  397. tool: "bash",
  398. state: {
  399. status: "error",
  400. input: { cmd: "ls" },
  401. error: "nope",
  402. time: { start: 0, end: 1 },
  403. metadata: {},
  404. },
  405. metadata: { openai: { tool: "meta" } },
  406. },
  407. ] as MessageV2.Part[],
  408. },
  409. ]
  410. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  411. {
  412. role: "user",
  413. content: [{ type: "text", text: "run tool" }],
  414. },
  415. {
  416. role: "assistant",
  417. content: [
  418. {
  419. type: "tool-call",
  420. toolCallId: "call-1",
  421. toolName: "bash",
  422. input: { cmd: "ls" },
  423. providerExecuted: undefined,
  424. providerOptions: { openai: { tool: "meta" } },
  425. },
  426. ],
  427. },
  428. {
  429. role: "tool",
  430. content: [
  431. {
  432. type: "tool-result",
  433. toolCallId: "call-1",
  434. toolName: "bash",
  435. output: { type: "error-text", value: "nope" },
  436. providerOptions: { openai: { tool: "meta" } },
  437. },
  438. ],
  439. },
  440. ])
  441. })
  442. test("filters assistant messages with non-abort errors", () => {
  443. const assistantID = "m-assistant"
  444. const input: MessageV2.WithParts[] = [
  445. {
  446. info: assistantInfo(
  447. assistantID,
  448. "m-parent",
  449. new MessageV2.APIError({ message: "boom", isRetryable: true }).toObject() as MessageV2.APIError,
  450. ),
  451. parts: [
  452. {
  453. ...basePart(assistantID, "a1"),
  454. type: "text",
  455. text: "should not render",
  456. },
  457. ] as MessageV2.Part[],
  458. },
  459. ]
  460. expect(MessageV2.toModelMessage(input)).toStrictEqual([])
  461. })
  462. test("includes aborted assistant messages only when they have non-step-start/reasoning content", () => {
  463. const assistantID1 = "m-assistant-1"
  464. const assistantID2 = "m-assistant-2"
  465. const aborted = new MessageV2.AbortedError({ message: "aborted" }).toObject() as MessageV2.Assistant["error"]
  466. const input: MessageV2.WithParts[] = [
  467. {
  468. info: assistantInfo(assistantID1, "m-parent", aborted),
  469. parts: [
  470. {
  471. ...basePart(assistantID1, "a1"),
  472. type: "reasoning",
  473. text: "thinking",
  474. time: { start: 0 },
  475. },
  476. {
  477. ...basePart(assistantID1, "a2"),
  478. type: "text",
  479. text: "partial answer",
  480. },
  481. ] as MessageV2.Part[],
  482. },
  483. {
  484. info: assistantInfo(assistantID2, "m-parent", aborted),
  485. parts: [
  486. {
  487. ...basePart(assistantID2, "b1"),
  488. type: "step-start",
  489. },
  490. {
  491. ...basePart(assistantID2, "b2"),
  492. type: "reasoning",
  493. text: "thinking",
  494. time: { start: 0 },
  495. },
  496. ] as MessageV2.Part[],
  497. },
  498. ]
  499. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  500. {
  501. role: "assistant",
  502. content: [
  503. { type: "reasoning", text: "thinking", providerOptions: undefined },
  504. { type: "text", text: "partial answer" },
  505. ],
  506. },
  507. ])
  508. })
  509. test("splits assistant messages on step-start boundaries", () => {
  510. const assistantID = "m-assistant"
  511. const input: MessageV2.WithParts[] = [
  512. {
  513. info: assistantInfo(assistantID, "m-parent"),
  514. parts: [
  515. {
  516. ...basePart(assistantID, "p1"),
  517. type: "text",
  518. text: "first",
  519. },
  520. {
  521. ...basePart(assistantID, "p2"),
  522. type: "step-start",
  523. },
  524. {
  525. ...basePart(assistantID, "p3"),
  526. type: "text",
  527. text: "second",
  528. },
  529. ] as MessageV2.Part[],
  530. },
  531. ]
  532. expect(MessageV2.toModelMessage(input)).toStrictEqual([
  533. {
  534. role: "assistant",
  535. content: [{ type: "text", text: "first" }],
  536. },
  537. {
  538. role: "assistant",
  539. content: [{ type: "text", text: "second" }],
  540. },
  541. ])
  542. })
  543. test("drops messages that only contain step-start parts", () => {
  544. const assistantID = "m-assistant"
  545. const input: MessageV2.WithParts[] = [
  546. {
  547. info: assistantInfo(assistantID, "m-parent"),
  548. parts: [
  549. {
  550. ...basePart(assistantID, "p1"),
  551. type: "step-start",
  552. },
  553. ] as MessageV2.Part[],
  554. },
  555. ]
  556. expect(MessageV2.toModelMessage(input)).toStrictEqual([])
  557. })
  558. test("converts pending/running tool calls to error results to prevent dangling tool_use", () => {
  559. const userID = "m-user"
  560. const assistantID = "m-assistant"
  561. const input: MessageV2.WithParts[] = [
  562. {
  563. info: userInfo(userID),
  564. parts: [
  565. {
  566. ...basePart(userID, "u1"),
  567. type: "text",
  568. text: "run tool",
  569. },
  570. ] as MessageV2.Part[],
  571. },
  572. {
  573. info: assistantInfo(assistantID, userID),
  574. parts: [
  575. {
  576. ...basePart(assistantID, "a1"),
  577. type: "tool",
  578. callID: "call-pending",
  579. tool: "bash",
  580. state: {
  581. status: "pending",
  582. input: { cmd: "ls" },
  583. raw: "",
  584. },
  585. },
  586. {
  587. ...basePart(assistantID, "a2"),
  588. type: "tool",
  589. callID: "call-running",
  590. tool: "read",
  591. state: {
  592. status: "running",
  593. input: { path: "/tmp" },
  594. time: { start: 0 },
  595. },
  596. },
  597. ] as MessageV2.Part[],
  598. },
  599. ]
  600. const result = MessageV2.toModelMessage(input)
  601. expect(result).toStrictEqual([
  602. {
  603. role: "user",
  604. content: [{ type: "text", text: "run tool" }],
  605. },
  606. {
  607. role: "assistant",
  608. content: [
  609. {
  610. type: "tool-call",
  611. toolCallId: "call-pending",
  612. toolName: "bash",
  613. input: { cmd: "ls" },
  614. providerExecuted: undefined,
  615. },
  616. {
  617. type: "tool-call",
  618. toolCallId: "call-running",
  619. toolName: "read",
  620. input: { path: "/tmp" },
  621. providerExecuted: undefined,
  622. },
  623. ],
  624. },
  625. {
  626. role: "tool",
  627. content: [
  628. {
  629. type: "tool-result",
  630. toolCallId: "call-pending",
  631. toolName: "bash",
  632. output: { type: "error-text", value: "[Tool execution was interrupted]" },
  633. },
  634. {
  635. type: "tool-result",
  636. toolCallId: "call-running",
  637. toolName: "read",
  638. output: { type: "error-text", value: "[Tool execution was interrupted]" },
  639. },
  640. ],
  641. },
  642. ])
  643. })
  644. })