experimental-a2 0.13.0 → 0.14.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 (99) hide show
  1. package/CHANGELOG.md +34 -0
  2. package/dist/{actor-DJi3RsNu.d.ts → actor-BfQSE0KC.d.ts} +4 -4
  3. package/dist/{actor-DJi3RsNu.d.ts.map → actor-BfQSE0KC.d.ts.map} +1 -1
  4. package/dist/actor-client.d.ts +1 -1
  5. package/dist/actor-client.js +1 -1
  6. package/dist/actor-react.d.ts +3 -3
  7. package/dist/actor-react.js +2 -2
  8. package/dist/{actor-shared-DI7J5upy.js → actor-shared-B5tJfzt-.js} +2 -2
  9. package/dist/{actor-shared-DI7J5upy.js.map → actor-shared-B5tJfzt-.js.map} +1 -1
  10. package/dist/actor.d.ts +1 -1
  11. package/dist/actor.js +3 -3
  12. package/dist/ai-Cai-lCbj.d.ts +580 -0
  13. package/dist/ai-Cai-lCbj.d.ts.map +1 -0
  14. package/dist/ai-control-CcD4hh3y.js +119 -0
  15. package/dist/ai-control-CcD4hh3y.js.map +1 -0
  16. package/dist/ai-server.d.ts +6 -6
  17. package/dist/ai-server.d.ts.map +1 -1
  18. package/dist/ai-server.js +1014 -501
  19. package/dist/ai-server.js.map +1 -1
  20. package/dist/ai.d.ts +2 -365
  21. package/dist/ai.js +801 -80
  22. package/dist/ai.js.map +1 -1
  23. package/dist/{client-P_NNNRM-.d.ts → client-BAEABRZB.d.ts} +2 -2
  24. package/dist/{client-P_NNNRM-.d.ts.map → client-BAEABRZB.d.ts.map} +1 -1
  25. package/dist/{client-Bf6uSEAk.js → client-BYzHjkwU.js} +21 -6
  26. package/dist/client-BYzHjkwU.js.map +1 -0
  27. package/dist/client.d.ts +1 -1
  28. package/dist/client.js +1 -1
  29. package/dist/{contract-48bUMgcL.js → contract-CKRg_E4q.js} +3 -26
  30. package/dist/contract-CKRg_E4q.js.map +1 -0
  31. package/dist/index.d.ts +2 -2
  32. package/dist/index.js +1 -1
  33. package/dist/react.d.ts +2 -2
  34. package/dist/react.js +1 -1
  35. package/dist/{reducer-DJKWm3cp.d.ts → reducer-BcS9VDKC.d.ts} +4 -1
  36. package/dist/{reducer-DJKWm3cp.d.ts.map → reducer-BcS9VDKC.d.ts.map} +1 -1
  37. package/dist/reducer-DEMjEY_O.js +29 -0
  38. package/dist/reducer-DEMjEY_O.js.map +1 -0
  39. package/dist/scheduler-qstash.d.ts +2 -2
  40. package/dist/scheduler-qstash.js +1 -1
  41. package/dist/scheduler-vercel.d.ts +2 -2
  42. package/dist/scheduler-vercel.js +1 -1
  43. package/dist/{server-DjZZa1wr.d.ts → server-Bp5Nd1pF.d.ts} +3 -3
  44. package/dist/{server-DjZZa1wr.d.ts.map → server-Bp5Nd1pF.d.ts.map} +1 -1
  45. package/dist/{server-BeNADlCI.js → server-CjJSGcF7.js} +4 -3
  46. package/dist/server-CjJSGcF7.js.map +1 -0
  47. package/dist/server.d.ts +3 -3
  48. package/dist/server.js +1 -1
  49. package/dist/{store-DtDOWLSn.d.ts → store-D_yhNdPz.d.ts} +7 -2
  50. package/dist/{store-DtDOWLSn.d.ts.map → store-D_yhNdPz.d.ts.map} +1 -1
  51. package/dist/store-N8PXxDAS.js.map +1 -1
  52. package/dist/store-memory.d.ts +1 -1
  53. package/dist/store-postgres.d.ts +1 -1
  54. package/dist/store-postgres.js +19 -0
  55. package/dist/store-postgres.js.map +1 -1
  56. package/dist/store-redis-http.d.ts +1 -1
  57. package/dist/store-redis-http.js +1 -1
  58. package/dist/{store-redis-notify-BUCyXOn0.js → store-redis-notify-D2EI6gwX.js} +27 -2
  59. package/dist/store-redis-notify-D2EI6gwX.js.map +1 -0
  60. package/dist/store-redis.d.ts +1 -1
  61. package/dist/store-redis.js +1 -1
  62. package/dist/store-sqlite.d.ts +1 -1
  63. package/docs/guides/06-ai-agents.mdx +265 -62
  64. package/docs/reference/01-api.mdx +122 -27
  65. package/examples/playground/app/agent/[agentId]/agent-client.tsx +118 -21
  66. package/examples/playground/app/agent/[agentId]/agent-queue.tsx +289 -0
  67. package/examples/playground/app/agent/compaction-settings.test.ts +22 -6
  68. package/examples/playground/app/agent/model.ts +11 -2
  69. package/examples/playground/app/agent/server.ts +8 -2
  70. package/examples/playground/app/chat/[chatId]/chat-client.tsx +3 -13
  71. package/examples/playground/app/chat/model.ts +2 -2
  72. package/examples/playground/app/chat/server.ts +24 -17
  73. package/examples/playground/app/globals.css +179 -0
  74. package/examples/playground/package.json +1 -1
  75. package/package.json +1 -1
  76. package/src/ai-client-state.ts +185 -0
  77. package/src/ai-control-server.ts +829 -0
  78. package/src/ai-control-state.ts +152 -0
  79. package/src/ai-control.ts +139 -0
  80. package/src/ai-coordinator.ts +99 -32
  81. package/src/ai-progress-batches.ts +68 -0
  82. package/src/ai-projector.ts +76 -15
  83. package/src/ai-sdk-step.ts +0 -1
  84. package/src/ai-server.ts +381 -608
  85. package/src/ai.ts +553 -108
  86. package/src/client.ts +31 -9
  87. package/src/licenses/Apache-2.0.txt +55 -0
  88. package/src/parse-partial-json.ts +441 -0
  89. package/src/reducer.ts +6 -0
  90. package/src/server.ts +8 -4
  91. package/src/store-postgres.ts +27 -0
  92. package/src/store-redis-core.ts +53 -1
  93. package/src/store-redis-notify.ts +1 -0
  94. package/src/store.ts +6 -0
  95. package/dist/ai.d.ts.map +0 -1
  96. package/dist/client-Bf6uSEAk.js.map +0 -1
  97. package/dist/contract-48bUMgcL.js.map +0 -1
  98. package/dist/server-BeNADlCI.js.map +0 -1
  99. package/dist/store-redis-notify-BUCyXOn0.js.map +0 -1
