@lota-sdk/core 0.4.8 → 0.4.9

package/src/services/thread/thread-turn-streaming.ts

@@ -0,0 +1,402 @@
+import type { ChatMessage } from '@lota-sdk/shared'
+import { withMessageCreatedAt } from '@lota-sdk/shared'
+import { convertToModelMessages, stepCountIs } from 'ai'
+import type { PrepareStepFunction, StopCondition, ToolSet, UIMessageStreamWriter } from 'ai'
+import { Effect, Ref, Schema, Stream } from 'effect'
+
+import { getAgentRuntimeConfig, getResolvedAgentFactoryConfig } from '../../config/agent-defaults'
+import { aiLogger } from '../../config/logger'
+import type { RecordIdRef } from '../../db/record-id'
+import { effectTryMaybeAsync, effectTryPromise as effectTryPromiseShared } from '../../effect/helpers'
+import { runPromise } from '../../effect/runtime'
+import {
+  readRuntimeAgentIdentityOverrides,
+  resolveRuntimeAgentDisplayName,
+} from '../../runtime/agent-identity-overrides'
+import { createAgentMessageMetadata } from '../../runtime/agent-stream-helpers'
+import { mergeInstructionSections } from '../../runtime/instruction-sections'
+import type { getTurnHooks } from '../../runtime/runtime-extensions'
+import {
+  asRecord,
+  collectToolOutputErrors,
+  extractMessageText,
+  readInstructionSections,
+  readOptionalString,
+} from '../../runtime/thread-chat-helpers'
+import { nowEpochMillis } from '../../utils/date-time'
+import { isRecord } from '../../utils/string'
+import { buildThreadTurnSpanAttributes, compactSpanAttributes } from './thread-turn-tracing'
+import type { NormalizedThread } from './thread.types'
+
+type ChatStreamChunk = Parameters<UIMessageStreamWriter<ChatMessage>['write']>[0]
+
+interface UIMessageStreamResult {
+  toUIMessageStream(options: Record<string, unknown>): ReadableStream<unknown>
+}
+
+interface ToolLoopGenerateResult {
+  text: string
+  steps?: Array<{ toolResults?: Array<{ toolCallId: string; toolName: string; input?: unknown; output?: unknown }> }>
+}
+
+export class ThreadTurnStreamingError extends Schema.TaggedErrorClass<ThreadTurnStreamingError>()(
+  'ThreadTurnStreamingError',
+  { message: Schema.String, cause: Schema.optional(Schema.Defect) },
+) {}
+
+function effectTryPromise<A>(
+  evaluate: () => PromiseLike<A>,
+  message: string,
+): Effect.Effect<A, ThreadTurnStreamingError> {
+  return effectTryPromiseShared(evaluate, (error) => new ThreadTurnStreamingError({ message, cause: error }))
+}
+
+function effectFromMaybeEffect<A>(
+  evaluate: () => A | PromiseLike<A> | Effect.Effect<A, ThreadTurnStreamingError>,
+  message: string,
+): Effect.Effect<A, ThreadTurnStreamingError> {
+  return Effect.try({ try: evaluate, catch: (error) => new ThreadTurnStreamingError({ message, cause: error }) }).pipe(
+    Effect.flatMap((result) =>
+      Effect.isEffect(result)
+        ? result.pipe(Effect.mapError((error) => new ThreadTurnStreamingError({ message, cause: error })))
+        : effectTryMaybeAsync(
+            () => result,
+            (error) => new ThreadTurnStreamingError({ message, cause: error }),
+          ),
+    ),
+  )
+}
+
+function hasUIMessageStream(value: unknown): value is UIMessageStreamResult {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    typeof (value as { toUIMessageStream?: unknown }).toUIMessageStream === 'function'
+  )
+}
+
+function optionalInstructionSection(value: unknown): string[] | undefined {
+  const section = readOptionalString(value)
+  return section ? [section] : undefined
+}
+
+function readStreamChunkType(value: ChatStreamChunk): string | undefined {
+  if (!isRecord(value)) return undefined
+  return typeof value.type === 'string' ? value.type : undefined
+}
+
+function isVisibleOutputChunkType(chunkType: string | undefined): boolean {
+  return chunkType === 'text-delta' || chunkType === 'reasoning-delta'
+}
+
+function isTextTokenChunkType(chunkType: string | undefined): boolean {
+  return chunkType === 'text-delta'
+}
+
+function buildFallbackResponseMessage(result: ToolLoopGenerateResult): ChatMessage {
+  const parts: ChatMessage['parts'] = []
+
+  for (const step of result.steps ?? []) {
+    for (const toolResult of step.toolResults ?? []) {
+      parts.push({
+        type: `tool-${toolResult.toolName}`,
+        toolCallId: toolResult.toolCallId,
+        state: 'output-available',
+        ...(toolResult.input !== undefined ? { input: toolResult.input } : {}),
+        ...(toolResult.output !== undefined ? { output: toolResult.output } : {}),
+      } as ChatMessage['parts'][number])
+    }
+  }
+
+  const text = result.text.trim()
+  if (text.length > 0) {
+    parts.push({ type: 'text', text })
+  }
+
+  if (parts.length === 0) {
+    throw new Error('Agent generate fallback did not produce any response parts.')
+  }
+
+  return { id: Bun.randomUUIDv7(), role: 'assistant', parts }
+}
+
+export interface StreamAgentResponseContext {
+  turnHooks: ReturnType<typeof getTurnHooks>
+  thread: NormalizedThread
+  threadRef: RecordIdRef
+  orgRef: RecordIdRef
+  userRef: RecordIdRef
+  userName?: string | null
+  onboardingActive: boolean
+  linearInstalled: boolean
+  githubInstalled: boolean
+  buildContextResult: Record<string, unknown> | null
+  getExecutionPlanInstructionSections: () => Effect.Effect<string[] | undefined, ThreadTurnStreamingError>
+  getPreSeededMemoriesSection: (agentId: string) => Effect.Effect<string | undefined, ThreadTurnStreamingError>
+  getLearnedSkillsSection: (
+    agentId: string,
+    queryText?: string,
+  ) => Effect.Effect<string | undefined, ThreadTurnStreamingError>
+  promptContext: { systemWorkspaceDetails?: string }
+  retrievedKnowledgeSection: string | undefined
+  memoryBlock: string
+  hookInstructionSections: string[]
+  runAbortSignal: AbortSignal
+}
+
+interface StreamAgentResponseParams {
+  agentId: string
+  mode: 'direct' | 'fixedThreadMode' | 'threadMode'
+  messages: ChatMessage[]
+  tools: ToolSet
+  observer: {
+    run: <T>(fn: () => T | Promise<T>) => Promise<T>
+    recordError: (error: unknown) => void
+    recordAbort: (error: unknown) => void
+  }
+  skills?: string[]
+  additionalInstructionSections?: string[]
+  includeExecutionPlanTools?: boolean
+  writer?: UIMessageStreamWriter<ChatMessage>
+  stopWhen?: StopCondition<ToolSet> | Array<StopCondition<ToolSet>>
+  prepareStep?: PrepareStepFunction<ToolSet>
+  abortSignal?: AbortSignal
+}
+
+const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResponse')(function* (
+  ctx: StreamAgentResponseContext,
+  streamParams: StreamAgentResponseParams,
+) {
+  const agentIdentityOverrides = readRuntimeAgentIdentityOverrides(ctx.buildContextResult)
+  const executionPlanInstructionSections =
+    streamParams.includeExecutionPlanTools === false
+      ? undefined
+      : yield* effectFromMaybeEffect(
+          () => ctx.getExecutionPlanInstructionSections(),
+          'Failed to load execution plan instructions.',
+        ).pipe(Effect.withSpan('ThreadTurnStreaming.loadExecutionPlanInstructions'))
+
+  const agentResolution = asRecord(
+    yield* effectTryMaybeAsync(
+      () =>
+        ctx.turnHooks.resolveAgent?.({
+          agentId: streamParams.agentId,
+          mode: streamParams.mode,
+          thread: ctx.thread,
+          threadRef: ctx.threadRef,
+          orgRef: ctx.orgRef,
+          userRef: ctx.userRef,
+          userName: ctx.userName,
+          onboardingActive: ctx.onboardingActive,
+          linearInstalled: ctx.linearInstalled,
+          githubInstalled: ctx.githubInstalled,
+          skills: streamParams.skills,
+          additionalInstructionSections: streamParams.additionalInstructionSections,
+          context: ctx.buildContextResult,
+        }),
+      (error) => new ThreadTurnStreamingError({ message: 'Failed to resolve the runtime agent.', cause: error }),
+    ).pipe(Effect.withSpan('ThreadTurnStreaming.resolveAgent')),
+  )
+
+  const resolvedAgentId = readOptionalString(agentResolution?.agentId) ?? streamParams.agentId
+  const latestUserMessage = [...streamParams.messages].reverse().find((message) => message.role === 'user')
+  const latestUserMessageText = latestUserMessage ? extractMessageText(latestUserMessage).trim() : undefined
+  const [preSeededMemoriesSection, learnedSkillsSection] = yield* Effect.all([
+    effectFromMaybeEffect(
+      () => ctx.getPreSeededMemoriesSection(resolvedAgentId),
+      `Failed to load pre-seeded memories for ${resolvedAgentId}.`,
+    ),
+    effectFromMaybeEffect(
+      () => ctx.getLearnedSkillsSection(resolvedAgentId, latestUserMessageText),
+      `Failed to load learned skills for ${resolvedAgentId}.`,
+    ),
+  ]).pipe(Effect.withSpan('ThreadTurnStreaming.loadMemoriesAndSkills'))
+
+  const toolNames = new Set(Object.keys(streamParams.tools))
+  const hasRetrievalTools = [
+    'memorySearch',
+    'conversationSearch',
+    'queryKnowledge',
+    'researchTopic',
+    'fetchWebpage',
+    'inspectWebsite',
+  ].some((toolName) => toolNames.has(toolName))
+  const hasDomainRoutingSkills =
+    (streamParams.skills ?? []).some((skill) => skill.startsWith('cpo-') || skill.startsWith('mentor-')) ||
+    resolvedAgentId === 'cpo' ||
+    resolvedAgentId === 'mentor'
+  const agentFactoryConfig = getResolvedAgentFactoryConfig()
+  const config = getAgentRuntimeConfig({
+    agentId: resolvedAgentId,
+    threadType: ctx.thread.type,
+    mode: streamParams.mode,
+    skills: streamParams.skills,
+    onboardingActive: ctx.onboardingActive,
+    linearInstalled: ctx.linearInstalled,
+    systemWorkspaceDetails: ctx.promptContext.systemWorkspaceDetails,
+    preSeededMemoriesSection,
+    retrievedKnowledgeSection: ctx.retrievedKnowledgeSection,
+    threadMemoryBlock: ctx.memoryBlock,
+    learnedSkillsSection,
+    userMessageText: latestUserMessageText,
+    ruleOptions: { includeMemr3Rule: hasRetrievalTools, includeDomainReasoningFallbackRule: hasDomainRoutingSkills },
+    additionalInstructionSections: mergeInstructionSections(
+      executionPlanInstructionSections,
+      streamParams.additionalInstructionSections,
+      ctx.hookInstructionSections,
+      readInstructionSections(agentResolution?.additionalInstructionSections),
+      optionalInstructionSection(agentResolution?.extraInstructions),
+    ),
+    context: ctx.buildContextResult,
+  })
+  yield* Effect.annotateCurrentSpan({
+    resolvedAgentId,
+    inputMessageCount: streamParams.messages.length,
+    toolCount: toolNames.size,
+    skillCount: streamParams.skills?.length ?? 0,
+    hasRetrievalTools,
+  })
+
+  const modelMessages = yield* effectTryPromise(
+    () => convertToModelMessages(streamParams.messages, { ignoreIncompleteToolCalls: true }),
+    'Failed to convert UI messages to model messages.',
+  ).pipe(Effect.withSpan('ThreadTurnStreaming.convertModelMessages'))
+
+  const agentFactory = agentFactoryConfig.createAgent[config.id]
+  if (!agentFactory) {
+    return yield* new ThreadTurnStreamingError({ message: `Agent factory "${config.id}" is not registered.` })
+  }
+  const agent = agentFactory({
+    mode: streamParams.mode,
+    tools: streamParams.tools,
+    extraInstructions: config.extraInstructions,
+    maxRetries: 3,
+    stopWhen: (agentResolution?.stopWhen as StopCondition<ToolSet> | Array<StopCondition<ToolSet>> | undefined) ??
+      streamParams.stopWhen ?? [stepCountIs(config.maxSteps)],
+    prepareStep: (agentResolution?.prepareStep as PrepareStepFunction<ToolSet> | undefined) ?? streamParams.prepareStep,
+  })
+  const agentAbortSignal = streamParams.abortSignal ?? ctx.runAbortSignal
+
+  const generateFallback = (cause: ThreadTurnStreamingError) =>
+    effectTryPromise(
+      () => streamParams.observer.run(() => agent.generate({ messages: modelMessages, abortSignal: agentAbortSignal })),
+      `Agent generate fallback failed for ${resolvedAgentId}.`,
+    ).pipe(
+      Effect.tap(() =>
+        Effect.sync(() => {
+          aiLogger.warn`Agent stream failed for ${resolvedAgentId}; falling back to generate: ${cause.message}`
+        }),
+      ),
+      Effect.map((result) => buildFallbackResponseMessage(result as ToolLoopGenerateResult)),
+    )
+
+  const result = yield* effectTryPromise(
+    () => streamParams.observer.run(() => agent.stream({ messages: modelMessages, abortSignal: agentAbortSignal })),
+    `Agent stream failed for ${resolvedAgentId}.`,
+  ).pipe(
+    Effect.tapError((error) =>
+      Effect.sync(() => {
+        if (agentAbortSignal.aborted) {
+          streamParams.observer.recordAbort(error)
+          return
+        }
+
+        streamParams.observer.recordError(error)
+      }),
+    ),
+    Effect.withSpan('ThreadTurnStreaming.startAgentStream'),
+  )
+
+  if (!hasUIMessageStream(result)) {
+    return yield* generateFallback(
+      new ThreadTurnStreamingError({ message: `Agent run for ${resolvedAgentId} did not expose a UI message stream.` }),
+    )
+  }
+
+  const { promise: finishedStream, resolve: resolveFinishedStream } = Promise.withResolvers<ChatMessage>()
+  const uiStream = result.toUIMessageStream({
+    generateMessageId: () => Bun.randomUUIDv7(),
+    originalMessages: streamParams.messages,
+    sendReasoning: true,
+    sendSources: true,
+    messageMetadata: createAgentMessageMetadata({
+      agentId: resolvedAgentId,
+      agentName: resolveRuntimeAgentDisplayName(agentIdentityOverrides, resolvedAgentId),
+    }),
+    onFinish: ({ responseMessage: finishedResponseMessage }: { responseMessage: ChatMessage }) => {
+      resolveFinishedStream(withMessageCreatedAt(finishedResponseMessage, nowEpochMillis()))
+    },
+  }) as ReadableStream<ChatStreamChunk>
+  const streamStartedAt = performance.now()
+  const firstVisibleOutputRecorded = yield* Ref.make(false)
+  const firstTextTokenRecorded = yield* Ref.make(false)
+
+  const streamedResponse = yield* Stream.fromReadableStream({
+    evaluate: () => uiStream,
+    onError: (cause) => new ThreadTurnStreamingError({ message: 'Failed to read chat stream chunk.', cause }),
+    releaseLockOnEnd: true,
+  }).pipe(
+    Stream.runForEach((value) =>
+      Effect.gen(function* () {
+        const chunkType = readStreamChunkType(value)
+        const elapsedMs = Number((performance.now() - streamStartedAt).toFixed(1))
+        const recordFirstVisibleOutput = yield* Ref.modify(firstVisibleOutputRecorded, (recorded) => [
+          !recorded && isVisibleOutputChunkType(chunkType),
+          recorded || isVisibleOutputChunkType(chunkType),
+        ])
+        const recordFirstTextToken = yield* Ref.modify(firstTextTokenRecorded, (recorded) => [
+          !recorded && isTextTokenChunkType(chunkType),
+          recorded || isTextTokenChunkType(chunkType),
+        ])
+
+        if (recordFirstVisibleOutput) {
+          yield* Effect.annotateCurrentSpan({ firstVisibleOutputMs: elapsedMs, firstVisibleOutputChunkType: chunkType })
+        }
+
+        if (recordFirstTextToken) {
+          yield* Effect.annotateCurrentSpan({ ttftMs: elapsedMs, ttftChunkType: chunkType })
+        }
+
+        if (streamParams.writer) {
+          yield* Effect.sync(() => {
+            streamParams.writer?.write(value)
+          })
+        }
+      }),
+    ),
+    Effect.withSpan('ThreadTurnStreaming.consumeUiStream'),
+    Effect.andThen(
+      effectTryPromise(() => finishedStream, `Agent run for ${resolvedAgentId} did not produce a response message.`),
+    ),
+    Effect.catchTag('ThreadTurnStreamingError', generateFallback),
+  )
+
+  for (const toolError of collectToolOutputErrors({ responseMessage: streamedResponse })) {
+    aiLogger.error`Tool execution failed (agent=${resolvedAgentId}, tool=${toolError.toolName}, toolCallId=${toolError.toolCallId}): ${toolError.errorText}`
+  }
+
+  return streamedResponse
+})
+
+export function streamAgentResponse(
+  ctx: StreamAgentResponseContext,
+  streamParams: StreamAgentResponseParams,
+): Promise<ChatMessage> {
+  return runPromise(
+    streamAgentResponseEffect(ctx, streamParams).pipe(
+      Effect.annotateSpans(
+        compactSpanAttributes({
+          ...buildThreadTurnSpanAttributes({
+            threadRef: ctx.threadRef,
+            orgRef: ctx.orgRef,
+            userRef: ctx.userRef,
+            agentId: streamParams.agentId,
+            threadType: ctx.thread.type,
+            mode: streamParams.mode,
+          }),
+          includeExecutionPlanTools: streamParams.includeExecutionPlanTools ?? true,
+        }),
+      ),
+    ),
+  )
+}

