simstudioai--sim/apps/sim/lib/copilot/async-runs/repository.ts

import { trace } from '@opentelemetry/api'
import { db } from '@sim/db'
import {
  type CopilotAsyncToolStatus,
  type CopilotRunStatus,
  copilotAsyncToolCalls,
  copilotRunCheckpoints,
  copilotRuns,
} from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { filterUndefined } from '@sim/utils/object'
import { and, desc, eq, inArray, isNull } from 'drizzle-orm'
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
import { markSpanForError } from '@/lib/copilot/request/otel'
import {
  ASYNC_TOOL_STATUS,
  type AsyncCompletionData,
  isDeliveredAsyncStatus,
  isTerminalAsyncStatus,
} from './lifecycle'

const logger = createLogger('CopilotAsyncRunsRepo')
// Resolve the tracer lazily per-call to avoid capturing the NoOp tracer
// before NodeSDK installs the global TracerProvider (Next.js 16/Turbopack
// can evaluate modules before instrumentation-node.ts finishes).
const getAsyncRunsTracer = () => trace.getTracer('sim-copilot-async-runs', '1.0.0')

// Wrap an async DB op in a client-kind span with canonical `db.*` attrs.
// Cancellation is routed through `markSpanForError` so aborts record the
// exception event but don't paint spans red.
async function withDbSpan<T>(
  name: string,
  op: string,
  table: string,
  attrs: Record<string, string | number | boolean | undefined>,
  fn: () => Promise<T>
): Promise<T> {
  const span = getAsyncRunsTracer().startSpan(name, {
    attributes: {
      [TraceAttr.DbSystem]: 'postgresql',
      [TraceAttr.DbOperation]: op,
      [TraceAttr.DbSqlTable]: table,
      ...filterUndefined(attrs),
    },
  })
  try {
    return await fn()
  } catch (error) {
    markSpanForError(span, error)
    throw error
  } finally {
    span.end()
  }
}

export interface CreateRunSegmentInput {
  id?: string
  executionId: string
  parentRunId?: string | null
  chatId: string
  userId: string
  workflowId?: string | null
  workspaceId?: string | null
  streamId: string
  agent?: string | null
  model?: string | null
  provider?: string | null
  requestContext?: Record<string, unknown>
  status?: CopilotRunStatus
}

export async function createRunSegment(input: CreateRunSegmentInput) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsCreateRunSegment,
    'INSERT',
    'copilot_runs',
    {
      [TraceAttr.CopilotExecutionId]: input.executionId,
      [TraceAttr.ChatId]: input.chatId,
      [TraceAttr.StreamId]: input.streamId,
      [TraceAttr.UserId]: input.userId,
      [TraceAttr.CopilotRunParentId]: input.parentRunId ?? undefined,
      [TraceAttr.CopilotRunAgent]: input.agent ?? undefined,
      [TraceAttr.CopilotRunModel]: input.model ?? undefined,
      [TraceAttr.CopilotRunProvider]: input.provider ?? undefined,
      [TraceAttr.CopilotRunStatus]: input.status ?? 'active',
    },
    async () => {
      const [run] = await db
        .insert(copilotRuns)
        .values({
          ...(input.id ? { id: input.id } : {}),
          executionId: input.executionId,
          parentRunId: input.parentRunId ?? null,
          chatId: input.chatId,
          userId: input.userId,
          workflowId: input.workflowId ?? null,
          workspaceId: input.workspaceId ?? null,
          streamId: input.streamId,
          agent: input.agent ?? null,
          model: input.model ?? null,
          provider: input.provider ?? null,
          requestContext: input.requestContext ?? {},
          status: input.status ?? 'active',
        })
        .returning()
      return run
    }
  )
}

export async function updateRunStatus(
  runId: string,
  status: CopilotRunStatus,
  updates: {
    completedAt?: Date | null
    error?: string | null
    requestContext?: Record<string, unknown>
  } = {}
) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsUpdateRunStatus,
    'UPDATE',
    'copilot_runs',
    {
      [TraceAttr.RunId]: runId,
      [TraceAttr.CopilotRunStatus]: status,
      [TraceAttr.CopilotRunHasError]: !!updates.error,
      [TraceAttr.CopilotRunHasCompletedAt]: !!updates.completedAt,
    },
    async () => {
      const [run] = await db
        .update(copilotRuns)
        .set({
          status,
          completedAt: updates.completedAt,
          error: updates.error,
          requestContext: updates.requestContext,
          updatedAt: new Date(),
        })
        .where(eq(copilotRuns.id, runId))
        .returning()
      return run ?? null
    }
  )
}

async function getLatestRunForExecution(executionId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsGetLatestForExecution,
    'SELECT',
    'copilot_runs',
    { [TraceAttr.CopilotExecutionId]: executionId },
    async () => {
      const [run] = await db
        .select()
        .from(copilotRuns)
        .where(eq(copilotRuns.executionId, executionId))
        .orderBy(desc(copilotRuns.startedAt))
        .limit(1)
      return run ?? null
    }
  )
}

