runtime.queue.test.ts 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465
  1. import { describe, expect, test } from "bun:test"
  2. import { runPromptQueue } from "@/cli/cmd/run/runtime.queue"
  3. import type { FooterApi, FooterEvent, RunPrompt, StreamCommit } from "@/cli/cmd/run/types"
  4. function footer() {
  5. const prompts = new Set<(input: RunPrompt) => void>()
  6. const queuedRemoves = new Set<(messageID: string) => void>()
  7. const closes = new Set<() => void>()
  8. const events: FooterEvent[] = []
  9. const commits: StreamCommit[] = []
  10. let closed = false
  11. const api: FooterApi = {
  12. get isClosed() {
  13. return closed
  14. },
  15. onPrompt(fn) {
  16. prompts.add(fn)
  17. return () => {
  18. prompts.delete(fn)
  19. }
  20. },
  21. onQueuedRemove(fn) {
  22. queuedRemoves.add(fn)
  23. return () => {
  24. queuedRemoves.delete(fn)
  25. }
  26. },
  27. onClose(fn) {
  28. if (closed) {
  29. fn()
  30. return () => {}
  31. }
  32. closes.add(fn)
  33. return () => {
  34. closes.delete(fn)
  35. }
  36. },
  37. event(next) {
  38. events.push(next)
  39. },
  40. append(next) {
  41. commits.push(next)
  42. },
  43. idle() {
  44. return Promise.resolve()
  45. },
  46. close() {
  47. if (closed) {
  48. return
  49. }
  50. closed = true
  51. for (const fn of [...closes]) {
  52. fn()
  53. }
  54. },
  55. destroy() {
  56. api.close()
  57. prompts.clear()
  58. closes.clear()
  59. },
  60. }
  61. return {
  62. api,
  63. events,
  64. commits,
  65. submit(text: string, mode?: RunPrompt["mode"]) {
  66. const next = mode ? { text, parts: [] as RunPrompt["parts"], mode } : { text, parts: [] as RunPrompt["parts"] }
  67. for (const fn of [...prompts]) {
  68. fn(next)
  69. }
  70. },
  71. removeQueued(messageID: string) {
  72. for (const fn of [...queuedRemoves]) fn(messageID)
  73. },
  74. }
  75. }
  76. describe("run runtime queue", () => {
  77. test("ignores empty prompts", async () => {
  78. const ui = footer()
  79. let calls = 0
  80. const task = runPromptQueue({
  81. footer: ui.api,
  82. run: async () => {
  83. calls += 1
  84. },
  85. })
  86. ui.submit(" ")
  87. ui.api.close()
  88. await task
  89. expect(calls).toBe(0)
  90. })
  91. test("treats /exit as a close command", async () => {
  92. const ui = footer()
  93. let calls = 0
  94. const task = runPromptQueue({
  95. footer: ui.api,
  96. run: async () => {
  97. calls += 1
  98. },
  99. })
  100. ui.submit("/exit")
  101. await task
  102. expect(calls).toBe(0)
  103. })
  104. test("treats /new as a local session command", async () => {
  105. const ui = footer()
  106. const seen: string[] = []
  107. let created = 0
  108. const task = runPromptQueue({
  109. footer: ui.api,
  110. onNewSession: async () => {
  111. created += 1
  112. },
  113. run: async (input) => {
  114. seen.push(input.text)
  115. ui.api.close()
  116. },
  117. })
  118. ui.submit("/new")
  119. ui.submit("hello")
  120. await task
  121. expect(created).toBe(1)
  122. expect(seen).toEqual(["hello"])
  123. expect(ui.commits).toEqual([
  124. {
  125. kind: "user",
  126. text: "hello",
  127. phase: "start",
  128. source: "system",
  129. messageID: expect.any(String),
  130. },
  131. ])
  132. })
  133. test("shell mode submits /exit as a shell command", async () => {
  134. const ui = footer()
  135. const seen: RunPrompt[] = []
  136. const task = runPromptQueue({
  137. footer: ui.api,
  138. run: async (input) => {
  139. seen.push(input)
  140. ui.api.close()
  141. },
  142. })
  143. ui.submit("/exit", "shell")
  144. await task
  145. expect(seen).toEqual([{ text: "/exit", parts: [], mode: "shell" }])
  146. expect(ui.commits).toEqual([])
  147. })
  148. test("shell mode submits /new instead of creating a session", async () => {
  149. const ui = footer()
  150. const seen: RunPrompt[] = []
  151. let created = 0
  152. const task = runPromptQueue({
  153. footer: ui.api,
  154. onNewSession: async () => {
  155. created += 1
  156. },
  157. run: async (input) => {
  158. seen.push(input)
  159. ui.api.close()
  160. },
  161. })
  162. ui.submit("/new", "shell")
  163. await task
  164. expect(created).toBe(0)
  165. expect(seen).toEqual([{ text: "/new", parts: [], mode: "shell" }])
  166. expect(ui.commits).toEqual([])
  167. })
  168. test("shell mode does not append a synthetic user row", async () => {
  169. const ui = footer()
  170. const task = runPromptQueue({
  171. footer: ui.api,
  172. run: async () => {
  173. expect(ui.commits).toEqual([])
  174. ui.api.close()
  175. },
  176. })
  177. ui.submit("ls", "shell")
  178. await task
  179. })
  180. test("preserves whitespace for initial input", async () => {
  181. const ui = footer()
  182. const seen: string[] = []
  183. await runPromptQueue({
  184. footer: ui.api,
  185. initialInput: " hello ",
  186. run: async (input) => {
  187. seen.push(input.text)
  188. ui.api.close()
  189. },
  190. })
  191. expect(seen).toEqual([" hello "])
  192. expect(ui.commits).toEqual([
  193. {
  194. kind: "user",
  195. text: " hello ",
  196. phase: "start",
  197. source: "system",
  198. messageID: expect.any(String),
  199. },
  200. ])
  201. })
  202. test("passes prompts to onSend", async () => {
  203. const ui = footer()
  204. const seen: string[] = []
  205. await runPromptQueue({
  206. footer: ui.api,
  207. initialInput: " hello ",
  208. onSend: (input) => {
  209. seen.push(input.text)
  210. },
  211. run: async () => {
  212. ui.api.close()
  213. },
  214. })
  215. expect(seen).toEqual([" hello "])
  216. })
  217. test("appends the user row before the turn starts", async () => {
  218. const ui = footer()
  219. await runPromptQueue({
  220. footer: ui.api,
  221. initialInput: "/fmt bash",
  222. run: async () => {
  223. expect(ui.commits).toEqual([
  224. {
  225. kind: "user",
  226. text: "/fmt bash",
  227. phase: "start",
  228. source: "system",
  229. messageID: expect.any(String),
  230. },
  231. ])
  232. ui.api.close()
  233. },
  234. })
  235. })
  236. test("runs queued prompts in order", async () => {
  237. const ui = footer()
  238. const seen: string[] = []
  239. let wake: (() => void) | undefined
  240. const gate = new Promise<void>((resolve) => {
  241. wake = resolve
  242. })
  243. const task = runPromptQueue({
  244. footer: ui.api,
  245. run: async (input) => {
  246. seen.push(input.text)
  247. if (seen.length === 1) {
  248. await gate
  249. return
  250. }
  251. ui.api.close()
  252. },
  253. })
  254. ui.submit("one")
  255. ui.submit("two")
  256. await Promise.resolve()
  257. expect(seen).toEqual(["one"])
  258. wake?.()
  259. await task
  260. expect(seen).toEqual(["one", "two"])
  261. })
  262. test("exposes ordinary in-flight prompts for removal before sending", async () => {
  263. const ui = footer()
  264. const turns: RunPrompt[] = []
  265. let wake: (() => void) | undefined
  266. const gate = new Promise<void>((resolve) => {
  267. wake = resolve
  268. })
  269. const task = runPromptQueue({
  270. footer: ui.api,
  271. run: async (input) => {
  272. turns.push(input)
  273. await gate
  274. },
  275. })
  276. ui.submit("one")
  277. ui.submit("two")
  278. await Promise.resolve()
  279. await Promise.resolve()
  280. expect(turns.map((item) => item.text)).toEqual(["one"])
  281. expect(turns[0]?.messageID).toEqual(expect.any(String))
  282. expect(ui.commits.map((item) => item.text)).toEqual(["one"])
  283. const first = ui.events.find((item) => item.type === "queued.prompts")
  284. const event = ui.events.findLast((item) => item.type === "queued.prompts")
  285. expect(first?.type === "queued.prompts" ? first.prompts : []).toEqual([])
  286. expect(
  287. first?.type === "queued.prompts" && event?.type === "queued.prompts" ? first.prompts === event.prompts : true,
  288. ).toBe(false)
  289. expect(ui.events.findLast((item) => item.type === "queue")).toEqual({ type: "queue", queue: 1 })
  290. expect(event?.type === "queued.prompts" ? event.prompts.map((item) => item.prompt.text) : []).toEqual(["two"])
  291. if (event?.type === "queued.prompts") ui.removeQueued(event.prompts[0]!.messageID)
  292. await Promise.resolve()
  293. wake?.()
  294. ui.api.close()
  295. await task
  296. expect(turns.map((item) => item.text)).toEqual(["one"])
  297. })
  298. test("removing one managed queued prompt preserves the others", async () => {
  299. const ui = footer()
  300. const turns: string[] = []
  301. let wake: (() => void) | undefined
  302. const gate = new Promise<void>((resolve) => {
  303. wake = resolve
  304. })
  305. const task = runPromptQueue({
  306. footer: ui.api,
  307. run: async (input) => {
  308. turns.push(input.text)
  309. if (input.text === "active") await gate
  310. if (input.text === "queued three") ui.api.close()
  311. },
  312. })
  313. ui.submit("active")
  314. ui.submit("queued one")
  315. ui.submit("queued two")
  316. ui.submit("queued three")
  317. await Promise.resolve()
  318. await Promise.resolve()
  319. const event = ui.events.findLast((item) => item.type === "queued.prompts")
  320. if (event?.type === "queued.prompts") {
  321. const second = event.prompts.find((item) => item.prompt.text === "queued two")
  322. if (second) ui.removeQueued(second.messageID)
  323. }
  324. wake?.()
  325. await task
  326. expect(turns).toEqual(["active", "queued one", "queued three"])
  327. })
  328. test("drains a prompt queued during an in-flight turn", async () => {
  329. const ui = footer()
  330. const seen: string[] = []
  331. let wake: (() => void) | undefined
  332. const gate = new Promise<void>((resolve) => {
  333. wake = resolve
  334. })
  335. const task = runPromptQueue({
  336. footer: ui.api,
  337. run: async (input) => {
  338. seen.push(input.text)
  339. if (seen.length === 1) {
  340. await gate
  341. return
  342. }
  343. ui.api.close()
  344. },
  345. })
  346. ui.submit("one")
  347. await Promise.resolve()
  348. expect(seen).toEqual(["one"])
  349. wake?.()
  350. await Promise.resolve()
  351. ui.submit("two")
  352. await task
  353. expect(seen).toEqual(["one", "two"])
  354. })
  355. test("close aborts the active run and drops pending queued work", async () => {
  356. const ui = footer()
  357. const seen: string[] = []
  358. let hit = false
  359. const task = runPromptQueue({
  360. footer: ui.api,
  361. run: async (input, signal) => {
  362. seen.push(input.text)
  363. await new Promise<void>((resolve) => {
  364. if (signal.aborted) {
  365. hit = true
  366. resolve()
  367. return
  368. }
  369. signal.addEventListener(
  370. "abort",
  371. () => {
  372. hit = true
  373. resolve()
  374. },
  375. { once: true },
  376. )
  377. })
  378. },
  379. })
  380. ui.submit("one")
  381. await Promise.resolve()
  382. ui.submit("two")
  383. ui.api.close()
  384. await task
  385. expect(hit).toBe(true)
  386. expect(seen).toEqual(["one"])
  387. })
  388. test("propagates run errors", async () => {
  389. const ui = footer()
  390. const task = runPromptQueue({
  391. footer: ui.api,
  392. run: async () => {
  393. throw new Error("boom")
  394. },
  395. })
  396. ui.submit("one")
  397. await expect(task).rejects.toThrow("boom")
  398. })
  399. })