@lota-sdk/core 0.4.8 → 0.4.10

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

@@ -0,0 +1,376 @@
+import type { ChatMessage } from '@lota-sdk/shared'
+import { createUIMessageStream } from 'ai'
+import { Context, Schema, Effect, Layer } from 'effect'
+
+import { ensureRecordId, recordIdToString } from '../../db/record-id'
+import { TABLES } from '../../db/tables'
+import { BadRequestError } from '../../effect/errors'
+import { runPromise } from '../../effect/runtime'
+import { hasApprovalRespondedParts, isApprovalContinuationRequest } from '../../runtime/approval-continuation'
+import { shouldPlanNodeUseVisibleTurn } from '../../runtime/execution-plan-visibility'
+import { wrapResponseWithKeepalive } from '../../utils/sse-keepalive'
+import type { makePlanExecutorService } from '../plan/plan-executor.service'
+import { PlanExecutorServiceTag } from '../plan/plan-executor.service'
+import type { makePlanRunService } from '../plan/plan-run.service'
+import { PlanRunServiceTag } from '../plan/plan-run.service'
+import type { makeUserService } from '../user.service'
+import { UserServiceTag } from '../user.service'
+import type {
+  PreparedThreadTurnResult,
+  ThreadApprovalContinuationParams,
+  ThreadPlanTurnParams,
+  ThreadTurnParams,
+  makeThreadTurnPreparationService,
+} from './thread-turn-preparation.service'
+import { ThreadTurnPreparationServiceTag } from './thread-turn-preparation.service'
+import { buildThreadTurnSpanAttributes, compactSpanAttributes } from './thread-turn-tracing'
+import type { makeThreadService } from './thread.service'
+import { ThreadServiceTag } from './thread.service'
+
+export { hasApprovalRespondedParts, isApprovalContinuationRequest }
+export { wrapResponseWithKeepalive }
+export type { PreparedThreadTurnResult }
+export type { ThreadPlanTurnParams }
+
+class ThreadTurnServiceError extends Schema.TaggedErrorClass<ThreadTurnServiceError>()('ThreadTurnServiceError', {
+  message: Schema.String,
+  cause: Schema.optional(Schema.Defect),
+}) {}
+
+interface ThreadTurnDeps {
+  planExecutor: ReturnType<typeof makePlanExecutorService>
+  planRun: ReturnType<typeof makePlanRunService>
+  thread: ReturnType<typeof makeThreadService>
+  threadTurnPreparation: ReturnType<typeof makeThreadTurnPreparationService>
+  user: ReturnType<typeof makeUserService>
+}
+
+const createThreadApprovalContinuationStreamEffect = Effect.fn('ThreadTurn.createApprovalContinuationStream')(
+  function* (deps: ThreadTurnDeps, params: ThreadApprovalContinuationParams) {
+    const currentContext = yield* Effect.context()
+    const runPromiseWithCurrentContext = Effect.runPromiseWith(currentContext)
+    const prepared = yield* deps.threadTurnPreparation.prepareThreadRunCoreEffect({
+      ...params,
+      kind: 'approvalContinuation',
+    })
+    return createUIMessageStream<ChatMessage>({
+      originalMessages: prepared.originalMessages,
+      onError: (error) => (error instanceof Error ? error.message : 'Approval continuation stream failed.'),
+      execute: ({ writer }) =>
+        runPromiseWithCurrentContext(
+          prepared.run(writer).pipe(Effect.withSpan('ThreadTurn.executeApprovalContinuationStream'), Effect.asVoid),
+        ),
+    })
+  },
+)
+
+function createThreadApprovalContinuationStreamWith(deps: ThreadTurnDeps, params: ThreadApprovalContinuationParams) {
+  return createThreadApprovalContinuationStreamEffect(deps, params).pipe(
+    Effect.annotateSpans(
+      buildThreadTurnSpanAttributes({
+        threadRef: params.threadRef,
+        orgRef: params.orgRef,
+        userRef: params.userRef,
+        kind: 'approvalContinuation',
+        streamId: params.streamId,
+        threadType: params.thread.type,
+      }),
+    ),
+  )
+}
+
+const createThreadNativeToolApprovalStreamEffect = Effect.fn('ThreadTurn.createNativeToolApprovalStream')(function* (
+  deps: ThreadTurnDeps,
+  params: ThreadApprovalContinuationParams,
+) {
+  const currentContext = yield* Effect.context()
+  const runPromiseWithCurrentContext = Effect.runPromiseWith(currentContext)
+  const prepared = yield* deps.threadTurnPreparation.prepareThreadRunCoreEffect({
+    ...params,
+    kind: 'nativeToolApprovalTurn',
+  })
+  return createUIMessageStream<ChatMessage>({
+    originalMessages: prepared.originalMessages,
+    onError: (error) => (error instanceof Error ? error.message : 'Native tool approval stream failed.'),
+    execute: ({ writer }) =>
+      runPromiseWithCurrentContext(
+        prepared.run(writer).pipe(Effect.withSpan('ThreadTurn.executeNativeToolApprovalStream'), Effect.asVoid),
+      ),
+  })
+})
+
+function createThreadNativeToolApprovalStreamWith(deps: ThreadTurnDeps, params: ThreadApprovalContinuationParams) {
+  return createThreadNativeToolApprovalStreamEffect(deps, params).pipe(
+    Effect.annotateSpans(
+      buildThreadTurnSpanAttributes({
+        threadRef: params.threadRef,
+        orgRef: params.orgRef,
+        userRef: params.userRef,
+        kind: 'nativeToolApprovalTurn',
+        streamId: params.streamId,
+        threadType: params.thread.type,
+      }),
+    ),
+  )
+}
+
+const createThreadTurnStreamEffect = Effect.fn('ThreadTurn.createThreadTurnStream')(function* (
+  deps: ThreadTurnDeps,
+  params: ThreadTurnParams,
+) {
+  const currentContext = yield* Effect.context()
+  const runPromiseWithCurrentContext = Effect.runPromiseWith(currentContext)
+  const prepared = yield* deps.threadTurnPreparation.prepareThreadRunCoreEffect({ ...params, kind: 'userTurn' })
+  return createUIMessageStream<ChatMessage>({
+    originalMessages: prepared.originalMessages,
+    onError: (error) => (error instanceof Error ? error.message : 'Chat stream failed.'),
+    execute: ({ writer }) =>
+      runPromiseWithCurrentContext(
+        prepared.run(writer).pipe(Effect.withSpan('ThreadTurn.executeThreadTurnStream'), Effect.asVoid),
+      ),
+  })
+})
+
+function createThreadTurnStreamWith(deps: ThreadTurnDeps, params: ThreadTurnParams) {
+  return createThreadTurnStreamEffect(deps, params).pipe(
+    Effect.annotateSpans(
+      buildThreadTurnSpanAttributes({
+        threadRef: params.threadRef,
+        orgRef: params.orgRef,
+        userRef: params.userRef,
+        kind: 'userTurn',
+        streamId: params.streamId,
+        threadType: params.thread.type,
+      }),
+    ),
+  )
+}
+
+const runThreadTurnInBackgroundEffect = Effect.fn('ThreadTurn.runThreadTurnInBackground')(function* (
+  deps: ThreadTurnDeps,
+  params: ThreadTurnParams,
+) {
+  const prepared = yield* deps.threadTurnPreparation.prepareThreadRunCoreEffect({ ...params, kind: 'userTurn' })
+  return yield* prepared.run().pipe(Effect.withSpan('ThreadTurn.executeBackgroundTurn'))
+})
+
+function runThreadTurnInBackgroundWith(deps: ThreadTurnDeps, params: ThreadTurnParams) {
+  return runThreadTurnInBackgroundEffect(deps, params).pipe(
+    Effect.annotateSpans(
+      buildThreadTurnSpanAttributes({
+        threadRef: params.threadRef,
+        orgRef: params.orgRef,
+        userRef: params.userRef,
+        kind: 'background',
+        streamId: params.streamId,
+        threadType: params.thread.type,
+        agentId: params.agentIdOverride,
+      }),
+    ),
+  )
+}
+
+const triggerPlanNodeTurnEffect = Effect.fn('ThreadTurn.triggerPlanNodeTurn')(function* (
+  deps: ThreadTurnDeps,
+  params: { runId: string; nodeId: string; abortSignal?: AbortSignal; streamId?: string },
+) {
+  const run = yield* deps.planRun.getRunById(params.runId)
+  const spec = yield* deps.planRun.getPlanSpecById(run.planSpecId)
+  const nodeSpec = yield* deps.planRun.getNodeSpecByNodeId(spec.id, params.nodeId)
+
+  if (!shouldPlanNodeUseVisibleTurn(spec, nodeSpec) || nodeSpec.owner.executorType !== 'agent') {
+    return yield* new BadRequestError({
+      message: `Plan node "${params.nodeId}" is not eligible for a visible plan turn.`,
+    })
+  }
+
+  let activeRun = run
+  let nodeRun = yield* deps.planRun.getNodeRunByNodeId(run.id, params.nodeId)
+  if (nodeRun.status === 'ready') {
+    yield* Effect.tryPromise({
+      try: () => deps.planExecutor.transitionNodeToRunning({ runId: params.runId, nodeId: params.nodeId }),
+      catch: (cause) => new ThreadTurnServiceError({ message: 'Failed to transition plan node to running.', cause }),
+    })
+    activeRun = yield* deps.planRun.getRunById(params.runId)
+    nodeRun = yield* deps.planRun.getNodeRunByNodeId(run.id, params.nodeId)
+  }
+
+  if (nodeRun.status !== 'running') {
+    return { assistantMessages: [], inputMessageId: undefined }
+  }
+
+  const thread = yield* deps.thread.getThread(activeRun.threadId)
+  const userRef = ensureRecordId(thread.userId, TABLES.USER)
+  const [artifacts, nodeRuns, nodeSpecs, userName] = yield* Effect.all([
+    deps.planRun.listArtifacts(activeRun.id),
+    deps.planRun.listNodeRuns(activeRun.id),
+    deps.planRun.listNodeSpecs(spec.id),
+    deps.user.getUser(recordIdToString(userRef, TABLES.USER)).pipe(
+      Effect.map((user) => user.name),
+      Effect.catch(() => Effect.succeed(undefined as string | undefined)),
+    ),
+  ]).pipe(Effect.withSpan('ThreadTurn.loadPlanTurnContext'))
+
+  const inputArtifacts = artifacts
+    .filter((artifact) => nodeSpec.upstreamNodeIds.includes(artifact.nodeId))
+    .map((artifact) => ({
+      name: artifact.name,
+      kind: artifact.kind,
+      ...(artifact.description ? { description: artifact.description } : {}),
+      ...(artifact.content !== undefined ? { content: artifact.content } : {}),
+      ...(artifact.payload !== undefined ? { payload: artifact.payload } : {}),
+      ...(artifact.publishedArtifactId
+        ? { publishedArtifactId: recordIdToString(artifact.publishedArtifactId, TABLES.ARTIFACT) }
+        : {}),
+    }))
+  const upstreamNodeSpecs = new Map(nodeSpecs.map((upstreamNodeSpec) => [upstreamNodeSpec.nodeId, upstreamNodeSpec]))
+  const upstreamHandoffs = nodeRuns
+    .filter((candidate) => nodeSpec.upstreamNodeIds.includes(candidate.nodeId) && candidate.latestNotes)
+    .map((candidate) => ({
+      nodeId: candidate.nodeId,
+      label: upstreamNodeSpecs.get(candidate.nodeId)?.label ?? candidate.nodeId,
+      ownerRef: upstreamNodeSpecs.get(candidate.nodeId)?.owner.ref ?? 'system',
+      ownerType: upstreamNodeSpecs.get(candidate.nodeId)?.owner.executorType ?? 'system',
+      summary: candidate.latestNotes ?? undefined,
+    }))
+
+  const threadRef = ensureRecordId(activeRun.threadId, TABLES.THREAD)
+  const orgRef = ensureRecordId(activeRun.organizationId, TABLES.ORGANIZATION)
+  const prepared = yield* deps.threadTurnPreparation.prepareThreadRunCoreEffect({
+    kind: 'planTurn',
+    thread,
+    threadRef,
+    orgRef,
+    userRef,
+    userName,
+    abortSignal: params.abortSignal,
+    streamId: params.streamId,
+    planTurn: {
+      runId: params.runId,
+      nodeId: params.nodeId,
+      planTitle: spec.title,
+      nodeSpec,
+      nodeRun,
+      resolvedInput: nodeRun.resolvedInput ?? {},
+      inputArtifacts,
+      upstreamHandoffs,
+    },
+  })
+  return yield* prepared.run().pipe(Effect.withSpan('ThreadTurn.executePlanTurn'))
+})
+
+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,
+      }),
+    ),
+  )
+}
+
+export function makeThreadTurnService(deps: ThreadTurnDeps) {
+  return {
+    createThreadApprovalContinuationStream(params: ThreadApprovalContinuationParams) {
+      return createThreadApprovalContinuationStreamWith(deps, params)
+    },
+    createThreadNativeToolApprovalStream(params: ThreadApprovalContinuationParams) {
+      return createThreadNativeToolApprovalStreamWith(deps, params)
+    },
+    createThreadTurnStream(params: ThreadTurnParams) {
+      return createThreadTurnStreamWith(deps, params)
+    },
+    runThreadTurnInBackground(params: ThreadTurnParams) {
+      return runThreadTurnInBackgroundWith(deps, params)
+    },
+    triggerPlanNodeTurn(params: { runId: string; nodeId: string; abortSignal?: AbortSignal; streamId?: string }) {
+      return triggerPlanNodeTurnWith(deps, params)
+    },
+  } as const
+}
+
+export class ThreadTurnServiceTag extends Context.Service<
+  ThreadTurnServiceTag,
+  ReturnType<typeof makeThreadTurnService>
+>()('@lota-sdk/core/ThreadTurnService') {}
+
+export const ThreadTurnServiceLive = Layer.effect(
+  ThreadTurnServiceTag,
+  Effect.gen(function* () {
+    const planExecutor = yield* PlanExecutorServiceTag
+    const planRun = yield* PlanRunServiceTag
+    const thread = yield* ThreadServiceTag
+    const threadTurnPreparation = yield* ThreadTurnPreparationServiceTag
+    const user = yield* UserServiceTag
+    return makeThreadTurnService({ planExecutor, planRun, thread, threadTurnPreparation, user })
+  }),
+)
+
+const createThreadApprovalContinuationStreamWithRuntime = Effect.fn(
+  'ThreadTurn.createApprovalContinuationStreamWithRuntime',
+)(function* (params: ThreadApprovalContinuationParams) {
+  const threadTurnService = yield* ThreadTurnServiceTag
+  return yield* threadTurnService.createThreadApprovalContinuationStream(params)
+})
+
+const createThreadNativeToolApprovalStreamWithRuntime = Effect.fn(
+  'ThreadTurn.createNativeToolApprovalStreamWithRuntime',
+)(function* (params: ThreadApprovalContinuationParams) {
+  const threadTurnService = yield* ThreadTurnServiceTag
+  return yield* threadTurnService.createThreadNativeToolApprovalStream(params)
+})
+
+const createThreadTurnStreamWithRuntime = Effect.fn('ThreadTurn.createThreadTurnStreamWithRuntime')(function* (
+  params: ThreadTurnParams,
+) {
+  const threadTurnService = yield* ThreadTurnServiceTag
+  return yield* threadTurnService.createThreadTurnStream(params)
+})
+
+const runThreadTurnInBackgroundWithRuntime = Effect.fn('ThreadTurn.runThreadTurnInBackgroundWithRuntime')(function* (
+  params: ThreadTurnParams,
+) {
+  const threadTurnService = yield* ThreadTurnServiceTag
+  return yield* threadTurnService.runThreadTurnInBackground(params)
+})
+
+const triggerPlanNodeTurnWithRuntime = Effect.fn('ThreadTurn.triggerPlanNodeTurnWithRuntime')(function* (params: {
+  runId: string
+  nodeId: string
+  abortSignal?: AbortSignal
+  streamId?: string
+}) {
+  const threadTurnService = yield* ThreadTurnServiceTag
+  return yield* threadTurnService.triggerPlanNodeTurn(params)
+})
+
+export function createThreadApprovalContinuationStream(params: ThreadApprovalContinuationParams) {
+  return runPromise(createThreadApprovalContinuationStreamWithRuntime(params))
+}
+
+export function createThreadNativeToolApprovalStream(params: ThreadApprovalContinuationParams) {
+  return runPromise(createThreadNativeToolApprovalStreamWithRuntime(params))
+}
+
+export function createThreadTurnStream(params: ThreadTurnParams) {
+  return runPromise(createThreadTurnStreamWithRuntime(params))
+}
+
+export function runThreadTurnInBackground(params: ThreadTurnParams): Promise<PreparedThreadTurnResult> {
+  return runPromise(runThreadTurnInBackgroundWithRuntime(params))
+}
+
+export function triggerPlanNodeTurn(params: {
+  runId: string
+  nodeId: string
+  abortSignal?: AbortSignal
+  streamId?: string
+}): Promise<PreparedThreadTurnResult> {
+  return runPromise(triggerPlanNodeTurnWithRuntime(params))
+}