@lota-sdk/core 0.4.13 → 0.4.15

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

@@ -6,7 +6,7 @@ import { Effect, Ref, Schema, Stream } from 'effect'
 
 import { aiLogger } from '../../config/logger'
 import type { RecordIdRef } from '../../db/record-id'
-import { effectTryMaybeAsync, makeEffectTryPromiseWithMessage } from '../../effect/helpers'
+import { ERROR_TAGS } from '../../effect/errors'
 import { AgentConfigServiceTag, AgentFactoryServiceTag } from '../../effect/services'
 import {
   readRuntimeAgentIdentityOverrides,
@@ -40,28 +40,12 @@ interface ToolLoopGenerateResult {
 }
 
 export class ThreadTurnStreamingError extends Schema.TaggedErrorClass<ThreadTurnStreamingError>()(
-  'ThreadTurnStreamingError',
+  '@lota-sdk/core/ThreadTurnStreamingError',
   { message: Schema.String, cause: Schema.optional(Schema.Defect) },
 ) {}
 
-const effectTryPromise = makeEffectTryPromiseWithMessage(
-  (message, cause) => new ThreadTurnStreamingError({ message, cause }),
-)
-
-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 toThreadTurnStreamingError(message: string, cause: unknown): ThreadTurnStreamingError {
+  return new ThreadTurnStreamingError({ message, cause })
 }
 
 function hasUIMessageStream(value: unknown): value is UIMessageStreamResult {
@@ -90,7 +74,7 @@ function isTextTokenChunkType(chunkType: string | undefined): boolean {
   return chunkType === 'text-delta'
 }
 
-function buildFallbackResponseMessage(
+export function buildFallbackResponseMessage(
   result: ToolLoopGenerateResult,
 ): Effect.Effect<ChatMessage, ThreadTurnStreamingError> {
   const parts: ChatMessage['parts'] = []
@@ -112,12 +96,6 @@ function buildFallbackResponseMessage(
     parts.push({ type: 'text', text })
   }
 
-  if (parts.length === 0) {
-    return Effect.fail(
-      new ThreadTurnStreamingError({ message: 'Agent generate fallback did not produce any response parts.' }),
-    )
-  }
-
   return Effect.succeed({ id: Bun.randomUUIDv7(), role: 'assistant', parts })
 }
 
@@ -151,7 +129,7 @@ interface StreamAgentResponseParams {
   messages: ChatMessage[]
   tools: ToolSet
   observer: {
-    run: <T>(fn: () => T | Promise<T>) => Promise<T>
+    run: <T>(fn: () => Promise<T>) => Promise<T>
     recordError: (error: unknown) => void
     recordAbort: (error: unknown) => void
   }
@@ -172,45 +150,54 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
   const executionPlanInstructionSections =
     streamParams.includeExecutionPlanTools === false
       ? undefined
-      : yield* effectFromMaybeEffect(
-          () => ctx.getExecutionPlanInstructionSections(),
-          'Failed to load execution plan instructions.',
-        ).pipe(Effect.withSpan('ThreadTurnStreaming.loadExecutionPlanInstructions'))
+      : yield* ctx.getExecutionPlanInstructionSections().pipe(
+          Effect.mapError((cause) => toThreadTurnStreamingError('Failed to load execution plan instructions.', cause)),
+          Effect.withSpan('ThreadTurnStreaming.loadExecutionPlanInstructions'),
+        )
 
+  const resolveAgent = ctx.turnHooks.resolveAgent
   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')),
+    resolveAgent
+      ? yield* Effect.tryPromise({
+          try: () =>
+            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,
+            }),
+          catch: (cause) => toThreadTurnStreamingError('Failed to resolve the runtime agent.', cause),
+        }).pipe(Effect.withSpan('ThreadTurnStreaming.resolveAgent'))
+      : undefined,
   )
 
   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}.`,
-    ),
+    ctx
+      .getPreSeededMemoriesSection(resolvedAgentId)
+      .pipe(
+        Effect.mapError((cause) =>
+          toThreadTurnStreamingError(`Failed to load pre-seeded memories for ${resolvedAgentId}.`, cause),
+        ),
+      ),
+    ctx
+      .getLearnedSkillsSection(resolvedAgentId, latestUserMessageText)
+      .pipe(
+        Effect.mapError((cause) =>
+          toThreadTurnStreamingError(`Failed to load learned skills for ${resolvedAgentId}.`, cause),
+        ),
+      ),
   ]).pipe(Effect.withSpan('ThreadTurnStreaming.loadMemoriesAndSkills'))
 
   const toolNames = new Set(Object.keys(streamParams.tools))
@@ -259,10 +246,10 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
     hasRetrievalTools,
   })
 
-  const modelMessages = yield* effectTryPromise(
-    () => convertToModelMessages(streamParams.messages, { ignoreIncompleteToolCalls: true }),
-    'Failed to convert UI messages to model messages.',
-  ).pipe(Effect.withSpan('ThreadTurnStreaming.convertModelMessages'))
+  const modelMessages = yield* Effect.tryPromise({
+    try: () => convertToModelMessages(streamParams.messages, { ignoreIncompleteToolCalls: true }),
+    catch: (cause) => toThreadTurnStreamingError('Failed to convert UI messages to model messages.', cause),
+  }).pipe(Effect.withSpan('ThreadTurnStreaming.convertModelMessages'))
 
   const agentFactory = agentFactoryConfig.createAgent[config.id]
   if (!agentFactory) {
@@ -281,10 +268,11 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
   const resolvedAgentName = resolveRuntimeAgentDisplayName(agentConfig, agentIdentityOverrides, resolvedAgentId)
 
   const generateFallback = (cause: ThreadTurnStreamingError) =>
-    effectTryPromise(
-      () => streamParams.observer.run(() => agent.generate({ messages: modelMessages, abortSignal: agentAbortSignal })),
-      `Agent generate fallback failed for ${resolvedAgentId}.`,
-    ).pipe(
+    Effect.tryPromise({
+      try: () =>
+        streamParams.observer.run(() => agent.generate({ messages: modelMessages, abortSignal: agentAbortSignal })),
+      catch: (error) => toThreadTurnStreamingError(`Agent generate fallback failed for ${resolvedAgentId}.`, error),
+    }).pipe(
       Effect.tap(() =>
         Effect.sync(() => {
           aiLogger.warn`Agent stream failed for ${resolvedAgentId}; falling back to generate: ${cause.message}`
@@ -293,10 +281,11 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
       Effect.flatMap((result) => buildFallbackResponseMessage(result as ToolLoopGenerateResult)),
     )
 
-  const generateWithoutUiStream = effectTryPromise(
-    () => streamParams.observer.run(() => agent.generate({ messages: modelMessages, abortSignal: agentAbortSignal })),
-    `Agent generate failed for ${resolvedAgentId}.`,
-  ).pipe(
+  const generateWithoutUiStream = Effect.tryPromise({
+    try: () =>
+      streamParams.observer.run(() => agent.generate({ messages: modelMessages, abortSignal: agentAbortSignal })),
+    catch: (cause) => toThreadTurnStreamingError(`Agent generate failed for ${resolvedAgentId}.`, cause),
+  }).pipe(
     Effect.tapError((error) =>
       Effect.sync(() => {
         if (agentAbortSignal.aborted) {
@@ -321,10 +310,12 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
     return generatedResponse
   }
 
-  const result = yield* effectTryPromise(
-    () => streamParams.observer.run(() => agent.stream({ messages: modelMessages, abortSignal: agentAbortSignal })),
-    `Agent stream failed for ${resolvedAgentId}.`,
-  ).pipe(
+  const writer = streamParams.writer
+  const result = yield* Effect.tryPromise({
+    try: () =>
+      streamParams.observer.run(() => agent.stream({ messages: modelMessages, abortSignal: agentAbortSignal })),
+    catch: (cause) => toThreadTurnStreamingError(`Agent stream failed for ${resolvedAgentId}.`, cause),
+  }).pipe(
     Effect.tapError((error) =>
       Effect.sync(() => {
         if (agentAbortSignal.aborted) {
@@ -355,14 +346,12 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
       resolveFinishedStream(withMessageCreatedAt(finishedResponseMessage, nowEpochMillis()))
     },
   }) as ReadableStream<ChatStreamChunk>
-  const liveTurnTrace = streamParams.writer
-    ? createLiveTurnTraceStreamObserver({
-        traceId: `trace:${Bun.randomUUIDv7()}`,
-        writer: streamParams.writer,
-        agentId: resolvedAgentId,
-        agentName: resolvedAgentName,
-      })
-    : null
+  const liveTurnTrace = createLiveTurnTraceStreamObserver({
+    traceId: `trace:${Bun.randomUUIDv7()}`,
+    writer,
+    agentId: resolvedAgentId,
+    agentName: resolvedAgentName,
+  })
   const streamStartedAt = performance.now()
   const firstVisibleOutputRecorded = yield* Ref.make(false)
   const firstTextTokenRecorded = yield* Ref.make(false)
@@ -393,21 +382,23 @@ const streamAgentResponseEffect = Effect.fn('ThreadTurnStreaming.streamAgentResp
           yield* Effect.annotateCurrentSpan({ ttftMs: elapsedMs, ttftChunkType: chunkType })
         }
 
-        if (streamParams.writer) {
-          yield* Effect.sync(() => {
-            streamParams.writer?.write(value)
-            liveTurnTrace?.observeChunk(value)
-          })
-        }
+        yield* Effect.sync(() => {
+          writer.write(value)
+          liveTurnTrace.observeChunk(value)
+        })
       }),
     ),
     Effect.withSpan('ThreadTurnStreaming.consumeUiStream'),
     Effect.andThen(
-      effectTryPromise(() => finishedStream, `Agent run for ${resolvedAgentId} did not produce a response message.`),
+      Effect.tryPromise({
+        try: () => finishedStream,
+        catch: (cause) =>
+          toThreadTurnStreamingError(`Agent run for ${resolvedAgentId} did not produce a response message.`, cause),
+      }),
     ),
-    Effect.catchTag('ThreadTurnStreamingError', generateFallback),
+    Effect.catchTag(ERROR_TAGS.ThreadTurnStreamingError, generateFallback),
   )
-  liveTurnTrace?.finish()
+  liveTurnTrace.finish()
 
   for (const toolError of collectToolOutputErrors({ responseMessage: streamedResponse })) {
     aiLogger.error`Tool execution failed (agent=${resolvedAgentId}, tool=${toolError.toolName}, toolCallId=${toolError.toolCallId}): ${toolError.errorText}`

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

@@ -5,12 +5,12 @@ import { Context, Schema, Effect, Layer } from 'effect'
 import type { ResolvedAgentConfig } from '../../config/agent-defaults'
 import { ensureRecordId, recordIdToString } from '../../db/record-id'
 import { TABLES } from '../../db/tables'
-import { BadRequestError, ForbiddenError } from '../../effect/errors'
+import { ERROR_TAGS, BadRequestError, ForbiddenError } from '../../effect/errors'
 import { AgentConfigServiceTag } from '../../effect/services'
 import { hasApprovalRespondedParts, isApprovalContinuationRequest } from '../../runtime/approval-continuation'
 import { shouldPlanNodeUseVisibleTurn } from '../../runtime/execution-plan-visibility'
 import { wrapResponseWithKeepalive } from '../../utils/sse-keepalive'
-import { BackgroundWorkService } from '../background-work.service'
+import { BackgroundWorkServiceTag } from '../background-work.service'
 import type { makePlanExecutorService } from '../plan/plan-executor.service'
 import { PlanExecutorServiceTag } from '../plan/plan-executor.service'
 import type { makePlanRunService } from '../plan/plan-run.service'
@@ -67,14 +67,14 @@ export interface LaunchBackgroundThreadWorkResult {
   message: string
 }
 
-class ThreadTurnServiceError extends Schema.TaggedErrorClass<ThreadTurnServiceError>()('ThreadTurnServiceError', {
-  message: Schema.String,
-  cause: Schema.optional(Schema.Defect),
-}) {}
+class ThreadTurnServiceError extends Schema.TaggedErrorClass<ThreadTurnServiceError>()(
+  ERROR_TAGS.ThreadTurnServiceError,
+  { message: Schema.String, cause: Schema.optional(Schema.Defect) },
+) {}
 
 interface ThreadTurnDeps {
   agentConfig: ResolvedAgentConfig
-  background: Context.Service.Shape<typeof BackgroundWorkService>
+  background: Context.Service.Shape<typeof BackgroundWorkServiceTag>
   planExecutor: ReturnType<typeof makePlanExecutorService>
   planRun: ReturnType<typeof makePlanRunService>
   provideCurrentContext: <A, E, R>(effect: Effect.Effect<A, E, R>) => Effect.Effect<A, E, never>
@@ -428,14 +428,16 @@ function triggerPlanNodeTurnWith(
   deps: ThreadTurnDeps,
   params: { runId: string; nodeId: string; abortSignal?: AbortSignal; streamId?: string },
 ) {
-  return triggerPlanNodeTurnEffect(deps, params).pipe(
-    Effect.annotateSpans(
-      compactSpanAttributes({
-        turnKind: 'planTurn',
-        streamId: params.streamId,
-        planRunId: params.runId,
-        planNodeId: params.nodeId,
-      }),
+  return deps.provideCurrentContext(
+    triggerPlanNodeTurnEffect(deps, params).pipe(
+      Effect.annotateSpans(
+        compactSpanAttributes({
+          turnKind: 'planTurn',
+          streamId: params.streamId,
+          planRunId: params.runId,
+          planNodeId: params.nodeId,
+        }),
+      ),
     ),
   )
 }
@@ -475,7 +477,7 @@ export const ThreadTurnServiceLive = Layer.effect(
     const provideCurrentContext = <A, E, R>(effect: Effect.Effect<A, E, R>): Effect.Effect<A, E, never> =>
       effect.pipe(Effect.provide(currentContext)) as Effect.Effect<A, E, never>
     const agentConfig = yield* AgentConfigServiceTag
-    const background = yield* BackgroundWorkService
+    const background = yield* BackgroundWorkServiceTag
     const planExecutor = yield* PlanExecutorServiceTag
     const planRun = yield* PlanRunServiceTag
     const thread = yield* ThreadServiceTag