@lota-sdk/core 0.1.20 → 0.1.22

package/src/services/workstream-turn.ts

@@ -1,19 +1,28 @@
-import type { ChatMessage } from '@lota-sdk/shared'
+import type { ChatMessage, PlanNodeHandoffContext } from '@lota-sdk/shared'
 import { createUIMessageStream } from 'ai'
 
 import { lotaDebugLogger } from '../config/debug-logger'
+import { ensureRecordId, recordIdToString } from '../db/record-id'
+import { TABLES } from '../db/tables'
 import { hasApprovalRespondedParts, isApprovalContinuationRequest } from '../runtime/approval-continuation'
+import { shouldPlanNodeUseVisibleTurn } from '../runtime/execution-plan-visibility'
 import { wrapResponseWithKeepalive } from '../utils/sse-keepalive'
+import { planExecutorService } from './plan-executor.service'
+import { planRunService } from './plan-run.service'
+import { userService } from './user.service'
 import { prepareWorkstreamRunCore } from './workstream-turn-preparation.service'
 import type {
   PreparedWorkstreamTurnResult,
   WorkstreamApprovalContinuationParams,
+  WorkstreamPlanTurnParams,
   WorkstreamTurnParams,
 } from './workstream-turn-preparation.service'
+import { workstreamService } from './workstream.service'
 
 export { hasApprovalRespondedParts, isApprovalContinuationRequest }
 export { wrapResponseWithKeepalive }
 export type { PreparedWorkstreamTurnResult }
+export type { WorkstreamPlanTurnParams }
 
 export async function createWorkstreamApprovalContinuationStream(params: WorkstreamApprovalContinuationParams) {
   const timer = lotaDebugLogger.timer('turn:approval-continuation')
@@ -70,3 +79,94 @@ export async function runWorkstreamTurnInBackground(
   timer.step('run')
   return result
 }
+
+export async function triggerPlanNodeTurn(params: {
+  runId: string
+  nodeId: string
+  abortSignal?: AbortSignal
+  streamId?: string
+}): Promise<PreparedWorkstreamTurnResult> {
+  const timer = lotaDebugLogger.timer('turn:plan-turn')
+  let run = await planRunService.getRunById(params.runId)
+  const spec = await planRunService.getPlanSpecById(run.planSpecId)
+  const nodeSpec = await planRunService.getNodeSpecByNodeId(spec.id, params.nodeId)
+
+  if (!shouldPlanNodeUseVisibleTurn(spec, nodeSpec) || nodeSpec.owner.executorType !== 'agent') {
+    throw new Error(`Plan node "${params.nodeId}" is not eligible for a visible plan turn.`)
+  }
+
+  let nodeRun = await planRunService.getNodeRunByNodeId(run.id, params.nodeId)
+  if (nodeRun.status === 'ready') {
+    await planExecutorService.transitionNodeToRunning({ runId: params.runId, nodeId: params.nodeId })
+    run = await planRunService.getRunById(params.runId)
+    nodeRun = await planRunService.getNodeRunByNodeId(run.id, params.nodeId)
+  }
+
+  if (nodeRun.status !== 'running') {
+    return { assistantMessages: [], inputMessageId: undefined }
+  }
+
+  const workstream = await workstreamService.getWorkstream(run.workstreamId)
+  const userRef = ensureRecordId(workstream.userId, TABLES.USER)
+
+  const [artifacts, nodeRuns, userName] = await Promise.all([
+    planRunService.listArtifacts(run.id),
+    planRunService.listNodeRuns(run.id),
+    userService
+      .getUser(recordIdToString(userRef, TABLES.USER))
+      .then((user) => user.name)
+      .catch(() => undefined),
+  ])
+  const inputArtifacts = artifacts
+    .filter((artifact) => nodeSpec.upstreamNodeIds.includes(artifact.nodeId))
+    .map((artifact) => ({
+      name: artifact.name,
+      kind: artifact.kind,
+      pointer: artifact.pointer,
+      ...(artifact.schemaRef ? { schemaRef: artifact.schemaRef } : {}),
+      ...(artifact.description ? { description: artifact.description } : {}),
+      ...(artifact.payload !== undefined ? { payload: artifact.payload } : {}),
+    }))
+  const upstreamNodeSpecs = new Map(
+    (await planRunService.listNodeSpecs(spec.id)).map((upstreamNodeSpec) => [
+      upstreamNodeSpec.nodeId,
+      upstreamNodeSpec,
+    ]),
+  )
+  const upstreamHandoffs = nodeRuns
+    .filter((candidate) => nodeSpec.upstreamNodeIds.includes(candidate.nodeId) && candidate.handoffContext)
+    .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',
+      handoffContext: candidate.handoffContext as PlanNodeHandoffContext,
+    }))
+  timer.step('load-plan-turn-context')
+
+  const prepared = await prepareWorkstreamRunCore({
+    kind: 'planTurn',
+    workstream,
+    workstreamRef: ensureRecordId(run.workstreamId, TABLES.WORKSTREAM),
+    orgRef: ensureRecordId(run.organizationId, TABLES.ORGANIZATION),
+    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,
+    },
+  })
+  timer.step('prepare')
+
+  const result = await prepared.run()
+  timer.step('run')
+  return result
+}

