experimental-a2 0.19.0 → 0.21.1

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 (320) hide show
  1. package/CHANGELOG.md +50 -0
  2. package/README.md +1 -1
  3. package/dist/{actor-C7b-1arn.d.ts → actor-HnGmrOZZ.d.ts} +5 -5
  4. package/dist/{actor-C7b-1arn.d.ts.map → actor-HnGmrOZZ.d.ts.map} +1 -1
  5. package/dist/actor-client.d.ts +1 -1
  6. package/dist/actor-client.js +2 -2
  7. package/dist/actor-react.d.ts +3 -3
  8. package/dist/actor-react.js +5 -5
  9. package/dist/{actor-shared-t7w8xIKP.js → actor-shared-DFET5q16.js} +3 -3
  10. package/dist/{actor-shared-t7w8xIKP.js.map → actor-shared-DFET5q16.js.map} +1 -1
  11. package/dist/actor.d.ts +1 -1
  12. package/dist/actor.js +54 -43
  13. package/dist/actor.js.map +1 -1
  14. package/dist/{ai-Cz3hqVhk.d.ts → ai-CvEnqVfo.d.ts} +8 -8
  15. package/dist/{ai-Cz3hqVhk.d.ts.map → ai-CvEnqVfo.d.ts.map} +1 -1
  16. package/dist/{ai-mMiWul1T.js → ai-CzE672Q3.js} +9 -9
  17. package/dist/ai-CzE672Q3.js.map +1 -0
  18. package/dist/ai-client.d.ts +2 -2
  19. package/dist/ai-client.js +2 -2
  20. package/dist/ai-fx.d.ts +1 -1
  21. package/dist/ai-fx.d.ts.map +1 -1
  22. package/dist/ai-fx.js +36 -9
  23. package/dist/ai-fx.js.map +1 -1
  24. package/dist/{ai-harness-CV6rXbB3.d.ts → ai-harness-Bfk-Lycn.d.ts} +2 -2
  25. package/dist/ai-harness-Bfk-Lycn.d.ts.map +1 -0
  26. package/dist/{ai-harness-BuuC-gJ3.js → ai-harness-dQeJaR1t.js} +17 -4
  27. package/dist/ai-harness-dQeJaR1t.js.map +1 -0
  28. package/dist/ai-pi.d.ts +1 -1
  29. package/dist/ai-pi.js +4 -4
  30. package/dist/ai-pi.js.map +1 -1
  31. package/dist/ai-react.d.ts +5 -5
  32. package/dist/ai-react.js +2 -2
  33. package/dist/ai-server.d.ts +14 -17
  34. package/dist/ai-server.d.ts.map +1 -1
  35. package/dist/ai-server.js +84 -75
  36. package/dist/ai-server.js.map +1 -1
  37. package/dist/ai.d.ts +2 -2
  38. package/dist/ai.js +2 -2
  39. package/dist/{client-BhHLoNKF.js → client-9IeTsiAs.js} +5 -5
  40. package/dist/{client-BhHLoNKF.js.map → client-9IeTsiAs.js.map} +1 -1
  41. package/dist/{client-h_DKFAjj.d.ts → client-CUMbS_38.d.ts} +3 -3
  42. package/dist/{client-h_DKFAjj.d.ts.map → client-CUMbS_38.d.ts.map} +1 -1
  43. package/dist/client.d.ts +3 -3
  44. package/dist/client.js +1 -1
  45. package/dist/{contract-BYKASlCV.js → contract-CWukt01d.js} +2 -2
  46. package/dist/contract-CWukt01d.js.map +1 -0
  47. package/dist/index.d.ts +5 -5
  48. package/dist/index.js +3 -3
  49. package/dist/{initial-snapshot-Bx-c0pDP.d.ts → initial-snapshot-D3Lve43o.d.ts} +4 -4
  50. package/dist/initial-snapshot-D3Lve43o.d.ts.map +1 -0
  51. package/dist/{internal-CscSY7K_.js → internal-Bc4PvfHk.js} +10 -3
  52. package/dist/internal-Bc4PvfHk.js.map +1 -0
  53. package/dist/otel.d.ts +1 -1
  54. package/dist/otel.js +1 -1
  55. package/dist/otel.js.map +1 -1
  56. package/dist/{react-BBXgHWFc.d.ts → react-BoHvC3v9.d.ts} +4 -4
  57. package/dist/{react-BBXgHWFc.d.ts.map → react-BoHvC3v9.d.ts.map} +1 -1
  58. package/dist/react.d.ts +1 -1
  59. package/dist/react.js +1 -1
  60. package/dist/{reducer-CqA59Q6U.d.ts → reducer-CYPHsSmd.d.ts} +2 -2
  61. package/dist/{reducer-CqA59Q6U.d.ts.map → reducer-CYPHsSmd.d.ts.map} +1 -1
  62. package/dist/{reducer-BAXc8JyF.js → reducer-CjmvZSAb.js} +2 -2
  63. package/dist/{reducer-BAXc8JyF.js.map → reducer-CjmvZSAb.js.map} +1 -1
  64. package/dist/runtime-consumer-BWmht60z.js +12 -0
  65. package/dist/runtime-consumer-BWmht60z.js.map +1 -0
  66. package/dist/scheduler-qstash.d.ts +2 -2
  67. package/dist/scheduler-qstash.d.ts.map +1 -1
  68. package/dist/scheduler-qstash.js +7 -7
  69. package/dist/scheduler-qstash.js.map +1 -1
  70. package/dist/scheduler-task-X1SToDjP.js.map +1 -1
  71. package/dist/scheduler-vercel.d.ts +2 -2
  72. package/dist/scheduler-vercel.d.ts.map +1 -1
  73. package/dist/scheduler-vercel.js +5 -4
  74. package/dist/scheduler-vercel.js.map +1 -1
  75. package/dist/{server-DYAATFuT.d.ts → server-BshFZ0eL.d.ts} +108 -41
  76. package/dist/server-BshFZ0eL.d.ts.map +1 -0
  77. package/dist/{server-D-Y4pqwu.js → server-DvIB8kj5.js} +423 -135
  78. package/dist/server-DvIB8kj5.js.map +1 -0
  79. package/dist/server.browser.js +1 -1
  80. package/dist/server.browser.js.map +1 -1
  81. package/dist/server.d.ts +4 -4
  82. package/dist/server.js +2 -2
  83. package/dist/{session-clear-C-QwTXV1.d.ts → session-clear-Bk_29eEw.d.ts} +2 -2
  84. package/dist/session-clear-Bk_29eEw.d.ts.map +1 -0
  85. package/dist/{session-clear-DGI4jTQ5.js → session-clear-wHd5GCpZ.js} +5 -4
  86. package/dist/session-clear-wHd5GCpZ.js.map +1 -0
  87. package/dist/store-N8PXxDAS.js.map +1 -1
  88. package/dist/store-Z2qjw50n.d.ts +523 -0
  89. package/dist/store-Z2qjw50n.d.ts.map +1 -0
  90. package/dist/store-codec-DwrWlgQx.js +65 -0
  91. package/dist/store-codec-DwrWlgQx.js.map +1 -0
  92. package/dist/store-memory.d.ts +1 -1
  93. package/dist/store-memory.d.ts.map +1 -1
  94. package/dist/store-memory.js +393 -210
  95. package/dist/store-memory.js.map +1 -1
  96. package/dist/{store-polling-7WJV0YWP.js → store-polling-DYNyxCev.js} +2 -2
  97. package/dist/{store-polling-7WJV0YWP.js.map → store-polling-DYNyxCev.js.map} +1 -1
  98. package/dist/store-postgres.d.ts +1 -1
  99. package/dist/store-postgres.d.ts.map +1 -1
  100. package/dist/store-postgres.js +1595 -673
  101. package/dist/store-postgres.js.map +1 -1
  102. package/dist/{store-presence-polling-B1ixeqIS.js → store-presence-polling-CzQ0wdtM.js} +3 -3
  103. package/dist/{store-presence-polling-B1ixeqIS.js.map → store-presence-polling-CzQ0wdtM.js.map} +1 -1
  104. package/dist/store-redis-http.d.ts +1 -1
  105. package/dist/store-redis-http.js +4 -4
  106. package/dist/store-redis-notify-IiVdlYHc.js +2919 -0
  107. package/dist/store-redis-notify-IiVdlYHc.js.map +1 -0
  108. package/dist/store-redis.d.ts +1 -1
  109. package/dist/store-redis.js +2 -2
  110. package/dist/store-sqlite.d.ts +1 -1
  111. package/dist/store-sqlite.d.ts.map +1 -1
  112. package/dist/store-sqlite.js +517 -287
  113. package/dist/store-sqlite.js.map +1 -1
  114. package/dist/{telemetry-Cyfqo6Qt.d.ts → telemetry-CWwsh-Gx.d.ts} +3 -3
  115. package/dist/telemetry-CWwsh-Gx.d.ts.map +1 -0
  116. package/dist/{wire-BKxH1vbP.js → wire-CZGqNZRx.js} +5 -4
  117. package/dist/wire-CZGqNZRx.js.map +1 -0
  118. package/docs/01-quickstart.mdx +28 -19
  119. package/docs/02-examples.mdx +1 -1
  120. package/docs/actors/01-introduction.mdx +9 -1
  121. package/docs/actors/02-concurrency.mdx +1 -1
  122. package/docs/actors/meta.ts +1 -1
  123. package/docs/agents/01-introduction.mdx +1 -1
  124. package/docs/agents/02-tools.mdx +13 -12
  125. package/docs/agents/03-conversations.mdx +31 -21
  126. package/docs/agents/04-harnesses.mdx +7 -3
  127. package/docs/agents/meta.ts +1 -1
  128. package/docs/client/03-presence.mdx +18 -14
  129. package/docs/client/meta.ts +1 -1
  130. package/docs/contract/01-events.mdx +2 -2
  131. package/docs/contract/02-state.mdx +18 -14
  132. package/docs/contract/03-sessions.mdx +32 -25
  133. package/docs/index.mdx +31 -26
  134. package/docs/reference/01-core.mdx +12 -12
  135. package/docs/reference/02-server.mdx +247 -192
  136. package/docs/reference/04-ai.mdx +29 -20
  137. package/docs/reference/05-telemetry.mdx +146 -0
  138. package/docs/reference/06-errors.mdx +24 -18
  139. package/docs/reference/meta.ts +1 -1
  140. package/docs/schedulers/01-overview.mdx +112 -0
  141. package/docs/schedulers/02-vercel-queues.mdx +87 -0
  142. package/docs/schedulers/03-qstash.mdx +123 -0
  143. package/docs/schedulers/meta.ts +1 -0
  144. package/docs/server/01-consumers.mdx +415 -0
  145. package/docs/server/02-timers.mdx +33 -23
  146. package/docs/server/03-cancellation.mdx +52 -44
  147. package/docs/server/04-routes.mdx +4 -1
  148. package/docs/server/05-sql-projections.mdx +391 -0
  149. package/docs/server/06-guarantees.mdx +60 -50
  150. package/docs/stores/01-overview.mdx +37 -0
  151. package/docs/stores/02-redis.mdx +57 -0
  152. package/docs/stores/03-redis-http.mdx +56 -0
  153. package/docs/stores/04-postgres.mdx +172 -0
  154. package/docs/stores/05-sqlite.mdx +36 -0
  155. package/docs/stores/06-memory.mdx +22 -0
  156. package/docs/stores/meta.ts +1 -0
  157. package/docs/upgrades/01-overview.mdx +30 -0
  158. package/docs/upgrades/79-0-21.mdx +165 -0
  159. package/docs/upgrades/80-0-20.mdx +43 -0
  160. package/docs/upgrades/81-0-19.mdx +33 -0
  161. package/docs/upgrades/82-0-18.mdx +33 -0
  162. package/docs/upgrades/83-0-17.mdx +69 -0
  163. package/docs/upgrades/meta.ts +3 -0
  164. package/examples/playground/DEPLOY.md +48 -11
  165. package/examples/playground/app/a0/apps/[appId]/route.ts +6 -10
  166. package/examples/playground/app/a0/chat.test.tsx +421 -330
  167. package/examples/playground/app/a0/sandbox-log.ts +17 -15
  168. package/examples/playground/app/a0/server.ts +4 -4
  169. package/examples/playground/app/a0/shell.tsx +1 -1
  170. package/examples/playground/app/a0/uploads/route.test.ts +59 -0
  171. package/examples/playground/app/a2slack/a2slack.css +1067 -0
  172. package/examples/playground/app/a2slack/c/[channel]/page.tsx +20 -0
  173. package/examples/playground/app/a2slack/consumers/agent.ts +63 -0
  174. package/examples/playground/app/a2slack/consumers/feed.ts +164 -0
  175. package/examples/playground/app/a2slack/consumers/notify.ts +116 -0
  176. package/examples/playground/app/a2slack/consumers/search.ts +37 -0
  177. package/examples/playground/app/a2slack/consumers/triage.ts +31 -0
  178. package/examples/playground/app/a2slack/consumers/welcome.ts +25 -0
  179. package/examples/playground/app/a2slack/db/a2.ts +206 -0
  180. package/examples/playground/app/a2slack/db/index.ts +73 -0
  181. package/examples/playground/app/a2slack/db/migrations/0000_extensions.sql +4 -0
  182. package/examples/playground/app/a2slack/db/migrations/0001_workspace.sql +74 -0
  183. package/examples/playground/app/a2slack/db/migrations/0002_a2.sql +215 -0
  184. package/examples/playground/app/a2slack/db/migrations/0003_a2_batch.sql +67 -0
  185. package/examples/playground/app/a2slack/db/migrations/meta/0000_snapshot.json +18 -0
  186. package/examples/playground/app/a2slack/db/migrations/meta/0001_snapshot.json +500 -0
  187. package/examples/playground/app/a2slack/db/migrations/meta/0002_snapshot.json +1159 -0
  188. package/examples/playground/app/a2slack/db/migrations/meta/0003_snapshot.json +1159 -0
  189. package/examples/playground/app/a2slack/db/migrations/meta/_journal.json +34 -0
  190. package/examples/playground/app/a2slack/db/queries.ts +263 -0
  191. package/examples/playground/app/a2slack/db/schema.ts +145 -0
  192. package/examples/playground/app/a2slack/drills/a2slack.drill.ts +708 -0
  193. package/examples/playground/app/a2slack/drills/harness.ts +244 -0
  194. package/examples/playground/app/a2slack/events/route.ts +36 -0
  195. package/examples/playground/app/a2slack/identity.ts +44 -0
  196. package/examples/playground/app/a2slack/inbox/route.ts +8 -0
  197. package/examples/playground/app/a2slack/lag/data/route.ts +8 -0
  198. package/examples/playground/app/a2slack/lag/page.tsx +101 -0
  199. package/examples/playground/app/a2slack/lag.ts +73 -0
  200. package/examples/playground/app/a2slack/layout.tsx +43 -0
  201. package/examples/playground/app/a2slack/messages/route.ts +17 -0
  202. package/examples/playground/app/a2slack/model.test.ts +208 -0
  203. package/examples/playground/app/a2slack/model.ts +369 -0
  204. package/examples/playground/app/a2slack/models.test.ts +72 -0
  205. package/examples/playground/app/a2slack/page.tsx +10 -0
  206. package/examples/playground/app/a2slack/scheduler/route.ts +5 -0
  207. package/examples/playground/app/a2slack/search/route.ts +40 -0
  208. package/examples/playground/app/a2slack/search-index.ts +108 -0
  209. package/examples/playground/app/a2slack/server.test.ts +315 -0
  210. package/examples/playground/app/a2slack/server.ts +55 -0
  211. package/examples/playground/app/a2slack/sessions.ts +50 -0
  212. package/examples/playground/app/a2slack/thread/route.ts +8 -0
  213. package/examples/playground/app/a2slack/ui/ChannelView.tsx +87 -0
  214. package/examples/playground/app/a2slack/ui/Composer.tsx +117 -0
  215. package/examples/playground/app/a2slack/ui/DotsShadow.tsx +42 -0
  216. package/examples/playground/app/a2slack/ui/Feed.tsx +157 -0
  217. package/examples/playground/app/a2slack/ui/InboxContext.tsx +63 -0
  218. package/examples/playground/app/a2slack/ui/InboxView.tsx +82 -0
  219. package/examples/playground/app/a2slack/ui/Menu.tsx +128 -0
  220. package/examples/playground/app/a2slack/ui/MessageRow.tsx +314 -0
  221. package/examples/playground/app/a2slack/ui/NameGate.tsx +35 -0
  222. package/examples/playground/app/a2slack/ui/SearchContext.tsx +146 -0
  223. package/examples/playground/app/a2slack/ui/SearchResults.tsx +97 -0
  224. package/examples/playground/app/a2slack/ui/Shell.tsx +34 -0
  225. package/examples/playground/app/a2slack/ui/Sidebar.tsx +199 -0
  226. package/examples/playground/app/a2slack/ui/ThreadPanel.tsx +134 -0
  227. package/examples/playground/app/a2slack/ui/useInbox.ts +71 -0
  228. package/examples/playground/app/agent/[agentId]/agent-queue.test.tsx +116 -90
  229. package/examples/playground/app/agent/[agentId]/compaction-panel.test.tsx +107 -76
  230. package/examples/playground/app/agent/ingress.test.ts +41 -17
  231. package/examples/playground/app/agent/model-input.test.tsx +92 -43
  232. package/examples/playground/app/agent/server.ts +3 -3
  233. package/examples/playground/app/chat/[chatId]/agent-stream-drawer.test.tsx +2 -2
  234. package/examples/playground/app/chat/[chatId]/chat-client.tsx +1 -1
  235. package/examples/playground/app/chat/server.test.ts +38 -65
  236. package/examples/playground/app/chat/server.ts +45 -35
  237. package/examples/playground/app/components/activity-history.test.tsx +60 -18
  238. package/examples/playground/app/components/demo-icons.tsx +7 -0
  239. package/examples/playground/app/components/demos.tsx +7 -0
  240. package/examples/playground/app/components/fast-link.test.tsx +3 -5
  241. package/examples/playground/app/initial-state.test.tsx +86 -45
  242. package/examples/playground/app/orders/[orderId]/order-client.tsx +2 -2
  243. package/examples/playground/app/orders/page.tsx +1 -1
  244. package/examples/playground/app/orders/server.ts +49 -30
  245. package/examples/playground/app/recovery/[recoveryId]/recovery-client.tsx +14 -15
  246. package/examples/playground/app/recovery/page.tsx +4 -4
  247. package/examples/playground/app/recovery/server.ts +23 -32
  248. package/examples/playground/drizzle.config.ts +14 -0
  249. package/examples/playground/next.config.ts +6 -1
  250. package/examples/playground/package.json +13 -1
  251. package/examples/playground/scripts/migrate.mjs +31 -0
  252. package/examples/playground/test/agent.ts +35 -0
  253. package/examples/playground/test/model.ts +149 -0
  254. package/examples/playground/test/serve.ts +184 -0
  255. package/examples/playground/test/server-only.ts +1 -0
  256. package/examples/playground/vercel.json +9 -0
  257. package/examples/playground/vitest.config.ts +27 -0
  258. package/examples/playground/vitest.drills.config.ts +32 -0
  259. package/package.json +7 -2
  260. package/sql/postgres/drizzle-schema.ts +206 -0
  261. package/sql/postgres/notifications.sql +24 -0
  262. package/sql/postgres/schema.sql +276 -0
  263. package/src/actor.ts +101 -71
  264. package/src/ai-agent.ts +154 -115
  265. package/src/ai-control-server.ts +8 -8
  266. package/src/ai-fx.ts +41 -9
  267. package/src/ai-harness.ts +28 -2
  268. package/src/ai-pi.ts +2 -2
  269. package/src/ai-server.ts +2 -2
  270. package/src/ai.ts +7 -7
  271. package/src/contract.ts +1 -1
  272. package/src/deterministic-id.ts +26 -3
  273. package/src/errors.ts +9 -6
  274. package/src/idempotent-replay.ts +0 -1
  275. package/src/index.ts +1 -1
  276. package/src/internal.ts +17 -6
  277. package/src/otel.ts +1 -1
  278. package/src/postgres-notifications.ts +0 -22
  279. package/src/postgres-resources.ts +6 -9
  280. package/src/postgres-schema.ts +121 -0
  281. package/src/push-envelope.ts +20 -0
  282. package/src/runtime-consumer.ts +10 -0
  283. package/src/scheduler-qstash.ts +12 -5
  284. package/src/scheduler-task.ts +2 -2
  285. package/src/scheduler-vercel.ts +10 -4
  286. package/src/server.browser.ts +1 -1
  287. package/src/server.ts +839 -238
  288. package/src/session-socket.ts +1 -1
  289. package/src/store-codec.ts +66 -0
  290. package/src/store-memory.ts +581 -255
  291. package/src/store-postgres.ts +1984 -943
  292. package/src/store-redis-core.ts +1624 -1152
  293. package/src/store-redis-notify.ts +2 -0
  294. package/src/store-sqlite.ts +891 -486
  295. package/src/store.ts +227 -97
  296. package/src/telemetry.ts +2 -1
  297. package/src/wire.ts +1 -0
  298. package/dist/ai-harness-BuuC-gJ3.js.map +0 -1
  299. package/dist/ai-harness-CV6rXbB3.d.ts.map +0 -1
  300. package/dist/ai-mMiWul1T.js.map +0 -1
  301. package/dist/contract-BYKASlCV.js.map +0 -1
  302. package/dist/idempotent-replay-JdgMQRRm.js +0 -21
  303. package/dist/idempotent-replay-JdgMQRRm.js.map +0 -1
  304. package/dist/initial-snapshot-Bx-c0pDP.d.ts.map +0 -1
  305. package/dist/internal-CscSY7K_.js.map +0 -1
  306. package/dist/server-D-Y4pqwu.js.map +0 -1
  307. package/dist/server-DYAATFuT.d.ts.map +0 -1
  308. package/dist/session-clear-C-QwTXV1.d.ts.map +0 -1
  309. package/dist/session-clear-DGI4jTQ5.js.map +0 -1
  310. package/dist/store-CyLUy86v.d.ts +0 -404
  311. package/dist/store-CyLUy86v.d.ts.map +0 -1
  312. package/dist/store-codec-DTG0Ftek.js +0 -10
  313. package/dist/store-codec-DTG0Ftek.js.map +0 -1
  314. package/dist/store-redis-notify-CM0AfLY0.js +0 -2544
  315. package/dist/store-redis-notify-CM0AfLY0.js.map +0 -1
  316. package/dist/telemetry-Cyfqo6Qt.d.ts.map +0 -1
  317. package/dist/wire-BKxH1vbP.js.map +0 -1
  318. package/docs/reference/05-adapters.mdx +0 -351
  319. package/docs/server/01-handlers.mdx +0 -200
  320. package/docs/server/05-production.mdx +0 -323
