experimental-a2 0.16.0 → 0.17.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (334) hide show
  1. package/CHANGELOG.md +360 -0
  2. package/dist/{actor-ohPC-81x.d.ts → actor-DBbe5xeE.d.ts} +21 -11
  3. package/dist/actor-DBbe5xeE.d.ts.map +1 -0
  4. package/dist/actor-client.d.ts +2 -5
  5. package/dist/actor-client.d.ts.map +1 -1
  6. package/dist/actor-client.js +13 -4
  7. package/dist/actor-client.js.map +1 -1
  8. package/dist/actor-react.d.ts +5 -3
  9. package/dist/actor-react.d.ts.map +1 -1
  10. package/dist/actor-react.js +20 -10
  11. package/dist/actor-react.js.map +1 -1
  12. package/dist/{actor-shared-B5tJfzt-.js → actor-shared-t7w8xIKP.js} +12 -7
  13. package/dist/actor-shared-t7w8xIKP.js.map +1 -0
  14. package/dist/actor.d.ts +1 -1
  15. package/dist/actor.js +42 -23
  16. package/dist/actor.js.map +1 -1
  17. package/dist/{ai-CrLYwNEx.js → ai-UrWfKeuv.js} +966 -541
  18. package/dist/ai-UrWfKeuv.js.map +1 -0
  19. package/dist/ai-client.d.ts +27 -0
  20. package/dist/ai-client.d.ts.map +1 -0
  21. package/dist/ai-client.js +29 -0
  22. package/dist/ai-client.js.map +1 -0
  23. package/dist/ai-react.d.ts +96 -0
  24. package/dist/ai-react.d.ts.map +1 -0
  25. package/dist/ai-react.js +142 -0
  26. package/dist/ai-react.js.map +1 -0
  27. package/dist/ai-server.browser.js +2 -2
  28. package/dist/ai-server.browser.js.map +1 -1
  29. package/dist/ai-server.d.ts +80 -36
  30. package/dist/ai-server.d.ts.map +1 -1
  31. package/dist/ai-server.js +466 -265
  32. package/dist/ai-server.js.map +1 -1
  33. package/dist/{ai-HJ9fHfYI.d.ts → ai-vpGtcaCB.d.ts} +115 -101
  34. package/dist/ai-vpGtcaCB.d.ts.map +1 -0
  35. package/dist/ai.d.ts +2 -2
  36. package/dist/ai.js +2 -2
  37. package/dist/{client-BYzHjkwU.js → client-COaiaDub.js} +447 -80
  38. package/dist/client-COaiaDub.js.map +1 -0
  39. package/dist/{client-CzyacQpJ.d.ts → client-qRnRUzBS.d.ts} +14 -13
  40. package/dist/client-qRnRUzBS.d.ts.map +1 -0
  41. package/dist/client.d.ts +4 -2
  42. package/dist/client.js +1 -1
  43. package/dist/{contract-CKRg_E4q.js → contract-BYKASlCV.js} +12 -21
  44. package/dist/contract-BYKASlCV.js.map +1 -0
  45. package/dist/{idempotent-replay-DVOlyYbx.js → idempotent-replay-JdgMQRRm.js} +2 -2
  46. package/dist/{idempotent-replay-DVOlyYbx.js.map → idempotent-replay-JdgMQRRm.js.map} +1 -1
  47. package/dist/index.d.ts +6 -5
  48. package/dist/index.d.ts.map +1 -1
  49. package/dist/index.js +4 -3
  50. package/dist/{errors-DvhSXnxk.d.ts → initial-snapshot-BrLso5x_.d.ts} +7 -3
  51. package/dist/initial-snapshot-BrLso5x_.d.ts.map +1 -0
  52. package/dist/{internal-Dq2qYxou.js → internal-CbVvEwad.js} +23 -2
  53. package/dist/internal-CbVvEwad.js.map +1 -0
  54. package/dist/otel.d.ts +1 -1
  55. package/dist/react-BbJ-g6Uu.d.ts +86 -0
  56. package/dist/react-BbJ-g6Uu.d.ts.map +1 -0
  57. package/dist/react.d.ts +2 -122
  58. package/dist/react.js +107 -100
  59. package/dist/react.js.map +1 -1
  60. package/dist/reducer-BAXc8JyF.js +87 -0
  61. package/dist/reducer-BAXc8JyF.js.map +1 -0
  62. package/dist/{reducer-otzHjuJj.d.ts → reducer-BYu6DT7h.d.ts} +45 -29
  63. package/dist/reducer-BYu6DT7h.d.ts.map +1 -0
  64. package/dist/reducer-runtime-DQwz_J59.js +28 -0
  65. package/dist/reducer-runtime-DQwz_J59.js.map +1 -0
  66. package/dist/scheduler-qstash.d.ts +3 -3
  67. package/dist/scheduler-qstash.d.ts.map +1 -1
  68. package/dist/scheduler-qstash.js +33 -10
  69. package/dist/scheduler-qstash.js.map +1 -1
  70. package/dist/{scheduler-task-BpzhPnRS.js → scheduler-task-X1SToDjP.js} +18 -3
  71. package/dist/scheduler-task-X1SToDjP.js.map +1 -0
  72. package/dist/scheduler-vercel.d.ts +4 -9
  73. package/dist/scheduler-vercel.d.ts.map +1 -1
  74. package/dist/scheduler-vercel.js +25 -44
  75. package/dist/scheduler-vercel.js.map +1 -1
  76. package/dist/{server-DbD7IJVK.d.ts → server-CGM8WzRG.d.ts} +87 -26
  77. package/dist/server-CGM8WzRG.d.ts.map +1 -0
  78. package/dist/{server-BD5ckxZb.js → server-Rj64fNvX.js} +854 -202
  79. package/dist/server-Rj64fNvX.js.map +1 -0
  80. package/dist/server.d.ts +4 -4
  81. package/dist/server.js +2 -2
  82. package/dist/{errors-DCk6ch5n.js → session-clear-DGI4jTQ5.js} +17 -2
  83. package/dist/session-clear-DGI4jTQ5.js.map +1 -0
  84. package/dist/session-clear-Dshpuhff.d.ts +12 -0
  85. package/dist/session-clear-Dshpuhff.d.ts.map +1 -0
  86. package/dist/{store-DyZM6fS5.d.ts → store-Gpkq7WBo.d.ts} +44 -3
  87. package/dist/store-Gpkq7WBo.d.ts.map +1 -0
  88. package/dist/store-N8PXxDAS.js.map +1 -1
  89. package/dist/store-memory.d.ts +1 -1
  90. package/dist/store-memory.d.ts.map +1 -1
  91. package/dist/store-memory.js +182 -69
  92. package/dist/store-memory.js.map +1 -1
  93. package/dist/{store-polling-CmxUbV93.js → store-polling-4_y8lETm.js} +21 -4
  94. package/dist/store-polling-4_y8lETm.js.map +1 -0
  95. package/dist/store-postgres.d.ts +1 -1
  96. package/dist/store-postgres.d.ts.map +1 -1
  97. package/dist/store-postgres.js +399 -108
  98. package/dist/store-postgres.js.map +1 -1
  99. package/dist/{store-presence-polling-C7-XZyW9.js → store-presence-polling-Q1C-k2dd.js} +3 -3
  100. package/dist/{store-presence-polling-C7-XZyW9.js.map → store-presence-polling-Q1C-k2dd.js.map} +1 -1
  101. package/dist/store-redis-http.d.ts +1 -1
  102. package/dist/store-redis-http.d.ts.map +1 -1
  103. package/dist/store-redis-http.js +293 -11
  104. package/dist/store-redis-http.js.map +1 -1
  105. package/dist/store-redis-notify-tqWNnMRu.js +2374 -0
  106. package/dist/store-redis-notify-tqWNnMRu.js.map +1 -0
  107. package/dist/store-redis.d.ts +1 -1
  108. package/dist/store-redis.d.ts.map +1 -1
  109. package/dist/store-redis.js +11 -11
  110. package/dist/store-redis.js.map +1 -1
  111. package/dist/store-sqlite.d.ts +1 -1
  112. package/dist/store-sqlite.d.ts.map +1 -1
  113. package/dist/store-sqlite.js +166 -32
  114. package/dist/store-sqlite.js.map +1 -1
  115. package/dist/stream-buffer-Cx7AAPGq.js +23 -0
  116. package/dist/stream-buffer-Cx7AAPGq.js.map +1 -0
  117. package/dist/{telemetry-B5jzpy6y.d.ts → telemetry-cZH-cibj.d.ts} +2 -2
  118. package/dist/telemetry-cZH-cibj.d.ts.map +1 -0
  119. package/dist/{wire-BO5wWCb1.js → wire-DjfHOvdW.js} +132 -10
  120. package/dist/wire-DjfHOvdW.js.map +1 -0
  121. package/docs/01-quickstart.mdx +2 -2
  122. package/docs/actors/01-introduction.mdx +4 -0
  123. package/docs/actors/04-routes.mdx +14 -1
  124. package/docs/concepts/01-contracts.mdx +95 -0
  125. package/docs/concepts/04-state.mdx +3 -3
  126. package/docs/guides/01-timers.mdx +22 -14
  127. package/docs/guides/02-cancellation.mdx +1 -1
  128. package/docs/guides/03-react.mdx +268 -328
  129. package/docs/guides/04-local-first.mdx +32 -73
  130. package/docs/guides/05-production.mdx +40 -27
  131. package/docs/guides/06-ai-agents.mdx +680 -580
  132. package/docs/guides/09-presence.mdx +3 -7
  133. package/docs/guides/10-transports.mdx +5 -5
  134. package/docs/index.mdx +10 -18
  135. package/docs/reference/01-api.mdx +662 -195
  136. package/docs/reference/02-errors.mdx +2 -1
  137. package/examples/playground/DEPLOY.md +52 -6
  138. package/examples/playground/app/a0/[appId]/loading.tsx +16 -0
  139. package/examples/playground/app/a0/[appId]/page.tsx +48 -0
  140. package/examples/playground/app/a0/a0.css +818 -0
  141. package/examples/playground/app/a0/app-list.tsx +211 -0
  142. package/examples/playground/app/a0/app-view.tsx +55 -0
  143. package/examples/playground/app/a0/apps/[appId]/preview/route.ts +27 -0
  144. package/examples/playground/app/a0/apps/[appId]/route.ts +23 -0
  145. package/examples/playground/app/a0/apps.test.ts +116 -0
  146. package/examples/playground/app/a0/apps.ts +87 -0
  147. package/examples/playground/app/a0/chat.tsx +364 -0
  148. package/examples/playground/app/a0/columns.tsx +109 -0
  149. package/examples/playground/app/a0/context-menu.tsx +116 -0
  150. package/examples/playground/app/a0/cookies.ts +46 -0
  151. package/examples/playground/app/a0/crash-on-request.ts +72 -0
  152. package/examples/playground/app/a0/crashes/route.ts +3 -0
  153. package/examples/playground/app/a0/crashes.test.ts +33 -0
  154. package/examples/playground/app/a0/crashes.ts +28 -0
  155. package/examples/playground/app/a0/devtools.tsx +164 -0
  156. package/examples/playground/app/a0/events/route.ts +6 -0
  157. package/examples/playground/app/a0/flags.ts +30 -0
  158. package/examples/playground/app/a0/layout.tsx +38 -0
  159. package/examples/playground/app/a0/list/route.ts +10 -0
  160. package/examples/playground/app/a0/model.test.ts +56 -0
  161. package/examples/playground/app/a0/model.ts +58 -0
  162. package/examples/playground/app/a0/page.tsx +14 -0
  163. package/examples/playground/app/a0/preview-pane.tsx +145 -0
  164. package/examples/playground/app/a0/preview.test.ts +53 -0
  165. package/examples/playground/app/a0/preview.ts +26 -0
  166. package/examples/playground/app/a0/sandbox-log.ts +69 -0
  167. package/examples/playground/app/a0/sandbox.ts +191 -0
  168. package/examples/playground/app/a0/sandboxes/scheduler/route.ts +6 -0
  169. package/examples/playground/app/a0/scheduler/route.ts +7 -0
  170. package/examples/playground/app/a0/server.ts +158 -0
  171. package/examples/playground/app/a0/session.ts +47 -0
  172. package/examples/playground/app/a0/shell.tsx +150 -0
  173. package/examples/playground/app/a0/short/events/route.ts +9 -0
  174. package/examples/playground/app/a0/splitter.tsx +120 -0
  175. package/examples/playground/app/agent/[agentId]/agent-client.tsx +144 -174
  176. package/examples/playground/app/agent/[agentId]/command-receipt.test.tsx +65 -0
  177. package/examples/playground/app/agent/[agentId]/command-receipt.tsx +27 -0
  178. package/examples/playground/app/agent/[agentId]/compaction-event.tsx +25 -24
  179. package/examples/playground/app/agent/[agentId]/compaction-panel.test.tsx +119 -0
  180. package/examples/playground/app/agent/[agentId]/compaction-panel.tsx +208 -165
  181. package/examples/playground/app/agent/[agentId]/page.tsx +2 -16
  182. package/examples/playground/app/agent/events/route.ts +5 -2
  183. package/examples/playground/app/agent/ingress.test.ts +98 -0
  184. package/examples/playground/app/agent/model-input.test.tsx +58 -0
  185. package/examples/playground/app/agent/model-input.tsx +115 -0
  186. package/examples/playground/app/agent/model.test.ts +73 -0
  187. package/examples/playground/app/agent/model.ts +58 -26
  188. package/examples/playground/app/agent/new-agent-session.tsx +37 -19
  189. package/examples/playground/app/agent/page.tsx +4 -4
  190. package/examples/playground/app/agent/server.ts +15 -14
  191. package/examples/playground/app/agent/session.ts +5 -6
  192. package/examples/playground/app/agent/telemetry.test.ts +168 -0
  193. package/examples/playground/app/agent/telemetry.ts +30 -0
  194. package/examples/playground/app/agent/trace.test.ts +23 -0
  195. package/examples/playground/app/agent/trace.ts +18 -0
  196. package/examples/playground/app/canvas/[canvasId]/canvas-client.tsx +11 -11
  197. package/examples/playground/app/canvas/[canvasId]/page.tsx +1 -2
  198. package/examples/playground/app/canvas/[canvasId]/session.ts +2 -2
  199. package/examples/playground/app/canvas/model.ts +53 -53
  200. package/examples/playground/app/chat/[chatId]/agent-state.test.ts +176 -0
  201. package/examples/playground/app/chat/[chatId]/agent-stream-drawer.tsx +16 -1
  202. package/examples/playground/app/chat/[chatId]/chat-client.tsx +55 -77
  203. package/examples/playground/app/chat/[chatId]/chat-view.ts +37 -3
  204. package/examples/playground/app/chat/[chatId]/page.tsx +7 -15
  205. package/examples/playground/app/chat/[chatId]/session.ts +2 -19
  206. package/examples/playground/app/chat/agent/route.ts +9 -0
  207. package/examples/playground/app/chat/model.test.ts +46 -119
  208. package/examples/playground/app/chat/model.ts +82 -128
  209. package/examples/playground/app/chat/scheduler/route.ts +2 -2
  210. package/examples/playground/app/chat/server.test.ts +108 -0
  211. package/examples/playground/app/chat/server.ts +69 -73
  212. package/examples/playground/app/components/activity-feed.tsx +30 -1
  213. package/examples/playground/app/components/activity-history.test.tsx +35 -0
  214. package/examples/playground/app/components/demo-icons.tsx +89 -0
  215. package/examples/playground/app/components/demos.tsx +33 -0
  216. package/examples/playground/app/components/fast-link.test.tsx +58 -0
  217. package/examples/playground/app/components/fast-link.tsx +55 -0
  218. package/examples/playground/app/components/site-header.tsx +96 -0
  219. package/examples/playground/app/counter/counter-client.tsx +3 -3
  220. package/examples/playground/app/counter/model.test.ts +3 -3
  221. package/examples/playground/app/counter/model.ts +11 -21
  222. package/examples/playground/app/counter/page.tsx +2 -15
  223. package/examples/playground/app/counter/session.ts +2 -2
  224. package/examples/playground/app/documents/[documentId]/code-editor.tsx +12 -1
  225. package/examples/playground/app/documents/[documentId]/document-client.tsx +10 -11
  226. package/examples/playground/app/documents/[documentId]/page.tsx +1 -2
  227. package/examples/playground/app/documents/[documentId]/session.ts +2 -2
  228. package/examples/playground/app/documents/model.ts +27 -27
  229. package/examples/playground/app/globals.css +420 -207
  230. package/examples/playground/app/initial-state.test.tsx +54 -0
  231. package/examples/playground/app/layout.tsx +7 -23
  232. package/examples/playground/app/orders/[orderId]/page.tsx +1 -2
  233. package/examples/playground/app/orders/[orderId]/session.ts +2 -2
  234. package/examples/playground/app/orders/model.ts +49 -49
  235. package/examples/playground/app/page.tsx +13 -102
  236. package/examples/playground/app/recovery/[recoveryId]/page.tsx +1 -2
  237. package/examples/playground/app/recovery/[recoveryId]/session.ts +2 -2
  238. package/examples/playground/app/recovery/model.ts +30 -30
  239. package/examples/playground/instrumentation.ts +5 -0
  240. package/examples/playground/package.json +10 -2
  241. package/examples/playground/scripts/a0-template.ts +77 -0
  242. package/examples/playground/vercel.json +18 -0
  243. package/package.json +3 -1
  244. package/src/actor-client.ts +16 -4
  245. package/src/actor-react.ts +25 -11
  246. package/src/actor-shared.ts +21 -6
  247. package/src/actor.ts +76 -31
  248. package/src/ai-agent-actions.ts +219 -0
  249. package/src/ai-agent.ts +2343 -0
  250. package/src/ai-app-state.ts +91 -0
  251. package/src/ai-client-state.ts +39 -2
  252. package/src/ai-client.ts +62 -0
  253. package/src/ai-context-schema.ts +25 -5
  254. package/src/ai-context.ts +90 -40
  255. package/src/ai-control-server.ts +94 -55
  256. package/src/ai-control.ts +2 -0
  257. package/src/ai-coordinator.ts +23 -22
  258. package/src/ai-message-base.ts +209 -0
  259. package/src/ai-message-projection.ts +382 -58
  260. package/src/ai-progress-batches.ts +26 -0
  261. package/src/ai-projector.ts +95 -14
  262. package/src/ai-react.ts +146 -0
  263. package/src/ai-server.browser.ts +2 -2
  264. package/src/ai-server.ts +20 -1951
  265. package/src/ai-stored-state.ts +48 -163
  266. package/src/ai.ts +343 -276
  267. package/src/bootstrap-stream.ts +72 -0
  268. package/src/bounded-cache.ts +55 -0
  269. package/src/client-render-session.ts +49 -0
  270. package/src/client.ts +688 -151
  271. package/src/contract.ts +92 -42
  272. package/src/deterministic-id.ts +37 -0
  273. package/src/equal-json.ts +32 -0
  274. package/src/errors.ts +13 -1
  275. package/src/index.ts +5 -1
  276. package/src/initial-snapshot.ts +39 -0
  277. package/src/internal.ts +45 -0
  278. package/src/postgres-notification-scope.ts +14 -3
  279. package/src/react.ts +171 -206
  280. package/src/redis-http-reply.ts +256 -0
  281. package/src/redis-http-stream.ts +73 -0
  282. package/src/redis-http-subscriptions.ts +13 -5
  283. package/src/reducer-runtime.ts +52 -0
  284. package/src/reducer.ts +92 -10
  285. package/src/scheduler-qstash.ts +49 -20
  286. package/src/scheduler-task.ts +34 -2
  287. package/src/scheduler-vercel.ts +32 -84
  288. package/src/server-fetch.ts +139 -35
  289. package/src/server.ts +957 -237
  290. package/src/session-clear.ts +23 -0
  291. package/src/session-socket.ts +154 -49
  292. package/src/sse.ts +71 -4
  293. package/src/store-memory.ts +220 -78
  294. package/src/store-polling.ts +18 -4
  295. package/src/store-postgres.ts +461 -114
  296. package/src/store-redis-core.ts +1260 -434
  297. package/src/store-redis-http.ts +18 -2
  298. package/src/store-redis-notify.ts +260 -85
  299. package/src/store-redis.ts +19 -9
  300. package/src/store-sqlite.ts +259 -45
  301. package/src/store.ts +46 -1
  302. package/src/stream-buffer.ts +43 -0
  303. package/src/telemetry.ts +13 -1
  304. package/src/wire.ts +241 -13
  305. package/dist/actor-ohPC-81x.d.ts.map +0 -1
  306. package/dist/actor-shared-B5tJfzt-.js.map +0 -1
  307. package/dist/ai-CrLYwNEx.js.map +0 -1
  308. package/dist/ai-HJ9fHfYI.d.ts.map +0 -1
  309. package/dist/client-BYzHjkwU.js.map +0 -1
  310. package/dist/client-CzyacQpJ.d.ts.map +0 -1
  311. package/dist/contract-CKRg_E4q.js.map +0 -1
  312. package/dist/errors-DCk6ch5n.js.map +0 -1
  313. package/dist/errors-DvhSXnxk.d.ts.map +0 -1
  314. package/dist/internal-Dq2qYxou.js.map +0 -1
  315. package/dist/react.d.ts.map +0 -1
  316. package/dist/reducer-DEMjEY_O.js +0 -29
  317. package/dist/reducer-DEMjEY_O.js.map +0 -1
  318. package/dist/reducer-otzHjuJj.d.ts.map +0 -1
  319. package/dist/scheduler-task-BpzhPnRS.js.map +0 -1
  320. package/dist/server-BD5ckxZb.js.map +0 -1
  321. package/dist/server-DbD7IJVK.d.ts.map +0 -1
  322. package/dist/store-DyZM6fS5.d.ts.map +0 -1
  323. package/dist/store-polling-CmxUbV93.js.map +0 -1
  324. package/dist/store-redis-notify-BVLUfI3j.js +0 -1583
  325. package/dist/store-redis-notify-BVLUfI3j.js.map +0 -1
  326. package/dist/telemetry-B5jzpy6y.d.ts.map +0 -1
  327. package/dist/validate-XKT4FSNn.js +0 -30
  328. package/dist/validate-XKT4FSNn.js.map +0 -1
  329. package/dist/wire-BO5wWCb1.js.map +0 -1
  330. package/examples/playground/app/agent/[agentId]/compaction/route.ts +0 -14
  331. package/examples/playground/app/agent/compaction-settings.test.ts +0 -144
  332. package/examples/playground/app/agent/compaction-settings.ts +0 -49
  333. package/examples/playground/app/agent/compaction-timeline.test.ts +0 -377
  334. package/examples/playground/app/agent/compaction-timeline.ts +0 -208
