prompt-effect.test.ts 38 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202
  1. import { NodeFileSystem } from "@effect/platform-node"
  2. import { expect, spyOn } from "bun:test"
  3. import { Cause, Effect, Exit, Fiber, Layer } from "effect"
  4. import path from "path"
  5. import z from "zod"
  6. import type { Agent } from "../../src/agent/agent"
  7. import { Agent as AgentSvc } from "../../src/agent/agent"
  8. import { Bus } from "../../src/bus"
  9. import { Command } from "../../src/command"
  10. import { Config } from "../../src/config/config"
  11. import { FileTime } from "../../src/file/time"
  12. import { LSP } from "../../src/lsp"
  13. import { MCP } from "../../src/mcp"
  14. import { Permission } from "../../src/permission"
  15. import { Plugin } from "../../src/plugin"
  16. import { Provider as ProviderSvc } from "../../src/provider/provider"
  17. import type { Provider } from "../../src/provider/provider"
  18. import { ModelID, ProviderID } from "../../src/provider/schema"
  19. import { Session } from "../../src/session"
  20. import { LLM } from "../../src/session/llm"
  21. import { MessageV2 } from "../../src/session/message-v2"
  22. import { AppFileSystem } from "../../src/filesystem"
  23. import { SessionCompaction } from "../../src/session/compaction"
  24. import { SessionProcessor } from "../../src/session/processor"
  25. import { SessionPrompt } from "../../src/session/prompt"
  26. import { MessageID, PartID, SessionID } from "../../src/session/schema"
  27. import { SessionStatus } from "../../src/session/status"
  28. import { Shell } from "../../src/shell/shell"
  29. import { Snapshot } from "../../src/snapshot"
  30. import { TaskTool } from "../../src/tool/task"
  31. import { ToolRegistry } from "../../src/tool/registry"
  32. import { Truncate } from "../../src/tool/truncate"
  33. import { Log } from "../../src/util/log"
  34. import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
  35. import { provideTmpdirInstance, provideTmpdirServer } from "../fixture/fixture"
  36. import { testEffect } from "../lib/effect"
  37. import { TestLLMServer } from "../lib/llm-server"
  38. Log.init({ print: false })
  39. const ref = {
  40. providerID: ProviderID.make("test"),
  41. modelID: ModelID.make("test-model"),
  42. }
  43. function defer<T>() {
  44. let resolve!: (value: T | PromiseLike<T>) => void
  45. const promise = new Promise<T>((done) => {
  46. resolve = done
  47. })
  48. return { promise, resolve }
  49. }
  50. function withSh<A, E, R>(fx: () => Effect.Effect<A, E, R>) {
  51. return Effect.acquireUseRelease(
  52. Effect.sync(() => {
  53. const prev = process.env.SHELL
  54. process.env.SHELL = "/bin/sh"
  55. Shell.preferred.reset()
  56. return prev
  57. }),
  58. () => fx(),
  59. (prev) =>
  60. Effect.sync(() => {
  61. if (prev === undefined) delete process.env.SHELL
  62. else process.env.SHELL = prev
  63. Shell.preferred.reset()
  64. }),
  65. )
  66. }
  67. function toolPart(parts: MessageV2.Part[]) {
  68. return parts.find((part): part is MessageV2.ToolPart => part.type === "tool")
  69. }
  70. type CompletedToolPart = MessageV2.ToolPart & { state: MessageV2.ToolStateCompleted }
  71. type ErrorToolPart = MessageV2.ToolPart & { state: MessageV2.ToolStateError }
  72. function completedTool(parts: MessageV2.Part[]) {
  73. const part = toolPart(parts)
  74. expect(part?.state.status).toBe("completed")
  75. return part?.state.status === "completed" ? (part as CompletedToolPart) : undefined
  76. }
  77. function errorTool(parts: MessageV2.Part[]) {
  78. const part = toolPart(parts)
  79. expect(part?.state.status).toBe("error")
  80. return part?.state.status === "error" ? (part as ErrorToolPart) : undefined
  81. }
  82. const mcp = Layer.succeed(
  83. MCP.Service,
  84. MCP.Service.of({
  85. status: () => Effect.succeed({}),
  86. clients: () => Effect.succeed({}),
  87. tools: () => Effect.succeed({}),
  88. prompts: () => Effect.succeed({}),
  89. resources: () => Effect.succeed({}),
  90. add: () => Effect.succeed({ status: { status: "disabled" as const } }),
  91. connect: () => Effect.void,
  92. disconnect: () => Effect.void,
  93. getPrompt: () => Effect.succeed(undefined),
  94. readResource: () => Effect.succeed(undefined),
  95. startAuth: () => Effect.die("unexpected MCP auth in prompt-effect tests"),
  96. authenticate: () => Effect.die("unexpected MCP auth in prompt-effect tests"),
  97. finishAuth: () => Effect.die("unexpected MCP auth in prompt-effect tests"),
  98. removeAuth: () => Effect.void,
  99. supportsOAuth: () => Effect.succeed(false),
  100. hasStoredTokens: () => Effect.succeed(false),
  101. getAuthStatus: () => Effect.succeed("not_authenticated" as const),
  102. }),
  103. )
  104. const lsp = Layer.succeed(
  105. LSP.Service,
  106. LSP.Service.of({
  107. init: () => Effect.void,
  108. status: () => Effect.succeed([]),
  109. hasClients: () => Effect.succeed(false),
  110. touchFile: () => Effect.void,
  111. diagnostics: () => Effect.succeed({}),
  112. hover: () => Effect.succeed(undefined),
  113. definition: () => Effect.succeed([]),
  114. references: () => Effect.succeed([]),
  115. implementation: () => Effect.succeed([]),
  116. documentSymbol: () => Effect.succeed([]),
  117. workspaceSymbol: () => Effect.succeed([]),
  118. prepareCallHierarchy: () => Effect.succeed([]),
  119. incomingCalls: () => Effect.succeed([]),
  120. outgoingCalls: () => Effect.succeed([]),
  121. }),
  122. )
  123. const filetime = Layer.succeed(
  124. FileTime.Service,
  125. FileTime.Service.of({
  126. read: () => Effect.void,
  127. get: () => Effect.succeed(undefined),
  128. assert: () => Effect.void,
  129. withLock: (_filepath, fn) => Effect.promise(fn),
  130. }),
  131. )
  132. const status = SessionStatus.layer.pipe(Layer.provideMerge(Bus.layer))
  133. const infra = Layer.mergeAll(NodeFileSystem.layer, CrossSpawnSpawner.defaultLayer)
  134. function makeHttp() {
  135. const deps = Layer.mergeAll(
  136. Session.defaultLayer,
  137. Snapshot.defaultLayer,
  138. LLM.defaultLayer,
  139. AgentSvc.defaultLayer,
  140. Command.defaultLayer,
  141. Permission.layer,
  142. Plugin.defaultLayer,
  143. Config.defaultLayer,
  144. ProviderSvc.defaultLayer,
  145. filetime,
  146. lsp,
  147. mcp,
  148. AppFileSystem.defaultLayer,
  149. status,
  150. ).pipe(Layer.provideMerge(infra))
  151. const registry = ToolRegistry.layer.pipe(Layer.provideMerge(deps))
  152. const trunc = Truncate.layer.pipe(Layer.provideMerge(deps))
  153. const proc = SessionProcessor.layer.pipe(Layer.provideMerge(deps))
  154. const compact = SessionCompaction.layer.pipe(Layer.provideMerge(proc), Layer.provideMerge(deps))
  155. return Layer.mergeAll(
  156. TestLLMServer.layer,
  157. SessionPrompt.layer.pipe(
  158. Layer.provideMerge(compact),
  159. Layer.provideMerge(proc),
  160. Layer.provideMerge(registry),
  161. Layer.provideMerge(trunc),
  162. Layer.provideMerge(deps),
  163. ),
  164. )
  165. }
  166. const it = testEffect(makeHttp())
  167. const unix = process.platform !== "win32" ? it.live : it.live.skip
  168. // Config that registers a custom "test" provider with a "test-model" model
  169. // so Provider.getModel("test", "test-model") succeeds inside the loop.
  170. const cfg = {
  171. provider: {
  172. test: {
  173. name: "Test",
  174. id: "test",
  175. env: [],
  176. npm: "@ai-sdk/openai-compatible",
  177. models: {
  178. "test-model": {
  179. id: "test-model",
  180. name: "Test Model",
  181. attachment: false,
  182. reasoning: false,
  183. temperature: false,
  184. tool_call: true,
  185. release_date: "2025-01-01",
  186. limit: { context: 100000, output: 10000 },
  187. cost: { input: 0, output: 0 },
  188. options: {},
  189. },
  190. },
  191. options: {
  192. apiKey: "test-key",
  193. baseURL: "http://localhost:1/v1",
  194. },
  195. },
  196. },
  197. }
  198. function providerCfg(url: string) {
  199. return {
  200. ...cfg,
  201. provider: {
  202. ...cfg.provider,
  203. test: {
  204. ...cfg.provider.test,
  205. options: {
  206. ...cfg.provider.test.options,
  207. baseURL: url,
  208. },
  209. },
  210. },
  211. }
  212. }
  213. const user = Effect.fn("test.user")(function* (sessionID: SessionID, text: string) {
  214. const session = yield* Session.Service
  215. const msg = yield* session.updateMessage({
  216. id: MessageID.ascending(),
  217. role: "user",
  218. sessionID,
  219. agent: "build",
  220. model: ref,
  221. time: { created: Date.now() },
  222. })
  223. yield* session.updatePart({
  224. id: PartID.ascending(),
  225. messageID: msg.id,
  226. sessionID,
  227. type: "text",
  228. text,
  229. })
  230. return msg
  231. })
  232. const seed = Effect.fn("test.seed")(function* (sessionID: SessionID, opts?: { finish?: string }) {
  233. const session = yield* Session.Service
  234. const msg = yield* user(sessionID, "hello")
  235. const assistant: MessageV2.Assistant = {
  236. id: MessageID.ascending(),
  237. role: "assistant",
  238. parentID: msg.id,
  239. sessionID,
  240. mode: "build",
  241. agent: "build",
  242. cost: 0,
  243. path: { cwd: "/tmp", root: "/tmp" },
  244. tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
  245. modelID: ref.modelID,
  246. providerID: ref.providerID,
  247. time: { created: Date.now() },
  248. ...(opts?.finish ? { finish: opts.finish } : {}),
  249. }
  250. yield* session.updateMessage(assistant)
  251. yield* session.updatePart({
  252. id: PartID.ascending(),
  253. messageID: assistant.id,
  254. sessionID,
  255. type: "text",
  256. text: "hi there",
  257. })
  258. return { user: msg, assistant }
  259. })
  260. const addSubtask = (sessionID: SessionID, messageID: MessageID, model = ref) =>
  261. Effect.gen(function* () {
  262. const session = yield* Session.Service
  263. yield* session.updatePart({
  264. id: PartID.ascending(),
  265. messageID,
  266. sessionID,
  267. type: "subtask",
  268. prompt: "look into the cache key path",
  269. description: "inspect bug",
  270. agent: "general",
  271. model,
  272. })
  273. })
  274. const boot = Effect.fn("test.boot")(function* (input?: { title?: string }) {
  275. const prompt = yield* SessionPrompt.Service
  276. const sessions = yield* Session.Service
  277. const chat = yield* sessions.create(input ?? { title: "Pinned" })
  278. return { prompt, sessions, chat }
  279. })
  280. // Loop semantics
  281. it.live("loop exits immediately when last assistant has stop finish", () =>
  282. provideTmpdirServer(
  283. Effect.fnUntraced(function* ({ llm }) {
  284. const prompt = yield* SessionPrompt.Service
  285. const sessions = yield* Session.Service
  286. const chat = yield* sessions.create({ title: "Pinned" })
  287. yield* seed(chat.id, { finish: "stop" })
  288. const result = yield* prompt.loop({ sessionID: chat.id })
  289. expect(result.info.role).toBe("assistant")
  290. if (result.info.role === "assistant") expect(result.info.finish).toBe("stop")
  291. expect(yield* llm.calls).toBe(0)
  292. }),
  293. { git: true, config: providerCfg },
  294. ),
  295. )
  296. it.live("loop calls LLM and returns assistant message", () =>
  297. provideTmpdirServer(
  298. Effect.fnUntraced(function* ({ llm }) {
  299. const prompt = yield* SessionPrompt.Service
  300. const sessions = yield* Session.Service
  301. const chat = yield* sessions.create({
  302. title: "Pinned",
  303. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  304. })
  305. yield* prompt.prompt({
  306. sessionID: chat.id,
  307. agent: "build",
  308. noReply: true,
  309. parts: [{ type: "text", text: "hello" }],
  310. })
  311. yield* llm.text("world")
  312. const result = yield* prompt.loop({ sessionID: chat.id })
  313. expect(result.info.role).toBe("assistant")
  314. const parts = result.parts.filter((p) => p.type === "text")
  315. expect(parts.some((p) => p.type === "text" && p.text === "world")).toBe(true)
  316. expect(yield* llm.hits).toHaveLength(1)
  317. }),
  318. { git: true, config: providerCfg },
  319. ),
  320. )
  321. it.live("static loop returns assistant text through local provider", () =>
  322. provideTmpdirServer(
  323. Effect.fnUntraced(function* ({ llm }) {
  324. const session = yield* Effect.promise(() =>
  325. Session.create({
  326. title: "Prompt provider",
  327. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  328. }),
  329. )
  330. yield* Effect.promise(() =>
  331. SessionPrompt.prompt({
  332. sessionID: session.id,
  333. agent: "build",
  334. noReply: true,
  335. parts: [{ type: "text", text: "hello" }],
  336. }),
  337. )
  338. yield* llm.text("world")
  339. const result = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
  340. expect(result.info.role).toBe("assistant")
  341. expect(result.parts.some((part) => part.type === "text" && part.text === "world")).toBe(true)
  342. expect(yield* llm.hits).toHaveLength(1)
  343. expect(yield* llm.pending).toBe(0)
  344. }),
  345. { git: true, config: providerCfg },
  346. ),
  347. )
  348. it.live("static loop consumes queued replies across turns", () =>
  349. provideTmpdirServer(
  350. Effect.fnUntraced(function* ({ llm }) {
  351. const session = yield* Effect.promise(() =>
  352. Session.create({
  353. title: "Prompt provider turns",
  354. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  355. }),
  356. )
  357. yield* Effect.promise(() =>
  358. SessionPrompt.prompt({
  359. sessionID: session.id,
  360. agent: "build",
  361. noReply: true,
  362. parts: [{ type: "text", text: "hello one" }],
  363. }),
  364. )
  365. yield* llm.text("world one")
  366. const first = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
  367. expect(first.info.role).toBe("assistant")
  368. expect(first.parts.some((part) => part.type === "text" && part.text === "world one")).toBe(true)
  369. yield* Effect.promise(() =>
  370. SessionPrompt.prompt({
  371. sessionID: session.id,
  372. agent: "build",
  373. noReply: true,
  374. parts: [{ type: "text", text: "hello two" }],
  375. }),
  376. )
  377. yield* llm.text("world two")
  378. const second = yield* Effect.promise(() => SessionPrompt.loop({ sessionID: session.id }))
  379. expect(second.info.role).toBe("assistant")
  380. expect(second.parts.some((part) => part.type === "text" && part.text === "world two")).toBe(true)
  381. expect(yield* llm.hits).toHaveLength(2)
  382. expect(yield* llm.pending).toBe(0)
  383. }),
  384. { git: true, config: providerCfg },
  385. ),
  386. )
  387. it.live("loop continues when finish is tool-calls", () =>
  388. provideTmpdirServer(
  389. Effect.fnUntraced(function* ({ llm }) {
  390. const prompt = yield* SessionPrompt.Service
  391. const sessions = yield* Session.Service
  392. const session = yield* sessions.create({
  393. title: "Pinned",
  394. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  395. })
  396. yield* prompt.prompt({
  397. sessionID: session.id,
  398. agent: "build",
  399. noReply: true,
  400. parts: [{ type: "text", text: "hello" }],
  401. })
  402. yield* llm.tool("first", { value: "first" })
  403. yield* llm.text("second")
  404. const result = yield* prompt.loop({ sessionID: session.id })
  405. expect(yield* llm.calls).toBe(2)
  406. expect(result.info.role).toBe("assistant")
  407. if (result.info.role === "assistant") {
  408. expect(result.parts.some((part) => part.type === "text" && part.text === "second")).toBe(true)
  409. expect(result.info.finish).toBe("stop")
  410. }
  411. }),
  412. { git: true, config: providerCfg },
  413. ),
  414. )
  415. it.live("failed subtask preserves metadata on error tool state", () =>
  416. provideTmpdirServer(
  417. Effect.fnUntraced(function* ({ llm }) {
  418. const prompt = yield* SessionPrompt.Service
  419. const sessions = yield* Session.Service
  420. const chat = yield* sessions.create({ title: "Pinned" })
  421. yield* llm.tool("task", {
  422. description: "inspect bug",
  423. prompt: "look into the cache key path",
  424. subagent_type: "general",
  425. })
  426. yield* llm.text("done")
  427. const msg = yield* user(chat.id, "hello")
  428. yield* addSubtask(chat.id, msg.id)
  429. const result = yield* prompt.loop({ sessionID: chat.id })
  430. expect(result.info.role).toBe("assistant")
  431. expect(yield* llm.calls).toBe(2)
  432. const msgs = yield* MessageV2.filterCompactedEffect(chat.id)
  433. const taskMsg = msgs.find((item) => item.info.role === "assistant" && item.info.agent === "general")
  434. expect(taskMsg?.info.role).toBe("assistant")
  435. if (!taskMsg || taskMsg.info.role !== "assistant") return
  436. const tool = errorTool(taskMsg.parts)
  437. if (!tool) return
  438. expect(tool.state.error).toContain("Tool execution failed")
  439. expect(tool.state.metadata).toBeDefined()
  440. expect(tool.state.metadata?.sessionId).toBeDefined()
  441. expect(tool.state.metadata?.model).toEqual({
  442. providerID: ProviderID.make("test"),
  443. modelID: ModelID.make("missing-model"),
  444. })
  445. }),
  446. {
  447. git: true,
  448. config: (url) => ({
  449. ...providerCfg(url),
  450. agent: {
  451. general: {
  452. model: "test/missing-model",
  453. },
  454. },
  455. }),
  456. },
  457. ),
  458. )
  459. it.live(
  460. "loop sets status to busy then idle",
  461. () =>
  462. provideTmpdirServer(
  463. Effect.fnUntraced(function* ({ llm }) {
  464. const prompt = yield* SessionPrompt.Service
  465. const sessions = yield* Session.Service
  466. const status = yield* SessionStatus.Service
  467. yield* llm.hang
  468. const chat = yield* sessions.create({})
  469. yield* user(chat.id, "hi")
  470. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  471. yield* llm.wait(1)
  472. expect((yield* status.get(chat.id)).type).toBe("busy")
  473. yield* prompt.cancel(chat.id)
  474. yield* Fiber.await(fiber)
  475. expect((yield* status.get(chat.id)).type).toBe("idle")
  476. }),
  477. { git: true, config: providerCfg },
  478. ),
  479. 3_000,
  480. )
  481. // Cancel semantics
  482. it.live(
  483. "cancel interrupts loop and resolves with an assistant message",
  484. () =>
  485. provideTmpdirServer(
  486. Effect.fnUntraced(function* ({ llm }) {
  487. const prompt = yield* SessionPrompt.Service
  488. const sessions = yield* Session.Service
  489. const chat = yield* sessions.create({ title: "Pinned" })
  490. yield* seed(chat.id)
  491. yield* llm.hang
  492. yield* user(chat.id, "more")
  493. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  494. yield* llm.wait(1)
  495. yield* prompt.cancel(chat.id)
  496. const exit = yield* Fiber.await(fiber)
  497. expect(Exit.isSuccess(exit)).toBe(true)
  498. if (Exit.isSuccess(exit)) {
  499. expect(exit.value.info.role).toBe("assistant")
  500. }
  501. }),
  502. { git: true, config: providerCfg },
  503. ),
  504. 3_000,
  505. )
  506. it.live(
  507. "cancel records MessageAbortedError on interrupted process",
  508. () =>
  509. provideTmpdirServer(
  510. Effect.fnUntraced(function* ({ llm }) {
  511. const prompt = yield* SessionPrompt.Service
  512. const sessions = yield* Session.Service
  513. const chat = yield* sessions.create({ title: "Pinned" })
  514. yield* llm.hang
  515. yield* user(chat.id, "hello")
  516. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  517. yield* llm.wait(1)
  518. yield* prompt.cancel(chat.id)
  519. const exit = yield* Fiber.await(fiber)
  520. expect(Exit.isSuccess(exit)).toBe(true)
  521. if (Exit.isSuccess(exit)) {
  522. const info = exit.value.info
  523. if (info.role === "assistant") {
  524. expect(info.error?.name).toBe("MessageAbortedError")
  525. }
  526. }
  527. }),
  528. { git: true, config: providerCfg },
  529. ),
  530. 3_000,
  531. )
  532. it.live(
  533. "cancel finalizes subtask tool state",
  534. () =>
  535. provideTmpdirInstance(
  536. (dir) =>
  537. Effect.gen(function* () {
  538. const ready = defer<void>()
  539. const aborted = defer<void>()
  540. const init = spyOn(TaskTool, "init").mockImplementation(async () => ({
  541. description: "task",
  542. parameters: z.object({
  543. description: z.string(),
  544. prompt: z.string(),
  545. subagent_type: z.string(),
  546. task_id: z.string().optional(),
  547. command: z.string().optional(),
  548. }),
  549. execute: async (_args, ctx) => {
  550. ready.resolve()
  551. ctx.abort.addEventListener("abort", () => aborted.resolve(), { once: true })
  552. await new Promise<void>(() => {})
  553. return {
  554. title: "",
  555. metadata: {
  556. sessionId: SessionID.make("task"),
  557. model: ref,
  558. },
  559. output: "",
  560. }
  561. },
  562. }))
  563. yield* Effect.addFinalizer(() => Effect.sync(() => init.mockRestore()))
  564. const { prompt, chat } = yield* boot()
  565. const msg = yield* user(chat.id, "hello")
  566. yield* addSubtask(chat.id, msg.id)
  567. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  568. yield* Effect.promise(() => ready.promise)
  569. yield* prompt.cancel(chat.id)
  570. yield* Effect.promise(() => aborted.promise)
  571. const exit = yield* Fiber.await(fiber)
  572. expect(Exit.isSuccess(exit)).toBe(true)
  573. const msgs = yield* MessageV2.filterCompactedEffect(chat.id)
  574. const taskMsg = msgs.find((item) => item.info.role === "assistant" && item.info.agent === "general")
  575. expect(taskMsg?.info.role).toBe("assistant")
  576. if (!taskMsg || taskMsg.info.role !== "assistant") return
  577. const tool = toolPart(taskMsg.parts)
  578. expect(tool?.type).toBe("tool")
  579. if (!tool) return
  580. expect(tool.state.status).not.toBe("running")
  581. expect(taskMsg.info.time.completed).toBeDefined()
  582. expect(taskMsg.info.finish).toBeDefined()
  583. }),
  584. { git: true, config: cfg },
  585. ),
  586. 30_000,
  587. )
  588. it.live(
  589. "cancel with queued callers resolves all cleanly",
  590. () =>
  591. provideTmpdirServer(
  592. Effect.fnUntraced(function* ({ llm }) {
  593. const prompt = yield* SessionPrompt.Service
  594. const sessions = yield* Session.Service
  595. const chat = yield* sessions.create({ title: "Pinned" })
  596. yield* llm.hang
  597. yield* user(chat.id, "hello")
  598. const a = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  599. yield* llm.wait(1)
  600. const b = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  601. yield* Effect.sleep(50)
  602. yield* prompt.cancel(chat.id)
  603. const [exitA, exitB] = yield* Effect.all([Fiber.await(a), Fiber.await(b)])
  604. expect(Exit.isSuccess(exitA)).toBe(true)
  605. expect(Exit.isSuccess(exitB)).toBe(true)
  606. if (Exit.isSuccess(exitA) && Exit.isSuccess(exitB)) {
  607. expect(exitA.value.info.id).toBe(exitB.value.info.id)
  608. }
  609. }),
  610. { git: true, config: providerCfg },
  611. ),
  612. 3_000,
  613. )
  614. // Queue semantics
  615. it.live("concurrent loop callers get same result", () =>
  616. provideTmpdirInstance(
  617. (dir) =>
  618. Effect.gen(function* () {
  619. const { prompt, chat } = yield* boot()
  620. yield* seed(chat.id, { finish: "stop" })
  621. const [a, b] = yield* Effect.all([prompt.loop({ sessionID: chat.id }), prompt.loop({ sessionID: chat.id })], {
  622. concurrency: "unbounded",
  623. })
  624. expect(a.info.id).toBe(b.info.id)
  625. expect(a.info.role).toBe("assistant")
  626. yield* prompt.assertNotBusy(chat.id)
  627. }),
  628. { git: true },
  629. ),
  630. )
  631. it.live(
  632. "concurrent loop callers all receive same error result",
  633. () =>
  634. provideTmpdirServer(
  635. Effect.fnUntraced(function* ({ llm }) {
  636. const prompt = yield* SessionPrompt.Service
  637. const sessions = yield* Session.Service
  638. const chat = yield* sessions.create({ title: "Pinned" })
  639. yield* llm.fail("boom")
  640. yield* user(chat.id, "hello")
  641. const [a, b] = yield* Effect.all([prompt.loop({ sessionID: chat.id }), prompt.loop({ sessionID: chat.id })], {
  642. concurrency: "unbounded",
  643. })
  644. expect(a.info.id).toBe(b.info.id)
  645. expect(a.info.role).toBe("assistant")
  646. }),
  647. { git: true, config: providerCfg },
  648. ),
  649. 3_000,
  650. )
  651. it.live(
  652. "prompt submitted during an active run is included in the next LLM input",
  653. () =>
  654. provideTmpdirServer(
  655. Effect.fnUntraced(function* ({ llm }) {
  656. const gate = defer<void>()
  657. const prompt = yield* SessionPrompt.Service
  658. const sessions = yield* Session.Service
  659. const chat = yield* sessions.create({ title: "Pinned" })
  660. yield* llm.hold("first", gate.promise)
  661. yield* llm.text("second")
  662. const a = yield* prompt
  663. .prompt({
  664. sessionID: chat.id,
  665. agent: "build",
  666. model: ref,
  667. parts: [{ type: "text", text: "first" }],
  668. })
  669. .pipe(Effect.forkChild)
  670. yield* llm.wait(1)
  671. const id = MessageID.ascending()
  672. const b = yield* prompt
  673. .prompt({
  674. sessionID: chat.id,
  675. messageID: id,
  676. agent: "build",
  677. model: ref,
  678. parts: [{ type: "text", text: "second" }],
  679. })
  680. .pipe(Effect.forkChild)
  681. yield* Effect.promise(async () => {
  682. const end = Date.now() + 5000
  683. while (Date.now() < end) {
  684. const msgs = await Effect.runPromise(sessions.messages({ sessionID: chat.id }))
  685. if (msgs.some((msg) => msg.info.role === "user" && msg.info.id === id)) return
  686. await new Promise((done) => setTimeout(done, 20))
  687. }
  688. throw new Error("timed out waiting for second prompt to save")
  689. })
  690. gate.resolve()
  691. const [ea, eb] = yield* Effect.all([Fiber.await(a), Fiber.await(b)])
  692. expect(Exit.isSuccess(ea)).toBe(true)
  693. expect(Exit.isSuccess(eb)).toBe(true)
  694. expect(yield* llm.calls).toBe(2)
  695. const msgs = yield* sessions.messages({ sessionID: chat.id })
  696. const assistants = msgs.filter((msg) => msg.info.role === "assistant")
  697. expect(assistants).toHaveLength(2)
  698. const last = assistants.at(-1)
  699. if (!last || last.info.role !== "assistant") throw new Error("expected second assistant")
  700. expect(last.info.parentID).toBe(id)
  701. expect(last.parts.some((part) => part.type === "text" && part.text === "second")).toBe(true)
  702. const inputs = yield* llm.inputs
  703. expect(inputs).toHaveLength(2)
  704. expect(JSON.stringify(inputs.at(-1)?.messages)).toContain("second")
  705. }),
  706. { git: true, config: providerCfg },
  707. ),
  708. 3_000,
  709. )
  710. it.live(
  711. "assertNotBusy throws BusyError when loop running",
  712. () =>
  713. provideTmpdirServer(
  714. Effect.fnUntraced(function* ({ llm }) {
  715. const prompt = yield* SessionPrompt.Service
  716. const sessions = yield* Session.Service
  717. yield* llm.hang
  718. const chat = yield* sessions.create({})
  719. yield* user(chat.id, "hi")
  720. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  721. yield* llm.wait(1)
  722. const exit = yield* prompt.assertNotBusy(chat.id).pipe(Effect.exit)
  723. expect(Exit.isFailure(exit)).toBe(true)
  724. if (Exit.isFailure(exit)) {
  725. expect(Cause.squash(exit.cause)).toBeInstanceOf(Session.BusyError)
  726. }
  727. yield* prompt.cancel(chat.id)
  728. yield* Fiber.await(fiber)
  729. }),
  730. { git: true, config: providerCfg },
  731. ),
  732. 3_000,
  733. )
  734. it.live("assertNotBusy succeeds when idle", () =>
  735. provideTmpdirInstance(
  736. (dir) =>
  737. Effect.gen(function* () {
  738. const prompt = yield* SessionPrompt.Service
  739. const sessions = yield* Session.Service
  740. const chat = yield* sessions.create({})
  741. const exit = yield* prompt.assertNotBusy(chat.id).pipe(Effect.exit)
  742. expect(Exit.isSuccess(exit)).toBe(true)
  743. }),
  744. { git: true },
  745. ),
  746. )
  747. // Shell semantics
  748. it.live(
  749. "shell rejects with BusyError when loop running",
  750. () =>
  751. provideTmpdirServer(
  752. Effect.fnUntraced(function* ({ llm }) {
  753. const prompt = yield* SessionPrompt.Service
  754. const sessions = yield* Session.Service
  755. const chat = yield* sessions.create({ title: "Pinned" })
  756. yield* llm.hang
  757. yield* user(chat.id, "hi")
  758. const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  759. yield* llm.wait(1)
  760. const exit = yield* prompt.shell({ sessionID: chat.id, agent: "build", command: "echo hi" }).pipe(Effect.exit)
  761. expect(Exit.isFailure(exit)).toBe(true)
  762. if (Exit.isFailure(exit)) {
  763. expect(Cause.squash(exit.cause)).toBeInstanceOf(Session.BusyError)
  764. }
  765. yield* prompt.cancel(chat.id)
  766. yield* Fiber.await(fiber)
  767. }),
  768. { git: true, config: providerCfg },
  769. ),
  770. 3_000,
  771. )
  772. unix("shell captures stdout and stderr in completed tool output", () =>
  773. provideTmpdirInstance(
  774. (dir) =>
  775. Effect.gen(function* () {
  776. const { prompt, chat } = yield* boot()
  777. const result = yield* prompt.shell({
  778. sessionID: chat.id,
  779. agent: "build",
  780. command: "printf out && printf err >&2",
  781. })
  782. expect(result.info.role).toBe("assistant")
  783. const tool = completedTool(result.parts)
  784. if (!tool) return
  785. expect(tool.state.output).toContain("out")
  786. expect(tool.state.output).toContain("err")
  787. expect(tool.state.metadata.output).toContain("out")
  788. expect(tool.state.metadata.output).toContain("err")
  789. yield* prompt.assertNotBusy(chat.id)
  790. }),
  791. { git: true, config: cfg },
  792. ),
  793. )
  794. unix("shell completes a fast command on the preferred shell", () =>
  795. provideTmpdirInstance(
  796. (dir) =>
  797. Effect.gen(function* () {
  798. const { prompt, chat } = yield* boot()
  799. const result = yield* prompt.shell({
  800. sessionID: chat.id,
  801. agent: "build",
  802. command: "pwd",
  803. })
  804. expect(result.info.role).toBe("assistant")
  805. const tool = completedTool(result.parts)
  806. if (!tool) return
  807. expect(tool.state.input.command).toBe("pwd")
  808. expect(tool.state.output).toContain(dir)
  809. expect(tool.state.metadata.output).toContain(dir)
  810. yield* prompt.assertNotBusy(chat.id)
  811. }),
  812. { git: true, config: cfg },
  813. ),
  814. )
  815. unix("shell lists files from the project directory", () =>
  816. provideTmpdirInstance(
  817. (dir) =>
  818. Effect.gen(function* () {
  819. const { prompt, chat } = yield* boot()
  820. yield* Effect.promise(() => Bun.write(path.join(dir, "README.md"), "# e2e\n"))
  821. const result = yield* prompt.shell({
  822. sessionID: chat.id,
  823. agent: "build",
  824. command: "command ls",
  825. })
  826. expect(result.info.role).toBe("assistant")
  827. const tool = completedTool(result.parts)
  828. if (!tool) return
  829. expect(tool.state.input.command).toBe("command ls")
  830. expect(tool.state.output).toContain("README.md")
  831. expect(tool.state.metadata.output).toContain("README.md")
  832. yield* prompt.assertNotBusy(chat.id)
  833. }),
  834. { git: true, config: cfg },
  835. ),
  836. )
  837. unix("shell captures stderr from a failing command", () =>
  838. provideTmpdirInstance(
  839. (dir) =>
  840. Effect.gen(function* () {
  841. const { prompt, chat } = yield* boot()
  842. const result = yield* prompt.shell({
  843. sessionID: chat.id,
  844. agent: "build",
  845. command: "command -v __nonexistent_cmd_e2e__ || echo 'not found' >&2; exit 1",
  846. })
  847. expect(result.info.role).toBe("assistant")
  848. const tool = completedTool(result.parts)
  849. if (!tool) return
  850. expect(tool.state.output).toContain("not found")
  851. expect(tool.state.metadata.output).toContain("not found")
  852. yield* prompt.assertNotBusy(chat.id)
  853. }),
  854. { git: true, config: cfg },
  855. ),
  856. )
  857. unix(
  858. "shell updates running metadata before process exit",
  859. () =>
  860. withSh(() =>
  861. provideTmpdirInstance(
  862. (dir) =>
  863. Effect.gen(function* () {
  864. const { prompt, chat } = yield* boot()
  865. const fiber = yield* prompt
  866. .shell({ sessionID: chat.id, agent: "build", command: "printf first && sleep 0.2 && printf second" })
  867. .pipe(Effect.forkChild)
  868. yield* Effect.promise(async () => {
  869. const start = Date.now()
  870. while (Date.now() - start < 5000) {
  871. const msgs = await MessageV2.filterCompacted(MessageV2.stream(chat.id))
  872. const taskMsg = msgs.find((item) => item.info.role === "assistant")
  873. const tool = taskMsg ? toolPart(taskMsg.parts) : undefined
  874. if (tool?.state.status === "running" && tool.state.metadata?.output.includes("first")) return
  875. await new Promise((done) => setTimeout(done, 20))
  876. }
  877. throw new Error("timed out waiting for running shell metadata")
  878. })
  879. const exit = yield* Fiber.await(fiber)
  880. expect(Exit.isSuccess(exit)).toBe(true)
  881. }),
  882. { git: true, config: cfg },
  883. ),
  884. ),
  885. 30_000,
  886. )
  887. it.live(
  888. "loop waits while shell runs and starts after shell exits",
  889. () =>
  890. provideTmpdirServer(
  891. Effect.fnUntraced(function* ({ llm }) {
  892. const prompt = yield* SessionPrompt.Service
  893. const sessions = yield* Session.Service
  894. const chat = yield* sessions.create({
  895. title: "Pinned",
  896. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  897. })
  898. yield* llm.text("after-shell")
  899. const sh = yield* prompt
  900. .shell({ sessionID: chat.id, agent: "build", command: "sleep 0.2" })
  901. .pipe(Effect.forkChild)
  902. yield* Effect.sleep(50)
  903. const loop = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  904. yield* Effect.sleep(50)
  905. expect(yield* llm.calls).toBe(0)
  906. yield* Fiber.await(sh)
  907. const exit = yield* Fiber.await(loop)
  908. expect(Exit.isSuccess(exit)).toBe(true)
  909. if (Exit.isSuccess(exit)) {
  910. expect(exit.value.info.role).toBe("assistant")
  911. expect(exit.value.parts.some((part) => part.type === "text" && part.text === "after-shell")).toBe(true)
  912. }
  913. expect(yield* llm.calls).toBe(1)
  914. }),
  915. { git: true, config: providerCfg },
  916. ),
  917. 3_000,
  918. )
  919. it.live(
  920. "shell completion resumes queued loop callers",
  921. () =>
  922. provideTmpdirServer(
  923. Effect.fnUntraced(function* ({ llm }) {
  924. const prompt = yield* SessionPrompt.Service
  925. const sessions = yield* Session.Service
  926. const chat = yield* sessions.create({
  927. title: "Pinned",
  928. permission: [{ permission: "*", pattern: "*", action: "allow" }],
  929. })
  930. yield* llm.text("done")
  931. const sh = yield* prompt
  932. .shell({ sessionID: chat.id, agent: "build", command: "sleep 0.2" })
  933. .pipe(Effect.forkChild)
  934. yield* Effect.sleep(50)
  935. const a = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  936. const b = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  937. yield* Effect.sleep(50)
  938. expect(yield* llm.calls).toBe(0)
  939. yield* Fiber.await(sh)
  940. const [ea, eb] = yield* Effect.all([Fiber.await(a), Fiber.await(b)])
  941. expect(Exit.isSuccess(ea)).toBe(true)
  942. expect(Exit.isSuccess(eb)).toBe(true)
  943. if (Exit.isSuccess(ea) && Exit.isSuccess(eb)) {
  944. expect(ea.value.info.id).toBe(eb.value.info.id)
  945. expect(ea.value.info.role).toBe("assistant")
  946. }
  947. expect(yield* llm.calls).toBe(1)
  948. }),
  949. { git: true, config: providerCfg },
  950. ),
  951. 3_000,
  952. )
  953. unix(
  954. "cancel interrupts shell and resolves cleanly",
  955. () =>
  956. withSh(() =>
  957. provideTmpdirInstance(
  958. (dir) =>
  959. Effect.gen(function* () {
  960. const { prompt, chat } = yield* boot()
  961. const sh = yield* prompt
  962. .shell({ sessionID: chat.id, agent: "build", command: "sleep 30" })
  963. .pipe(Effect.forkChild)
  964. yield* Effect.sleep(50)
  965. yield* prompt.cancel(chat.id)
  966. const status = yield* SessionStatus.Service
  967. expect((yield* status.get(chat.id)).type).toBe("idle")
  968. const busy = yield* prompt.assertNotBusy(chat.id).pipe(Effect.exit)
  969. expect(Exit.isSuccess(busy)).toBe(true)
  970. const exit = yield* Fiber.await(sh)
  971. expect(Exit.isSuccess(exit)).toBe(true)
  972. if (Exit.isSuccess(exit)) {
  973. expect(exit.value.info.role).toBe("assistant")
  974. const tool = completedTool(exit.value.parts)
  975. if (tool) {
  976. expect(tool.state.output).toContain("User aborted the command")
  977. }
  978. }
  979. }),
  980. { git: true, config: cfg },
  981. ),
  982. ),
  983. 30_000,
  984. )
  985. unix(
  986. "cancel persists aborted shell result when shell ignores TERM",
  987. () =>
  988. withSh(() =>
  989. provideTmpdirInstance(
  990. (dir) =>
  991. Effect.gen(function* () {
  992. const { prompt, chat } = yield* boot()
  993. const sh = yield* prompt
  994. .shell({ sessionID: chat.id, agent: "build", command: "trap '' TERM; sleep 30" })
  995. .pipe(Effect.forkChild)
  996. yield* Effect.sleep(50)
  997. yield* prompt.cancel(chat.id)
  998. const exit = yield* Fiber.await(sh)
  999. expect(Exit.isSuccess(exit)).toBe(true)
  1000. if (Exit.isSuccess(exit)) {
  1001. expect(exit.value.info.role).toBe("assistant")
  1002. const tool = completedTool(exit.value.parts)
  1003. if (tool) {
  1004. expect(tool.state.output).toContain("User aborted the command")
  1005. }
  1006. }
  1007. }),
  1008. { git: true, config: cfg },
  1009. ),
  1010. ),
  1011. 30_000,
  1012. )
  1013. unix(
  1014. "cancel interrupts loop queued behind shell",
  1015. () =>
  1016. provideTmpdirInstance(
  1017. (dir) =>
  1018. Effect.gen(function* () {
  1019. const { prompt, chat } = yield* boot()
  1020. const sh = yield* prompt
  1021. .shell({ sessionID: chat.id, agent: "build", command: "sleep 30" })
  1022. .pipe(Effect.forkChild)
  1023. yield* Effect.sleep(50)
  1024. const loop = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
  1025. yield* Effect.sleep(50)
  1026. yield* prompt.cancel(chat.id)
  1027. const exit = yield* Fiber.await(loop)
  1028. expect(Exit.isSuccess(exit)).toBe(true)
  1029. yield* Fiber.await(sh)
  1030. }),
  1031. { git: true, config: cfg },
  1032. ),
  1033. 30_000,
  1034. )
  1035. unix(
  1036. "shell rejects when another shell is already running",
  1037. () =>
  1038. withSh(() =>
  1039. provideTmpdirInstance(
  1040. (dir) =>
  1041. Effect.gen(function* () {
  1042. const { prompt, chat } = yield* boot()
  1043. const a = yield* prompt
  1044. .shell({ sessionID: chat.id, agent: "build", command: "sleep 30" })
  1045. .pipe(Effect.forkChild)
  1046. yield* Effect.sleep(50)
  1047. const exit = yield* prompt
  1048. .shell({ sessionID: chat.id, agent: "build", command: "echo hi" })
  1049. .pipe(Effect.exit)
  1050. expect(Exit.isFailure(exit)).toBe(true)
  1051. if (Exit.isFailure(exit)) {
  1052. expect(Cause.squash(exit.cause)).toBeInstanceOf(Session.BusyError)
  1053. }
  1054. yield* prompt.cancel(chat.id)
  1055. yield* Fiber.await(a)
  1056. }),
  1057. { git: true, config: cfg },
  1058. ),
  1059. ),
  1060. 30_000,
  1061. )