package/src/services/thread/thread-turn-tracing.ts

@@ -0,0 +1,35 @@
+import type { RecordIdRef } from '../../db/record-id'
+import { recordIdToString } from '../../db/record-id'
+import { TABLES } from '../../db/tables'
+
+export function compactSpanAttributes(attributes: Record<string, unknown>): Record<string, unknown> {
+  return Object.fromEntries(
+    Object.entries(attributes).filter(([, value]) => value !== undefined && value !== null && value !== ''),
+  )
+}
+
+export function buildThreadTurnSpanAttributes(params: {
+  threadRef?: RecordIdRef
+  orgRef?: RecordIdRef
+  userRef?: RecordIdRef
+  kind?: string
+  streamId?: string
+  agentId?: string
+  threadType?: string
+  mode?: string
+  planRunId?: string
+  planNodeId?: string
+}) {
+  return compactSpanAttributes({
+    threadId: params.threadRef ? recordIdToString(params.threadRef, TABLES.THREAD) : undefined,
+    orgId: params.orgRef ? recordIdToString(params.orgRef, TABLES.ORGANIZATION) : undefined,
+    userId: params.userRef ? recordIdToString(params.userRef, TABLES.USER) : undefined,
+    turnKind: params.kind,
+    streamId: params.streamId,
+    agentId: params.agentId,
+    threadType: params.threadType,
+    threadMode: params.mode,
+    planRunId: params.planRunId,
+    planNodeId: params.planNodeId,
+  })
+}