package/src/ai.ts CHANGED
@@ -13,19 +13,54 @@ import type {
13
13
  UIMessage,
14
14
  UIMessageChunk,
15
15
  } from 'ai'
16
+ import {
17
+ isQueueCommand,
18
+ projectAIClientState,
19
+ type PendingQueueCommand,
20
+ } from './ai-client-state.ts'
21
+ export type { PendingQueueCommand } from './ai-client-state.ts'
22
+ import { applyControlChanges } from './ai-control-state.ts'
23
+ import { initialControlState } from './ai-control.ts'
24
+ import type {
25
+ ControlCommit,
26
+ ActiveTurn,
27
+ ControlCommand,
28
+ ControlReceipt,
29
+ InboxItem,
30
+ ToolExecutionRequest,
31
+ WorkReport,
32
+ } from './ai-control.ts'
33
+ export type {
34
+ ActiveTurn,
35
+ ControlCommand,
36
+ ControlReceipt,
37
+ InboxItem,
38
+ } from './ai-control.ts'
16
39
  import { contract as createContract } from './contract.ts'
17
- import type { AppendInput, Contract, EventDefs } from './contract.ts'
18
- import { interruptUIMessage, projectUIMessage } from './ai-projector.ts'
40
+ import type { Contract, EventDefs } from './contract.ts'
41
+ import {
42
+ interruptUIMessage,
43
+ pauseUIMessage,
44
+ projectUIMessage,
45
+ projectUIMessageStream,
46
+ type UIMessageCursor,
47
+ } from './ai-projector.ts'
19
48
  import { ambientToolScopeStorage } from './internal.ts'
20
- import type { Reducer } from './reducer.ts'
49
+ import { clientStateProjection, type Reducer } from './reducer.ts'
50
+ import {
51
+ appendProgressBatch,
52
+ flattenProgressBatches,
53
+ progressBatches,
54
+ type ProgressBatches,
55
+ } from './ai-progress-batches.ts'
21
56
  import type { HandlerContext } from './server.ts'
22
57
  import type { StandardSchemaV1 } from './standard-schema.ts'
23
58
  import { assertSyncSchema, validateSync } from './validate.ts'
24
59
 
25
60
  export type AgentStatus =
26
- 'idle' | 'generating' | 'waiting' | 'failed' | 'closed'
61
+ 'idle' | 'generating' | 'waiting' | 'failed' | 'closed' | 'paused'
27
62
 
28
- export type GenerationReason = 'message' | 'tool' | 'input' | 'retry'
63
+ export type GenerationReason = 'message' | 'tool' | 'input' | 'retry' | 'resume'
29
64
 
30
65
  export type SessionCreatedPayload = { metadata?: JSONValue }
31
66
  export type SessionClosedPayload = { reason?: string }
@@ -47,6 +82,7 @@ export type MessageInterruptedPayload = MessageInterruption &
47
82
  | { requestId: string; generationId?: never }
48
83
  )
