prompt-effect.test.ts 36 KB

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