experimental-a2 0.3.0 → 0.5.0

package/src/ai-sdk-step.ts

@@ -0,0 +1,261 @@
+import {
+  convertToModelMessages,
+  stepCountIs,
+  streamText,
+  toUIMessageStream,
+} from 'ai'
+import type {
+  FinishReason,
+  InferUIMessageChunk,
+  Instructions,
+  LanguageModel,
+  LanguageModelUsage,
+  TextStreamPart,
+  ToolSet,
+  UIMessage,
+} from 'ai'
+
+type StreamTextOptions<T extends ToolSet> = Parameters<typeof streamText<T>>[0]
+
+type UIMessageMetadata<M extends UIMessage> =
+  M extends UIMessage<infer Metadata> ? Metadata : unknown
+
+type ModelStepTimeout<T extends ToolSet> =
+  | number
+  | Omit<
+      Exclude<StreamTextOptions<T>['timeout'], number | undefined>,
+      'toolMs' | 'tools'
+    >
+
+type ControlledStreamTextSetting =
+  | '_internal'
+  | 'abortSignal'
+  | 'instructions'
+  | 'messages'
+  | 'model'
+  | 'onToolExecutionEnd'
+  | 'onToolExecutionStart'
+  | 'experimental_onToolCallFinish'
+  | 'experimental_onToolCallStart'
+  | 'experimental_sandbox'
+  | 'experimental_toolCallers'
+  | 'experimental_toolApprovalSecret'
+  | 'prompt'
+  | 'prepareStep'
+  | 'stopWhen'
+  | 'system'
+  | 'tools'
+
+type RetypedStreamTextSetting =
+  | 'experimental_onStart'
+  | 'experimental_onStepStart'
+  | 'experimental_repairToolCall'
+  | 'experimental_transform'
+  | 'onStart'
+  | 'onStepStart'
+  | 'repairToolCall'
+
+type ModelFacingTools<T extends ToolSet> = {
+  readonly [K in keyof T]: Omit<T[K], 'execute'>
+}
+
+type SettingEvent<T extends ToolSet, K extends keyof StreamTextOptions<T>> =
+  NonNullable<StreamTextOptions<T>[K]> extends (event: infer EVENT) => unknown
+    ? EVENT
+    : never
+
+type ModelPhaseStartEvent<T extends ToolSet> = Omit<
+  SettingEvent<T, 'onStart'>,
+  'tools'
+> & {
+  readonly tools: ModelFacingTools<T> | undefined
+}
+
+type ModelPhaseStepStartEvent<T extends ToolSet> = Omit<
+  SettingEvent<T, 'onStepStart'>,
+  'tools'
+> & {
+  readonly tools: ModelFacingTools<T> | undefined
+}
+
+type ModelPhaseCallback<EVENT> = (event: EVENT) => void | PromiseLike<void>
+
+type StreamTextRepairToolCall<T extends ToolSet> = NonNullable<
+  StreamTextOptions<T>['repairToolCall']
+>
+
+type ModelPhaseRepairToolCall<T extends ToolSet> = (
+  options: Omit<Parameters<StreamTextRepairToolCall<T>>[0], 'tools'> & {
+    tools: ModelFacingTools<T>
+  },
+) => ReturnType<StreamTextRepairToolCall<T>>
+
+type StreamTextTransformSetting<T extends ToolSet> = NonNullable<
+  StreamTextOptions<T>['experimental_transform']
+>
+
+type StreamTextTransform<T extends ToolSet> = Extract<
+  StreamTextTransformSetting<T>,
+  (...args: never[]) => unknown
+>
+
+type ModelPhaseTransform<T extends ToolSet> = (
+  options: Omit<Parameters<StreamTextTransform<T>>[0], 'tools'> & {
+    tools: ModelFacingTools<T>
+  },
+) => ReturnType<StreamTextTransform<T>>
+
+/** AI SDK streamText settings that remain meaningful for one durable step. */
+export type AISDKStepSettings<T extends ToolSet> = Partial<
+  Omit<
+    StreamTextOptions<T>,
+    ControlledStreamTextSetting | RetypedStreamTextSetting | 'timeout'
+  >
+> & {
+  timeout?: ModelStepTimeout<T>
+  onStart?: ModelPhaseCallback<ModelPhaseStartEvent<T>>
+  experimental_onStart?: ModelPhaseCallback<ModelPhaseStartEvent<T>>
+  onStepStart?: ModelPhaseCallback<ModelPhaseStepStartEvent<T>>
+  experimental_onStepStart?: ModelPhaseCallback<ModelPhaseStepStartEvent<T>>
+  repairToolCall?: ModelPhaseRepairToolCall<T>
+  experimental_repairToolCall?: ModelPhaseRepairToolCall<T>
+  experimental_transform?:
+    ModelPhaseTransform<T> | Array<ModelPhaseTransform<T>>
+}
+
+export type AISDKStepInput<M extends UIMessage, T extends ToolSet> = {
+  model: LanguageModel
+  tools: T
+  messages: M[]
+  responseMessageId: string
+  abortSignal: AbortSignal
+  instructions?: Instructions
+  settings?: AISDKStepSettings<T>
+  messageMetadata?: (options: {
+    part: TextStreamPart<T>
+  }) => UIMessageMetadata<M> | undefined
+}
+
+export type AISDKStepCompletion = {
+  finishReason: FinishReason
+  usage: LanguageModelUsage
+}
+
+export type AISDKStepResult<M extends UIMessage> = {
+  stream: ReadableStream<InferUIMessageChunk<M>>
+  completion(): AISDKStepCompletion
+}
+
+const CONTROLLED_SETTINGS: readonly ControlledStreamTextSetting[] = [
+  '_internal',
+  'abortSignal',
+  'instructions',
+  'messages',
+  'model',
+  'onToolExecutionEnd',
+  'onToolExecutionStart',
+  'experimental_onToolCallFinish',
+  'experimental_onToolCallStart',
+  'experimental_sandbox',
+  'experimental_toolCallers',
+  'experimental_toolApprovalSecret',
+  'prompt',
+  'prepareStep',
+  'stopWhen',
+  'system',
+  'tools',
+]
+
+const safeSettings = <T extends ToolSet>(
+  settings: AISDKStepSettings<T> | undefined,
+): AISDKStepSettings<T> => {
+  const result = { ...settings } as AISDKStepSettings<T>
+  for (const key of CONTROLLED_SETTINGS) Reflect.deleteProperty(result, key)
+  if (typeof result.timeout === 'object' && result.timeout !== null) {
+    const timeout = { ...result.timeout } as Record<string, unknown>
+    Reflect.deleteProperty(timeout, 'toolMs')
+    Reflect.deleteProperty(timeout, 'tools')
+    result.timeout = timeout as ModelStepTimeout<T>
+  }
+  return result
+}
+
+const modelToolSet = <T extends ToolSet>(tools: T): T =>
+  Object.fromEntries(
+    Object.entries(tools).map(([name, tool]) => {
+      if (tool.type === 'provider' && tool.isProviderExecuted === true) {
+        return [name, tool]
+      }
+      const definition = { ...tool }
+      Reflect.deleteProperty(definition, 'execute')
+      return [name, definition]
+    }),
+  ) as T
+
+const errorMessage = (error: unknown): string =>
+  error instanceof Error ? error.message : String(error)
+
+/** Run exactly one AI SDK model step while leaving local tool execution to A2. */
+export async function generateAISDKStep<M extends UIMessage, T extends ToolSet>(
+  input: AISDKStepInput<M, T>,
+): Promise<AISDKStepResult<M>> {
+  const settings = safeSettings(input.settings)
+  const tools = modelToolSet(input.tools)
+  const messages = await convertToModelMessages(input.messages, { tools })
+  const streamOptions = {
+    ...settings,
+    model: input.model,
+    tools,
+    messages,
+    abortSignal: input.abortSignal,
+    ...(input.instructions === undefined
+      ? {}
+      : { instructions: input.instructions }),
+    stopWhen: stepCountIs(1),
+    // A2 records stream failures durably; the AI SDK default writes them to
+    // stderr in addition to yielding an error part.
+    onError: settings.onError ?? (() => {}),
+  } as unknown as StreamTextOptions<T>
+  const result = streamText<T>(streamOptions)
+
+  let completion: AISDKStepCompletion | undefined
+  let fatalError: { value: unknown } | undefined
+  const observed = result.stream.pipeThrough(
+    new TransformStream<TextStreamPart<T>, TextStreamPart<T>>({
+      transform(part, controller) {
+        if (part.type === 'error' && fatalError === undefined) {
+          fatalError = { value: part.error }
+        } else if (part.type === 'finish') {
+          completion = {
+            finishReason: part.finishReason,
+            usage: part.totalUsage,
+          }
+        }
+        controller.enqueue(part)
+      },
+    }),
+  )
+
+  return {
+    stream: toUIMessageStream<T, M>({
+      stream: observed,
+      tools,
+      originalMessages: input.messages,
+      generateMessageId: () => input.responseMessageId,
+      onError: errorMessage,
+      ...(input.messageMetadata === undefined
+        ? {}
+        : { messageMetadata: input.messageMetadata }),
+    }),
+    completion() {
+      if (fatalError !== undefined) throw fatalError.value
+      if (completion === undefined) {
+        throw new Error('AI SDK step ended without a finish part')
+      }
+      if (completion.finishReason === 'error') {
+        throw new Error('AI SDK step finished with an error')
+      }
+      return completion
+    },
+  }
+}

package/src/ai-server.browser.ts

@@ -0,0 +1,5 @@
+/** experimental-a2/ai/server is server-only; agent definitions live in experimental-a2/ai. */
+throw new Error(
+  'experimental-a2/ai/server is server-only — a client bundle imported it. Import agent from experimental-a2/ai in browser code instead.',
+)
+export {}