// Un-instrumented: called from a 4 Hz resume poll; per-call spans
// swamped traces. Use Prom histograms if latency visibility is needed.
export async function getLatestRunForStream(streamId: string, userId?: string) {
  const conditions = userId
    ? and(eq(copilotRuns.streamId, streamId), eq(copilotRuns.userId, userId))
    : eq(copilotRuns.streamId, streamId)
  const [run] = await db
    .select()
    .from(copilotRuns)
    .where(conditions)
    .orderBy(desc(copilotRuns.startedAt))
    .limit(1)
  return run ?? null
}

export async function getRunSegment(runId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsGetRunSegment,
    'SELECT',
    'copilot_runs',
    { [TraceAttr.RunId]: runId },
    async () => {
      const [run] = await db
        .select({ id: copilotRuns.id, userId: copilotRuns.userId })
        .from(copilotRuns)
        .where(eq(copilotRuns.id, runId))
        .limit(1)
      return run ?? null
    }
  )
}

async function createRunCheckpoint(input: {
  runId: string
  pendingToolCallId: string
  conversationSnapshot: Record<string, unknown>
  agentState: Record<string, unknown>
  providerRequest: Record<string, unknown>
}) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsCreateRunCheckpoint,
    'INSERT',
    'copilot_run_checkpoints',
    {
      [TraceAttr.RunId]: input.runId,
      [TraceAttr.CopilotCheckpointPendingToolCallId]: input.pendingToolCallId,
    },
    async () => {
      const [checkpoint] = await db
        .insert(copilotRunCheckpoints)
        .values({
          runId: input.runId,
          pendingToolCallId: input.pendingToolCallId,
          conversationSnapshot: input.conversationSnapshot,
          agentState: input.agentState,
          providerRequest: input.providerRequest,
        })
        .returning()

      return checkpoint
    }
  )
}

export async function upsertAsyncToolCall(input: {
  runId?: string | null
  checkpointId?: string | null
  toolCallId: string
  toolName: string
  args?: Record<string, unknown>
  status?: CopilotAsyncToolStatus
}) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsUpsertAsyncToolCall,
    'UPSERT',
    'copilot_async_tool_calls',
    {
      [TraceAttr.ToolCallId]: input.toolCallId,
      [TraceAttr.ToolName]: input.toolName,
      [TraceAttr.CopilotAsyncToolStatus]: input.status ?? 'pending',
      [TraceAttr.RunId]: input.runId ?? undefined,
    },
    async () => {
      const existing = await getAsyncToolCall(input.toolCallId)
      const incomingStatus = input.status ?? 'pending'
      if (
        existing &&
        (isTerminalAsyncStatus(existing.status) || isDeliveredAsyncStatus(existing.status)) &&
        !isTerminalAsyncStatus(incomingStatus) &&
        !isDeliveredAsyncStatus(incomingStatus)
      ) {
        logger.info('Ignoring async tool upsert that would downgrade terminal state', {
          toolCallId: input.toolCallId,
          existingStatus: existing.status,
          incomingStatus,
        })
        return existing
      }
      const effectiveRunId = input.runId ?? existing?.runId ?? null
      if (!effectiveRunId) {
        logger.warn('upsertAsyncToolCall missing runId and no existing row', {
          toolCallId: input.toolCallId,
          toolName: input.toolName,
          status: input.status ?? 'pending',
        })
        return null
      }

      const now = new Date()
      const [row] = await db
        .insert(copilotAsyncToolCalls)
        .values({
          runId: effectiveRunId,
          checkpointId: input.checkpointId ?? null,
          toolCallId: input.toolCallId,
          toolName: input.toolName,
          args: input.args ?? {},
          status: incomingStatus,
          updatedAt: now,
        })
        .onConflictDoUpdate({
          target: copilotAsyncToolCalls.toolCallId,
          set: {
            runId: effectiveRunId,
            checkpointId: input.checkpointId ?? null,
            toolName: input.toolName,
            args: input.args ?? {},
            status: incomingStatus,
            updatedAt: now,
          },
        })
        .returning()

      return row
    }
  )
}

export async function getAsyncToolCall(toolCallId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsGetAsyncToolCall,
    'SELECT',
    'copilot_async_tool_calls',
    { [TraceAttr.ToolCallId]: toolCallId },
    async () => {
      const [row] = await db
        .select()
        .from(copilotAsyncToolCalls)
        .where(eq(copilotAsyncToolCalls.toolCallId, toolCallId))
        .limit(1)
      return row ?? null
    }
  )
}