package/src/services/workstream.service.ts

@@ -10,6 +10,7 @@ import type { RecordIdInput, RecordIdRef } from '../db/record-id'
 import { databaseService } from '../db/service'
 import type { DatabaseTable } from '../db/tables'
 import { TABLES } from '../db/tables'
+import { getRedisConnection, withRedisLeaseLock } from '../redis'
 import {
   MEMORY_BLOCK_COMPACTION_CHUNK_ENTRIES,
   MEMORY_BLOCK_COMPACTION_TRIGGER_ENTRIES,
@@ -62,6 +63,10 @@ WHERE userId = $userId
   AND status = "regular"
 ORDER BY updatedAt DESC`
 
+const WORKSTREAM_ACTIVE_RUN_LOCK_TTL_MS = 90_000
+const WORKSTREAM_ACTIVE_RUN_LOCK_MAX_WAIT_MS = 750
+const WORKSTREAM_ACTIVE_RUN_LOCK_RETRY_DELAY_MS = 75
+
 function toSafeDirectIdSegment(value: string): string {
   return value.replace(/[^a-zA-Z0-9_-]/g, '_')
 }
@@ -107,6 +112,17 @@ function requireString(coreType: string | undefined): string {
   return coreType
 }
 
+function buildActiveRunLockKey(workstreamId: RecordIdRef): string {
+  return `workstream-active-run:${recordIdToString(ensureRecordId(workstreamId, TABLES.WORKSTREAM), TABLES.WORKSTREAM)}`
+}
+
+export class ActiveWorkstreamRunConflictError extends Error {
+  constructor() {
+    super('A chat run is already active.')
+    this.name = 'ActiveWorkstreamRunConflictError'
+  }
+}
+
 function buildDirectWorkstreamId({
   userId,
   orgId,
@@ -188,7 +204,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
 
       const existing = await this.findById(directWorkstreamId)
       if (existing) {
-        return this.normalizeWorkstream(existing)
+        return await this.toNormalizedWorkstream(existing)
       }
 
       let createError: unknown = null
@@ -216,7 +232,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
         throw new Error('Failed to create or load direct workstream')
       }
 
-      return this.normalizeWorkstream(workstream)
+      return await this.toNormalizedWorkstream(workstream)
     }
 
     if (core) {
@@ -225,7 +241,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
       const coreWorkstreamId = buildCoreWorkstreamId({ userId, orgId, coreType: resolvedCoreType })
       const existing = await this.findById(coreWorkstreamId)
       if (existing) {
-        return this.normalizeWorkstream(existing)
+        return await this.toNormalizedWorkstream(existing)
       }
 
       let createError: unknown = null
@@ -263,7 +279,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
         throw new Error('Failed to create or load core workstream')
       }
 
-      return this.normalizeWorkstream(workstream)
+      return await this.toNormalizedWorkstream(workstream)
     }
 
     const groupWorkstream = await this.create({
@@ -276,7 +292,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
       nameGenerated: options?.title !== undefined && options.title !== WORKSTREAM.DEFAULT_TITLE,
     })
 
-    return this.normalizeWorkstream(groupWorkstream)
+    return await this.toNormalizedWorkstream(groupWorkstream)
   }
 
   async ensureBootstrapWorkstreams(
@@ -387,7 +403,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
         WorkstreamSchema,
       )
 
-      return { workstreams: workstreams.map((workstream) => this.normalizeWorkstream(workstream)), hasMore: false }
+      return { workstreams: await this.toNormalizedWorkstreams(workstreams), hasMore: false }
     }
 
     const take = options.take ?? WORKSTREAM.DEFAULT_PAGE_LIMIT
@@ -408,7 +424,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
     const hasMore = workstreams.length > take
     const sliced = hasMore ? workstreams.slice(0, take) : workstreams
 
-    return { workstreams: sliced.map((workstream) => this.normalizeWorkstream(workstream)), hasMore }
+    return { workstreams: await this.toNormalizedWorkstreams(sliced), hasMore }
   }
 
   async listOrganizationWorkstreams(params: {
@@ -450,19 +466,19 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
       WorkstreamSchema,
     )
 
-    return workstreams.map((workstream) => this.normalizeWorkstream(workstream))
+    return await this.toNormalizedWorkstreams(workstreams)
   }
 
   async getWorkstream(workstreamId: RecordIdRef): Promise<NormalizedWorkstream> {
     const workstream = await this.getById(workstreamId)
-    return this.normalizeWorkstream(workstream)
+    return await this.toNormalizedWorkstream(workstream)
   }
 
   async updateTitle(workstreamId: RecordIdRef, title: string): Promise<NormalizedWorkstream> {
     const existing = await this.getById(workstreamId)
     this.assertMutableWorkstream(existing, 'rename')
     const workstream = await this.update(workstreamId, { title, nameGenerated: true })
-    return this.normalizeWorkstream(workstream)
+    return await this.toNormalizedWorkstream(workstream)
   }
 
   async updateStatus(workstreamId: RecordIdRef, status: string): Promise<NormalizedWorkstream> {
@@ -470,7 +486,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
     const existing = await this.getById(workstreamId)
     this.assertMutableWorkstream(existing, validStatus === 'archived' ? 'archive' : 'unarchive')
     const workstream = await this.update(workstreamId, { status: validStatus })
-    return this.normalizeWorkstream(workstream)
+    return await this.toNormalizedWorkstream(workstream)
   }
 
   async setActiveRunId(workstreamId: RecordIdRef, runId: string | null): Promise<void> {
@@ -497,6 +513,33 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
     return normalized.length > 0 ? normalized : null
   }
 
+  async hasActiveRunLease(workstreamId: RecordIdRef): Promise<boolean> {
+    const count = await getRedisConnection().exists(buildActiveRunLockKey(workstreamId))
+    return count > 0
+  }
+
+  async withActiveRunLease<T>(workstreamId: RecordIdRef, fn: (signal: AbortSignal) => Promise<T>): Promise<T> {
+    try {
+      return await withRedisLeaseLock(
+        {
+          redis: getRedisConnection(),
+          lockKey: buildActiveRunLockKey(workstreamId),
+          lockTtlMs: WORKSTREAM_ACTIVE_RUN_LOCK_TTL_MS,
+          retryDelayMs: WORKSTREAM_ACTIVE_RUN_LOCK_RETRY_DELAY_MS,
+          maxWaitMs: WORKSTREAM_ACTIVE_RUN_LOCK_MAX_WAIT_MS,
+          label: 'workstream active run',
+          logger: serverLogger,
+        },
+        fn,
+      )
+    } catch (error) {
+      if (error instanceof Error && error.message.startsWith('Timed out waiting for workstream active run')) {
+        throw new ActiveWorkstreamRunConflictError()
+      }
+      throw error
+    }
+  }
+
   async clearActiveRunIdIfMatches(workstreamId: RecordIdRef, runId: string): Promise<void> {
     const workstreamRef = ensureRecordId(workstreamId, TABLES.WORKSTREAM)
     await databaseService.query(surql`UPDATE ONLY ${workstreamRef} SET activeRunId = NONE WHERE activeRunId = ${runId}`)
@@ -527,7 +570,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
 
   async clearStaleActiveRunIfMissingFromRegistry(workstreamId: RecordIdRef): Promise<boolean> {
     const activeRunId = await this.getActiveRunId(workstreamId)
-    if (!activeRunId || chatRunRegistry.has(activeRunId)) {
+    if (!activeRunId || (await this.hasActiveRunLease(workstreamId))) {
       return false
     }
 
@@ -537,7 +580,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
       activeStreamId ? this.clearActiveStreamIdIfMatches(workstreamId, activeStreamId) : Promise.resolve(),
     ])
 
-    serverLogger.warn`Cleared stale workstream run after process restart: workstream=${recordIdToString(ensureRecordId(workstreamId, TABLES.WORKSTREAM), TABLES.WORKSTREAM)} run=${activeRunId}`
+    serverLogger.warn`Cleared stale workstream run after lease expired: workstream=${recordIdToString(ensureRecordId(workstreamId, TABLES.WORKSTREAM), TABLES.WORKSTREAM)} run=${activeRunId}`
     return true
   }
 
@@ -656,7 +699,7 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
       WorkstreamSchema,
     )
 
-    return workstreams.map((workstream) => this.normalizeWorkstream(workstream))
+    return await this.toNormalizedWorkstreams(workstreams)
   }
 
   private normalizeWorkstreamId(id: unknown): string {
@@ -686,12 +729,25 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
     return WORKSTREAM.DEFAULT_TITLE
   }
 
-  normalizeWorkstream(workstream: WorkstreamRecord): NormalizedWorkstream {
+  private async computeIsRunning(workstream: Pick<WorkstreamRecord, 'id' | 'activeRunId'>): Promise<boolean> {
     const activeRunId =
       typeof workstream.activeRunId === 'string' && workstream.activeRunId.trim().length > 0
         ? workstream.activeRunId
         : null
-    const isRunning = activeRunId !== null && chatRunRegistry.has(activeRunId)
+
+    if (activeRunId === null) {
+      return false
+    }
+
+    if (chatRunRegistry.has(activeRunId)) {
+      return true
+    }
+
+    return await this.hasActiveRunLease(ensureRecordId(workstream.id, TABLES.WORKSTREAM))
+  }
+
+  private async toNormalizedWorkstream(workstream: WorkstreamRecord): Promise<NormalizedWorkstream> {
+    const isRunning = await this.computeIsRunning(workstream)
     const isCompacting = workstream.isCompacting === true
     const mode = workstream.mode
     const core = workstream.core
@@ -716,6 +772,10 @@ class WorkstreamService extends BaseService<typeof WorkstreamSchema> {
     }
   }
 
+  private async toNormalizedWorkstreams(workstreams: WorkstreamRecord[]): Promise<NormalizedWorkstream[]> {
+    return await Promise.all(workstreams.map(async (workstream) => await this.toNormalizedWorkstream(workstream)))
+  }
+
   toPublicWorkstream(workstream: NormalizedWorkstream) {
     return {
       id: workstream.id,

package/src/tools/execution-plan.tool.ts

@@ -27,7 +27,6 @@ export function createCreateExecutionPlanTool(params: {
         organizationId: params.orgId,
         workstreamId: params.workstreamId,
         leadAgentId: params.agentId,
-        dispatchMode: 'deferred',
         input,
       })
       params.onPlanChanged?.()
@@ -51,7 +50,6 @@ export function createReplaceExecutionPlanTool(params: {
         organizationId: params.orgId,
         workstreamId: params.workstreamId,
         leadAgentId: params.agentId,
-        dispatchMode: 'deferred',
         input,
       })
       params.onPlanChanged?.()