@lota-sdk/core 0.4.8 → 0.4.10

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

@@ -0,0 +1,280 @@
+import type { ChatMessage, MessageMetadata, PlanNodeSpecRecord } from '@lota-sdk/shared'
+import { SUBMIT_PLAN_TURN_RESULT_TOOL_NAME, toTimestamp, withMessageCreatedAt } from '@lota-sdk/shared'
+import type { ToolSet, UIMessageStreamWriter } from 'ai'
+import { Schema, Effect } from 'effect'
+
+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 type { readRuntimeAgentIdentityOverrides } from '../../runtime/agent-identity-overrides'
+import { resolveRuntimeAgentDisplayName } from '../../runtime/agent-identity-overrides'
+import { OWNERSHIP_DISPATCH_BLOCKED_TOOL_NAMES } from '../../runtime/agent-runtime-policy'
+import { collectCompletedConsultTeamMessages, readOptionalString } from '../../runtime/thread-chat-helpers'
+import { nowEpochMillis } from '../../utils/date-time'
+import type { makeThreadMessageService } from './thread-message.service'
+import type { ThreadTurnPreparationError } from './thread-turn-preparation.service'
+import { streamAgentResponse } from './thread-turn-streaming'
+import type { StreamAgentResponseContext } from './thread-turn-streaming'
+import { buildThreadTurnSpanAttributes, compactSpanAttributes } from './thread-turn-tracing'
+
+type ChatStreamChunk = Parameters<UIMessageStreamWriter<ChatMessage>['write']>[0]
+
+export function writeMultiAgentEvent(
+  writer: UIMessageStreamWriter<ChatMessage> | undefined,
+  event: {
+    phase: 'routing' | 'waiting-for-agent' | 'agent-message-persisted' | 'complete'
+    agentId?: string
+    agentName?: string
+    messageId?: string
+    note?: string
+  },
+): void {
+  if (!writer) return
+
+  const chunk: ChatStreamChunk = {
+    type: 'data-multi-agent-event',
+    id: `multi-agent-${Bun.randomUUIDv7()}`,
+    data: event,
+    transient: true,
+  }
+  writer.write(chunk)
+}
+
+export function applyPlanTurnToolPolicy(tools: ToolSet, nodeSpec: PlanNodeSpecRecord): ToolSet {
+  const blockedToolNames = new Set([...OWNERSHIP_DISPATCH_BLOCKED_TOOL_NAMES, ...nodeSpec.toolPolicy.deny])
+  const allowList = nodeSpec.toolPolicy.allow.length > 0 ? new Set(nodeSpec.toolPolicy.allow) : null
+
+  return Object.fromEntries(
+    Object.entries(tools).filter(
+      ([toolName]) =>
+        !blockedToolNames.has(toolName) &&
+        (toolName === SUBMIT_PLAN_TURN_RESULT_TOOL_NAME || allowList === null || allowList.has(toolName)),
+    ),
+  )
+}
+
+interface VisibleAgentRunParams {
+  agentId: string
+  mode: 'direct' | 'fixedThreadMode' | 'threadMode'
+  skills?: string[]
+  additionalInstructionSections?: string[]
+  extraMessages?: ChatMessage[]
+  extraTools?: ToolSet
+  filterTools?: (tools: ToolSet) => ToolSet
+  includeExecutionPlanTools?: boolean
+  metadataPatch?: NonNullable<MessageMetadata>
+}
+
+interface VisibleAgentRunnerDeps<TBuildTurnToolParams> {
+  agentFactoryConfig: { buildAgentTools: (params: TBuildTurnToolParams) => Promise<ToolSet> | ToolSet }
+  buildTurnToolParams: (params: {
+    agentId: string
+    mode: 'direct' | 'fixedThreadMode' | 'threadMode'
+    memoryBlock: string
+    onAppendMemoryBlock: (value: string) => void
+    extraMessages?: ChatMessage[]
+    skills?: string[]
+    includeExecutionPlanTools: boolean
+  }) => TBuildTurnToolParams
+  threadTurnMessageContext: {
+    buildRunInputMessages: (extraMessages?: ChatMessage[]) => ChatMessage[]
+    appendMessages: (messages: ChatMessage[]) => void
+  }
+  threadMessageService: Pick<ReturnType<typeof makeThreadMessageService>, 'upsertMessagesEffect'>
+  threadRef: RecordIdRef
+  writer?: UIMessageStreamWriter<ChatMessage>
+  streamCtx: StreamAgentResponseContext
+  agentIdentityOverrides: ReturnType<typeof readRuntimeAgentIdentityOverrides>
+  getMemoryBlock: () => string
+  setMemoryBlock: (value: string) => void
+  onPersistedMessage: (message: ChatMessage) => void
+  failIfRunAborted: () => Effect.Effect<void, ThreadTurnPreparationError>
+}
+
+export function createThreadTurnVisibleAgentRunner<TBuildTurnToolParams>(
+  deps: VisibleAgentRunnerDeps<TBuildTurnToolParams>,
+) {
+  class ThreadTurnExecutionError extends Schema.TaggedErrorClass<ThreadTurnExecutionError>()(
+    'ThreadTurnExecutionError',
+    { message: Schema.String, cause: Schema.optional(Schema.Defect) },
+  ) {}
+
+  const effectTryPromise = <A>(
+    evaluate: () => PromiseLike<A> | Effect.Effect<A, unknown>,
+    message: string,
+  ): Effect.Effect<A, ThreadTurnExecutionError> =>
+    effectTryPromiseShared(evaluate, (error) => new ThreadTurnExecutionError({ message, cause: error }))
+
+  function effectFromMaybePromise<A>(
+    evaluate: () => A | PromiseLike<A>,
+    message: string,
+  ): Effect.Effect<A, ThreadTurnExecutionError> {
+    return effectTryMaybeAsync(evaluate, (error) => new ThreadTurnExecutionError({ message, cause: error }))
+  }
+
+  const buildAgentMetadataPatch = (agentId: string, agentName: string): NonNullable<MessageMetadata> => ({
+    agentId,
+    agentName,
+    semanticTerminationReason: 'none',
+  })
+
+  const failIfRunAborted = () =>
+    deps
+      .failIfRunAborted()
+      .pipe(
+        Effect.mapError(
+          (error) =>
+            new ThreadTurnExecutionError({ message: 'Thread turn was aborted during agent execution.', cause: error }),
+        ),
+      )
+
+  const createObserver = (agentId: string) => ({
+    run: <T>(fn: () => T | Promise<T>) =>
+      runPromise(
+        effectTryMaybeAsync(
+          fn,
+          (error) =>
+            new ThreadTurnExecutionError({ message: `Agent observer run failed (agent=${agentId}).`, cause: error }),
+        ),
+      ),
+    recordError: (error: unknown) => {
+      aiLogger.error`Agent run failed (agent=${agentId}): ${error}`
+    },
+    recordAbort: (error: unknown) => {
+      aiLogger.debug`Agent run aborted (agent=${agentId}): ${error instanceof Error ? error.message : String(error)}`
+    },
+  })
+
+  const commitAssistantResponse = (
+    response: ChatMessage,
+    agentId: string,
+    agentName: string,
+    metadataPatch?: NonNullable<MessageMetadata>,
+  ): Promise<ChatMessage> =>
+    runPromise(
+      Effect.gen(function* () {
+        yield* failIfRunAborted()
+
+        const toCommittedAssistantMessage = (
+          message: ChatMessage,
+          resolvedAgentId: string,
+          resolvedAgentName: string,
+          patch?: NonNullable<MessageMetadata>,
+        ) =>
+          withMessageCreatedAt(
+            {
+              ...message,
+              metadata: {
+                ...message.metadata,
+                ...buildAgentMetadataPatch(resolvedAgentId, resolvedAgentName),
+                ...patch,
+              },
+            },
+            toTimestamp(message.metadata?.createdAt) ?? nowEpochMillis(),
+          )
+
+        const committedConsultMessages = collectCompletedConsultTeamMessages({ responseMessage: response }).flatMap(
+          (consultMessage) => {
+            const consultAgentId = readOptionalString(consultMessage.metadata?.agentId)
+            const consultAgentName = readOptionalString(consultMessage.metadata?.agentName)
+            if (!consultAgentId || !consultAgentName) {
+              return []
+            }
+
+            return [toCommittedAssistantMessage(consultMessage, consultAgentId, consultAgentName)]
+          },
+        )
+
+        const committed = toCommittedAssistantMessage(response, agentId, agentName, metadataPatch)
+        const messagesToPersist = [...committedConsultMessages, committed]
+        yield* deps.threadMessageService.upsertMessagesEffect({ threadId: deps.threadRef, messages: messagesToPersist })
+
+        for (const persistedMessage of messagesToPersist) {
+          deps.threadTurnMessageContext.appendMessages([persistedMessage])
+          deps.onPersistedMessage(persistedMessage)
+        }
+        yield* failIfRunAborted()
+
+        return committed
+      }),
+    )
+
+  const runVisibleAgentEffect = Effect.fn('ThreadTurnExecution.runVisibleAgent')(function* (
+    runParams: VisibleAgentRunParams,
+  ) {
+    let runMemoryBlock = deps.getMemoryBlock()
+    const includeExecutionPlanTools = runParams.includeExecutionPlanTools ?? runParams.mode !== 'fixedThreadMode'
+    const builtTools = yield* effectFromMaybePromise(
+      () =>
+        deps.agentFactoryConfig.buildAgentTools(
+          deps.buildTurnToolParams({
+            agentId: runParams.agentId,
+            mode: runParams.mode,
+            skills: runParams.skills,
+            memoryBlock: runMemoryBlock,
+            onAppendMemoryBlock: (value: string) => {
+              runMemoryBlock = value
+            },
+            extraMessages: runParams.extraMessages,
+            includeExecutionPlanTools,
+          }),
+        ),
+      `Failed to build tools for ${runParams.agentId}.`,
+    ).pipe(Effect.withSpan('ThreadTurnExecution.buildAgentTools'))
+    const rawTools = { ...builtTools, ...runParams.extraTools }
+    const tools = runParams.filterTools ? runParams.filterTools(rawTools) : rawTools
+    deps.streamCtx.memoryBlock = runMemoryBlock
+    yield* failIfRunAborted()
+    const responseMessage = yield* effectTryPromise(
+      () =>
+        streamAgentResponse(deps.streamCtx, {
+          agentId: runParams.agentId,
+          mode: runParams.mode,
+          messages: deps.threadTurnMessageContext.buildRunInputMessages(runParams.extraMessages),
+          tools,
+          observer: createObserver(runParams.agentId),
+          skills: runParams.skills,
+          additionalInstructionSections: runParams.additionalInstructionSections,
+          includeExecutionPlanTools,
+          writer: deps.writer,
+        }),
+      `Failed to stream response for ${runParams.agentId}.`,
+    ).pipe(Effect.withSpan('ThreadTurnExecution.streamVisibleAgent'))
+    deps.setMemoryBlock(runMemoryBlock)
+
+    return yield* effectTryPromise(
+      () =>
+        commitAssistantResponse(
+          responseMessage,
+          runParams.agentId,
+          resolveRuntimeAgentDisplayName(deps.agentIdentityOverrides, runParams.agentId),
+          runParams.metadataPatch,
+        ),
+      `Failed to commit assistant response for ${runParams.agentId}.`,
+    ).pipe(Effect.withSpan('ThreadTurnExecution.commitAssistantResponse'))
+  })
+
+  const runVisibleAgent = (runParams: VisibleAgentRunParams): Promise<ChatMessage> =>
+    runPromise(
+      runVisibleAgentEffect(runParams).pipe(
+        Effect.annotateSpans(
+          compactSpanAttributes({
+            ...buildThreadTurnSpanAttributes({
+              threadRef: deps.threadRef,
+              orgRef: deps.streamCtx.orgRef,
+              userRef: deps.streamCtx.userRef,
+              agentId: runParams.agentId,
+              threadType: deps.streamCtx.thread.type,
+              mode: runParams.mode,
+            }),
+            includeExecutionPlanTools: runParams.includeExecutionPlanTools ?? runParams.mode !== 'fixedThreadMode',
+            extraMessageCount: runParams.extraMessages?.length ?? 0,
+            extraToolCount: runParams.extraTools ? Object.keys(runParams.extraTools).length : 0,
+          }),
+        ),
+      ),
+    )
+
+  return { runVisibleAgent }
+}

