service-usage.test.ts 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264
  1. import { describe, expect, test } from "bun:test"
  2. import { makeACPFixture, makeSession, secondModel, type FixtureContext, type FixtureRequest } from "./service-fixture"
  3. describe("acp service prompt routing and usage", () => {
  4. test("routes slash commands, skills, and compact through their session endpoints", async () => {
  5. await using fixture = makeACPFixture({
  6. fetch(request, context) {
  7. if (request.method === "POST" && request.path === "/api/session") {
  8. return Response.json({ data: makeSession("ses_routes") })
  9. }
  10. if (request.method === "GET" && request.path === "/api/session/ses_routes") {
  11. return Response.json({ data: makeSession("ses_routes") })
  12. }
  13. if (request.method === "POST" && request.path === "/api/session/ses_routes/command") {
  14. const id = requestID(request)
  15. completeTurn(context, "ses_routes", {
  16. id: `evt_${id}`,
  17. type: "session.input.promoted",
  18. data: { sessionID: "ses_routes", inputID: id },
  19. })
  20. return Response.json({ data: {} })
  21. }
  22. if (request.method === "POST" && request.path === "/api/session/ses_routes/skill") {
  23. const id = requestID(request)
  24. completeTurn(context, "ses_routes", {
  25. id: id.replace(/^msg_/, "evt_"),
  26. type: "session.skill.activated",
  27. data: { sessionID: "ses_routes", skill: "verify" },
  28. })
  29. return new Response(null, { status: 204 })
  30. }
  31. if (request.method === "POST" && request.path === "/api/session/ses_routes/compact") {
  32. const id = requestID(request)
  33. completeTurn(context, "ses_routes", {
  34. id: `evt_${id}`,
  35. type: "session.compaction.admitted",
  36. data: { sessionID: "ses_routes", inputID: id },
  37. })
  38. return Response.json({ data: {} })
  39. }
  40. return undefined
  41. },
  42. })
  43. const session = await fixture.service.newSession({ cwd: "/workspace", mcpServers: [] })
  44. const commandResult = await fixture.service.prompt({
  45. sessionId: session.sessionId,
  46. messageId: "client-command",
  47. prompt: [{ type: "text", text: "/review now" }],
  48. })
  49. const skillResult = await fixture.service.prompt({
  50. sessionId: session.sessionId,
  51. messageId: "client-skill",
  52. prompt: [{ type: "text", text: "/verify" }],
  53. })
  54. const compactResult = await fixture.service.prompt({
  55. sessionId: session.sessionId,
  56. messageId: "client-compact",
  57. prompt: [{ type: "text", text: "/compact" }],
  58. })
  59. expect([commandResult.stopReason, skillResult.stopReason, compactResult.stopReason]).toEqual([
  60. "end_turn",
  61. "end_turn",
  62. "end_turn",
  63. ])
  64. const command = fixture.requests.find((request) => request.path === "/api/session/ses_routes/command")
  65. const skill = fixture.requests.find((request) => request.path === "/api/session/ses_routes/skill")
  66. const compact = fixture.requests.find((request) => request.path === "/api/session/ses_routes/compact")
  67. expect(command?.body).toMatchObject({
  68. id: expect.any(String),
  69. command: "review",
  70. arguments: "now",
  71. files: [],
  72. delivery: "steer",
  73. })
  74. expect(skill?.body).toMatchObject({ id: expect.any(String), skill: "verify" })
  75. expect(compact?.body).toMatchObject({ id: expect.any(String) })
  76. expect(fixture.requests.some((request) => request.path === "/api/session/ses_routes/prompt")).toBe(false)
  77. })
  78. test("returns turn usage and publishes current context usage with cumulative session cost", async () => {
  79. const assistantTokens = {
  80. input: 100,
  81. output: 40,
  82. reasoning: 7,
  83. cache: { read: 11, write: 13 },
  84. }
  85. await using fixture = makeACPFixture({
  86. fetch(request, context) {
  87. if (request.method === "POST" && request.path === "/api/session") {
  88. return Response.json({ data: makeSession("ses_usage") })
  89. }
  90. if (request.method === "POST" && request.path === "/api/session/ses_usage/model") {
  91. return new Response(null, { status: 204 })
  92. }
  93. if (request.method === "POST" && request.path === "/api/session/ses_usage/prompt") {
  94. const id = requestID(request)
  95. context.send({
  96. id: `evt_${id}`,
  97. type: "session.input.promoted",
  98. data: { sessionID: "ses_usage", inputID: id },
  99. })
  100. context.send({
  101. id: "evt_step",
  102. type: "session.step.ended",
  103. data: {
  104. sessionID: "ses_usage",
  105. assistantMessageID: "msg_assistant",
  106. finish: "stop",
  107. cost: 0.5,
  108. tokens: assistantTokens,
  109. },
  110. })
  111. context.send({
  112. id: "evt_done",
  113. type: "session.execution.succeeded",
  114. data: { sessionID: "ses_usage" },
  115. })
  116. return Response.json({ data: {} })
  117. }
  118. if (request.method === "GET" && request.path === "/api/session/ses_usage/message/msg_assistant") {
  119. return Response.json({
  120. data: {
  121. id: "msg_assistant",
  122. type: "assistant",
  123. agent: "build",
  124. model: { providerID: "test", id: secondModel.id },
  125. content: [{ type: "text", text: "done" }],
  126. finish: "stop",
  127. tokens: assistantTokens,
  128. time: { created: 1, completed: 2 },
  129. },
  130. })
  131. }
  132. if (request.method === "GET" && request.path === "/api/session/ses_usage") {
  133. return Response.json({
  134. data: makeSession("ses_usage", {
  135. model: { providerID: "test", id: secondModel.id },
  136. cost: 3.5,
  137. tokens: { input: 120, output: 50, reasoning: 8, cache: { read: 30, write: 4 } },
  138. }),
  139. })
  140. }
  141. return undefined
  142. },
  143. })
  144. const session = await fixture.service.newSession({ cwd: "/workspace", mcpServers: [] })
  145. await fixture.service.setSessionConfigOption({
  146. sessionId: session.sessionId,
  147. configId: "model",
  148. value: "test/second-model",
  149. })
  150. const response = await fixture.service.prompt({
  151. sessionId: session.sessionId,
  152. messageId: "client-message",
  153. prompt: [{ type: "text", text: "hello" }],
  154. })
  155. expect(response).toEqual({
  156. stopReason: "end_turn",
  157. userMessageId: "client-message",
  158. usage: {
  159. inputTokens: 100,
  160. outputTokens: 40,
  161. thoughtTokens: 7,
  162. cachedReadTokens: 11,
  163. cachedWriteTokens: 13,
  164. totalTokens: 171,
  165. },
  166. _meta: {},
  167. })
  168. expect(fixture.updates.filter((item) => item.update.sessionUpdate === "usage_update")).toEqual([
  169. {
  170. sessionId: "ses_usage",
  171. update: {
  172. sessionUpdate: "usage_update",
  173. used: 171,
  174. size: 200_000,
  175. cost: { amount: 3.5, currency: "USD" },
  176. },
  177. },
  178. ])
  179. })
  180. test("does not fail a completed prompt when the usage refresh fails", async () => {
  181. await using fixture = makeACPFixture({
  182. fetch(request, context) {
  183. if (request.method === "POST" && request.path === "/api/session") {
  184. return Response.json({ data: makeSession("ses_usage_failure") })
  185. }
  186. if (request.method === "POST" && request.path === "/api/session/ses_usage_failure/prompt") {
  187. const id = requestID(request)
  188. context.send({
  189. id: `evt_${id}`,
  190. type: "session.input.promoted",
  191. data: { sessionID: "ses_usage_failure", inputID: id },
  192. })
  193. context.send({
  194. id: "evt_step_failure",
  195. type: "session.step.ended",
  196. data: {
  197. sessionID: "ses_usage_failure",
  198. assistantMessageID: "msg_usage_failure",
  199. finish: "stop",
  200. cost: 0,
  201. tokens: { input: 3, output: 2, reasoning: 0, cache: { read: 0, write: 0 } },
  202. },
  203. })
  204. context.send({
  205. id: "evt_done_failure",
  206. type: "session.execution.succeeded",
  207. data: { sessionID: "ses_usage_failure" },
  208. })
  209. return Response.json({ data: {} })
  210. }
  211. if (request.method === "GET" && request.path === "/api/session/ses_usage_failure/message/msg_usage_failure") {
  212. return Response.json({
  213. data: {
  214. id: "msg_usage_failure",
  215. type: "assistant",
  216. agent: "build",
  217. model: { providerID: "test", id: "test-model" },
  218. content: [],
  219. finish: "stop",
  220. tokens: { input: 3, output: 2, reasoning: 0, cache: { read: 0, write: 0 } },
  221. time: { created: 1, completed: 2 },
  222. },
  223. })
  224. }
  225. if (request.method === "GET" && request.path === "/api/session/ses_usage_failure") {
  226. return new Response(null, { status: 500 })
  227. }
  228. return undefined
  229. },
  230. })
  231. const session = await fixture.service.newSession({ cwd: "/workspace", mcpServers: [] })
  232. const response = await fixture.service.prompt({
  233. sessionId: session.sessionId,
  234. prompt: [{ type: "text", text: "hello" }],
  235. })
  236. expect(response.stopReason).toBe("end_turn")
  237. expect(fixture.updates.some((item) => item.update.sessionUpdate === "usage_update")).toBe(false)
  238. })
  239. })
  240. function requestID(request: FixtureRequest) {
  241. if (!request.body || typeof request.body !== "object") throw new Error(`missing body for ${request.path}`)
  242. const id = Reflect.get(request.body, "id")
  243. if (typeof id !== "string") throw new Error(`missing prompt id for ${request.path}`)
  244. return id
  245. }
  246. function completeTurn(context: FixtureContext, sessionID: string, start: unknown) {
  247. context.send(start)
  248. context.send({
  249. id: `evt_done_${sessionID}`,
  250. type: "session.execution.succeeded",
  251. data: { sessionID },
  252. })
  253. }