@@ -0,0 +1,2343 @@
1
+ /**
2
+ * The server-only implementation of the experimental-a2/ai protocol: the
3
+ * handler table and `agent()`, its assembly as an ordinary A2 server.
4
+ */
5
+
6
+ // oxlint-disable no-await-in-loop -- stream chunks and durable appends are
7
+ // ordered; parallel consumption would corrupt progress sequence and state.
8
+
9
+ import { createControlRuntime, ControlCancelled } from './ai-control-server.ts'
10
+ import {
11
+ createContextReducer,
12
+ contextForAttempt,
13
+ generationBase,
14
+ type AIContextState,
15
+ type AIModelContext,
16
+ } from './ai-context.ts'
17
+ import { isAgentContract } from './ai-app-state.ts'
18
+ import { boundedCache } from './bounded-cache.ts'
19
+ import type { Reducer } from './reducer.ts'
20
+ import { aiEventId, AI_EVENT_ID_PREFIX } from './ai-id.ts'
21
+ import type { ControlCommit, ControlCommand } from './ai-control.ts'
22
+ import { AsyncLocalStorage } from 'node:async_hooks'
23
+ import { createHash } from 'node:crypto'
24
+ import { convertToModelMessages, isToolUIPart } from 'ai'
25
+ import type {
26
+ FinishReason,
27
+ Instructions,
28
+ LanguageModel,
29
+ LanguageModelUsage,
30
+ ModelMessage,
31
+ TextStreamPart,
32
+ ToolSet,
33
+ UIMessage,
34
+ UIMessageChunk,
35
+ } from 'ai'
36
+ import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
37
+ import { gatewayModelId, readModelLimits } from './ai-model-metadata.ts'
38
+ import { estimateInputTokens, measuredInputTokens } from './ai-input-tokens.ts'
39
+ import { AI_STATE_REDUCER, agentContract, events as aiEvents } from './ai.ts'
40
+ import type {
41
+ AIEventDefs,
42
+ AIMessageOf,
43
+ AIState,
44
+ AgentToolContext,
45
+ AgentToolEventDefs,
46
+ ApprovalRequestedPayload,
47
+ CompactionCompletedPayload,
48
+ CompactionRequestedPayload,
49
+ GenerationCompletedPayload,
50
+ GenerationFailedPayload,
51
+ GenerationProgressPayload,
52
+ GenerationStartedPayload,
53
+ MessageCompletedPayload,
54
+ ReasoningLevel,
55
+ ToolCalledPayload,
56
+ ToolResultPayload,
57
+ } from './ai.ts'
58
+ import type { StandardSchemaV1 } from './standard-schema.ts'
59
+ import type { AICoordinatorState } from './ai-coordinator.ts'
60
+ import type {
61
+ AppendInput,
62
+ Contract,
63
+ ContractEvent,
64
+ EventDefs,
65
+ } from './contract.ts'
66
+ import { A2Error } from './errors.ts'
67
+ import {
68
+ consumeSchedulerSendFailure,
69
+ installAmbientToolScopeStorage,
70
+ type AmbientToolScope,
71
+ } from './internal.ts'
72
+ import {
73
+ createServer,
74
+ type A2Server,
75
+ type HandlerContext,
76
+ type ServerOptions,
77
+ } from './server.ts'
78
+ import {
79
+ setServerFetchHooks,
80
+ type ServerIngressContext,
81
+ } from './server-fetch.ts'
82
+ import { CONSOLE_TELEMETRY, type A2Telemetry } from './telemetry.ts'
83
+
84
+ const isAgentIngressEvent = (event: {
85
+ type: string
86
+ payload: unknown
87
+ }): boolean => {
88
+ if (event.type !== 'ai.control.requested') return false
89
+ const command = event.payload
90
+ if (typeof command !== 'object' || command === null) return false
91
+ const action = (command as { action?: unknown }).action
92
+ if (action === 'request-input' || action === 'tool-result') return false
93
+ if (action === 'send' || action === 'edit' || action === 'steer')
94
+ return (
95
+ (command as { message?: { role?: unknown } }).message?.role === 'user'
96
+ )
97
+ return true
98
+ }
99
+
100
+ // With a shared contract, browsers push application events like any A2
101
+ // client, under the route's `authorize`. Events declared with `name` have no
102
+ // client that can push them, so they stay server-only.
103
+ const validateAgentIngress = ({
104
+ context,
105
+ applicationEvents,
106
+ }: {
107
+ context: ServerIngressContext
108
+ applicationEvents: boolean
109
+ }): void => {
110
+ const rejected = context.events.find((event) =>
111
+ Object.hasOwn(aiEvents, event.type)
112
+ ? !isAgentIngressEvent(event)
113
+ : !applicationEvents,
114
+ )
115
+ if (rejected !== undefined) {
116
+ throw new A2Error(
117
+ 'INVALID_PAYLOAD',
118
+ `event '${rejected.type}' is server-only for an AI agent`,
119
+ )
120
+ }
121
+ }
122
+
123
+ type AgentGenerationProgress = {
124
+ type: 'progress'
125
+ chunks: UIMessageChunk[]
126
+ }
127
+
128
+ type AgentGenerationFinish = {
129
+ type: 'finish'
130
+ finishReason?: FinishReason
131
+ usage?: LanguageModelUsage
132
+ }
133
+
134
+ type AgentGenerationUpdate = AgentGenerationProgress | AgentGenerationFinish
135
+
136
+ export type AgentGenerateContext<
137
+ M extends UIMessage,
138
+ D extends AIEventDefs<M> & EventDefs,
139
+ T extends ToolSet,
140
+ X extends object = Record<never, never>,
141
+ > = AgentResolverContext<M, D, X> & {
142
+ generationId: string
143
+ modelMessages: ModelMessage[]
144
+ model: LanguageModel
145
+ tools: T
146
+ instructions?: Instructions
147
+ generation: AgentGenerationSettings<T>
148
+ }
149
+
150
+ export type AgentGenerate<
151
+ M extends UIMessage,
152
+ D extends AIEventDefs<M> & EventDefs,
153
+ T extends ToolSet,
154
+ X extends object = Record<never, never>,
155
+ > = (
156
+ context: AgentGenerateContext<M, D, T, X>,
157
+ ) => ReadableStream<UIMessageChunk> | Promise<ReadableStream<UIMessageChunk>>
158
+
159
+ type UIMessageMetadata<M extends UIMessage> =
160
+ M extends UIMessage<infer Metadata> ? Metadata : unknown
161
+
162
+ export type AgentMessageMetadata<
163
+ M extends UIMessage,
164
+ D extends AIEventDefs<M> & EventDefs,
165
+ T extends ToolSet,
166
+ X extends object = Record<never, never>,
167
+ > = (
168
+ context: AgentGenerateContext<M, D, T, X> & { part: TextStreamPart<T> },
169
+ ) => UIMessageMetadata<M> | undefined
170
+
171
+ export type AgentResolverContext<
172
+ M extends UIMessage,
173
+ D extends AIEventDefs<M> & EventDefs,
174
+ X extends object = Record<never, never>,
175
+ > = {
176
+ event: ContractEvent<D, 'ai.generation.requested'>
177
+ messages: M[]
178
+ /**
179
+ * The contract's application state, as of the captured snapshot. It rides
180
+ * in the checkpoint the generation already reads, so it costs no I/O.
181
+ */
182
+ state: X
183
+ session: Pick<HandlerContext<D>['session'], 'state'>
184
+ signal: AbortSignal
185
+ }
186
+
187
+ export type Resolvable<T, C> = T | ((context: C) => T | Promise<T>)
188
+
189
+ /** A model with its reasoning level and speed, chosen together. */
190
+ export type ModelSettings = {
191
+ model: LanguageModel
192
+ /** Overrides `generation.reasoning`. */
193
+ reasoning?: ReasoningLevel
194
+ /** `fast` prefers AI Gateway's faster tier when the model offers one. */
195
+ speed?: 'fast' | 'standard'
196
+ }
197
+
198
+ export type CompactionPolicy<
199
+ M extends UIMessage,
200
+ D extends AIEventDefs<M> & EventDefs,
201
+ X extends object = Record<never, never>,
202
+ > = {
203
+ shouldCompact(
204
+ context: AgentResolverContext<M, D, X> & {
205
+ messages: M[]
206
+ modelMessages: ModelMessage[]
207
+ },
208
+ ): boolean | Promise<boolean>
209
+ compact(
210
+ context: AgentResolverContext<M, D, X> & {
211
+ messages: M[]
212
+ modelMessages: ModelMessage[]
213
+ },
214
+ ): M[] | Promise<M[]>
215
+ }
216
+
217
+ export type AutomaticCompaction = {
218
+ thresholdTokens?: number
219
+ instructions?: string
220
+ }
221
+
222
+ export type CompactionOptions<
223
+ M extends UIMessage,
224
+ D extends AIEventDefs<M> & EventDefs,
225
+ X extends object = Record<never, never>,
226
+ > = false | AutomaticCompaction | CompactionPolicy<M, D, X>
227
+
228
+ export type AgentGenerationSettings<T extends ToolSet> = AISDKStepSettings<T>
229
+
230
+ export type HandlersOptions<
231
+ M extends UIMessage,
232
+ D extends AIEventDefs<M> & EventDefs,
233
+ T extends ToolSet = ToolSet,
234
+ X extends object = Record<never, never>,
235
+ > = AgentModelOptions<M, D, T, X> & {
236
+ contract: Contract<D, Record<never, never>, unknown>
237
+ telemetry?: A2Telemetry
238
+ }
239
+
240
+ export type AgentModelOptions<
241
+ M extends UIMessage,
242
+ D extends AIEventDefs<M> & EventDefs,
243
+ T extends ToolSet = ToolSet,
244
+ X extends object = Record<never, never>,
245
+ > = {
246
+ /**
247
+ * AI SDK model string or provider model, or `ModelSettings` to choose its
248
+ * reasoning level and speed with it, optionally resolved per generation.
249
+ */
250
+ model: Resolvable<
251
+ LanguageModel | ModelSettings,
252
+ AgentResolverContext<M, D, X>
253
+ >
254
+ tools?: T
255
+ instructions?: Resolvable<
256
+ Instructions | undefined,
257
+ AgentResolverContext<M, D, X>
258
+ >
259
+ /** Settings for each one-step AI SDK streamText call, resolved per generation. */
260
+ generation?: Resolvable<
261
+ AgentGenerationSettings<T>,
262
+ AgentResolverContext<M, D, X>
263
+ >
264
+ /**
265
+ * Maximum model steps in one assistant response. Defaults to Infinity.
266
+ */
267
+ maxSteps?: number
268
+ /** Map default AI SDK stream parts to typed assistant message metadata. */
269
+ messageMetadata?: AgentMessageMetadata<M, D, T, X>
270
+ /**
271
+ * Replace model generation with an AI SDK UI message chunk stream.
272
+ * The custom stream owns its metadata chunks.
273
+ */
274
+ generate?: AgentGenerate<M, D, T, X>
275
+ compaction?: Resolvable<
276
+ CompactionOptions<M, D, X>,
277
+ AgentResolverContext<M, D, X>
278
+ >
279
+ /** Progress is durably flushed at either limit, whichever is reached first. */
280
+ progress?: { maxChunks?: number; maxDelayMs?: number }
281
+ /**
282
+ * Run local tool executions one at a time on this lane. By default,
283
+ * independent tools run concurrently.
284
+ */
285
+ toolLane?: string
286
+ }
287
+
288
+ type AgentServerOptions<
289
+ M extends UIMessage,
290
+ D extends AIEventDefs<M> & EventDefs,
291
+ S,
292
+ T extends ToolSet,
293
+ X extends object = Record<never, never>,
294
+ > = AgentModelOptions<M, D, T, X> &
295
+ Omit<
296
+ ServerOptions<D, Record<never, never>, Reducer<D, S>>,
297
+ 'contract' | 'handlers'
298
+ > & {
299
+ /** Handlers for application events. */
300
+ handlers?: ServerOptions<D, Record<never, never>, Reducer<D, S>>['handlers']
301
+ }
302
+
303
+ export type AgentOptions<
304
+ M extends UIMessage,
305
+ E extends EventDefs,
306
+ T extends ToolSet,
307
+ > = AgentServerOptions<
308
+ NoInfer<M>,
309
+ AIEventDefs<NoInfer<M>> & E,
310
+ AIState<NoInfer<M>>,
311
+ T
312
+ > & {
313
+ /** Identity: prefixes storage keys and queue messages. */
314
+ name: string
315
+ /**
316
+ * Server-only application events. Browsers push application events through
317
+ * a shared `agentContract()`.
318
+ */
319
+ events?: E
320
+ /** Optional stricter schema for custom UIMessage metadata/data/tool types. */
321
+ messageSchema?: StandardSchemaV1<M>
322
+ }
323
+
324
+ export type AgentFromContractOptions<
325
+ D extends AIEventDefs<AIMessageOf<D>> & EventDefs,
326
+ S,
327
+ T extends ToolSet,
328
+ Name extends string = string,
329
+ X extends object = Record<never, never>,
330
+ > = AgentServerOptions<AIMessageOf<D>, D, S, T, X> & {
331
+ /**
332
+ * A shared `agentContract()`. Its application state reaches every resolver
333
+ * as `state`; a fully custom contract is core A2, served by `createServer()`.
334
+ */
335
+ contract: Contract<
336
+ D,
337
+ Record<never, never>,
338
+ Reducer<D, S> & { readonly name: Name; readonly '~appState': X }
339
+ >
340
+ }
341
+
342
+ /**
343
+ * An A2 server for the AI protocol; `typeof` it types `useAgent`. `Name` is
344
+ * its reducer's name: the built-in one lets contract-free clients fold state.
345
+ */
346
+ export type Agent<
347
+ D extends EventDefs,
348
+ S,
349
+ Name extends string = string,
350
+ > = A2Server<D, Record<never, never>, Reducer<D, S> & { readonly name: Name }>
351
+
352
+ const checkAbort = (signal: AbortSignal): boolean => {
353
+ if (!signal.aborted) return false
354
+ const reason: unknown = signal.reason
355
+ if (
356
+ reason instanceof A2Error &&
357
+ (reason.code === 'CLAIM_EXPIRED' || reason.code === 'SUPERSEDED_ATTEMPT')
358
+ ) {
359
+ throw reason
360
+ }
361
+ return true
362
+ }
363
+
364
+ const resolve = async <T, C extends { signal: AbortSignal }>(
365
+ value: Resolvable<T, C>,
366
+ context: C,
367
+ ): Promise<T> => {
368
+ let result: T
369
+ try {
370
+ result =
371
+ typeof value === 'function'
372
+ ? await (value as (context: C) => T | Promise<T>)(context)
373
+ : value
374
+ } catch (error) {
375
+ checkAbort(context.signal)
376
+ throw error
377
+ }
378
+ checkAbort(context.signal)
379
+ return result
380
+ }
381
+
382
+ const withModelSettings = <T extends ToolSet>(options: {
383
+ generation: AgentGenerationSettings<T>
384
+ settings: ModelSettings
385
+ }): AgentGenerationSettings<T> => {
386
+ const { generation, settings } = options
387
+ if (settings.reasoning === undefined && settings.speed === undefined)
388
+ return generation
389
+ const { speed: _speed, ...gateway } =
390
+ generation.providerOptions?.['gateway'] ?? {}
391
+ return {
392
+ ...generation,
393
+ ...(settings.reasoning === undefined
394
+ ? {}
395
+ : { reasoning: settings.reasoning }),
396
+ ...(settings.speed === undefined
397
+ ? {}
398
+ : {
399
+ providerOptions: {
400
+ ...generation.providerOptions,
401
+ gateway:
402
+ settings.speed === 'fast'
403
+ ? { ...gateway, speed: 'fast' }
404
+ : gateway,
405
+ },
406
+ }),
407
+ }
408
+ }
409
+
410
+ const modelName = (model: LanguageModel): string => {
411
+ if (typeof model === 'string') return model
412
+ if ('modelId' in model && typeof model.modelId === 'string') {
413
+ return model.modelId
414
+ }
415
+ return 'custom'
416
+ }
417
+
418
+ const errorMessage = (error: unknown): string =>
419
+ error instanceof Error ? error.message : String(error)
420
+
421
+ const toolScopeStorage = installAmbientToolScopeStorage(
422
+ () => new AsyncLocalStorage<AmbientToolScope>(),
423
+ )
424
+
425
+ /** Run `fn` with `ctx` as the ambient tool handler context, for direct tool tests. */
426
+ export function runWithHandlerContext<D extends AgentToolEventDefs, R>(
427
+ ctx: AgentToolContext<D>,
428
+ fn: () => R,
429
+ ): R {
430
+ return toolScopeStorage.run({ context: ctx }, fn)
431
+ }
432
+
433
+ /** Sessions whose latest step prompt one process keeps for that step's tools. */
434
+ const PROMPT_SESSIONS = 64
435
+
436
+ const asyncIterator = (value: unknown): AsyncIterator<unknown> | undefined => {
437
+ if (typeof value !== 'object' || value === null) return undefined
438
+ const factory = Reflect.get(value, Symbol.asyncIterator)
439
+ if (typeof factory !== 'function') return undefined
440
+ return (factory as (this: object) => AsyncIterator<unknown>).call(value)
441
+ }
442
+
443
+ const isBoundaryChunk = (chunk: UIMessageChunk | undefined): boolean =>
444
+ chunk !== undefined &&
445
+ (chunk.type === 'text-end' ||
446
+ chunk.type === 'reasoning-end' ||
447
+ chunk.type === 'tool-input-available' ||
448
+ chunk.type === 'tool-input-error' ||
449
+ chunk.type === 'tool-output-available' ||
450
+ chunk.type === 'tool-output-error' ||
451
+ chunk.type === 'tool-output-denied' ||
452
+ chunk.type === 'tool-approval-request' ||
453
+ chunk.type === 'finish' ||
454
+ chunk.type === 'abort' ||
455
+ chunk.type === 'error')
456
+
457
+ type DeltaChunk = Extract<
458
+ UIMessageChunk,
459
+ { type: 'text-delta' | 'reasoning-delta' | 'tool-input-delta' }
460
+ >
461
+
462
+ const DELTA_KEYS: Record<DeltaChunk['type'], ReadonlySet<string>> = {
463
+ 'text-delta': new Set(['type', 'id', 'delta', 'providerMetadata']),
464
+ 'reasoning-delta': new Set(['type', 'id', 'delta', 'providerMetadata']),
465
+ 'tool-input-delta': new Set(['type', 'toolCallId', 'inputTextDelta']),
466
+ }
467
+
468
+ const isDeltaChunk = (chunk: UIMessageChunk): chunk is DeltaChunk =>
469
+ Object.hasOwn(DELTA_KEYS, chunk.type) &&
470
+ Object.keys(chunk).every((key) =>
471
+ DELTA_KEYS[chunk.type as DeltaChunk['type']].has(key),
472
+ ) &&
473
+ typeof (chunk.type === 'tool-input-delta'
474
+ ? chunk.inputTextDelta
475
+ : (chunk as { delta?: unknown }).delta) === 'string'
476
+
477
+ // Adjacent deltas of one part project exactly as their concatenation: the text
478
+ // appends and the last defined `providerMetadata` wins.
479
+ const joinDeltas = (
480
+ left: DeltaChunk,
481
+ right: DeltaChunk,
482
+ ): DeltaChunk | undefined => {
483
+ if (left.type === 'tool-input-delta')
484
+ return right.type === 'tool-input-delta' &&
485
+ right.toolCallId === left.toolCallId
486
+ ? { ...left, inputTextDelta: left.inputTextDelta + right.inputTextDelta }
487
+ : undefined
488
+ if (
489
+ right.type === 'tool-input-delta' ||
490
+ right.type !== left.type ||
491
+ right.id !== left.id
492
+ )
493
+ return undefined
494
+ const providerMetadata =
495
+ right.providerMetadata !== undefined
496
+ ? right.providerMetadata
497
+ : left.providerMetadata
498
+ return {
499
+ ...left,
500
+ delta: left.delta + right.delta,
501
+ ...(providerMetadata === undefined ? {} : { providerMetadata }),
502
+ }
503
+ }
504
+
505
+ const mergeDeltas = (chunks: UIMessageChunk[]): UIMessageChunk[] => {
506
+ const merged: UIMessageChunk[] = []
507
+ for (const chunk of chunks) {
508
+ const previous = merged.at(-1)
509
+ const joined =
510
+ previous !== undefined && isDeltaChunk(previous) && isDeltaChunk(chunk)
511
+ ? joinDeltas(previous, chunk)
512
+ : undefined
513
+ if (joined === undefined) merged.push(chunk)
514
+ else merged[merged.length - 1] = joined
515
+ }
516
+ return merged
517
+ }
518
+
519
+ type StreamRead<T> = Awaited<ReturnType<ReadableStreamDefaultReader<T>['read']>>
520
+
521
+ const nextOrTimer = async <T>(
522
+ next: Promise<StreamRead<T>>,
523
+ delayMs: number,
524
+ ): Promise<{ type: 'next'; result: StreamRead<T> } | { type: 'timer' }> => {
525
+ let timer: ReturnType<typeof setTimeout> | undefined
526
+ try {
527
+ return await Promise.race([
528
+ next.then((result) => ({ type: 'next' as const, result })),
529
+ new Promise<{ type: 'timer' }>((finishTimer) => {
530
+ timer = setTimeout(() => finishTimer({ type: 'timer' }), delayMs)
531
+ ;(timer as { unref?: () => void }).unref?.()
532
+ }),
533
+ ])
534
+ } finally {
535
+ if (timer !== undefined) clearTimeout(timer)
536
+ }
537
+ }
538
+
539
+ type GenerationCompletion = {
540
+ finishReason?: FinishReason
541
+ usage?: LanguageModelUsage
542
+ }
543
+
544
+ type GenerationSource = {
545
+ stream: ReadableStream<UIMessageChunk>
546
+ completion?: () => GenerationCompletion
547
+ }
548
+
549
+ async function generateWithAISDK<
550
+ M extends UIMessage,
551
+ D extends AIEventDefs<M> & EventDefs,
552
+ T extends ToolSet,
553
+ X extends object,
554
+ >(options: {
555
+ agentName: string
556
+ context: AgentGenerateContext<M, D, T, X>
557
+ messageMetadata: AgentMessageMetadata<M, D, T, X> | undefined
558
+ }): Promise<GenerationSource> {
559
+ const { agentName, context, messageMetadata } = options
560
+ const headers = context.generation.headers ?? {}
561
+ const generation =
562
+ gatewayModelId(context.model) === undefined
563
+ ? context.generation
564
+ : {
565
+ ...context.generation,
566
+ headers: Object.keys(headers).some(
567
+ (key) => key.toLowerCase() === 'x-session-affinity',
568
+ )
569
+ ? headers
570
+ : {
571
+ ...headers,
572
+ 'x-session-affinity': createHash('sha256')
573
+ .update(JSON.stringify([agentName, context.event.sessionId]))
574
+ .digest('hex'),
575
+ },
576
+ providerOptions: {
577
+ ...context.generation.providerOptions,
578
+ gateway: {
579
+ caching: 'auto',
580
+ ...context.generation.providerOptions?.['gateway'],
581
+ },
582
+ },
583
+ }
584
+ const step = await generateAISDKStep<M, T>({
585
+ model: context.model,
586
+ tools: context.tools,
587
+ messages: context.messages,
588
+ modelMessages: context.modelMessages,
589
+ responseMessageId: context.event.payload.responseMessageId,
590
+ abortSignal: context.signal,
591
+ ...(context.instructions === undefined
592
+ ? {}
593
+ : { instructions: context.instructions }),
594
+ settings: generation,
595
+ ...(messageMetadata === undefined
596
+ ? {}
597
+ : {
598
+ messageMetadata: ({ part }) =>
599
+ messageMetadata({ ...context, generation, part }),
600
+ }),
601
+ })
602
+ return {
603
+ stream: step.stream as ReadableStream<UIMessageChunk>,
604
+ completion: step.completion,
605
+ }
606
+ }
607
+
608
+ async function* consumeGeneration(options: {
609
+ source: GenerationSource
610
+ progress?: { maxChunks?: number; maxDelayMs?: number }
611
+ }): AsyncGenerator<AgentGenerationUpdate> {
612
+ const pending: UIMessageChunk[] = []
613
+ let streamedFinishReason: FinishReason | undefined
614
+ let streamedError: string | undefined
615
+ // Flush as soon as the previous write finishes: chunks that arrive while
616
+ // it is in flight wait in the stream and share the next batch.
617
+ const maxChunks = options.progress?.maxChunks ?? 128
618
+ const maxDelayMs = options.progress?.maxDelayMs ?? 0
619
+ let lastFlush = Date.now()
620
+ const batch = (): AgentGenerationProgress => ({
621
+ type: 'progress',
622
+ chunks: mergeDeltas(pending.splice(0)),
623
+ })
624
+ const reader = options.source.stream.getReader()
625
+ let pendingChunk = reader.read()
626
+ try {
627
+ for (;;) {
628
+ let result: StreamRead<UIMessageChunk>
629
+ if (pending.length > 0) {
630
+ const remaining = Math.max(0, maxDelayMs - (Date.now() - lastFlush))
631
+ const outcome = await nextOrTimer(pendingChunk, remaining)
632
+ if (outcome.type === 'timer') {
633
+ yield batch()
634
+ lastFlush = Date.now()
635
+ continue
636
+ }
637
+ result = outcome.result
638
+ } else {
639
+ result = await pendingChunk
640
+ }
641
+ if (result.done) break
642
+ const chunk = result.value
643
+ pending.push(chunk)
644
+ if (chunk.type === 'finish') streamedFinishReason = chunk.finishReason
645
+ if (chunk.type === 'error') streamedError = chunk.errorText
646
+ pendingChunk = reader.read()
647
+
648
+ if (pending.length >= maxChunks || isBoundaryChunk(pending.at(-1))) {
649
+ yield batch()
650
+ lastFlush = Date.now()
651
+ }
652
+ }
653
+ } catch (error) {
654
+ if (pending.length > 0) yield batch()
655
+ throw error
656
+ } finally {
657
+ void pendingChunk.catch(() => {})
658
+ try {
659
+ await reader.cancel()
660
+ } catch {
661
+ // Cleanup must not replace the error that ended generation.
662
+ } finally {
663
+ reader.releaseLock()
664
+ }
665
+ }
666
+
667
+ let completion: GenerationCompletion | undefined
668
+ try {
669
+ completion = options.source.completion?.()
670
+ } catch (error) {
671
+ if (pending.length > 0) {
672
+ yield batch()
673
+ }
674
+ throw error
675
+ }
676
+ if (pending.length > 0) {
677
+ yield batch()
678
+ }
679
+ if (streamedError !== undefined) throw new Error(streamedError)
680
+ const finishReason = completion?.finishReason ?? streamedFinishReason
681
+ yield {
682
+ type: 'finish',
683
+ ...(finishReason === undefined ? {} : { finishReason }),
684
+ ...(completion?.usage === undefined ? {} : { usage: completion.usage }),
685
+ }
686
+ }
687
+
688
+ const foldAIEvents = <
689
+ M extends UIMessage,
690
+ D extends AIEventDefs<M> & EventDefs,
691
+ >(options: {
692
+ reducer: Reducer<D, AIContextState<M>>
693
+ state: AIContextState<M>
694
+ events: ContractEvent<D>[]
695
+ }): AIContextState<M> => {
696
+ let state = options.state
697
+ for (const event of options.events) state = options.reducer.fold(state, event)
698
+ return state
699
+ }
700
+
701
+ const summarize = async <
702
+ M extends UIMessage,
703
+ D extends AIEventDefs<M> & EventDefs,
704
+ T extends ToolSet,
705
+ >(options: {
706
+ agentName: string
707
+ context: AgentGenerateContext<M, D, T>
708
+ instructions: string | undefined
709
+ }): Promise<{ summary: string; usage?: LanguageModelUsage }> => {
710
+ const instruction = [
711
+ 'Summarize this conversation so you can continue the task after the earlier context is replaced by your summary.',
712
+ 'Preserve the objective, constraints, decisions, exact identifiers, completed actions and their results, unresolved questions, and next steps. Distinguish observations from assumptions.',
713
+ 'Respond only with the summary as text. Do not call any tools or continue working on the task.',
714
+ options.instructions,
715
+ ]
716
+ .filter(Boolean)
717
+ .join('\n')
718
+ const generation = { ...options.context.generation }
719
+ for (const key of Object.keys(generation)) {
720
+ if (
721
+ key.startsWith('on') ||
722
+ key.startsWith('experimental_on') ||
723
+ [
724
+ 'repairToolCall',
725
+ 'experimental_repairToolCall',
726
+ 'experimental_refineToolInput',
727
+ 'experimental_transform',
728
+ 'toolApproval',
729
+ ].includes(key)
730
+ ) {
731
+ Reflect.deleteProperty(generation, key)
732
+ }
733
+ }
734
+ const tools = Object.fromEntries(
735
+ Object.entries(options.context.tools).map(([name, tool]) => {
736
+ const definition = { ...tool, needsApproval: false }
737
+ for (const key of [
738
+ 'execute',
739
+ 'onInputStart',
740
+ 'onInputDelta',
741
+ 'onInputAvailable',
742
+ ]) {
743
+ Reflect.deleteProperty(definition, key)
744
+ }
745
+ return [name, definition]
746
+ }),
747
+ ) as T
748
+ const source = await generateWithAISDK({
749
+ agentName: options.agentName,
750
+ context: {
751
+ ...options.context,
752
+ generation,
753
+ tools,
754
+ modelMessages: [
755
+ ...options.context.modelMessages,
756
+ { role: 'user', content: instruction },
757
+ ],
758
+ },
759
+ messageMetadata: undefined,
760
+ })
761
+ let summary = ''
762
+ let calledTool = false
763
+ let finish: AgentGenerationFinish | undefined
764
+ for await (const update of consumeGeneration({ source })) {
765
+ options.context.signal.throwIfAborted()
766
+ if (update.type === 'finish') {
767
+ finish = update
768
+ continue
769
+ }
770
+ for (const chunk of update.chunks) {
771
+ if (chunk.type === 'text-delta') summary += chunk.delta
772
+ if (chunk.type.startsWith('tool-')) calledTool = true
773
+ }
774
+ }
775
+ if (
776
+ calledTool ||
777
+ summary.trim().length === 0 ||
778
+ finish?.finishReason !== 'stop'
779
+ ) {
780
+ throw new Error(
781
+ 'compaction must produce a complete text summary without calling tools',
782
+ )
783
+ }
784
+ return {
785
+ summary,
786
+ ...(finish.usage === undefined ? {} : { usage: finish.usage }),
787
+ }
788
+ }
789
+
790
+ const activeCompaction = <M extends UIMessage>(
791
+ state: AIModelContext<M>,
792
+ ): AIModelContext<M>['compaction'] => {
793
+ const compaction = state.compaction
794
+ return compaction?.status === 'completed' ? compaction : null
795
+ }
796
+
797
+ const contextMessages = <M extends UIMessage>({
798
+ state,
799
+ coordinator,
800
+ }: {
801
+ state: AIModelContext<M>
802
+ coordinator: AICoordinatorState
803
+ }): M[] => {
804
+ const queued = new Set(coordinator.queued.map((item) => item.messageId))
805
+ return state.messages.filter((message) => !queued.has(message.id))
806
+ }
807
+
808
+ const modelContext = async <M extends UIMessage>(options: {
809
+ messages: M[]
810
+ state: AIModelContext<M>
811
+ tools: ToolSet
812
+ }): Promise<ModelMessage[]> => {
813
+ const messages = await convertToModelMessages(options.messages, {
814
+ tools: options.tools,
815
+ })
816
+ const summary = activeCompaction(options.state)?.summary
817
+ return summary === undefined
818
+ ? messages
819
+ : [{ role: 'user', content: summary }, ...messages]
820
+ }
821
+
822
+ type ToolCallClassification = 'automatic' | 'approval' | 'unknown'
823
+
824
+ type PendingToolCall = {
825
+ call: ToolCalledPayload
826
+ classification: ToolCallClassification
827
+ }
828
+
829
+ const toolCalledEvent = async <D extends EventDefs>(
830
+ payload: ToolCalledPayload,
831
+ ): Promise<AppendInput<D>> =>
832
+ ({
833
+ type: 'ai.tool.called',
834
+ id: await aiEventId(
835
+ 'a2.ai.tool.called',
836
+ payload.generationId,
837
+ payload.toolCallId,
838
+ ),
839
+ payload,
840
+ }) as AppendInput<D>
841
+
842
+ const approvalClassification = <T extends ToolSet>(
843
+ tools: T,
844
+ generation: AgentGenerationSettings<T>,
845
+ toolName: string,
846
+ ): ToolCallClassification => {
847
+ const configured = generation.toolApproval
848
+ if (typeof configured === 'function') return 'unknown'
849
+ const policy = configured?.[toolName]
850
+ if (typeof policy === 'function') return 'unknown'
851
+ if (
852
+ policy === 'user-approval' ||
853
+ (typeof policy === 'object' && policy?.type === 'user-approval')
854
+ ) {
855
+ return 'approval'
856
+ }
857
+ if (
858
+ policy === 'denied' ||
859
+ (typeof policy === 'object' && policy?.type === 'denied')
860
+ ) {
861
+ return 'approval'
862
+ }
863
+ const tool = tools[toolName] as { needsApproval?: unknown } | undefined
864
+ if (typeof tool?.needsApproval === 'function') return 'unknown'
865
+ return tool?.needsApproval === true ? 'approval' : 'automatic'
866
+ }
867
+
868
+ const lifecycleEvents = async <
869
+ D extends EventDefs,
870
+ T extends ToolSet,
871
+ >(options: {
872
+ requestId: string
873
+ messageId: string
874
+ generationId: string
875
+ sequence: number
876
+ chunks: UIMessageChunk[]
877
+ pending: PendingToolCall[]
878
+ tools: T
879
+ generation: AgentGenerationSettings<T>
880
+ custom: boolean
881
+ }): Promise<{ events: AppendInput<D>[]; pending: PendingToolCall[] }> => {
882
+ const result: AppendInput<D>[] = []
883
+ const pending = [...options.pending]
884
+ const take = (toolCallId: string): PendingToolCall | undefined => {
885
+ const index = pending.findIndex(
886
+ (candidate) => candidate.call.toolCallId === toolCallId,
887
+ )
888
+ if (index === -1) return undefined
889
+ return pending.splice(index, 1)[0]
890
+ }
891
+ const flush = async (toolCallId: string): Promise<void> => {
892
+ const pendingCall = take(toolCallId)
893
+ if (pendingCall) result.push(await toolCalledEvent<D>(pendingCall.call))
894
+ }
895
+ for (const [chunkIndex, chunk] of options.chunks.entries()) {
896
+ const resultId = (): Promise<string> =>
897
+ aiEventId(
898
+ 'a2.ai.tool.stream-result',
899
+ options.generationId,
900
+ String(options.sequence),
901
+ String(chunkIndex),
902
+ )
903
+ if (chunk.type === 'tool-input-available') {
904
+ const payload: ToolCalledPayload = {
905
+ requestId: options.requestId,
906
+ messageId: options.messageId,
907
+ generationId: options.generationId,
908
+ toolCallId: chunk.toolCallId,
909
+ toolName: chunk.toolName,
910
+ input: chunk.input,
911
+ ...(chunk.dynamic === undefined ? {} : { dynamic: chunk.dynamic }),
912
+ ...(chunk.providerExecuted === undefined
913
+ ? {}
914
+ : { providerExecuted: chunk.providerExecuted }),
915
+ ...((
916
+ options.tools[chunk.toolName] as
917
+ { supportsDeferredResults?: boolean } | undefined
918
+ )?.supportsDeferredResults === true
919
+ ? { supportsDeferredResults: true }
920
+ : {}),
921
+ ...(chunk.providerMetadata === undefined
922
+ ? {}
923
+ : { providerMetadata: chunk.providerMetadata }),
924
+ ...(chunk.toolMetadata === undefined
925
+ ? {}
926
+ : { toolMetadata: chunk.toolMetadata }),
927
+ ...(chunk.title === undefined ? {} : { title: chunk.title }),
928
+ }
929
+ take(payload.toolCallId)
930
+ const classification = approvalClassification(
931
+ options.tools,
932
+ options.generation,
933
+ chunk.toolName,
934
+ )
935
+ if (classification === 'automatic') {
936
+ result.push(await toolCalledEvent<D>(payload))
937
+ } else {
938
+ pending.push({ call: payload, classification })
939
+ }
940
+ } else if (chunk.type === 'tool-input-error') {
941
+ const payload: ToolResultPayload = {
942
+ requestId: options.requestId,
943
+ messageId: options.messageId,
944
+ generationId: options.generationId,
945
+ toolCallId: chunk.toolCallId,
946
+ toolName: chunk.toolName,
947
+ ...(chunk.dynamic ? { input: chunk.input } : { rawInput: chunk.input }),
948
+ error: chunk.errorText,
949
+ phase: 'input',
950
+ ...(chunk.dynamic === undefined ? {} : { dynamic: chunk.dynamic }),
951
+ ...(chunk.providerExecuted === undefined
952
+ ? {}
953
+ : { providerExecuted: chunk.providerExecuted }),
954
+ ...(chunk.providerMetadata === undefined
955
+ ? {}
956
+ : { providerMetadata: chunk.providerMetadata }),
957
+ ...(chunk.toolMetadata === undefined
958
+ ? {}
959
+ : { toolMetadata: chunk.toolMetadata }),
960
+ }
961
+ const called: ToolCalledPayload = {
962
+ requestId: options.requestId,
963
+ messageId: options.messageId,
964
+ generationId: options.generationId,
965
+ toolCallId: chunk.toolCallId,
966
+ toolName: chunk.toolName,
967
+ input: chunk.input,
968
+ }
969
+ take(chunk.toolCallId)
970
+ result.push(await toolCalledEvent<D>(called), {
971
+ type: 'ai.tool.result',
972
+ id: await resultId(),
973
+ payload,
974
+ } as AppendInput<D>)
975
+ } else if (chunk.type === 'tool-output-available') {
976
+ await flush(chunk.toolCallId)
977
+ const payload: ToolResultPayload = {
978
+ requestId: options.requestId,
979
+ messageId: options.messageId,
980
+ generationId: options.generationId,
981
+ toolCallId: chunk.toolCallId,
982
+ output: chunk.output,
983
+ phase: 'execution',
984
+ ...(chunk.preliminary === undefined
985
+ ? {}
986
+ : { preliminary: chunk.preliminary }),
987
+ ...(chunk.dynamic === undefined ? {} : { dynamic: chunk.dynamic }),
988
+ ...(chunk.providerExecuted === undefined
989
+ ? {}
990
+ : { providerExecuted: chunk.providerExecuted }),
991
+ ...(chunk.providerMetadata === undefined
992
+ ? {}
993
+ : { providerMetadata: chunk.providerMetadata }),
994
+ ...(chunk.toolMetadata === undefined
995
+ ? {}
996
+ : { toolMetadata: chunk.toolMetadata }),
997
+ }
998
+ result.push({
999
+ type: 'ai.tool.result',
1000
+ id: await resultId(),
1001
+ payload,
1002
+ } as AppendInput<D>)
1003
+ } else if (chunk.type === 'tool-output-error') {
1004
+ await flush(chunk.toolCallId)
1005
+ const payload: ToolResultPayload = {
1006
+ requestId: options.requestId,
1007
+ messageId: options.messageId,
1008
+ generationId: options.generationId,
1009
+ toolCallId: chunk.toolCallId,
1010
+ error: chunk.errorText,
1011
+ phase: 'execution',
1012
+ ...(chunk.dynamic === undefined ? {} : { dynamic: chunk.dynamic }),
1013
+ ...(chunk.providerExecuted === undefined
1014
+ ? {}
1015
+ : { providerExecuted: chunk.providerExecuted }),
1016
+ ...(chunk.providerMetadata === undefined
1017
+ ? {}
1018
+ : { providerMetadata: chunk.providerMetadata }),
1019
+ ...(chunk.toolMetadata === undefined
1020
+ ? {}
1021
+ : { toolMetadata: chunk.toolMetadata }),
1022
+ }
1023
+ result.push({
1024
+ type: 'ai.tool.result',
1025
+ id: await resultId(),
1026
+ payload,
1027
+ } as AppendInput<D>)
1028
+ } else if (chunk.type === 'tool-output-denied') {
1029
+ await flush(chunk.toolCallId)
1030
+ const payload: ToolResultPayload = {
1031
+ requestId: options.requestId,
1032
+ messageId: options.messageId,
1033
+ generationId: options.generationId,
1034
+ toolCallId: chunk.toolCallId,
1035
+ denied: true,
1036
+ phase: 'execution',
1037
+ }
1038
+ result.push({
1039
+ type: 'ai.tool.result',
1040
+ id: await resultId(),
1041
+ payload,
1042
+ } as AppendInput<D>)
1043
+ } else if (chunk.type === 'tool-approval-request') {
1044
+ await flush(chunk.toolCallId)
1045
+ const payload: ApprovalRequestedPayload = {
1046
+ requestId: options.requestId,
1047
+ messageId: options.messageId,
1048
+ generationId: options.generationId,
1049
+ approvalId: chunk.approvalId,
1050
+ toolCallId: chunk.toolCallId,
1051
+ ...(chunk.isAutomatic === undefined
1052
+ ? {}
1053
+ : { isAutomatic: chunk.isAutomatic }),
1054
+ ...(chunk.signature === undefined
1055
+ ? {}
1056
+ : { signature: chunk.signature }),
1057
+ }
1058
+ result.push({
1059
+ type: 'ai.approval.requested',
1060
+ id: await aiEventId(
1061
+ 'a2.ai.approval.requested',
1062
+ options.generationId,
1063
+ chunk.approvalId,
1064
+ ),
1065
+ payload,
1066
+ } as AppendInput<D>)
1067
+ } else if (chunk.type === 'tool-approval-response') {
1068
+ result.push({
1069
+ type: 'ai.approval.responded',
1070
+ id: await aiEventId(
1071
+ 'a2.ai.approval.automatic-response',
1072
+ options.generationId,
1073
+ chunk.approvalId,
1074
+ ),
1075
+ payload: {
1076
+ requestId: options.requestId,
1077
+ messageId: options.messageId,
1078
+ generationId: options.generationId,
1079
+ approvalId: chunk.approvalId,
1080
+ approved: chunk.approved,
1081
+ ...(chunk.reason === undefined ? {} : { reason: chunk.reason }),
1082
+ },
1083
+ } as AppendInput<D>)
1084
+ }
1085
+ }
1086
+ return { events: result, pending }
1087
+ }
1088
+
1089
+ const validateCompaction = <
1090
+ M extends UIMessage,
1091
+ D extends AIEventDefs<M> & EventDefs,
1092
+ T extends ToolSet,
1093
+ >(options: {
1094
+ compaction: CompactionOptions<M, D>
1095
+ explicit: boolean
1096
+ generation: AgentGenerationSettings<T> | undefined
1097
+ tools: T | undefined
1098
+ }): CompactionOptions<M, D> => {
1099
+ let compaction = options.compaction
1100
+ if (
1101
+ compaction !== false &&
1102
+ (typeof compaction !== 'object' || compaction === null)
1103
+ ) {
1104
+ throw new TypeError('compaction must resolve to false or a policy object')
1105
+ }
1106
+ if (compaction !== false && 'then' in compaction) {
1107
+ void Promise.resolve(compaction).catch(() => {})
1108
+ throw new TypeError('compaction options must be a policy or a resolver')
1109
+ }
1110
+ if (compaction !== false && !('shouldCompact' in compaction)) {
1111
+ if (
1112
+ compaction.thresholdTokens !== undefined &&
1113
+ (!Number.isSafeInteger(compaction.thresholdTokens) ||
1114
+ compaction.thresholdTokens < 1)
1115
+ ) {
1116
+ throw new TypeError(
1117
+ 'compaction.thresholdTokens must be a positive safe integer',
1118
+ )
1119
+ }
1120
+ const choice = options.generation?.toolChoice
1121
+ const unsupported =
1122
+ options.generation?.output !== undefined
1123
+ ? 'structured output'
1124
+ : choice !== undefined && choice !== 'auto' && choice !== 'none'
1125
+ ? 'forced tool choice'
1126
+ : Object.values(options.tools ?? {}).some(
1127
+ (tool) =>
1128
+ tool.type === 'provider' && tool.isProviderExecuted === true,
1129
+ )
1130
+ ? 'provider-executed tools'
1131
+ : undefined
1132
+ if (unsupported !== undefined) {
1133
+ if (options.explicit) {
1134
+ throw new TypeError(
1135
+ `automatic compaction does not support ${unsupported}; use a custom policy or compaction: false`,
1136
+ )
1137
+ }
1138
+ compaction = false
1139
+ }
1140
+ }
1141
+
1142
+ return compaction
1143
+ }
1144
+
1145
+ export function createHandlers<
1146
+ M extends UIMessage,
1147
+ D extends AIEventDefs<M> & EventDefs,
1148
+ T extends ToolSet = ToolSet,
1149
+ X extends object = Record<never, never>,
1150
+ >(
1151
+ options: HandlersOptions<M, D, T, X>,
1152
+ ): NonNullable<ServerOptions<D>['handlers']> {
1153
+ if (options.model === undefined) {
1154
+ throw new TypeError('agent requires model')
1155
+ }
1156
+ if (
1157
+ options.progress?.maxChunks !== undefined &&
1158
+ (!Number.isInteger(options.progress.maxChunks) ||
1159
+ options.progress.maxChunks < 1)
1160
+ ) {
1161
+ throw new TypeError('progress.maxChunks must be a positive integer')
1162
+ }
1163
+ if (
1164
+ options.progress?.maxDelayMs !== undefined &&
1165
+ (!Number.isFinite(options.progress.maxDelayMs) ||
1166
+ options.progress.maxDelayMs < 0)
1167
+ ) {
1168
+ throw new TypeError('progress.maxDelayMs must be a non-negative number')
1169
+ }
1170
+ if (
1171
+ options.maxSteps !== undefined &&
1172
+ options.maxSteps !== Number.POSITIVE_INFINITY &&
1173
+ (!Number.isInteger(options.maxSteps) || options.maxSteps < 1)
1174
+ ) {
1175
+ throw new TypeError('maxSteps must be a positive integer or Infinity')
1176
+ }
1177
+
1178
+ const configuredCompaction =
1179
+ options.compaction ?? (options.generate === undefined ? {} : false)
1180
+ const staticCompaction =
1181
+ typeof configuredCompaction === 'function'
1182
+ ? undefined
1183
+ : validateCompaction({
1184
+ compaction: configuredCompaction,
1185
+ explicit: options.compaction !== undefined,
1186
+ generation:
1187
+ typeof options.generation === 'function'
1188
+ ? undefined
1189
+ : options.generation,
1190
+ tools: options.tools,
1191
+ })
1192
+
1193
+ const tools = options.tools ?? ({} as T)
1194
+ const telemetry = options.telemetry ?? CONSOLE_TELEMETRY
1195
+ const maxSteps = options.maxSteps ?? Number.POSITIVE_INFINITY
1196
+ const control = createControlRuntime({
1197
+ contract: options.contract,
1198
+ telemetry,
1199
+ })
1200
+ const controlHandler: typeof control.handler = (ctx) =>
1201
+ telemetry.span(
1202
+ 'a2.ai.control',
1203
+ {
1204
+ 'a2.contract': options.contract.name,
1205
+ 'a2.session_id': ctx.event.sessionId,
1206
+ 'a2.ai.control.type': ctx.event.type,
1207
+ 'a2.ai.control.id': ctx.event.id,
1208
+ 'a2.ai.control.attempt': ctx.attempt,
1209
+ },
1210
+ () => control.handler(ctx),
1211
+ )
1212
+ const coordinator = control.coordinator
1213
+ const contextReducer = createContextReducer<M, D>({
1214
+ contract: options.contract,
1215
+ })
1216
+ const readContext = async (
1217
+ session: HandlerContext<D>['session'],
1218
+ through: number | 'latest',
1219
+ ) => {
1220
+ const snapshot = await session.state(contextReducer, { through })
1221
+ return snapshot
1222
+ }
1223
+ // A session's next step starts only after its tools settle, so only its
1224
+ // latest step's prompt is ever reused. A miss rebuilds the prompt from the
1225
+ // context checkpoint, so eviction is memory policy, never correctness.
1226
+ const prompts = boundedCache<
1227
+ string,
1228
+ { generationId: string; messages: Promise<ModelMessage[]> }
1229
+ >({ maxEntries: PROMPT_SESSIONS, maxBytes: PROMPT_SESSIONS })
1230
+ const cachePrompt = ({
1231
+ sessionId,
1232
+ generationId,
1233
+ messages,
1234
+ }: {
1235
+ sessionId: string
1236
+ generationId: string
1237
+ messages: Promise<ModelMessage[]>
1238
+ }): void =>
1239
+ // Entries count as one byte, so the bound is the session count.
1240
+ prompts.set(sessionId, { generationId, messages }, 1)
1241
+ const forgetPrompt = (sessionId: string, generationId: string): void => {
1242
+ if (prompts.get(sessionId)?.generationId === generationId)
1243
+ prompts.delete(sessionId)
1244
+ }
1245
+
1246
+ type RuntimeTool = {
1247
+ execute?: (
1248
+ input: unknown,
1249
+ options: {
1250
+ toolCallId: string
1251
+ messages: ModelMessage[]
1252
+ abortSignal?: AbortSignal
1253
+ },
1254
+ ) => unknown
1255
+ }
1256
+
1257
+ type ToolHandlerContext = HandlerContext<D, 'ai.tool.execution.requested'>
1258
+
1259
+ const resultEvent = async (
1260
+ call: ToolCalledPayload,
1261
+ suffix: string,
1262
+ result: {
1263
+ output?: unknown
1264
+ error?: string
1265
+ denied?: boolean
1266
+ preliminary?: boolean
1267
+ },
1268
+ ): Promise<AppendInput<D>> =>
1269
+ ({
1270
+ type: 'ai.tool.result',
1271
+ id: await aiEventId(
1272
+ 'a2.ai.tool.result',
1273
+ call.generationId,
1274
+ call.toolCallId,
1275
+ suffix,
1276
+ ),
1277
+ payload: {
1278
+ requestId: call.requestId,
1279
+ messageId: call.messageId,
1280
+ generationId: call.generationId,
1281
+ toolCallId: call.toolCallId,
1282
+ toolName: call.toolName,
1283
+ input: call.input,
1284
+ phase: 'execution',
1285
+ ...(result.output === undefined ? {} : { output: result.output }),
1286
+ ...(result.error === undefined ? {} : { error: result.error }),
1287
+ ...(result.denied === undefined ? {} : { denied: result.denied }),
1288
+ ...(result.preliminary === undefined
1289
+ ? {}
1290
+ : { preliminary: result.preliminary }),
1291
+ ...(call.dynamic === undefined ? {} : { dynamic: call.dynamic }),
1292
+ ...(call.providerExecuted === undefined
1293
+ ? {}
1294
+ : { providerExecuted: call.providerExecuted }),
1295
+ ...(call.toolMetadata === undefined
1296
+ ? {}
1297
+ : { toolMetadata: call.toolMetadata }),
1298
+ },
1299
+ }) as AppendInput<D>
1300
+
1301
+ const promptMessages = (input: {
1302
+ ctx: ToolHandlerContext
1303
+ call: ToolCalledPayload
1304
+ generation: GenerationStartedPayload
1305
+ }): Promise<ModelMessage[]> => {
1306
+ const { ctx, call, generation: owner } = input
1307
+ const sessionId = ctx.event.sessionId
1308
+ const cached = prompts.get(sessionId)
1309
+ if (cached?.generationId === call.generationId) return cached.messages
1310
+ const computation = (async () => {
1311
+ const frontier = owner.promptThroughIndex
1312
+ if (frontier === undefined)
1313
+ throw new TypeError('AI work requires a prompt checkpoint')
1314
+ const snapshot = await readContext(ctx.session, ctx.event.index)
1315
+ if (
1316
+ snapshot.state.generation?.owner.generationId !== owner.generationId ||
1317
+ snapshot.state.generation.owner.requestId !== owner.requestId ||
1318
+ snapshot.state.activeProjection?.generationId !== owner.generationId
1319
+ )
1320
+ throw new TypeError(
1321
+ 'Tool prompt checkpoint does not own its generation',
1322
+ )
1323
+ const state = generationBase(snapshot.state)
1324
+ return telemetry.span(
1325
+ 'a2.ai.context',
1326
+ {
1327
+ 'a2.contract': options.contract.name,
1328
+ 'a2.session_id': ctx.event.sessionId,
1329
+ 'a2.ai.context.request_id': call.requestId,
1330
+ 'a2.ai.context.phase': 'tool',
1331
+ },
1332
+ () =>
1333
+ modelContext({
1334
+ messages: state.messages.filter(
1335
+ (message) => state.messageIndexes[message.id]! <= frontier,
1336
+ ),
1337
+ state,
1338
+ tools,
1339
+ }),
1340
+ )
1341
+ })()
1342
+ cachePrompt({
1343
+ sessionId,
1344
+ generationId: call.generationId,
1345
+ messages: computation,
1346
+ })
1347
+ void computation.catch(() => forgetPrompt(sessionId, call.generationId))
1348
+ return computation
1349
+ }
1350
+
1351
+ const toolExecutionFailure = async (
1352
+ ctx: ToolHandlerContext,
1353
+ call: ToolCalledPayload,
1354
+ error: unknown,
1355
+ ): Promise<AppendInput<D> | void> => {
1356
+ const schedulerFailure = consumeSchedulerSendFailure(error)
1357
+ if (checkAbort(ctx.signal)) return
1358
+ if (schedulerFailure === 'retryable') throw error
1359
+ return resultEvent(call, `${ctx.event.id}:execution:error`, {
1360
+ error: errorMessage(error),
1361
+ })
1362
+ }
1363
+
1364
+ const runToolExecution = async (
1365
+ ctx: ToolHandlerContext,
1366
+ call: ToolCalledPayload,
1367
+ execute: NonNullable<RuntimeTool['execute']>,
1368
+ messages: ModelMessage[],
1369
+ ): Promise<AppendInput<D> | void> => {
1370
+ let output: unknown
1371
+ try {
1372
+ output = await execute(call.input, {
1373
+ toolCallId: call.toolCallId,
1374
+ messages,
1375
+ abortSignal: ctx.signal,
1376
+ })
1377
+ } catch (error) {
1378
+ return toolExecutionFailure(ctx, call, error)
1379
+ }
1380
+ let iterator: AsyncIterator<unknown> | undefined
1381
+ try {
1382
+ iterator = asyncIterator(output)
1383
+ } catch (error) {
1384
+ return toolExecutionFailure(ctx, call, error)
1385
+ }
1386
+ if (iterator === undefined) {
1387
+ return resultEvent(call, `${ctx.event.id}:execution:0`, { output })
1388
+ }
1389
+ let last: unknown
1390
+ let sequence = 0
1391
+ let done = false
1392
+ try {
1393
+ for (;;) {
1394
+ let result: IteratorResult<unknown>
1395
+ try {
1396
+ result = await iterator.next()
1397
+ } catch (error) {
1398
+ return toolExecutionFailure(ctx, call, error)
1399
+ }
1400
+ if (result.done) {
1401
+ if (ctx.signal.reason instanceof A2Error) throw ctx.signal.reason
1402
+ done = true
1403
+ break
1404
+ }
1405
+ if (checkAbort(ctx.signal)) return
1406
+ last = result.value
1407
+ await control.append({
1408
+ ctx,
1409
+ name: `tool:${call.toolCallId}:preliminary:${sequence}`,
1410
+ events: [
1411
+ await resultEvent(
1412
+ call,
1413
+ `${ctx.event.id}:execution:${ctx.attempt}:${sequence}:preliminary`,
1414
+ {
1415
+ output: result.value,
1416
+ preliminary: true,
1417
+ },
1418
+ ),
1419
+ ],
1420
+ })
1421
+ sequence += 1
1422
+ }
1423
+ } finally {
1424
+ if (!done) {
1425
+ try {
1426
+ await iterator.return?.()
1427
+ } catch {
1428
+ // Preserve the failure or cancellation that ended execution.
1429
+ }
1430
+ }
1431
+ }
1432
+ return resultEvent(
1433
+ call,
1434
+ `${ctx.event.id}:execution:${sequence}:final`,
1435
+ sequence === 0 ? {} : { output: last },
1436
+ )
1437
+ }
1438
+
1439
+ const executeTool = async (input: {
1440
+ ctx: ToolHandlerContext
1441
+ call: ToolCalledPayload
1442
+ generation: GenerationStartedPayload
1443
+ }): Promise<AppendInput<D> | void> => {
1444
+ const { ctx, call } = input
1445
+ const tool = tools[call.toolName] as RuntimeTool | undefined
1446
+ const execute = tool?.execute
1447
+ if (execute === undefined) {
1448
+ return resultEvent(call, `${ctx.event.id}:execution:error`, {
1449
+ error: `Tool '${call.toolName}' has no server executor`,
1450
+ })
1451
+ }
1452
+ const messages = await promptMessages(input)
1453
+ if (checkAbort(ctx.signal)) return
1454
+ const scope: AmbientToolScope = {
1455
+ contract: options.contract,
1456
+ context: ctx,
1457
+ }
1458
+ // Async iteration resumes in the caller's context, so the whole
1459
+ // execution, not just the execute call, runs inside the scope.
1460
+ return toolScopeStorage.run(scope, () =>
1461
+ runToolExecution(ctx, call, execute, messages),
1462
+ )
1463
+ }
1464
+
1465
+ const handleToolExecution = async (
1466
+ ctx: ToolHandlerContext,
1467
+ ): Promise<AppendInput<D>> => {
1468
+ try {
1469
+ const state = (
1470
+ await ctx.session.state(control.reducer, { through: 'latest' })
1471
+ ).state
1472
+ const request = ctx.event.payload
1473
+ const call = state.coordinator.response?.calls.find(
1474
+ (candidate) => candidate.work?.id === ctx.event.id,
1475
+ )
1476
+ if (
1477
+ state.active?.turnId !== request.turnId ||
1478
+ Object.values(state.pendingStops).includes(request.turnId) ||
1479
+ state.active.suspended ||
1480
+ state.active.version !== request.version ||
1481
+ state.coordinator.response?.failure !== undefined ||
1482
+ !call ||
1483
+ call.terminal ||
1484
+ call.work?.settled
1485
+ )
1486
+ return control.settled({ ctx, events: undefined })
1487
+ return control.settled({
1488
+ ctx,
1489
+ events: await executeTool({
1490
+ ctx,
1491
+ call: request.call,
1492
+ generation: request.generation,
1493
+ }),
1494
+ })
1495
+ } catch (error) {
1496
+ if (control.cancelled({ error, signal: ctx.signal }))
1497
+ return control.settled({ ctx, events: undefined })
1498
+ throw error
1499
+ }
1500
+ }
1501
+
1502
+ const generationHandler = async (
1503
+ ctx: HandlerContext<D, 'ai.generation.requested'>,
1504
+ ): Promise<void | AppendInput<D> | readonly AppendInput<D>[]> => {
1505
+ if (checkAbort(ctx.signal)) return
1506
+ const requestId = ctx.event.id
1507
+ const request = ctx.event.payload
1508
+ const snapshot = await readContext(ctx.session, 'latest')
1509
+ if (
1510
+ snapshot.state.closed ||
1511
+ snapshot.state.authority?.requestId !== requestId ||
1512
+ snapshot.state.authority?.phase === 'paused' ||
1513
+ snapshot.state.authority?.phase === 'pausing'
1514
+ )
1515
+ return
1516
+ const controlState = (
1517
+ await ctx.session.state(control.reducer, { through: snapshot.index })
1518
+ ).state
1519
+ if (
1520
+ request.control !== undefined &&
1521
+ Object.values(controlState.pendingStops).includes(request.control.turnId)
1522
+ )
1523
+ return
1524
+ const coordinatorState = controlState.coordinator
1525
+ if (
1526
+ coordinatorState.closed ||
1527
+ coordinatorState.response?.activeRequestId !== requestId
1528
+ )
1529
+ return
1530
+ const response = coordinatorState.response
1531
+ const current =
1532
+ response.generation?.requestId === requestId
1533
+ ? response.generation
1534
+ : undefined
1535
+ if (
1536
+ current !== undefined &&
1537
+ (response.completion !== undefined ||
1538
+ (response.failure !== undefined &&
1539
+ response.failure.superseded !== true))
1540
+ )
1541
+ return
1542
+ if (current !== undefined && current.attempt >= ctx.attempt) return
1543
+ const attempt = ctx.attempt
1544
+ const generationId = await aiEventId(
1545
+ 'a2.ai.generation',
1546
+ requestId,
1547
+ String(attempt),
1548
+ )
1549
+ const responseMessageId = request.responseMessageId
1550
+ const responseStepCount = response.stepCount
1551
+ if (
1552
+ request.reason === 'tool' &&
1553
+ (request.sourceGenerationId !==
1554
+ response.source?.generation.generationId ||
1555
+ response.responseMessageId !== request.messageId ||
1556
+ response.responseMessageId !== request.responseMessageId)
1557
+ )
1558
+ return
1559
+ const previous = current
1560
+ const state =
1561
+ previous === undefined
1562
+ ? snapshot.state
1563
+ : contextForAttempt(snapshot.state)
1564
+ const resolverContext: AgentResolverContext<M, D, X> = {
1565
+ event: ctx.event,
1566
+ messages: contextMessages({ state, coordinator: coordinatorState }),
1567
+ // The contract registered this reducer with its declared state type.
1568
+ state: (snapshot.state.app ?? {}) as X,
1569
+ session: {
1570
+ state: (reducer, readOptions) =>
1571
+ ctx.session.state(reducer, {
1572
+ ...readOptions,
1573
+ through: readOptions?.through ?? snapshot.index,
1574
+ }),
1575
+ },
1576
+ signal: ctx.signal,
1577
+ }
1578
+
1579
+ if (request.reason === 'tool' && responseStepCount >= maxSteps) {
1580
+ const source = response.source?.generation
1581
+ if (source === undefined) return
1582
+ return {
1583
+ type: 'ai.generation.failed',
1584
+ id: await aiEventId('a2.ai.step-limit', requestId),
1585
+ payload: {
1586
+ requestId: source.requestId,
1587
+ messageId: source.messageId,
1588
+ generationId: source.generationId,
1589
+ responseMessageId,
1590
+ error: `agent exceeded the ${maxSteps}-step limit`,
1591
+ stepLimit: true,
1592
+ },
1593
+ } as AppendInput<D>
1594
+ }
1595
+
1596
+ const startEvents = async (
1597
+ started: GenerationStartedPayload,
1598
+ metadataModelId?: string,
1599
+ ): Promise<AppendInput<D>[]> => {
1600
+ const events: AppendInput<D>[] = []
1601
+ if (metadataModelId !== undefined) {
1602
+ events.push({
1603
+ type: 'ai.model.metadata.requested',
1604
+ id: await aiEventId(
1605
+ 'a2.ai.model.metadata',
1606
+ options.contract.name,
1607
+ ctx.event.sessionId,
1608
+ metadataModelId,
1609
+ ),
1610
+ payload: { modelId: metadataModelId },
1611
+ } as AppendInput<D>)
1612
+ }
1613
+ if (previous) {
1614
+ const superseded: GenerationFailedPayload = {
1615
+ requestId,
1616
+ messageId: previous.messageId,
1617
+ generationId: previous.generationId,
1618
+ responseMessageId: previous.responseMessageId,
1619
+ error: 'generation attempt was superseded after an incomplete run',
1620
+ superseded: true,
1621
+ }
1622
+ events.push({
1623
+ type: 'ai.generation.failed',
1624
+ id: await aiEventId('a2.ai.superseded', previous.generationId),
1625
+ payload: superseded,
1626
+ } as AppendInput<D>)
1627
+ }
1628
+ events.push({
1629
+ type: 'ai.generation.started',
1630
+ id: generationId,
1631
+ payload: started,
1632
+ } as AppendInput<D>)
1633
+ return events
1634
+ }
1635
+ const startedPayload = (
1636
+ model?: LanguageModel,
1637
+ ): GenerationStartedPayload => ({
1638
+ requestId,
1639
+ messageId: request.messageId,
1640
+ generationId,
1641
+ responseMessageId,
1642
+ attempt,
1643
+ ...(model === undefined ? {} : { model: modelName(model) }),
1644
+ promptThroughIndex: snapshot.index,
1645
+ })
1646
+
1647
+ let resolvedModel: LanguageModel | undefined
1648
+ let resolving = true
1649
+ let generationStarted = false
1650
+ try {
1651
+ const choice = await resolve(options.model, resolverContext)
1652
+ if (choice === undefined) {
1653
+ throw new TypeError('the model resolver returned undefined')
1654
+ }
1655
+ const settings: ModelSettings =
1656
+ typeof choice === 'string' || 'specificationVersion' in choice
1657
+ ? { model: choice }
1658
+ : choice
1659
+ const model = settings.model
1660
+ if (
1661
+ typeof model !== 'string' &&
1662
+ (typeof model !== 'object' || model === null)
1663
+ )
1664
+ throw new TypeError('the model resolver must return a model')
1665
+ resolvedModel = model
1666
+ if (checkAbort(ctx.signal)) return
1667
+ const resolvedInstructions =
1668
+ options.instructions === undefined
1669
+ ? undefined
1670
+ : await resolve(options.instructions, resolverContext)
1671
+ if (checkAbort(ctx.signal)) return
1672
+ const resolvedGeneration: AgentGenerationSettings<T> = await resolve(
1673
+ options.generation ?? {},
1674
+ resolverContext,
1675
+ )
1676
+ if (checkAbort(ctx.signal)) return
1677
+ if (
1678
+ resolvedGeneration === null ||
1679
+ typeof resolvedGeneration !== 'object' ||
1680
+ Array.isArray(resolvedGeneration)
1681
+ )
1682
+ throw new TypeError('the generation resolver must return settings')
1683
+ const generation = withModelSettings({
1684
+ generation: resolvedGeneration,
1685
+ settings,
1686
+ })
1687
+ const compaction =
1688
+ typeof configuredCompaction === 'function' ||
1689
+ typeof options.generation === 'function'
1690
+ ? validateCompaction({
1691
+ compaction:
1692
+ typeof configuredCompaction === 'function'
1693
+ ? await resolve(configuredCompaction, resolverContext)
1694
+ : staticCompaction!,
1695
+ explicit: options.compaction !== undefined,
1696
+ generation,
1697
+ tools: options.tools,
1698
+ })
1699
+ : staticCompaction!
1700
+ if (checkAbort(ctx.signal)) return
1701
+ resolving = false
1702
+
1703
+ const catalogModelId = gatewayModelId(model)
1704
+ const gatewayOptions = generation.providerOptions?.['gateway']
1705
+ const usesFallbackModels =
1706
+ Array.isArray(gatewayOptions?.['models']) &&
1707
+ gatewayOptions['models'].length > 0
1708
+ const discoversMetadata =
1709
+ compaction !== false &&
1710
+ !('shouldCompact' in compaction) &&
1711
+ compaction.thresholdTokens === undefined &&
1712
+ !usesFallbackModels &&
1713
+ catalogModelId !== undefined
1714
+ const generationEvents = await control.append({
1715
+ ctx,
1716
+ name: 'generation-start',
1717
+ events: await startEvents(
1718
+ startedPayload(model),
1719
+ discoversMetadata &&
1720
+ !Object.hasOwn(state.modelMetadata, catalogModelId)
1721
+ ? catalogModelId
1722
+ : undefined,
1723
+ ),
1724
+ })
1725
+ const startIndex = generationEvents.find(
1726
+ (event) => event.id === generationId,
1727
+ )!.index
1728
+ generationStarted = true
1729
+
1730
+ if (checkAbort(ctx.signal)) return
1731
+
1732
+ const promptCoordinatorState = coordinatorState
1733
+ const { messages } = resolverContext
1734
+ let generationMessages = messages
1735
+ let modelMessages = await telemetry.span(
1736
+ 'a2.ai.context',
1737
+ {
1738
+ 'a2.contract': options.contract.name,
1739
+ 'a2.session_id': ctx.event.sessionId,
1740
+ 'a2.ai.context.request_id': requestId,
1741
+ 'a2.ai.context.phase': 'generation',
1742
+ },
1743
+ () => modelContext({ messages, state, tools }),
1744
+ )
1745
+ const baseContext: AgentGenerateContext<M, D, T, X> = {
1746
+ ...resolverContext,
1747
+ generationId,
1748
+ modelMessages,
1749
+ model: resolvedModel,
1750
+ tools,
1751
+ ...(resolvedInstructions === undefined
1752
+ ? {}
1753
+ : { instructions: resolvedInstructions }),
1754
+ generation,
1755
+ }
1756
+ const policy = compaction
1757
+ const tracksInput =
1758
+ policy !== false &&
1759
+ !('shouldCompact' in policy) &&
1760
+ (policy.thresholdTokens !== undefined || discoversMetadata)
1761
+ const metadata =
1762
+ catalogModelId === undefined
1763
+ ? undefined
1764
+ : state.modelMetadata[catalogModelId]
1765
+ const limits =
1766
+ metadata?.status === 'resolved' && !usesFallbackModels
1767
+ ? metadata.limits
1768
+ : undefined
1769
+ let inputTokenEstimate: number | undefined
1770
+ let compacted = false
1771
+ let compactionEstimate:
1772
+ | Pick<
1773
+ CompactionRequestedPayload,
1774
+ 'inputTokenEstimate' | 'inputTokens' | 'thresholdTokens'
1775
+ >
1776
+ | undefined
1777
+ const hasPendingTools = messages.some((message) =>
1778
+ message.parts.some(
1779
+ (part) =>
1780
+ isToolUIPart(part) &&
1781
+ !(
1782
+ part.state === 'output-error' ||
1783
+ part.state === 'output-denied' ||
1784
+ (part.state === 'output-available' && part.preliminary !== true)
1785
+ ),
1786
+ ),
1787
+ )
1788
+ if (policy && !hasPendingTools) {
1789
+ const compactionContext = {
1790
+ ...resolverContext,
1791
+ messages,
1792
+ modelMessages,
1793
+ }
1794
+ let shouldCompact: boolean
1795
+ if ('shouldCompact' in policy) {
1796
+ shouldCompact = await policy.shouldCompact(compactionContext)
1797
+ } else if (tracksInput) {
1798
+ inputTokenEstimate = await estimateInputTokens({
1799
+ messages: modelMessages,
1800
+ instructions: resolvedInstructions,
1801
+ tools,
1802
+ })
1803
+ const inputTokens = measuredInputTokens({
1804
+ estimate: inputTokenEstimate,
1805
+ model: modelName(resolvedModel),
1806
+ calibration: coordinatorState.calibration,
1807
+ })
1808
+ const threshold =
1809
+ policy.thresholdTokens ??
1810
+ (limits === undefined
1811
+ ? undefined
1812
+ : Math.floor(
1813
+ Math.min(
1814
+ limits.contextWindow * 0.75,
1815
+ limits.contextWindow - (generation.maxOutputTokens ?? 0),
1816
+ ),
1817
+ ))
1818
+ if (
1819
+ limits !== undefined &&
1820
+ threshold !== undefined &&
1821
+ threshold <= 0
1822
+ ) {
1823
+ throw new Error(
1824
+ `output allowance exhausts the ${limits.contextWindow}-token context window`,
1825
+ )
1826
+ }
1827
+ shouldCompact = threshold !== undefined && inputTokens >= threshold
1828
+ if (threshold !== undefined && shouldCompact) {
1829
+ compactionEstimate = {
1830
+ inputTokenEstimate,
1831
+ inputTokens,
1832
+ thresholdTokens: threshold,
1833
+ }
1834
+ }
1835
+ } else {
1836
+ shouldCompact = false
1837
+ }
1838
+ if (checkAbort(ctx.signal)) return
1839
+ if (shouldCompact) {
1840
+ const latest = (
1841
+ await ctx.session.state(coordinator, { through: 'latest' })
1842
+ ).state
1843
+ shouldCompact = Object.keys(latest.unsettledTools ?? {}).length === 0
1844
+ }
1845
+ if (checkAbort(ctx.signal)) return
1846
+ if (shouldCompact) {
1847
+ const throughMessageId =
1848
+ state.messages.findLast(
1849
+ (message) =>
1850
+ !promptCoordinatorState.queued.some(
1851
+ (queued) => queued.messageId === message.id,
1852
+ ),
1853
+ )?.id ?? request.messageId
1854
+ const throughIndex = snapshot.index
1855
+ generationEvents.push(
1856
+ ...(await control.append({
1857
+ ctx,
1858
+ name: 'compaction-requested',
1859
+ events: [
1860
+ {
1861
+ type: 'ai.compaction.requested',
1862
+ id: await aiEventId(
1863
+ 'a2.ai.compaction.requested',
1864
+ generationId,
1865
+ ),
1866
+ payload: {
1867
+ generationId,
1868
+ throughMessageId,
1869
+ throughIndex,
1870
+ ...compactionEstimate,
1871
+ },
1872
+ } as AppendInput<D>,
1873
+ ],
1874
+ })),
1875
+ )
1876
+ const result = await telemetry.span(
1877
+ 'a2.ai.compaction',
1878
+ {
1879
+ 'a2.contract': options.contract.name,
1880
+ 'a2.session_id': ctx.event.sessionId,
1881
+ 'a2.ai.compaction.request_id': requestId,
1882
+ 'a2.ai.compaction.generation_id': generationId,
1883
+ 'a2.ai.compaction.policy':
1884
+ 'shouldCompact' in policy ? 'custom' : 'automatic',
1885
+ },
1886
+ async () =>
1887
+ !('shouldCompact' in policy)
1888
+ ? {
1889
+ messages: [] as M[],
1890
+ ...(await summarize({
1891
+ agentName: options.contract.name,
1892
+ context: baseContext,
1893
+ instructions: policy.instructions,
1894
+ })),
1895
+ }
1896
+ : { messages: await policy.compact(compactionContext) },
1897
+ )
1898
+ if (checkAbort(ctx.signal)) return
1899
+ const completed: CompactionCompletedPayload<M> = {
1900
+ generationId,
1901
+ throughMessageId,
1902
+ throughIndex,
1903
+ ...result,
1904
+ ...(promptCoordinatorState.queued.length === 0
1905
+ ? {}
1906
+ : {
1907
+ retainedMessageIds: promptCoordinatorState.queued.map(
1908
+ (item) => item.messageId,
1909
+ ),
1910
+ }),
1911
+ }
1912
+ generationEvents.push(
1913
+ ...(await control.append({
1914
+ ctx,
1915
+ name: 'compaction-completed',
1916
+ events: [
1917
+ {
1918
+ type: 'ai.compaction.completed',
1919
+ id: await aiEventId(
1920
+ 'a2.ai.compaction.completed',
1921
+ generationId,
1922
+ ),
1923
+ payload: completed,
1924
+ } as AppendInput<D>,
1925
+ ],
1926
+ })),
1927
+ )
1928
+ const queued = new Set(
1929
+ promptCoordinatorState.queued.map((item) => item.messageId),
1930
+ )
1931
+ generationMessages = result.messages.filter(
1932
+ (message) => !queued.has(message.id),
1933
+ )
1934
+ compacted = true
1935
+ }
1936
+ }
1937
+
1938
+ const currentState = foldAIEvents({
1939
+ reducer: contextReducer,
1940
+ events: generationEvents,
1941
+ state,
1942
+ })
1943
+ if (request.reason === 'retry' || previous !== undefined) {
1944
+ generationMessages = contextMessages({
1945
+ state: currentState,
1946
+ coordinator: promptCoordinatorState,
1947
+ })
1948
+ }
1949
+ if (
1950
+ (policy && 'shouldCompact' in policy) ||
1951
+ generationMessages !== messages ||
1952
+ activeCompaction(currentState)?.summary !==
1953
+ activeCompaction(state)?.summary
1954
+ ) {
1955
+ modelMessages = await telemetry.span(
1956
+ 'a2.ai.context',
1957
+ {
1958
+ 'a2.contract': options.contract.name,
1959
+ 'a2.session_id': ctx.event.sessionId,
1960
+ 'a2.ai.context.request_id': requestId,
1961
+ 'a2.ai.context.phase': 'refresh',
1962
+ },
1963
+ () =>
1964
+ modelContext({
1965
+ messages: generationMessages,
1966
+ state: currentState,
1967
+ tools,
1968
+ }),
1969
+ )
1970
+ }
1971
+ cachePrompt({
1972
+ sessionId: ctx.event.sessionId,
1973
+ generationId,
1974
+ messages: Promise.resolve(modelMessages),
1975
+ })
1976
+ if (tracksInput && (inputTokenEstimate === undefined || compacted)) {
1977
+ inputTokenEstimate = await estimateInputTokens({
1978
+ messages: modelMessages,
1979
+ instructions: resolvedInstructions,
1980
+ tools,
1981
+ })
1982
+ }
1983
+
1984
+ let pendingToolCalls: PendingToolCall[] = []
1985
+ if (checkAbort(ctx.signal)) return
1986
+ const generateContext: AgentGenerateContext<M, D, T, X> = {
1987
+ ...baseContext,
1988
+ messages: generationMessages,
1989
+ modelMessages,
1990
+ }
1991
+ const custom = options.generate !== undefined
1992
+ const source =
1993
+ options.generate === undefined
1994
+ ? await generateWithAISDK({
1995
+ agentName: options.contract.name,
1996
+ context: generateContext,
1997
+ messageMetadata: options.messageMetadata,
1998
+ })
1999
+ : { stream: await options.generate(generateContext) }
2000
+ const updates = consumeGeneration({
2001
+ source,
2002
+ ...(options.progress === undefined
2003
+ ? {}
2004
+ : { progress: options.progress }),
2005
+ })
2006
+
2007
+ let sequence = 0
2008
+ let finish: AgentGenerationFinish | undefined
2009
+ for await (const update of updates) {
2010
+ ctx.signal.throwIfAborted()
2011
+ if (update.type === 'finish') {
2012
+ finish = update
2013
+ continue
2014
+ }
2015
+ const progress: GenerationProgressPayload = {
2016
+ startIndex,
2017
+ chunks: update.chunks,
2018
+ }
2019
+ const lifecycle = await lifecycleEvents<D, T>({
2020
+ requestId,
2021
+ messageId: responseMessageId,
2022
+ generationId,
2023
+ sequence,
2024
+ chunks: update.chunks,
2025
+ pending: pendingToolCalls,
2026
+ tools,
2027
+ generation,
2028
+ custom,
2029
+ })
2030
+ pendingToolCalls = lifecycle.pending
2031
+ await control.append({
2032
+ ctx,
2033
+ name: `generation-progress:${sequence}`,
2034
+ events: [
2035
+ {
2036
+ type: 'ai.generation.progress',
2037
+ id: await aiEventId(
2038
+ 'a2.ai.progress',
2039
+ generationId,
2040
+ String(sequence),
2041
+ ),
2042
+ payload: progress,
2043
+ } as AppendInput<D>,
2044
+ ...lifecycle.events,
2045
+ ],
2046
+ })
2047
+ sequence += 1
2048
+ }
2049
+
2050
+ if (!finish) throw new Error('agent generation finished without output')
2051
+ const completed: GenerationCompletedPayload = {
2052
+ requestId,
2053
+ messageId: request.messageId,
2054
+ generationId,
2055
+ responseMessageId,
2056
+ ...(inputTokenEstimate === undefined ? {} : { inputTokenEstimate }),
2057
+ ...(finish.finishReason === undefined
2058
+ ? {}
2059
+ : { finishReason: finish.finishReason }),
2060
+ ...(finish.usage === undefined ? {} : { usage: finish.usage }),
2061
+ }
2062
+ const completionEvent = {
2063
+ type: 'ai.generation.completed',
2064
+ id: await aiEventId('a2.ai.completed', generationId),
2065
+ payload: completed,
2066
+ } as AppendInput<D>
2067
+ const unresolved = options.generate
2068
+ ? pendingToolCalls.filter(
2069
+ (pending) => pending.classification !== 'automatic',
2070
+ )
2071
+ : []
2072
+ if (unresolved.length > 0) {
2073
+ throw new Error(
2074
+ `generation ended with unresolved tool authorization for '${unresolved[0]!.call.toolName}'`,
2075
+ )
2076
+ }
2077
+ const pendingEvents = await Promise.all(
2078
+ pendingToolCalls
2079
+ .filter(
2080
+ (pending) =>
2081
+ pending.classification === 'automatic' ||
2082
+ (options.generate === undefined &&
2083
+ pending.classification === 'unknown'),
2084
+ )
2085
+ .map((pending) => toolCalledEvent<D>(pending.call)),
2086
+ )
2087
+ if (finish.finishReason === 'tool-calls') {
2088
+ return [...pendingEvents, completionEvent]
2089
+ }
2090
+ forgetPrompt(ctx.event.sessionId, generationId)
2091
+ return [
2092
+ ...pendingEvents,
2093
+ completionEvent,
2094
+ {
2095
+ type: 'ai.message.completed',
2096
+ id: await aiEventId('a2.ai.message.completed', generationId),
2097
+ payload: {
2098
+ messageId: responseMessageId,
2099
+ } satisfies MessageCompletedPayload,
2100
+ } as AppendInput<D>,
2101
+ ]
2102
+ } catch (error) {
2103
+ if (error instanceof ControlCancelled) throw error
2104
+ if (error instanceof A2Error) {
2105
+ checkAbort(ctx.signal)
2106
+ throw error
2107
+ }
2108
+ if (checkAbort(ctx.signal)) return
2109
+ if (!generationStarted && !resolving) throw error
2110
+ // A resolver or its result failed before the model ran: record the
2111
+ // attempt and fail it once, like a provider error, rather than retrying.
2112
+ if (!generationStarted)
2113
+ await control.append({
2114
+ ctx,
2115
+ name: 'generation-start',
2116
+ events: await startEvents(startedPayload(resolvedModel)),
2117
+ })
2118
+ forgetPrompt(ctx.event.sessionId, generationId)
2119
+ const failed: GenerationFailedPayload = {
2120
+ requestId,
2121
+ messageId: request.messageId,
2122
+ generationId,
2123
+ responseMessageId,
2124
+ error: errorMessage(error),
2125
+ }
2126
+ return {
2127
+ type: 'ai.generation.failed',
2128
+ id: await aiEventId('a2.ai.failed', generationId),
2129
+ payload: failed,
2130
+ } as AppendInput<D>
2131
+ }
2132
+ }
2133
+
2134
+ return {
2135
+ 'ai.model.metadata.requested': {
2136
+ handler: async (ctx) => {
2137
+ const limits = await readModelLimits({
2138
+ modelId: ctx.event.payload.modelId,
2139
+ signal: ctx.signal,
2140
+ })
2141
+ ctx.signal.throwIfAborted()
2142
+ return {
2143
+ type: 'ai.model.metadata.resolved',
2144
+ id: await aiEventId('a2.ai.model.metadata.resolved', ctx.event.id),
2145
+ payload: { modelId: ctx.event.payload.modelId, limits },
2146
+ } as AppendInput<D>
2147
+ },
2148
+ },
2149
+ 'ai.control.requested': {
2150
+ lane: ({ event }) => {
2151
+ const command = event.payload as ControlCommand<M>
2152
+ if (event.id === '')
2153
+ throw new A2Error(
2154
+ 'INVALID_PAYLOAD',
2155
+ 'AI command IDs must be nonempty',
2156
+ )
2157
+ if (event.id?.startsWith(AI_EVENT_ID_PREFIX))
2158
+ throw new A2Error(
2159
+ 'INVALID_PAYLOAD',
2160
+ `AI command IDs cannot use the reserved '${AI_EVENT_ID_PREFIX}' prefix`,
2161
+ )
2162
+ if (
2163
+ (command.action === 'send' || command.action === 'steer') &&
2164
+ event.id !== command.message.id
2165
+ )
2166
+ throw new A2Error(
2167
+ 'INVALID_PAYLOAD',
2168
+ 'AI input event ID must equal message.id',
2169
+ )
2170
+ return 'a2.ai.control'
2171
+ },
2172
+ handler: controlHandler,
2173
+ },
2174
+ 'ai.work.reported': { lane: 'a2.ai.control', handler: controlHandler },
2175
+ 'ai.session.closed': { lane: 'a2.ai.control', handler: controlHandler },
2176
+ 'ai.generation.requested': {
2177
+ lane: 'a2.ai.model',
2178
+ abortOn: {
2179
+ 'ai.control.requested': (event, trigger) =>
2180
+ event.payload.action === 'stop' &&
2181
+ event.payload.turnId === trigger.payload.control?.turnId,
2182
+ 'ai.control.committed': (event, trigger) => {
2183
+ const active = (event.payload as ControlCommit).view
2184
+ return (
2185
+ active?.turnId !== trigger.payload.control?.turnId ||
2186
+ active?.version !== trigger.payload.control?.version
2187
+ )
2188
+ },
2189
+ 'ai.session.closed': true,
2190
+ },
2191
+ handler: (ctx) =>
2192
+ telemetry.span(
2193
+ 'a2.ai.generate',
2194
+ {
2195
+ 'a2.contract': options.contract.name,
2196
+ 'a2.session_id': ctx.event.sessionId,
2197
+ 'a2.ai.generate.request_id': ctx.event.id,
2198
+ 'a2.ai.generate.attempt': ctx.attempt,
2199
+ },
2200
+ async () => {
2201
+ try {
2202
+ return control.settled({
2203
+ ctx,
2204
+ events: await generationHandler(ctx),
2205
+ })
2206
+ } catch (error) {
2207
+ if (control.cancelled({ error, signal: ctx.signal }))
2208
+ return control.settled({ ctx, events: undefined })
2209
+ throw error
2210
+ }
2211
+ },
2212
+ ),
2213
+ },
2214
+ 'ai.tool.execution.requested': {
2215
+ ...(options.toolLane === undefined ? {} : { lane: options.toolLane }),
2216
+ abortOn: {
2217
+ 'ai.control.requested': (event, trigger) =>
2218
+ event.payload.action === 'stop' &&
2219
+ event.payload.turnId === trigger.payload.turnId,
2220
+ 'ai.generation.failed': (event, trigger) =>
2221
+ event.payload.generationId === trigger.payload.call.generationId,
2222
+ 'ai.control.committed': (event, trigger) => {
2223
+ const active = (event.payload as ControlCommit).view
2224
+ return (
2225
+ active?.turnId !== trigger.payload.turnId ||
2226
+ active?.version !== trigger.payload.version
2227
+ )
2228
+ },
2229
+ 'ai.session.closed': true,
2230
+ },
2231
+ handler: (ctx) =>
2232
+ telemetry.span(
2233
+ 'a2.ai.tool',
2234
+ {
2235
+ 'a2.contract': options.contract.name,
2236
+ 'a2.session_id': ctx.event.sessionId,
2237
+ 'a2.ai.tool.request_id': ctx.event.payload.call.requestId,
2238
+ 'a2.ai.tool.call_id': ctx.event.payload.call.toolCallId,
2239
+ 'a2.ai.tool.name': ctx.event.payload.call.toolName,
2240
+ 'a2.ai.tool.attempt': ctx.attempt,
2241
+ },
2242
+ () => handleToolExecution(ctx),
2243
+ ),
2244
+ },
2245
+ } as NonNullable<ServerOptions<D>['handlers']>
2246
+ }
2247
+
2248
+ /**
2249
+ * Serve an AI agent: the built-in protocol, its handlers, and A2's store,
2250
+ * scheduler, and telemetry, as one ordinary A2 server. Pass `name` for the
2251
+ * built-in protocol, or a shared `agentContract()` when clients fold
2252
+ * application state.
2253
+ */
2254
+ export function agent<
2255
+ T extends ToolSet = ToolSet,
2256
+ M extends UIMessage = UIMessage,
2257
+ E extends EventDefs = Record<never, never>,
2258
+ >(
2259
+ options: AgentOptions<M, E, T>,
2260
+ ): Agent<AIEventDefs<M> & E, AIState<M>, typeof AI_STATE_REDUCER>
2261
+ export function agent<
2262
+ D extends AIEventDefs<AIMessageOf<D>> & EventDefs,
2263
+ S,
2264
+ T extends ToolSet = ToolSet,
2265
+ Name extends string = string,
2266
+ X extends object = Record<never, never>,
2267
+ >(options: AgentFromContractOptions<D, S, T, Name, X>): Agent<D, S, Name>
2268
+ export function agent(
2269
+ options:
2270
+ | AgentOptions<UIMessage, EventDefs, ToolSet>
2271
+ | AgentFromContractOptions<AIEventDefs & EventDefs, AIState, ToolSet>,
2272
+ ): Agent<AIEventDefs & EventDefs, AIState> {
2273
+ type D = AIEventDefs & EventDefs
2274
+ const {
2275
+ model,
2276
+ tools,
2277
+ instructions,
2278
+ generation,
2279
+ maxSteps,
2280
+ messageMetadata,
2281
+ generate,
2282
+ compaction,
2283
+ progress,
2284
+ toolLane,
2285
+ handlers,
2286
+ ...rest
2287
+ } = options
2288
+ let contract: Contract<D, Record<never, never>, Reducer<D, AIState>>
2289
+ let serverOptions: Omit<
2290
+ ServerOptions<D, Record<never, never>, Reducer<D, AIState>>,
2291
+ 'contract' | 'handlers'
2292
+ >
2293
+ if ('contract' in rest) {
2294
+ const { contract: shared, ...remaining } = rest
2295
+ if ('name' in remaining)
2296
+ throw new TypeError('agent takes either name or contract, not both')
2297
+ if (!isAgentContract(shared))
2298
+ throw new TypeError('agent({ contract }) takes an agentContract()')
2299
+ contract = shared
2300
+ serverOptions = remaining
2301
+ } else {
2302
+ const { name, events, messageSchema, ...remaining } = rest
2303
+ contract = agentContract({
2304
+ name,
2305
+ ...(events === undefined ? {} : { events }),
2306
+ ...(messageSchema === undefined ? {} : { messageSchema }),
2307
+ })
2308
+ serverOptions = remaining
2309
+ }
2310
+ const builtIns = createHandlers<UIMessage, D>({
2311
+ contract,
2312
+ model,
2313
+ ...(serverOptions.telemetry === undefined
2314
+ ? {}
2315
+ : { telemetry: serverOptions.telemetry }),
2316
+ ...(tools === undefined ? {} : { tools }),
2317
+ ...(instructions === undefined ? {} : { instructions }),
2318
+ ...(generation === undefined ? {} : { generation }),
2319
+ ...(maxSteps === undefined ? {} : { maxSteps }),
2320
+ ...(messageMetadata === undefined ? {} : { messageMetadata }),
2321
+ ...(generate === undefined ? {} : { generate }),
2322
+ ...(compaction === undefined ? {} : { compaction }),
2323
+ ...(progress === undefined ? {} : { progress }),
2324
+ ...(toolLane === undefined ? {} : { toolLane }),
2325
+ })
2326
+ for (const type of Object.keys(handlers ?? {})) {
2327
+ if (Object.hasOwn(builtIns, type))
2328
+ throw new TypeError(`agent handles '${type}' itself`)
2329
+ }
2330
+ const server = createServer({
2331
+ ...serverOptions,
2332
+ contract,
2333
+ handlers: { ...builtIns, ...handlers },
2334
+ })
2335
+ const applicationEvents = 'contract' in rest
2336
+ setServerFetchHooks(server, {
2337
+ validateIngress: (context) =>
2338
+ validateAgentIngress({ context, applicationEvents }),
2339
+ })
2340
+ return server
2341
+ }
2342
+
2343
+ export type { Instructions, LanguageModel, ToolSet }