package/src/ai-agent.ts CHANGED
@@ -1,6 +1,6 @@
1
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.
2
+ * The server-only implementation of the experimental-a2/ai protocol: its
3
+ * consumers and `agent()`, its assembly as an ordinary A2 server.
4
4
  */
5
5
 
6
6
  // oxlint-disable no-await-in-loop -- stream chunks and durable appends are
@@ -20,7 +20,6 @@ import type { Reducer } from './reducer.ts'
20
20
  import { aiEventId, AI_EVENT_ID_PREFIX } from './ai-id.ts'
21
21
  import type { ControlCommit, ControlCommand } from './ai-control.ts'
22
22
  import { AsyncLocalStorage } from 'node:async_hooks'
23
- import { createHash } from 'node:crypto'
24
23
  import { convertToModelMessages, isToolUIPart } from 'ai'
25
24
  import type {
26
25
  Experimental_SandboxSession,
@@ -35,6 +34,7 @@ import type {
35
34
  } from 'ai'
36
35
  import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
37
36
  import {
37
+ gatewaySessionAffinity,
38
38
  harnessDriver,
39
39
  harnessLog,
40
40
  harnessStates,
@@ -79,10 +79,12 @@ import {
79
79
  installAmbientToolScopeStorage,
80
80
  type AmbientToolScope,
81
81
  } from './internal.ts'
82
+ import { runtimeConsumer } from './runtime-consumer.ts'
82
83
  import {
83
84
  createServer,
84
85
  type A2Server,
85
- type HandlerContext,
86
+ type Consumer,
87
+ type ConsumerContext,
86
88
  type ServerOptions,
87
89
  } from './server.ts'
88
90
  import {
@@ -191,7 +193,7 @@ export type AgentResolverContext<
191
193
  * in the checkpoint the generation already reads, so it costs no I/O.
192
194
  */
193
195
  state: X
194
- session: Pick<HandlerContext<D>['session'], 'state'>
196
+ session: Pick<ConsumerContext<D>['session'], 'state'>
195
197
  signal: AbortSignal
196
198
  }
197
199
 
@@ -239,7 +241,7 @@ export type CompactionOptions<
239
241
 
240
242
  export type AgentGenerationSettings<T extends ToolSet> = AISDKStepSettings<T>
241
243
 
242
- export type HandlersOptions<
244
+ export type ConsumersOptions<
243
245
  M extends UIMessage,
244
246
  D extends AIEventDefs<M> & EventDefs,
245
247
  T extends ToolSet = ToolSet,
@@ -344,27 +346,24 @@ type AgentServerOptions<
344
346
  T extends ToolSet,
345
347
  X extends object = Record<never, never>,
346
348
  H extends string = string,
349
+ C extends readonly unknown[] = readonly Consumer<D>[],
347
350
  > = AgentModelOptions<M, D, T, X, H> &
348
- Omit<
349
- ServerOptions<D, Record<never, never>, Reducer<D, S>>,
350
- 'contract' | 'handlers'
351
- > & {
352
- /** Handlers for application events. */
353
- handlers?: ServerOptions<D, Record<never, never>, Reducer<D, S>>['handlers']
354
- }
351
+ Omit<ServerOptions<D, Record<never, never>, Reducer<D, S>, C>, 'contract'>
355
352
 
356
353
  export type AgentOptions<
357
354
  M extends UIMessage,
358
355
  E extends EventDefs,
359
356
  T extends ToolSet,
360
357
  H extends string = string,
358
+ C extends readonly unknown[] = readonly Consumer<AIEventDefs<M> & E>[],
361
359
  > = AgentServerOptions<
362
360
  NoInfer<M>,
363
361
  AIEventDefs<NoInfer<M>> & E,
364
362
  AIState<NoInfer<M>>,
365
363
  T,
366
364
  Record<never, never>,
367
- H
365
+ H,
366
+ C
368
367
  > & {
369
368
  /** Identity: prefixes storage keys and queue messages. */
370
369
  name: string
@@ -384,7 +383,8 @@ export type AgentFromContractOptions<
384
383
  Name extends string = string,
385
384
  X extends object = Record<never, never>,
386
385
  H extends string = string,
387
- > = AgentServerOptions<AIMessageOf<D>, D, S, T, X, H> & {
386
+ C extends readonly unknown[] = readonly Consumer<D>[],
387
+ > = AgentServerOptions<AIMessageOf<D>, D, S, T, X, H, C> & {
388
388
  /**
389
389
  * A shared `agentContract()`. Its application state reaches every resolver
390
390
  * as `state`; a fully custom contract is core A2, served by `createServer()`.
@@ -480,8 +480,8 @@ const toolScopeStorage = installAmbientToolScopeStorage(
480
480
  () => new AsyncLocalStorage<AmbientToolScope>(),
481
481
  )
482
482
 
483
- /** Run `fn` with `ctx` as the ambient tool handler context, for direct tool tests. */
484
- export function runWithHandlerContext<D extends AgentToolEventDefs, R>(
483
+ /** Run `fn` with `ctx` as the ambient tool consumer context, for direct tool tests. */
484
+ export function runWithConsumerContext<D extends AgentToolEventDefs, R>(
485
485
  ctx: AgentToolContext<D>,
486
486
  fn: () => R,
487
487
  ): R {
@@ -627,9 +627,10 @@ async function generateWithAISDK<
627
627
  ? headers
628
628
  : {
629
629
  ...headers,
630
- 'x-session-affinity': createHash('sha256')
631
- .update(JSON.stringify([agentName, context.event.sessionId]))
632
- .digest('hex'),
630
+ 'x-session-affinity': gatewaySessionAffinity({
631
+ agentName,
632
+ sessionId: context.event.sessionId,
633
+ }),
633
634
  },
634
635
  providerOptions: {
635
636
  ...context.generation.providerOptions,
@@ -682,7 +683,7 @@ async function* consumeGeneration(options: {
682
683
  const reader = options.source.stream.getReader()
683
684
  // The next read starts before a batch is yielded, and the consumer may
684
685
  // take a while to ask again: a stream that fails meanwhile must not leave
685
- // that read rejected with no handler, which ends the process.
686
+ // that read rejected with nothing to catch it, which ends the process.
686
687
  const readAhead = (): Promise<StreamRead<UIMessageChunk>> => {
687
688
  const read = reader.read()
688
689
  read.catch(() => {})
@@ -1231,14 +1232,12 @@ const assertNoLoopOptions = (options: { harness?: { type: string } }): void => {
1231
1232
  }
1232
1233
  }
1233
1234
 
1234
- export function createHandlers<
1235
+ export function createConsumers<
1235
1236
  M extends UIMessage,
1236
1237
  D extends AIEventDefs<M> & EventDefs,
1237
1238
  T extends ToolSet = ToolSet,
1238
1239
  X extends object = Record<never, never>,
1239
- >(
1240
- options: HandlersOptions<M, D, T, X>,
1241
- ): NonNullable<ServerOptions<D>['handlers']> {
1240
+ >(options: ConsumersOptions<M, D, T, X>): Consumer<D>[] {
1242
1241
  if (options.model === undefined) {
1243
1242
  throw new TypeError('agent requires model')
1244
1243
  }
@@ -1348,7 +1347,7 @@ export function createHandlers<
1348
1347
  contract: options.contract,
1349
1348
  })
1350
1349
  const readContext = async (
1351
- session: HandlerContext<D>['session'],
1350
+ session: ConsumerContext<D>['session'],
1352
1351
  through: number | 'latest',
1353
1352
  ) => {
1354
1353
  const snapshot = await session.state(contextReducer, { through })
@@ -1389,7 +1388,7 @@ export function createHandlers<
1389
1388
  ) => unknown
1390
1389
  }
1391
1390
 
1392
- type ToolHandlerContext = HandlerContext<D, 'ai.tool.execution.requested'>
1391
+ type ToolConsumerContext = ConsumerContext<D, 'ai.tool.execution.requested'>
1393
1392
 
1394
1393
  const resultEvent = async (
1395
1394
  call: ToolCalledPayload,
@@ -1434,7 +1433,7 @@ export function createHandlers<
1434
1433
  }) as AppendInput<D>
1435
1434
 
1436
1435
  const promptMessages = (input: {
1437
- ctx: ToolHandlerContext
1436
+ ctx: ToolConsumerContext
1438
1437
  call: ToolCalledPayload
1439
1438
  generation: GenerationStartedPayload
1440
1439
  }): Promise<ModelMessage[]> => {
@@ -1484,7 +1483,7 @@ export function createHandlers<
1484
1483
  }
1485
1484
 
1486
1485
  const toolExecutionFailure = async (
1487
- ctx: ToolHandlerContext,
1486
+ ctx: ToolConsumerContext,
1488
1487
  call: ToolCalledPayload,
1489
1488
  error: unknown,
1490
1489
  ): Promise<AppendInput<D> | void> => {
@@ -1497,7 +1496,7 @@ export function createHandlers<
1497
1496
  }
1498
1497
 
1499
1498
  const runToolExecution = async (
1500
- ctx: ToolHandlerContext,
1499
+ ctx: ToolConsumerContext,
1501
1500
  call: ToolCalledPayload,
1502
1501
  execute: NonNullable<RuntimeTool['execute']>,
1503
1502
  messages: ModelMessage[],
@@ -1574,7 +1573,7 @@ export function createHandlers<
1574
1573
  }
1575
1574
 
1576
1575
  const executeTool = async (input: {
1577
- ctx: ToolHandlerContext
1576
+ ctx: ToolConsumerContext
1578
1577
  call: ToolCalledPayload
1579
1578
  generation: GenerationStartedPayload
1580
1579
  }): Promise<AppendInput<D> | void> => {
@@ -1628,7 +1627,7 @@ export function createHandlers<
1628
1627
  }
1629
1628
 
1630
1629
  const handleToolExecution = async (
1631
- ctx: ToolHandlerContext,
1630
+ ctx: ToolConsumerContext,
1632
1631
  ): Promise<AppendInput<D>> => {
1633
1632
  try {
1634
1633
  const state = (
@@ -1665,7 +1664,7 @@ export function createHandlers<
1665
1664
  }
1666
1665
 
1667
1666
  const generationHandler = async (
1668
- ctx: HandlerContext<D, 'ai.generation.requested'>,
1667
+ ctx: ConsumerContext<D, 'ai.generation.requested'>,
1669
1668
  ): Promise<void | AppendInput<D> | readonly AppendInput<D>[]> => {
1670
1669
  if (checkAbort(ctx.signal)) return
1671
1670
  const requestId = ctx.event.id
@@ -2314,23 +2313,41 @@ export function createHandlers<
2314
2313
  }
2315
2314
  }
2316
2315
 
2317
- return {
2318
- 'ai.model.metadata.requested': {
2319
- handler: async (ctx) => {
2320
- const limits = await readModelLimits({
2321
- modelId: ctx.event.payload.modelId,
2322
- signal: ctx.signal,
2323
- })
2324
- ctx.signal.throwIfAborted()
2325
- return {
2326
- type: 'ai.model.metadata.resolved',
2327
- id: await aiEventId('a2.ai.model.metadata.resolved', ctx.event.id),
2328
- payload: { modelId: ctx.event.payload.modelId, limits },
2329
- } as AppendInput<D>
2330
- },
2331
- },
2332
- 'ai.control.requested': {
2316
+ const contract = options.contract as unknown as Contract<
2317
+ D,
2318
+ Record<never, never>,
2319
+ undefined
2320
+ >
2321
+ const { toolLane } = options
2322
+ type Of<K extends string> = ContractEvent<D, K & keyof D & string>
2323
+ type Generation = Of<'ai.generation.requested'>
2324
+ type ToolCall = Of<'ai.tool.execution.requested'>
2325
+ return [
2326
+ runtimeConsumer({
2327
+ contract,
2328
+ name: 'a2.ai.metadata',
2329
+ types: ['ai.model.metadata.requested'],
2330
+ }).run(async (ctx) => {
2331
+ const limits = await readModelLimits({
2332
+ modelId: ctx.event.payload.modelId,
2333
+ signal: ctx.signal,
2334
+ })
2335
+ ctx.signal.throwIfAborted()
2336
+ return {
2337
+ type: 'ai.model.metadata.resolved',
2338
+ id: await aiEventId('a2.ai.model.metadata.resolved', ctx.event.id),
2339
+ payload: { modelId: ctx.event.payload.modelId, limits },
2340
+ } as AppendInput<D>
2341
+ }),
2342
+ // One lane for every control event: commands, work reports, and closes
2343
+ // apply one at a time. Its lane function also rejects a bad command at
2344
+ // append, before anything is written.
2345
+ runtimeConsumer({
2346
+ contract,
2347
+ name: 'a2.ai.control',
2348
+ types: ['ai.control.requested', 'ai.work.reported', 'ai.session.closed'],
2333
2349
  lane: ({ event }) => {
2350
+ if (event.type !== 'ai.control.requested') return 'a2.ai.control'
2334
2351
  const command = event.payload as ControlCommand<M>
2335
2352
  if (event.id === '')
2336
2353
  throw new A2Error(
@@ -2352,17 +2369,25 @@ export function createHandlers<
2352
2369
  )
2353
2370
  return 'a2.ai.control'
2354
2371
  },
2355
- handler: controlHandler,
2356
- },
2357
- 'ai.work.reported': { lane: 'a2.ai.control', handler: controlHandler },
2358
- 'ai.session.closed': { lane: 'a2.ai.control', handler: controlHandler },
2359
- 'ai.generation.requested': {
2360
- lane: 'a2.ai.model',
2372
+ }).run(controlHandler as never),
2373
+ runtimeConsumer({
2374
+ contract,
2375
+ name: 'a2.ai.model',
2376
+ types: ['ai.generation.requested'],
2377
+ lane: () => 'a2.ai.model',
2378
+ // The contract is generic here, so the map's keys widen to string and
2379
+ // each predicate names its own event types.
2361
2380
  abortOn: {
2362
- 'ai.control.requested': (event, trigger) =>
2381
+ 'ai.control.requested': (
2382
+ event: Of<'ai.control.requested'>,
2383
+ trigger: Generation,
2384
+ ) =>
2363
2385
  event.payload.action === 'stop' &&
2364
2386
  event.payload.turnId === trigger.payload.control?.turnId,
2365
- 'ai.control.committed': (event, trigger) => {
2387
+ 'ai.control.committed': (
2388
+ event: Of<'ai.control.committed'>,
2389
+ trigger: Generation,
2390
+ ) => {
2366
2391
  const active = (event.payload as ControlCommit).view
2367
2392
  return (
2368
2393
  active?.turnId !== trigger.payload.control?.turnId ||
@@ -2370,39 +2395,50 @@ export function createHandlers<
2370
2395
  )
2371
2396
  },
2372
2397
  'ai.session.closed': true,
2373
- },
2374
- handler: (ctx) =>
2375
- telemetry.span(
2376
- 'a2.ai.generate',
2377
- {
2378
- 'a2.contract': options.contract.name,
2379
- 'a2.session_id': ctx.event.sessionId,
2380
- 'a2.ai.generate.request_id': ctx.event.id,
2381
- 'a2.ai.generate.attempt': ctx.attempt,
2382
- },
2383
- async () => {
2384
- try {
2385
- return control.settled({
2386
- ctx,
2387
- events: await generationHandler(ctx),
2388
- })
2389
- } catch (error) {
2390
- if (control.cancelled({ error, signal: ctx.signal }))
2391
- return control.settled({ ctx, events: undefined })
2392
- throw error
2393
- }
2394
- },
2395
- ),
2396
- },
2397
- 'ai.tool.execution.requested': {
2398
- ...(options.toolLane === undefined ? {} : { lane: options.toolLane }),
2398
+ } as never,
2399
+ }).run((ctx) =>
2400
+ telemetry.span(
2401
+ 'a2.ai.generate',
2402
+ {
2403
+ 'a2.contract': options.contract.name,
2404
+ 'a2.session_id': ctx.event.sessionId,
2405
+ 'a2.ai.generate.request_id': ctx.event.id,
2406
+ 'a2.ai.generate.attempt': ctx.attempt,
2407
+ },
2408
+ async () => {
2409
+ try {
2410
+ return control.settled({
2411
+ ctx,
2412
+ events: await generationHandler(ctx),
2413
+ })
2414
+ } catch (error) {
2415
+ if (control.cancelled({ error, signal: ctx.signal }))
2416
+ return control.settled({ ctx, events: undefined })
2417
+ throw error
2418
+ }
2419
+ },
2420
+ ),
2421
+ ),
2422
+ runtimeConsumer({
2423
+ contract,
2424
+ name: 'a2.ai.tools',
2425
+ types: ['ai.tool.execution.requested'],
2426
+ ...(toolLane === undefined ? {} : { lane: () => toolLane }),
2399
2427
  abortOn: {
2400
- 'ai.control.requested': (event, trigger) =>
2428
+ 'ai.control.requested': (
2429
+ event: Of<'ai.control.requested'>,
2430
+ trigger: ToolCall,
2431
+ ) =>
2401
2432
  event.payload.action === 'stop' &&
2402
2433
  event.payload.turnId === trigger.payload.turnId,
2403
- 'ai.generation.failed': (event, trigger) =>
2404
- event.payload.generationId === trigger.payload.call.generationId,
2405
- 'ai.control.committed': (event, trigger) => {
2434
+ 'ai.generation.failed': (
2435
+ event: Of<'ai.generation.failed'>,
2436
+ trigger: ToolCall,
2437
+ ) => event.payload.generationId === trigger.payload.call.generationId,
2438
+ 'ai.control.committed': (
2439
+ event: Of<'ai.control.committed'>,
2440
+ trigger: ToolCall,
2441
+ ) => {
2406
2442
  const active = (event.payload as ControlCommit).view
2407
2443
  return (
2408
2444
  active?.turnId !== trigger.payload.turnId ||
@@ -2410,26 +2446,26 @@ export function createHandlers<
2410
2446
  )
2411
2447
  },
2412
2448
  'ai.session.closed': true,
2413
- },
2414
- handler: (ctx) =>
2415
- telemetry.span(
2416
- 'a2.ai.tool',
2417
- {
2418
- 'a2.contract': options.contract.name,
2419
- 'a2.session_id': ctx.event.sessionId,
2420
- 'a2.ai.tool.request_id': ctx.event.payload.call.requestId,
2421
- 'a2.ai.tool.call_id': ctx.event.payload.call.toolCallId,
2422
- 'a2.ai.tool.name': ctx.event.payload.call.toolName,
2423
- 'a2.ai.tool.attempt': ctx.attempt,
2424
- },
2425
- () => handleToolExecution(ctx),
2426
- ),
2427
- },
2428
- } as NonNullable<ServerOptions<D>['handlers']>
2449
+ } as never,
2450
+ }).run((ctx) =>
2451
+ telemetry.span(
2452
+ 'a2.ai.tool',
2453
+ {
2454
+ 'a2.contract': options.contract.name,
2455
+ 'a2.session_id': ctx.event.sessionId,
2456
+ 'a2.ai.tool.request_id': ctx.event.payload.call.requestId,
2457
+ 'a2.ai.tool.call_id': ctx.event.payload.call.toolCallId,
2458
+ 'a2.ai.tool.name': ctx.event.payload.call.toolName,
2459
+ 'a2.ai.tool.attempt': ctx.attempt,
2460
+ },
2461
+ () => handleToolExecution(ctx),
2462
+ ),
2463
+ ),
2464
+ ] as Consumer<D>[]
2429
2465
  }
2430
2466
 
2431
2467
  /**
2432
- * Serve an AI agent: the built-in protocol, its handlers, and A2's store,
2468
+ * Serve an AI agent: the built-in protocol, its consumers, and A2's store,
2433
2469
  * scheduler, and telemetry, as one ordinary A2 server. Pass `name` for the
2434
2470
  * built-in protocol, or a shared `agentContract()` when clients fold
2435
2471
  * application state.
@@ -2439,8 +2475,9 @@ export function agent<
2439
2475
  M extends UIMessage = UIMessage,
2440
2476
  E extends EventDefs = Record<never, never>,
2441
2477
  H extends string = never,
2478
+ const C extends readonly unknown[] = readonly Consumer<AIEventDefs<M> & E>[],
2442
2479
  >(
2443
- options: AgentOptions<M, E, T, H>,
2480
+ options: AgentOptions<M, E, T, H, C>,
2444
2481
  ): Agent<AIEventDefs<M> & E, AIState<M>, typeof AI_STATE_REDUCER>
2445
2482
  export function agent<
2446
2483
  D extends AIEventDefs<AIMessageOf<D>> & EventDefs,
@@ -2449,7 +2486,8 @@ export function agent<
2449
2486
  Name extends string = string,
2450
2487
  X extends object = Record<never, never>,
2451
2488
  H extends string = never,
2452
- >(options: AgentFromContractOptions<D, S, T, Name, X, H>): Agent<D, S, Name>
2489
+ const C extends readonly unknown[] = readonly Consumer<D>[],
2490
+ >(options: AgentFromContractOptions<D, S, T, Name, X, H, C>): Agent<D, S, Name>
2453
2491
  export function agent(
2454
2492
  options:
2455
2493
  | AgentOptions<UIMessage, EventDefs, ToolSet>
@@ -2457,6 +2495,12 @@ export function agent(
2457
2495
  ): Agent<AIEventDefs & EventDefs, AIState> {
2458
2496
  type D = AIEventDefs & EventDefs
2459
2497
  assertNoLoopOptions(options)
2498
+ if (
2499
+ options.toolLane !== undefined &&
2500
+ (typeof options.toolLane !== 'string' || options.toolLane.length === 0)
2501
+ ) {
2502
+ throw new TypeError('toolLane must be a non-empty string')
2503
+ }
2460
2504
  const {
2461
2505
  model,
2462
2506
  tools,
@@ -2466,13 +2510,12 @@ export function agent(
2466
2510
  toolRecovery,
2467
2511
  progress,
2468
2512
  toolLane,
2469
- handlers,
2470
2513
  ...rest
2471
2514
  } = options
2472
2515
  let contract: Contract<D, Record<never, never>, Reducer<D, AIState>>
2473
2516
  let serverOptions: Omit<
2474
2517
  ServerOptions<D, Record<never, never>, Reducer<D, AIState>>,
2475
- 'contract' | 'handlers'
2518
+ 'contract'
2476
2519
  >
2477
2520
  if ('contract' in rest) {
2478
2521
  const { contract: shared, ...remaining } = rest
@@ -2492,7 +2535,7 @@ export function agent(
2492
2535
  serverOptions = remaining
2493
2536
  }
2494
2537
  const states = harnessStates()
2495
- const builtIns = createHandlers<UIMessage, D>({
2538
+ const builtIns = createConsumers<UIMessage, D>({
2496
2539
  contract,
2497
2540
  model,
2498
2541
  harnessStates: states,
@@ -2506,15 +2549,11 @@ export function agent(
2506
2549
  ...(toolRecovery === undefined ? {} : { toolRecovery }),
2507
2550
  ...(progress === undefined ? {} : { progress }),
2508
2551
  ...(toolLane === undefined ? {} : { toolLane }),
2509
- } as HandlersOptions<UIMessage, D>)
2510
- for (const type of Object.keys(handlers ?? {})) {
2511
- if (Object.hasOwn(builtIns, type))
2512
- throw new TypeError(`agent handles '${type}' itself`)
2513
- }
2552
+ } as ConsumersOptions<UIMessage, D>)
2514
2553
  const server = createServer({
2515
2554
  ...serverOptions,
2516
2555
  contract,
2517
- handlers: { ...builtIns, ...handlers },
2556
+ consumers: [...builtIns, ...(serverOptions.consumers ?? [])],
2518
2557
  })
2519
2558
  registerHarnessStates(server, states)
2520
2559
  const applicationEvents = 'contract' in rest
@@ -31,7 +31,7 @@ import type {
31
31
  ContractEvent,
32
32
  EventDefs,
33
33
  } from './contract.ts'
34
- import type { HandlerContext } from './server.ts'
34
+ import type { ConsumerContext } from './server.ts'
35
35
  import type { Reducer } from './reducer.ts'
36
36
  import type { AICoordinatorState, CoordinatorCall } from './ai-coordinator.ts'
37
37
  import { validateSync } from './validate.ts'
@@ -56,18 +56,18 @@ export type AIControlRuntime<
56
56
  D extends AIEventDefs<M> & EventDefs,
57
57
  > = {
58
58
  handler(
59
- ctx: HandlerContext<
59
+ ctx: ConsumerContext<
60
60
  D,
61
61
  'ai.control.requested' | 'ai.work.reported' | 'ai.session.closed'
62
62
  >,
63
63
  ): Promise<AppendInput<D>[]>
64
64
  append(options: {
65
- ctx: HandlerContext<D>
65
+ ctx: ConsumerContext<D>
66
66
  name: string
67
67
  events: AppendInput<D>[]
68
68
  }): Promise<ContractEvent<D>[]>
69
69
  settled(options: {
70
- ctx: HandlerContext<D>
70
+ ctx: ConsumerContext<D>
71
71
  events: void | AppendInput<D> | readonly AppendInput<D>[]
72
72
  }): Promise<AppendInput<D>>
73
73
  cancelled(options: { error: unknown; signal: AbortSignal }): boolean
@@ -117,7 +117,7 @@ export const createControlRuntime = <
117
117
  const coordinator = aiCoordinatorReducer(options.contract)
118
118
 
119
119
  const handler = async (
120
- ctx: HandlerContext<
120
+ ctx: ConsumerContext<
121
121
  D,
122
122
  'ai.control.requested' | 'ai.work.reported' | 'ai.session.closed'
123
123
  >,
@@ -761,7 +761,7 @@ export const createControlRuntime = <
761
761
  settled,
762
762
  name,
763
763
  }: {
764
- ctx: HandlerContext<D>
764
+ ctx: ConsumerContext<D>
765
765
  events: readonly AppendInput<D>[]
766
766
  settled: boolean
767
767
  name: string
@@ -828,7 +828,7 @@ export const createControlRuntime = <
828
828
  name,
829
829
  events,
830
830
  }: {
831
- ctx: HandlerContext<D>
831
+ ctx: ConsumerContext<D>
832
832
  name: string
833
833
  events: AppendInput<D>[]
834
834
  }): Promise<ContractEvent<D>[]> => {
@@ -896,7 +896,7 @@ export const createControlRuntime = <
896
896
  ctx,
897
897
  events,
898
898
  }: {
899
- ctx: HandlerContext<D>
899
+ ctx: ConsumerContext<D>
900
900
  events: void | AppendInput<D> | readonly AppendInput<D>[]
901
901
  }): Promise<AppendInput<D>> =>
902
902
  report({
package/src/ai-fx.ts CHANGED
@@ -20,6 +20,7 @@ import {
20
20
  Divergence,
21
21
  gatewayCredential,
22
22
  type GatewayCredential,
23
+ gatewaySessionAffinity,
23
24
  type Harness,
24
25
  type HarnessGeneration,
25
26
  harnessImages,
@@ -174,7 +175,7 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
174
175
  if (outcome.result.stopReason === 'end_turn' && !step.signal.aborted)
175
176
  await saveCheckpoint()
176
177
  } else {
177
- // FX now waits on tools that A2's tool handlers run.
178
+ // FX now waits on tools that A2's tool consumer run.
178
179
  fxTurn.cancel()
179
180
  }
180
181
  controller.enqueue({ type: 'finish-step' })
@@ -244,19 +245,27 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
244
245
  return pending(init.signal)
245
246
  }
246
247
  state.live = true
247
- return live(state, url, init, compaction)
248
+ return live(state, url, { ...init, body }, compaction)
248
249
  }
249
250
 
251
+ const affinity = gatewaySessionAffinity({
252
+ agentName: step.agentName,
253
+ sessionId: step.sessionId,
254
+ })
250
255
  let compactions = 0
251
256
  const live = async (
252
257
  state: TurnState,
253
258
  url: string,
254
- init: RequestInit,
259
+ init: RequestInit & { body: string },
255
260
  compaction: boolean,
256
261
  ): Promise<Response> => {
262
+ const headers = withAuth(init.headers, credential)
263
+ // The kernel's own affinity is new with every rebuild; A2's lasts the session.
264
+ headers.set('x-session-affinity', affinity)
257
265
  const response = await fetch(url, {
258
266
  ...init,
259
- headers: withAuth(init.headers, credential),
267
+ body: withGatewayCaching(init.body),
268
+ headers,
260
269
  signal: network,
261
270
  })
262
271
  // Abandoned before its headers came: the kernel has moved on without it.
@@ -408,14 +417,12 @@ async function turnPlan(
408
417
  ): Promise<TurnPlan> {
409
418
  const calls: RecordedCall[] = []
410
419
  const steers = new Map<number, string[]>()
411
- const steer = (at: number, inputs: HarnessInput[] | undefined) => {
412
- if (inputs !== undefined) steers.set(Math.max(0, at), inputs.map(steerText))
413
- }
414
- steer(0, turn.generations.length === 0 ? turn.steers.get(0) : undefined)
415
420
  turn.generations.forEach((generation, index) => {
416
421
  calls.push(...recordedCalls(generation.chunks, generation.completed))
422
+ const inputs = turn.steers.get(index)
417
423
  // A generation's step is its last call.
418
- steer(calls.length - 1, turn.steers.get(index))
424
+ if (inputs !== undefined)
425
+ steers.set(calls.length - 1, inputs.map(steerText))
419
426
  })
420
427
  const images = await harnessImages({
421
428
  harness: 'fx()',
@@ -774,6 +781,31 @@ async function loadLibfx(): Promise<typeof import('libfx')> {
774
781
  }
775
782
  }
776
783
 
784
+ /**
785
+ * A model request with Gateway's automatic caching, as the AI SDK harness
786
+ * asks for it; options the request already has take precedence. Every other
787
+ * byte stays as the kernel wrote it, a tool input's spacing included.
788
+ */
789
+ function withGatewayCaching(body: string): string {
790
+ const request = JSON.parse(body) as {
791
+ providerOptions?: Record<string, Record<string, unknown> | undefined>
792
+ }
793
+ const options = request.providerOptions
794
+ if (options === undefined)
795
+ return body.replace(
796
+ /^\{/,
797
+ '{"providerOptions":{"gateway":{"caching":"auto"}},',
798
+ )
799
+ // The kernel writes its own options last; any other shape goes unchanged.
800
+ const own = `"providerOptions":${JSON.stringify(options)}}`
801
+ if (!body.endsWith(own)) return body
802
+ const merged = {
803
+ ...options,
804
+ gateway: { caching: 'auto', ...options['gateway'] },
805
+ }
806
+ return `${body.slice(0, -own.length)}"providerOptions":${JSON.stringify(merged)}}`
807
+ }
808
+
777
809
  function withAuth(
778
810
  headers: HeadersInit | undefined,
779
811
  credential: GatewayCredential,