49
84
  export type GenerationRequestedPayload = {
85
+ control?: { turnId: string; version: number }
50
86
  messageId: string
51
87
  reason: GenerationReason
52
88
  responseMessageId?: string
@@ -184,6 +220,15 @@ export type CompactionCompletedPayload<M extends UIMessage> = {
184
220
  }
185
221
 
186
222
  export type AIEventDefs<M extends UIMessage = UIMessage> = {
223
+ 'ai.control.requested': StandardSchemaV1<ControlCommand<M>>
224
+ 'ai.control.committed': StandardSchemaV1<ControlCommit>
225
+ 'ai.control.decided': StandardSchemaV1<ControlReceipt>
226
+ 'ai.work.reported': StandardSchemaV1<WorkReport>
227
+ 'ai.tool.execution.requested': StandardSchemaV1<ToolExecutionRequest>
228
+ 'ai.generation.paused': StandardSchemaV1<{
229
+ generationId: string
230
+ messageId: string
231
+ }>
187
232
  'ai.session.created': StandardSchemaV1<SessionCreatedPayload>
188
233
  'ai.session.closed': StandardSchemaV1<SessionClosedPayload>
189
234
  'ai.message.created': StandardSchemaV1<MessageCreatedPayload<M>>
@@ -215,8 +260,10 @@ export type ActiveGenerationProjection<M extends UIMessage = UIMessage> = {
215
260
  /** Log index of the completed model step while its response remains interruptible. */
216
261
  completionIndex?: number
217
262
  baseMessage?: M
263
+ promptThroughIndex?: number
264
+ cursor?: { throughIndex: number; stream: UIMessageCursor } | null
218
265
  baseCompaction?: CompactionState<M> | null
219
- batches: Array<{ index: number; chunks: UIMessageChunk[] }>
266
+ batches: ProgressBatches
220
267
  baseTools: ToolActivity[]
221
268
  toolEvents: Array<
222
269
  | { index: number; type: 'ai.tool.called'; payload: ToolCalledPayload }
@@ -277,6 +324,12 @@ export type GenerationUsage = {
277
324
  }
278
325
 
279
326
  export type AIState<M extends UIMessage = UIMessage> = {
327
+ inbox: { paused: boolean; items: InboxItem<M>[] }
328
+ pendingQueueCommands?: PendingQueueCommand<M>[]
329
+ active: ActiveTurn<M> | null
330
+ starting: boolean
331
+ receipt: ControlReceipt | null
332
+ pausedGenerationIds?: Record<string, true>
280
333
  session: { created: boolean; metadata?: JSONValue }
281
334
  messages: M[]
282
335
  status: AgentStatus
@@ -308,12 +361,13 @@ export type AgentDefinition<
308
361
  }
309
362
 
310
363
  export type AgentToolEventDefs = EventDefs &
311
- Pick<AIEventDefs, 'ai.tool.called' | 'ai.approval.responded'>
364
+ Pick<AIEventDefs, 'ai.tool.execution.requested'>
312
365
 
313
366
  /** The durable handler context supplied when A2 executes a local AI tool. */
314
- export type AgentToolContext<D extends AgentToolEventDefs> =
315
- | HandlerContext<D, 'ai.tool.called'>
316
- | HandlerContext<D, 'ai.approval.responded'>
367
+ export type AgentToolContext<D extends AgentToolEventDefs> = HandlerContext<
368
+ D,
369
+ 'ai.tool.execution.requested'
370
+ >
317
371
 
318
372
  export type AgentOptions<M extends UIMessage, E extends EventDefs> = {
319
373
  name: string
@@ -545,6 +599,143 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
545
599
  const messageSchema = options?.messageSchema
546
600
  if (messageSchema) assertSyncSchema(messageSchema, 'the AI message schema')
547
601
  const definition: AIEventDefs<M> = {
602
+ 'ai.control.requested': schema('ai.control.requested', (value) => {
603
+ if (!isRecord(value) || !isJSONCompatible(value))
604
+ return issue('invalid AI command')
605
+ const action = value['action']
606
+ const forwarded = {
607
+ approval: ['ai.approval.responded', 'response'],
608
+ 'tool-result': ['ai.tool.result', 'result'],
609
+ input: ['ai.input.responded', 'response'],
610
+ 'request-input': ['ai.input.requested', 'request'],
611
+ retry: ['ai.retry.requested', 'request'],
612
+ interrupt: ['ai.message.interrupted', 'request'],
613
+ } as const
614
+ if (typeof action === 'string' && Object.hasOwn(forwarded, action)) {
615
+ const [type, key] = forwarded[action as keyof typeof forwarded]
616
+ const payload = validateSync(
617
+ definition[type] as StandardSchemaV1<unknown>,
618
+ value[key],
619
+ 'AI command',
620
+ )
621
+ if (payload.issues) return { issues: payload.issues }
622
+ return { value: { action, [key]: payload.value } as ControlCommand<M> }
623
+ }
624
+ if (action === 'send' || action === 'edit' || action === 'steer') {
625
+ const parsed = validateMessage(value['message'], messageSchema)
626
+ if (parsed.issues) return parsed
627
+ if (
628
+ (action === 'edit' &&
629
+ (!Number.isSafeInteger(value['expectedRevision']) ||
630
+ Number(value['expectedRevision']) < 0)) ||
631
+ (action === 'steer' &&
632
+ (!stringField(value, 'turnId') ||
633
+ !optionalIndexField(value, 'lastSeenIndex'))) ||
634
+ (action === 'send' &&
635
+ value['generate'] !== undefined &&
636
+ typeof value['generate'] !== 'boolean')
637
+ )
638
+ return issue('invalid AI command')
639
+ return {
640
+ value: { ...value, message: parsed.value } as ControlCommand<M>,
641
+ }
642
+ }
643
+ if (
644
+ (action === 'remove' && stringField(value, 'inputId')) ||
645
+ (action === 'send-now' &&
646
+ stringField(value, 'inputId') &&
647
+ (value['turnId'] === null || stringField(value, 'turnId')) &&
648
+ optionalIndexField(value, 'lastSeenIndex')) ||
649
+ (action === 'move' &&
650
+ stringField(value, 'inputId') &&
651
+ (value['beforeId'] === null || stringField(value, 'beforeId'))) ||
652
+ (action === 'stop' &&
653
+ stringField(value, 'turnId') &&
654
+ optionalIndexField(value, 'lastSeenIndex')) ||
655
+ (action === 'pause' &&
656
+ (value['when'] === 'now' || value['when'] === 'after-turn')) ||
657
+ action === 'resume'
658
+ )
659
+ return { value: value as ControlCommand<M> }
660
+ return issue('invalid AI command')
661
+ }),
662
+ 'ai.control.committed': jsonObject(
663
+ 'ai.control.committed',
664
+ (value) => Array.isArray(value['changes']) && isRecord(value['view']),
665
+ ),
666
+ 'ai.control.decided': jsonObject(
667
+ 'ai.control.decided',
668
+ (value) =>
669
+ stringField(value, 'commandId') &&
670
+ (value['outcome'] === 'applied' || value['outcome'] === 'rejected'),
671
+ ),
672
+ 'ai.work.reported': schema('ai.work.reported', (value) => {
673
+ if (
674
+ !isRecord(value) ||
675
+ !isJSONCompatible(value) ||
676
+ !stringField(value, 'workId') ||
677
+ !stringField(value, 'turnId') ||
678
+ typeof value['version'] !== 'number' ||
679
+ !optionalIndexField(value, 'version') ||
680
+ typeof value['attempt'] !== 'number' ||
681
+ !Number.isSafeInteger(value['attempt']) ||
682
+ value['attempt'] < 1 ||
683
+ typeof value['settled'] !== 'boolean' ||
684
+ (value['kind'] !== 'model' && value['kind'] !== 'tool') ||
685
+ !Array.isArray(value['events'])
686
+ )
687
+ return issue('invalid AI work report')
688
+ const allowed = new Set([
689
+ 'ai.generation.started',
690
+ 'ai.generation.completed',
691
+ 'ai.generation.failed',
692
+ 'ai.tool.called',
693
+ 'ai.tool.result',
694
+ 'ai.approval.requested',
695
+ 'ai.approval.responded',
696
+ 'ai.compaction.requested',
697
+ 'ai.compaction.completed',
698
+ 'ai.model.metadata.requested',
699
+ ])
700
+ const events: WorkReport['events'] = []
701
+ for (const event of value['events']) {
702
+ if (
703
+ !isRecord(event) ||
704
+ !stringField(event, 'id') ||
705
+ typeof event['type'] !== 'string' ||
706
+ !allowed.has(event['type']) ||
707
+ (value['kind'] === 'tool' && event['type'] !== 'ai.tool.result')
708
+ )
709
+ return issue('invalid AI work fact')
710
+ const parsed = validateSync(
711
+ definition[
712
+ event['type'] as keyof AIEventDefs<M>
713
+ ] as StandardSchemaV1<unknown>,
714
+ event['payload'],
715
+ 'AI work fact',
716
+ )
717
+ if (parsed.issues) return { issues: parsed.issues }
718
+ events.push({
719
+ id: event['id'] as string,
720
+ type: event['type'],
721
+ payload: parsed.value,
722
+ })
723
+ }
724
+ return { value: { ...value, events } as WorkReport }
725
+ }),
726
+ 'ai.tool.execution.requested': jsonObject(
727
+ 'ai.tool.execution.requested',
728
+ (value) =>
729
+ stringField(value, 'turnId') &&
730
+ optionalIndexField(value, 'version') &&
731
+ isRecord(value['call']) &&
732
+ isRecord(value['generation']),
733
+ ),
734
+ 'ai.generation.paused': jsonObject(
735
+ 'ai.generation.paused',
736
+ (value) =>
737
+ stringField(value, 'generationId') && stringField(value, 'messageId'),
738
+ ),
548
739
  'ai.session.created': jsonObject('ai.session.created', () => true),
549
740
  'ai.session.closed': jsonObject('ai.session.closed', (value) =>
550
741
  optionalStringField(value, 'reason'),
@@ -575,7 +766,8 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
575
766
  (value['reason'] === 'message' ||
576
767
  value['reason'] === 'tool' ||
577
768
  value['reason'] === 'input' ||
578
- value['reason'] === 'retry'),
769
+ value['reason'] === 'retry' ||
770
+ value['reason'] === 'resume'),
579
771
  ),
580
772
  'ai.retry.requested': jsonObject(
581
773
  'ai.retry.requested',
@@ -717,6 +909,11 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
717
909
  export const events: AIEventDefs<UIMessage> = createEvents()
718
910
 
719
911
  const initialState = <M extends UIMessage>(): AIState<M> => ({
912
+ inbox: initialControlState<M>().inbox,
913
+ pendingQueueCommands: [],
914
+ active: null,
915
+ starting: false,
916
+ receipt: null,
720
917
  session: { created: false },
721
918
  messages: [],
722
919
  status: 'idle',
@@ -742,6 +939,23 @@ const upsertMessage = <M extends UIMessage>(messages: M[], message: M): M[] => {
742
939
  return next
743
940
  }
744
941
 
942
+ const upsertResponse = <M extends UIMessage>(options: {
943
+ messages: M[]
944
+ message: M
945
+ inputId: string
946
+ }): M[] => {
947
+ const { messages, message, inputId } = options
948
+ if (messages.some((item) => item.id === message.id))
949
+ return upsertMessage(messages, message)
950
+ const inputIndex = messages.findIndex((item) => item.id === inputId)
951
+ if (inputIndex === -1) return [...messages, message]
952
+ return [
953
+ ...messages.slice(0, inputIndex + 1),
954
+ message,
955
+ ...messages.slice(inputIndex + 1),
956
+ ]
957
+ }
958
+
745
959
  const answerApproval = <M extends UIMessage>(
746
960
  messages: M[],
747
961
  response: ApprovalRespondedPayload,
@@ -893,12 +1107,10 @@ const isStaleGeneration = <M extends UIMessage>(
893
1107
  return current !== undefined && current !== generationId
894
1108
  }
895
1109
 
896
- const projectedMessage = <M extends UIMessage>(
1110
+ const projectedMessageStream = <M extends UIMessage>(
897
1111
  projection: ActiveGenerationProjection<M>,
898
- ): M => {
899
- const projectionEvents: Array<{ index: number; chunks: UIMessageChunk[] }> = [
900
- ...projection.batches,
901
- ]
1112
+ ): { message: M; cursor: UIMessageCursor } => {
1113
+ const projectionEvents = flattenProgressBatches(projection.batches)
902
1114
  for (const event of projection.toolEvents) {
903
1115
  if (event.type === 'ai.tool.result') {
904
1116
  projectionEvents.push({
@@ -917,7 +1129,7 @@ const projectedMessage = <M extends UIMessage>(
917
1129
  ],
918
1130
  })
919
1131
  }
920
- return projectUIMessage({
1132
+ return projectUIMessageStream({
921
1133
  messageId: projection.responseMessageId,
922
1134
  ...(projection.baseMessage === undefined
923
1135
  ? {}
@@ -928,6 +1140,25 @@ const projectedMessage = <M extends UIMessage>(
928
1140
  })
929
1141
  }
930
1142
 
1143
+ const projectedMessage = <M extends UIMessage>(
1144
+ projection: ActiveGenerationProjection<M>,
1145
+ ): M => projectedMessageStream(projection).message
1146
+
1147
+ const projectedResponse = <M extends UIMessage>(options: {
1148
+ state: AIState<M>
1149
+ projection: ActiveGenerationProjection<M>
1150
+ }): M => {
1151
+ const message = projectedMessage(options.projection)
1152
+ return options.state.pausedGenerationIds?.[options.projection.generationId]
1153
+ ? pauseUIMessage({
1154
+ message,
1155
+ toolCallIds: new Set(
1156
+ options.state.tools.map((tool) => tool.toolCallId),
1157
+ ),
1158
+ })
1159
+ : message
1160
+ }
1161
+
931
1162
  const reduceToolActivity = (
932
1163
  tools: ToolActivity[],
933
1164
  event:
@@ -1087,7 +1318,11 @@ const interruptGenerationProjection = <M extends UIMessage>(
1087
1318
  if (projection === null) return null
1088
1319
  const retained: ActiveGenerationProjection<M> = {
1089
1320
  ...projection,
1090
- batches: projection.batches.filter((batch) => batch.index <= cutoff),
1321
+ batches: progressBatches(
1322
+ flattenProgressBatches(projection.batches).filter(
1323
+ (batch) => batch.index <= cutoff,
1324
+ ),
1325
+ ),
1091
1326
  toolEvents: projection.toolEvents.filter(
1092
1327
  (toolEvent) => toolEvent.index <= cutoff,
1093
1328
  ),
@@ -1107,7 +1342,7 @@ const interruptGenerationProjection = <M extends UIMessage>(
1107
1342
  ...restored,
1108
1343
  messages: shouldProject
1109
1344
  ? upsertMessage(
1110
- restored.messages,
1345
+ state.messages,
1111
1346
  interruptUIMessage(projectedMessage(retained)),
1112
1347
  )
1113
1348
  : restored.messages,
@@ -1151,6 +1386,117 @@ export function reduceAIState<M extends UIMessage>(
1151
1386
  ): AIState<M> {
1152
1387
  if (state.status === 'closed') return state
1153
1388
  switch (event.type) {
1389
+ case 'ai.control.requested': {
1390
+ const command = event.payload as ControlCommand<M>
1391
+ if (!event.id || !isQueueCommand(command)) return state
1392
+ const pending = state.pendingQueueCommands ?? []
1393
+ if (pending.some((item) => item.commandId === event.id)) return state
1394
+ if (
1395
+ (command.action === 'send' || command.action === 'steer') &&
1396
+ (state.messages.some((message) => message.id === command.message.id) ||
1397
+ state.inbox.items.some((item) => item.id === command.message.id))
1398
+ )
1399
+ return state
1400
+ return {
1401
+ ...state,
1402
+ pendingQueueCommands: [...pending, { commandId: event.id, command }],
1403
+ }
1404
+ }
1405
+ case 'ai.control.committed': {
1406
+ const commit = event.payload as ControlCommit
1407
+ const visible = applyControlChanges({
1408
+ state: { inbox: state.inbox, active: state.active },
1409
+ changes: commit.changes.filter(
1410
+ (change) => change.path[0] === 'inbox' || change.path[0] === 'active',
1411
+ ),
1412
+ })
1413
+ const { phase } = commit.view
1414
+ const status: AgentStatus = commit.view.closed
1415
+ ? 'closed'
1416
+ : phase === 'paused' ||
1417
+ phase === 'pausing' ||
1418
+ (commit.view.paused && phase === null)
1419
+ ? 'paused'
1420
+ : phase === 'failed'
1421
+ ? 'failed'
1422
+ : phase === 'starting' || phase === 'generating'
1423
+ ? 'generating'
1424
+ : phase
1425
+ ? 'waiting'
1426
+ : 'idle'
1427
+ return {
1428
+ ...state,
1429
+ ...visible,
1430
+ status,
1431
+ activeRequestId: commit.view.requestId,
1432
+ activeResponseMessageId: commit.view.responseMessageId,
1433
+ }
1434
+ }
1435
+ case 'ai.control.decided': {
1436
+ const receipt = event.payload as ControlReceipt
1437
+ if (receipt.work) return state
1438
+ const pending = state.pendingQueueCommands ?? []
1439
+ const settled = pending.find(
1440
+ (item) => item.commandId === receipt.commandId,
1441
+ )
1442
+ return {
1443
+ ...state,
1444
+ receipt,
1445
+ ...(settled === undefined
1446
+ ? {}
1447
+ : {
1448
+ pendingQueueCommands: pending.filter(
1449
+ (item) => item.commandId !== receipt.commandId,
1450
+ ),
1451
+ }),
1452
+ }
1453
+ }
1454
+ case 'ai.generation.paused': {
1455
+ const payload = event.payload as {
1456
+ generationId: string
1457
+ messageId: string
1458
+ }
1459
+ const toolCallIds = new Set(
1460
+ state.tools
1461
+ .filter((tool) => tool.messageId === payload.messageId)
1462
+ .map((tool) => tool.toolCallId),
1463
+ )
1464
+ const projection = state.activeProjection
1465
+ return {
1466
+ ...state,
1467
+ activeGeneration: null,
1468
+ activeProjection:
1469
+ projection?.generationId === payload.generationId
1470
+ ? {
1471
+ ...projection,
1472
+ cursor: null,
1473
+ batches: progressBatches(
1474
+ flattenProgressBatches(projection.batches).map((batch) => ({
1475
+ index: batch.index,
1476
+ chunks: batch.chunks.filter(
1477
+ (chunk) =>
1478
+ !('toolCallId' in chunk) ||
1479
+ toolCallIds.has(chunk.toolCallId),
1480
+ ),
1481
+ })),
1482
+ ),
1483
+ }
1484
+ : projection,
1485
+ compaction:
1486
+ state.compaction?.status === 'running'
1487
+ ? (projection?.baseCompaction ?? null)
1488
+ : state.compaction,
1489
+ messages: state.messages.map((message) =>
1490
+ message.id === payload.messageId
1491
+ ? pauseUIMessage({ message, toolCallIds })
1492
+ : message,
1493
+ ),
1494
+ pausedGenerationIds: {
1495
+ ...state.pausedGenerationIds,
1496
+ [payload.generationId]: true,
1497
+ },
1498
+ }
1499
+ }
1154
1500
  case 'ai.model.metadata.requested': {
1155
1501
  const { modelId } = event.payload as ModelMetadataRequestedPayload
1156
1502
  if (Object.hasOwn(state.modelMetadata, modelId)) return state
@@ -1192,6 +1538,9 @@ export function reduceAIState<M extends UIMessage>(
1192
1538
  case 'ai.session.closed':
1193
1539
  return {
1194
1540
  ...state,
1541
+ active: null,
1542
+ inbox: { paused: true, items: [] },
1543
+ pendingQueueCommands: [],
1195
1544
  status: 'closed',
1196
1545
  compaction:
1197
1546
  state.compaction?.status === 'running'
@@ -1204,7 +1553,14 @@ export function reduceAIState<M extends UIMessage>(
1204
1553
  }
1205
1554
  case 'ai.message.created': {
1206
1555
  const { message } = event.payload as MessageCreatedPayload<M>
1207
- return { ...state, messages: upsertMessage(state.messages, message) }
1556
+ return {
1557
+ ...state,
1558
+ messages: upsertMessage(state.messages, message),
1559
+ activeProjection:
1560
+ state.activeProjection?.responseMessageId === message.id
1561
+ ? { ...state.activeProjection, cursor: null }
1562
+ : state.activeProjection,
1563
+ }
1208
1564
  }
1209
1565
  case 'ai.message.completed': {
1210
1566
  const { messageId } = event.payload as MessageCompletedPayload
@@ -1261,6 +1617,7 @@ export function reduceAIState<M extends UIMessage>(
1261
1617
  return state
1262
1618
  }
1263
1619
  const acceptsRequest =
1620
+ (payload.control !== undefined && state.activeRequestId === event.id) ||
1264
1621
  payload.reason === 'message' ||
1265
1622
  (payload.reason === 'tool'
1266
1623
  ? ownerGenerationId !== undefined &&
@@ -1361,9 +1718,10 @@ export function reduceAIState<M extends UIMessage>(
1361
1718
  activeProjection: {
1362
1719
  generationId: payload.generationId,
1363
1720
  responseMessageId: payload.responseMessageId,
1721
+ promptThroughIndex: payload.promptThroughIndex ?? event.index - 1,
1364
1722
  ...(baseMessage === undefined ? {} : { baseMessage }),
1365
1723
  baseCompaction: baseState.compaction,
1366
- batches: [],
1724
+ batches: progressBatches(),
1367
1725
  baseTools: baseState.tools,
1368
1726
  toolEvents: [],
1369
1727
  approvalEvents: [],
@@ -1373,6 +1731,7 @@ export function reduceAIState<M extends UIMessage>(
1373
1731
  }
1374
1732
  case 'ai.generation.progress': {
1375
1733
  const payload = event.payload as GenerationProgressPayload
1734
+ if (state.pausedGenerationIds?.[payload.generationId]) return state
1376
1735
  if (terminalGenerationReason(state, payload.generationId) !== undefined) {
1377
1736
  return state
1378
1737
  }
@@ -1383,18 +1742,63 @@ export function reduceAIState<M extends UIMessage>(
1383
1742
  ) {
1384
1743
  return state
1385
1744
  }
1745
+ const ordered =
1746
+ event.index > (activeProjection.batches.tail.at(-1)?.index ?? -1)
1386
1747
  const projection: ActiveGenerationProjection<M> = {
1387
1748
  ...activeProjection,
1388
- batches: [
1389
- ...activeProjection.batches.filter(
1390
- (batch) => batch.index !== event.index,
1391
- ),
1392
- { index: event.index, chunks: payload.chunks },
1393
- ].toSorted((left, right) => left.index - right.index),
1749
+ batches: ordered
1750
+ ? appendProgressBatch({
1751
+ batches: activeProjection.batches,
1752
+ batch: { index: event.index, chunks: payload.chunks },
1753
+ })
1754
+ : progressBatches(
1755
+ [
1756
+ ...flattenProgressBatches(activeProjection.batches).filter(
1757
+ (batch) => batch.index !== event.index,
1758
+ ),
1759
+ { index: event.index, chunks: payload.chunks },
1760
+ ].toSorted((left, right) => left.index - right.index),
1761
+ ),
1762
+ }
1763
+ const base = state.messages.find(
1764
+ (message) => message.id === projection.responseMessageId,
1765
+ )
1766
+ const incremental =
1767
+ base !== undefined &&
1768
+ activeProjection.cursor != null &&
1769
+ event.index > activeProjection.cursor.throughIndex
1770
+ const projected = incremental
1771
+ ? projectUIMessageStream({
1772
+ messageId: projection.responseMessageId,
1773
+ base,
1774
+ cursor: activeProjection.cursor!.stream,
1775
+ chunks: payload.chunks,
1776
+ })
1777
+ : projectedMessageStream(projection)
1778
+ projection.cursor = {
1779
+ throughIndex: incremental
1780
+ ? event.index
1781
+ : Math.max(
1782
+ event.index,
1783
+ projection.batches.tail.at(-1)?.index ?? 0,
1784
+ projection.toolEvents.reduce(
1785
+ (index, entry) => Math.max(index, entry.index),
1786
+ 0,
1787
+ ),
1788
+ projection.approvalEvents.reduce(
1789
+ (index, entry) => Math.max(index, entry.index),
1790
+ 0,
1791
+ ),
1792
+ ),
1793
+ stream: projected.cursor,
1394
1794
  }
1395
1795
  return {
1396
1796
  ...state,
1397
- messages: upsertMessage(state.messages, projectedMessage(projection)),
1797
+ messages: upsertResponse({
1798
+ messages: state.messages,
1799
+ message: projected.message,
1800
+ inputId: state.activeGeneration.messageId,
1801
+ }),
1398
1802
  activeProjection: projection,
1399
1803
  status: 'generating',
1400
1804
  }
@@ -1822,12 +2226,16 @@ export function reduceAIState<M extends UIMessage>(
1822
2226
  }
1823
2227
  const activeProjection: ActiveGenerationProjection<M> = {
1824
2228
  ...projection,
2229
+ cursor: null,
1825
2230
  toolEvents: [
1826
2231
  ...projection.toolEvents,
1827
2232
  { index: event.index, type: 'ai.tool.result', payload },
1828
2233
  ],
1829
2234
  }
1830
- const projected = projectedMessage(activeProjection)
2235
+ const projected = projectedResponse({
2236
+ state,
2237
+ projection: activeProjection,
2238
+ })
1831
2239
  const tools = projectedTools(activeProjection)
1832
2240
  const failureError = state.error ?? 'Generation failed.'
1833
2241
  const next = {
@@ -1878,6 +2286,7 @@ export function reduceAIState<M extends UIMessage>(
1878
2286
  projection?.generationId === payload.generationId
1879
2287
  ? {
1880
2288
  ...projection,
2289
+ cursor: null,
1881
2290
  approvalEvents: [
1882
2291
  ...projection.approvalEvents,
1883
2292
  {
@@ -1895,7 +2304,10 @@ export function reduceAIState<M extends UIMessage>(
1895
2304
  ? projectMessageChunks(state.messages, payload.messageId, [
1896
2305
  approvalRequestedChunk(payload),
1897
2306
  ])
1898
- : upsertMessage(state.messages, projectedMessage(activeProjection)),
2307
+ : upsertMessage(
2308
+ state.messages,
2309
+ projectedResponse({ state, projection: activeProjection }),
2310
+ ),
1899
2311
  activeProjection,
1900
2312
  pendingApprovals:
1901
2313
  payload.isAutomatic === true
@@ -1947,6 +2359,7 @@ export function reduceAIState<M extends UIMessage>(
1947
2359
  projection !== null && activeRequest !== undefined
1948
2360
  ? {
1949
2361
  ...projection,
2362
+ cursor: null,
1950
2363
  approvalEvents: [
1951
2364
  ...projection.approvalEvents,
1952
2365
  {
@@ -1962,8 +2375,14 @@ export function reduceAIState<M extends UIMessage>(
1962
2375
  messages:
1963
2376
  activeProjection === null || activeProjection === projection
1964
2377
  ? answerApproval(state.messages, payload)
1965
- : upsertMessage(state.messages, projectedMessage(activeProjection)),
1966
- activeProjection,
2378
+ : upsertMessage(
2379
+ state.messages,
2380
+ projectedResponse({ state, projection: activeProjection }),
2381
+ ),
2382
+ activeProjection:
2383
+ activeProjection === null
2384
+ ? null
2385
+ : { ...activeProjection, cursor: null },
1967
2386
  pendingApprovals: state.pendingApprovals.filter(
1968
2387
  (approval) =>
1969
2388
  approval.approvalId !== payload.approvalId ||
@@ -2038,7 +2457,12 @@ export function reduceAIState<M extends UIMessage>(
2038
2457
  }
2039
2458
  return {
2040
2459
  ...state,
2041
- compaction: { status: 'completed', ...payload },
2460
+ compaction: {
2461
+ status: 'completed',
2462
+ ...payload,
2463
+ throughIndex:
2464
+ payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
2465
+ },
2042
2466
  }
2043
2467
  }
2044
2468
  default:
@@ -2061,59 +2485,84 @@ export function createReducer<D extends EventDefs>(options: {
2061
2485
  name?: string
2062
2486
  }): Reducer<D, AIState<AIMessageFrom<D>>> {
2063
2487
  type M = AIMessageFrom<D>
2064
- return options.contract
2488
+ const reducer = options.contract
2065
2489
  .reducer({
2066
- name: options.name ?? 'a2.ai.state.v10',
2490
+ name: options.name ?? 'a2.ai.state.v17',
2067
2491
  initialState: initialState<M>(),
2068
2492
  })
2069
2493
  .fold((state, event) => reduceAIState(state, event))
2494
+ return { ...reducer, [clientStateProjection]: projectAIClientState }
2070
2495
  }
2071
2496
 
2072
- type ProtocolInput<
2497
+ type CommandInput<C> = { type: 'ai.control.requested'; payload: C; id?: string }
2498
+ type ControlInput<
2073
2499
  M extends UIMessage,
2074
- K extends keyof AIEventDefs<M> & string,
2075
- > = AppendInput<Pick<AIEventDefs<M>, K>>
2500
+ A extends ControlCommand<M>['action'],
2501
+ > = CommandInput<Extract<ControlCommand<M>, { action: A }>>
2076
2502
 
2077
2503
  export type AIInputs = {
2078
- /** Record a message. User messages schedule unless generate is false. */
2079
2504
  message<M extends UIMessage>(
2080
2505
  message: M,
2081
2506
  options?: { generate?: boolean },
2082
- ): ProtocolInput<M, 'ai.message.created'>[]
2083
- /** Record a message without requesting a generation. */
2084
- seed<M extends UIMessage>(
2085
- message: M,
2086
- ): ProtocolInput<M, 'ai.message.created'>[]
2087
- /** Answer an AI SDK tool approval. The tool barrier owns continuation. */
2507
+ ): ControlInput<M, 'send'>[]
2508
+ seed<M extends UIMessage>(message: M): ControlInput<M, 'send'>[]
2509
+ queue: {
2510
+ remove(options: { inputId: string }): ControlInput<UIMessage, 'remove'>[]
2511
+ sendNow(options: {
2512
+ inputId: string
2513
+ turnId: string | null
2514
+ lastSeenIndex?: number
2515
+ }): ControlInput<UIMessage, 'send-now'>[]
2516
+ edit<M extends UIMessage>(options: {
2517
+ message: M
2518
+ expectedRevision: number
2519
+ }): ControlInput<M, 'edit'>[]
2520
+ move(options: {
2521
+ inputId: string
2522
+ beforeId: string | null
2523
+ }): ControlInput<UIMessage, 'move'>[]
2524
+ }
2525
+ stop(options: {
2526
+ turnId: string
2527
+ lastSeenIndex?: number
2528
+ }): ControlInput<UIMessage, 'stop'>[]
2529
+ steer<M extends UIMessage>(options: {
2530
+ turnId: string
2531
+ message: M
2532
+ lastSeenIndex?: number
2533
+ }): ControlInput<M, 'steer'>[]
2534
+ pause(options: {
2535
+ when: 'now' | 'after-turn'
2536
+ }): ControlInput<UIMessage, 'pause'>[]
2537
+ resume(): ControlInput<UIMessage, 'resume'>[]
2088
2538
  approval(
2089
2539
  response: ApprovalRespondedPayload,
2090
- ): ProtocolInput<UIMessage, 'ai.approval.responded'>[]
2091
- /** Answer an application-defined input request. */
2092
- input(
2093
- response: InputRespondedPayload,
2094
- ): ProtocolInput<UIMessage, 'ai.input.responded'>[]
2095
- /** Record an application-defined request for input. */
2540
+ ): ControlInput<UIMessage, 'approval'>[]
2541
+ input(response: InputRespondedPayload): ControlInput<UIMessage, 'input'>[]
2096
2542
  requestInput(
2097
2543
  request: InputRequestedPayload,
2098
- ): ProtocolInput<UIMessage, 'ai.input.requested'>[]
2099
- /** Request a fresh attempt for a failed response. */
2100
- retry(
2101
- options: RetryGenerationOptions,
2102
- ): ProtocolInput<UIMessage, 'ai.retry.requested'>[]
2103
- /** Interrupt the active response for a message. */
2544
+ ): ControlInput<UIMessage, 'request-input'>[]
2545
+ retry(request: RetryGenerationOptions): ControlInput<UIMessage, 'retry'>[]
2104
2546
  interrupt(
2105
- interruption: MessageInterruptedPayload,
2106
- ): ProtocolInput<UIMessage, 'ai.message.interrupted'>[]
2547
+ request: MessageInterruptedPayload,
2548
+ ): ControlInput<UIMessage, 'interrupt'>[]
2549
+ toolResult(
2550
+ result: ToolResultPayload,
2551
+ ): ControlInput<UIMessage, 'tool-result'>[]
2107
2552
  }
2108
2553
 
2554
+ const commandInput = <C extends ControlCommand<UIMessage>>(
2555
+ payload: C,
2556
+ ): CommandInput<C>[] => [{ type: 'ai.control.requested', payload }]
2109
2557
  const messageInput = <M extends UIMessage>(
2110
2558
  message: M,
2111
2559
  generate?: boolean,
2112
- ): ProtocolInput<M, 'ai.message.created'>[] => [
2560
+ ): ControlInput<M, 'send'>[] => [
2113
2561
  {
2114
- type: 'ai.message.created',
2562
+ type: 'ai.control.requested',
2115
2563
  id: `ai.message:${message.id}`,
2116
2564
  payload: {
2565
+ action: 'send',
2117
2566
  message,
2118
2567
  ...(generate === false ? { generate: false } : {}),
2119
2568
  },
@@ -2122,57 +2571,53 @@ const messageInput = <M extends UIMessage>(
2122
2571
 
2123
2572
  /** Pure append inputs for the built-in AI protocol. */
2124
2573
  export const inputs: AIInputs = {
2574
+ toolResult: (result) => commandInput({ action: 'tool-result', result }),
2575
+ queue: {
2576
+ remove: (options) => commandInput({ action: 'remove', ...options }),
2577
+ sendNow: (options) => commandInput({ action: 'send-now', ...options }),
2578
+ edit: (options) => commandInput({ action: 'edit', ...options }),
2579
+ move: (options) => commandInput({ action: 'move', ...options }),
2580
+ },
2581
+ stop: (options) => commandInput({ action: 'stop', ...options }),
2582
+ steer: (options) => commandInput({ action: 'steer', ...options }),
2583
+ pause: (options) => commandInput({ action: 'pause', ...options }),
2584
+ resume: () => commandInput({ action: 'resume' }),
2125
2585
  message<M extends UIMessage>(message: M, options?: { generate?: boolean }) {
2126
2586
  return messageInput(message, options?.generate)
2127
2587
  },
2128
2588
  seed<M extends UIMessage>(message: M) {
2129
2589
  return messageInput(message, false)
2130
2590
  },
2131
- approval(response) {
2132
- return [
2133
- {
2134
- type: 'ai.approval.responded',
2135
- id: `ai.approval:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.approvalId)}:response`,
2136
- payload: response,
2137
- },
2138
- ]
2139
- },
2140
- input(response) {
2141
- return [
2142
- {
2143
- type: 'ai.input.responded',
2144
- id: `ai.input:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.inputId)}:response`,
2145
- payload: response,
2146
- },
2147
- ]
2148
- },
2149
- requestInput(request) {
2150
- return [
2151
- {
2152
- type: 'ai.input.requested',
2153
- id: `ai.input:${encodeURIComponent(request.messageId)}:${encodeURIComponent(request.generationId)}:${encodeURIComponent(request.inputId)}:request`,
2154
- payload: request,
2155
- },
2156
- ]
2157
- },
2158
- retry(retryOptions) {
2159
- return [
2160
- {
2161
- type: 'ai.retry.requested',
2162
- id: `ai.retry:${retryOptions.retryId}`,
2163
- payload: retryOptions,
2164
- },
2165
- ]
2166
- },
2167
- interrupt(interruption) {
2168
- return [
2169
- {
2170
- type: 'ai.message.interrupted',
2171
- id: `ai.interrupt:${interruption.generationId ?? interruption.requestId}`,
2172
- payload: interruption,
2173
- },
2174
- ]
2175
- },
2591
+ approval: (response) => [
2592
+ {
2593
+ ...commandInput({ action: 'approval', response })[0]!,
2594
+ id: `ai.approval:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.approvalId)}:response`,
2595
+ },
2596
+ ],
2597
+ input: (response) => [
2598
+ {
2599
+ ...commandInput({ action: 'input', response })[0]!,
2600
+ id: `ai.input:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.inputId)}:response`,
2601
+ },
2602
+ ],
2603
+ requestInput: (request) => [
2604
+ {
2605
+ ...commandInput({ action: 'request-input', request })[0]!,
2606
+ id: `ai.input:${encodeURIComponent(request.messageId)}:${encodeURIComponent(request.generationId)}:${encodeURIComponent(request.inputId)}:request`,
2607
+ },
2608
+ ],
2609
+ retry: (request) => [
2610
+ {
2611
+ ...commandInput({ action: 'retry', request })[0]!,
2612
+ id: `ai.retry:${request.retryId}`,
2613
+ },
2614
+ ],
2615
+ interrupt: (request) => [
2616
+ {
2617
+ ...commandInput({ action: 'interrupt', request })[0]!,
2618
+ id: `ai.interrupt:${request.generationId ?? request.requestId}`,
2619
+ },
2620
+ ],
2176
2621
  }
2177
2622
 
2178
2623
  /**
@@ -2229,14 +2674,14 @@ export function agent<
2229
2674
  })
2230
2675
  const reducer = agentContract
2231
2676
  .reducer({
2232
- name: options.reducerName ?? 'a2.ai.state.v10',
2677
+ name: options.reducerName ?? 'a2.ai.state.v17',
2233
2678
  initialState: initialState<M>(),
2234
2679
  })
2235
2680
  .fold((state, event) => reduceAIState(state, event))
2236
2681
 
2237
2682
  const definition = {
2238
2683
  contract: agentContract,
2239
- reducer,
2684
+ reducer: { ...reducer, [clientStateProjection]: projectAIClientState<M> },
2240
2685
  }
2241
2686
  return definition as unknown as AgentDefinition<M, AIEventDefs<M> & E>
2242
2687
  }