package/src/services/thread/thread-turn-message-context.ts

@@ -0,0 +1,73 @@
+import type { ChatMessage } from '@lota-sdk/shared'
+
+import type { RecordIdRef } from '../../db/record-id'
+import {
+  buildReadableUploadMetadataText,
+  buildModelInputMessagesWithUploadMetadata,
+} from '../../runtime/chat-attachments'
+import type { ReadableUploadMetadataLike } from '../../runtime/chat-types'
+import type { ContextCompactionRuntime } from '../../runtime/context-compaction/context-compaction'
+
+export function upsertChatHistoryMessage(messages: ChatMessage[], nextMessage: ChatMessage): ChatMessage[] {
+  const existingIndex = messages.findIndex((message) => message.id === nextMessage.id)
+  if (existingIndex === -1) {
+    return [...messages, nextMessage]
+  }
+
+  const nextMessages = [...messages]
+  nextMessages[existingIndex] = nextMessage
+  return nextMessages
+}
+
+interface ThreadTurnMessageContext {
+  readonly currentMessages: ChatMessage[]
+  appendMessages(messages: ChatMessage[]): void
+  listReadableUploads(extraMessages?: ChatMessage[]): ReadableUploadMetadataLike[]
+  buildRunInputMessages(extraMessages?: ChatMessage[]): ChatMessage[]
+}
+
+export function createThreadTurnMessageContext(params: {
+  contextCompactionRuntime: Pick<ContextCompactionRuntime, 'prependSummaryMessage'>
+  persistedCompactionSummary: string
+  messagesForContext: ChatMessage[]
+  orgRef: RecordIdRef
+  userRef: RecordIdRef
+  latestUserMessageId: string
+  listReadableUploadsFromMessages: (params: {
+    messages: ChatMessage[]
+    orgId: RecordIdRef
+    userId: RecordIdRef
+  }) => ReadableUploadMetadataLike[]
+}): ThreadTurnMessageContext {
+  let currentMessages = params.contextCompactionRuntime.prependSummaryMessage(
+    params.persistedCompactionSummary,
+    params.messagesForContext,
+  )
+
+  const listReadableUploads = (extraMessages: ChatMessage[] = []) =>
+    params.listReadableUploadsFromMessages({
+      messages: [...currentMessages, ...extraMessages],
+      orgId: params.orgRef,
+      userId: params.userRef,
+    })
+
+  const buildRunInputMessages = (extraMessages: ChatMessage[] = []): ChatMessage[] =>
+    buildModelInputMessagesWithUploadMetadata({
+      messages: [...currentMessages, ...extraMessages],
+      latestUserMessageId: params.latestUserMessageId,
+      uploadMetadataText: buildReadableUploadMetadataText(listReadableUploads(extraMessages)),
+    })
+
+  return {
+    get currentMessages() {
+      return currentMessages
+    },
+    appendMessages(messages: ChatMessage[]) {
+      for (const message of messages) {
+        currentMessages = upsertChatHistoryMessage(currentMessages, message)
+      }
+    },
+    listReadableUploads,
+    buildRunInputMessages,
+  }
+}