async function markAsyncToolStatus(
  toolCallId: string,
  status: CopilotAsyncToolStatus,
  updates: {
    claimedBy?: string | null
    claimedAt?: Date | null
    result?: AsyncCompletionData | null
    error?: string | null
    completedAt?: Date | null
  } = {}
) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsMarkAsyncToolStatus,
    'UPDATE',
    'copilot_async_tool_calls',
    {
      [TraceAttr.ToolCallId]: toolCallId,
      [TraceAttr.CopilotAsyncToolStatus]: status,
      [TraceAttr.CopilotAsyncToolHasError]: !!updates.error,
      [TraceAttr.CopilotAsyncToolClaimedBy]: updates.claimedBy ?? undefined,
    },
    async () => {
      const claimedAt =
        updates.claimedAt !== undefined
          ? updates.claimedAt
          : status === 'running' && updates.claimedBy
            ? new Date()
            : undefined

      const [row] = await db
        .update(copilotAsyncToolCalls)
        .set({
          status,
          claimedBy: updates.claimedBy,
          claimedAt,
          result: updates.result,
          error: updates.error,
          completedAt: updates.completedAt,
          updatedAt: new Date(),
        })
        .where(eq(copilotAsyncToolCalls.toolCallId, toolCallId))
        .returning()

      return row ?? null
    }
  )
}

export async function markAsyncToolRunning(toolCallId: string, claimedBy: string) {
  return markAsyncToolStatus(toolCallId, 'running', { claimedBy })
}

export async function completeAsyncToolCall(input: {
  toolCallId: string
  status: Extract<CopilotAsyncToolStatus, 'completed' | 'failed' | 'cancelled'>
  result?: AsyncCompletionData | null
  error?: string | null
}) {
  const existing = await getAsyncToolCall(input.toolCallId)

  if (!existing) {
    logger.warn('completeAsyncToolCall called before pending row existed', {
      toolCallId: input.toolCallId,
      status: input.status,
    })
    return null
  }

  if (isTerminalAsyncStatus(existing.status) || isDeliveredAsyncStatus(existing.status)) {
    return existing
  }

  return markAsyncToolStatus(input.toolCallId, input.status, {
    claimedBy: null,
    claimedAt: null,
    result: input.result ?? null,
    error: input.error ?? null,
    completedAt: new Date(),
  })
}

export async function markAsyncToolDelivered(toolCallId: string) {
  return markAsyncToolStatus(toolCallId, ASYNC_TOOL_STATUS.delivered, {
    claimedBy: null,
    claimedAt: null,
  })
}

async function listAsyncToolCallsForRun(runId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsListForRun,
    'SELECT',
    'copilot_async_tool_calls',
    { [TraceAttr.RunId]: runId },
    async () =>
      db
        .select()
        .from(copilotAsyncToolCalls)
        .where(eq(copilotAsyncToolCalls.runId, runId))
        .orderBy(desc(copilotAsyncToolCalls.createdAt))
  )
}

export async function getAsyncToolCalls(toolCallIds: string[]) {
  if (toolCallIds.length === 0) return []
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsGetMany,
    'SELECT',
    'copilot_async_tool_calls',
    { [TraceAttr.CopilotAsyncToolIdsCount]: toolCallIds.length },
    async () =>
      db
        .select()
        .from(copilotAsyncToolCalls)
        .where(inArray(copilotAsyncToolCalls.toolCallId, toolCallIds))
  )
}

export async function claimCompletedAsyncToolCall(toolCallId: string, workerId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsClaimCompleted,
    'UPDATE',
    'copilot_async_tool_calls',
    {
      [TraceAttr.ToolCallId]: toolCallId,
      [TraceAttr.CopilotAsyncToolWorkerId]: workerId,
    },
    async () => {
      const [row] = await db
        .update(copilotAsyncToolCalls)
        .set({
          claimedBy: workerId,
          claimedAt: new Date(),
          updatedAt: new Date(),
        })
        .where(
          and(
            eq(copilotAsyncToolCalls.toolCallId, toolCallId),
            inArray(copilotAsyncToolCalls.status, ['completed', 'failed', 'cancelled']),
            isNull(copilotAsyncToolCalls.claimedBy)
          )
        )
        .returning()
      return row ?? null
    }
  )
}

async function releaseCompletedAsyncToolClaim(toolCallId: string, workerId: string) {
  return withDbSpan(
    TraceSpan.CopilotAsyncRunsReleaseClaim,
    'UPDATE',
    'copilot_async_tool_calls',
    {
      [TraceAttr.ToolCallId]: toolCallId,
      [TraceAttr.CopilotAsyncToolWorkerId]: workerId,
    },
    async () => {
      const [row] = await db
        .update(copilotAsyncToolCalls)
        .set({
          claimedBy: null,
          claimedAt: null,
          updatedAt: new Date(),
        })
        .where(
          and(
            eq(copilotAsyncToolCalls.toolCallId, toolCallId),
            inArray(copilotAsyncToolCalls.status, ['completed', 'failed', 'cancelled']),
            eq(copilotAsyncToolCalls.claimedBy, workerId)
          )
        )
        .returning()
      return row ?? null
    }
  )
}