remove getStream function

Author: charleslien

packages/agent-runtime/src/prompt-agent-stream.ts

@@ -14,52 +14,54 @@ import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { OpenRouterProviderOptions } from '@codebuff/internal/openrouter-ai-sdk'
 
 export const getAgentStreamFromTemplate = (params: {
+  agentId?: string
   apiKey: string
-  runId: string
   clientSessionId: string
   fingerprintId: string
-  userInputId: string
-  userId: string | undefined
-  onCostCalculated?: (credits: number) => Promise<void>
-  agentId?: string
   includeCacheControl?: boolean
+  liveUserInputRecord: UserInputRecord
+  logger: Logger
+  messages: Message[]
+  runId: string
+  sessionConnections: SessionRecord
+  template: AgentTemplate
   textOverride: string | null
+  userId: string | undefined
+  userInputId: string
 
-  template: AgentTemplate
-  logger: Logger
-  sendAction: SendActionFn
+  onCostCalculated?: (credits: number) => Promise<void>
   promptAiSdkStream: PromptAiSdkStreamFn
-  liveUserInputRecord: UserInputRecord
-  sessionConnections: SessionRecord
+  sendAction: SendActionFn
   trackEvent: TrackEventFn
-}): { getStream: (messages: Message[]) => ReturnType<PromptAiSdkStreamFn> } => {
+}): ReturnType<PromptAiSdkStreamFn> => {
   const {
+    agentId,
     apiKey,
-    runId,
     clientSessionId,
     fingerprintId,
-    userInputId,
-    userId,
-    onCostCalculated,
-    agentId,
     includeCacheControl,
-    textOverride,
-    template,
+    liveUserInputRecord,
     logger,
+    messages,
+    runId,
+    sessionConnections,
+    template,
+    textOverride,
+    userId,
+    userInputId,
+
     sendAction,
+    onCostCalculated,
     promptAiSdkStream,
-    liveUserInputRecord,
-    sessionConnections,
     trackEvent,
   } = params
 
   if (textOverride !== null) {
-    return {
-      getStream: async function* stream(): ReturnType<PromptAiSdkStreamFn> {
-        yield { type: 'text', text: textOverride!, agentId }
-        return crypto.randomUUID()
-      },
+    async function* stream(): ReturnType<PromptAiSdkStreamFn> {
+      yield { type: 'text', text: textOverride!, agentId }
+      return crypto.randomUUID()
     }
+    return stream()
   }
 
   if (!template) {
@@ -68,46 +70,42 @@ export const getAgentStreamFromTemplate = (params: {
 
   const { model } = template
 
-  const getStream = (messages: Message[]): ReturnType<PromptAiSdkStreamFn> => {
-    const aiSdkStreamParams: ParamsOf<PromptAiSdkStreamFn> = {
-      apiKey,
-      runId,
-      messages,
-      model,
-      stopSequences: [globalStopSequence],
-      clientSessionId,
-      fingerprintId,
-      userInputId,
-      userId,
-      maxOutputTokens: 32_000,
-      onCostCalculated,
-      includeCacheControl,
-      agentId,
-      maxRetries: 3,
-      sendAction,
-      liveUserInputRecord,
-      sessionConnections,
-      logger,
-      trackEvent,
-    }
+  const aiSdkStreamParams: ParamsOf<PromptAiSdkStreamFn> = {
+    apiKey,
+    runId,
+    messages,
+    model,
+    stopSequences: [globalStopSequence],
+    clientSessionId,
+    fingerprintId,
+    userInputId,
+    userId,
+    maxOutputTokens: 32_000,
+    onCostCalculated,
+    includeCacheControl,
+    agentId,
+    maxRetries: 3,
+    sendAction,
+    liveUserInputRecord,
+    sessionConnections,
+    logger,
+    trackEvent,
+  }
 
-    if (!aiSdkStreamParams.providerOptions) {
-      aiSdkStreamParams.providerOptions = {}
-    }
-    for (const provider of ['openrouter', 'codebuff'] as const) {
-      if (!aiSdkStreamParams.providerOptions[provider]) {
-        aiSdkStreamParams.providerOptions[provider] = {}
-      }
-      ;(
-        aiSdkStreamParams.providerOptions[provider] as OpenRouterProviderOptions
-      ).reasoning = template.reasoningOptions
+  if (!aiSdkStreamParams.providerOptions) {
+    aiSdkStreamParams.providerOptions = {}
+  }
+  for (const provider of ['openrouter', 'codebuff'] as const) {
+    if (!aiSdkStreamParams.providerOptions[provider]) {
+      aiSdkStreamParams.providerOptions[provider] = {}
     }
-
-    // Pass agent's provider routing options to SDK
-    aiSdkStreamParams.agentProviderOptions = template.providerOptions
-
-    return promptAiSdkStream(aiSdkStreamParams)
+    ;(
+      aiSdkStreamParams.providerOptions[provider] as OpenRouterProviderOptions
+    ).reasoning = template.reasoningOptions
   }
 
-  return { getStream }
+  // Pass agent's provider routing options to SDK
+  aiSdkStreamParams.agentProviderOptions = template.providerOptions
+
+  return promptAiSdkStream(aiSdkStreamParams)
 }

packages/agent-runtime/src/run-agent-step.ts

@@ -90,7 +90,11 @@ export const runAgentStep = async (
   > &
     ParamsExcluding<
       typeof getAgentStreamFromTemplate,
-      'agentId' | 'template' | 'onCostCalculated' | 'includeCacheControl'
+      | 'agentId'
+      | 'includeCacheControl'
+      | 'messages'
+      | 'onCostCalculated'
+      | 'template'
     > &
     ParamsExcluding<typeof getAgentTemplate, 'agentId'> &
     ParamsExcluding<
@@ -327,19 +331,15 @@ export const runAgentStep = async (
   let fullResponse = ''
   const toolResults: ToolMessage[] = []
 
-  const { getStream } = getAgentStreamFromTemplate({
+  const stream = getAgentStreamFromTemplate({
     ...params,
     agentId: agentState.parentId ? agentState.agentId : undefined,
+    includeCacheControl: supportsCacheControl(agentTemplate.model),
+    messages: [systemMessage(system), ...agentState.messageHistory],
     template: agentTemplate,
     onCostCalculated,
-    includeCacheControl: supportsCacheControl(agentTemplate.model),
   })
 
-  const stream = getStream([
-    systemMessage(system),
-    ...agentState.messageHistory,
-  ])
-
   const {
     fullResponse: fullResponseAfterStream,
     fullResponseChunks,