@lota-sdk/core 0.4.12 → 0.4.14

package/src/services/agent-executor.service.ts

@@ -12,12 +12,12 @@ import { stepCountIs, tool } from 'ai'
 import type { ToolSet } from 'ai'
 import { Context, Schema, Effect, Layer } from 'effect'
 
+import { RuntimeBridgeTag } from '../ai-gateway/ai-gateway'
 import type { ResolvedAgentConfig, ResolvedAgentFactoryConfig } from '../config/agent-defaults'
 import { ensureRecordId } from '../db/record-id'
 import type { SurrealDBService } from '../db/service'
 import { TABLES } from '../db/tables'
-import { BadRequestError, NotFoundError } from '../effect/errors'
-import { makeEffectTryPromiseWithOperation } from '../effect/helpers'
+import { ERROR_TAGS, BadRequestError, NotFoundError } from '../effect/errors'
 import {
   AgentConfigServiceTag,
   AgentFactoryServiceTag,
@@ -111,22 +111,12 @@ function validateOwner(agentConfig: ResolvedAgentConfig, agentId: string, nodeId
   return []
 }
 
-class AgentExecutorError extends Schema.TaggedErrorClass<AgentExecutorError>()('AgentExecutorError', {
+class AgentExecutorError extends Schema.TaggedErrorClass<AgentExecutorError>()(ERROR_TAGS.AgentExecutorError, {
   operation: Schema.String,
   message: Schema.String,
   cause: Schema.Defect,
 }) {}
 
-const tryAgentExecutorPromise = makeEffectTryPromiseWithOperation(
-  (operation, message, cause) => new AgentExecutorError({ operation, message, cause }),
-)
-const tryAgentExecutorTask = <A>(
-  operation: string,
-  message: string,
-  task: () => PromiseLike<A>,
-): Effect.Effect<A, AgentExecutorError> =>
-  Effect.tryPromise({ try: task, catch: (cause) => new AgentExecutorError({ operation, message, cause }) })
-
 function isAgentExecutorOwner(
   owner: PlanNodeSpec['owner'],
 ): owner is Extract<PlanNodeSpec['owner'], { executorType: 'agent' }> {
@@ -138,13 +128,14 @@ function buildWriteIntentTool(params: {
   workspace: ReturnType<ReturnType<typeof makeNodeWorkspaceService>['initialize']>
   writeIntentValidatorService: ReturnType<typeof makeWriteIntentValidatorService>
   stageWrite: (intent: WriteIntent) => Effect.Effect<void, BadRequestError>
+  runPromise: <A, E>(effect: Effect.Effect<A, E, never>, options?: { signal?: AbortSignal }) => Promise<A>
 }) {
   return tool({
     description:
       'Write a validated artifact or structured output field. Call this for each deliverable. If validation fails, correct your payload and try again.',
     inputSchema: WriteIntentSchema,
-    execute: (intent: WriteIntent) =>
-      Effect.runPromise(
+    execute: (intent: WriteIntent, { abortSignal }: { abortSignal?: AbortSignal } = {}) =>
+      params.runPromise(
         Effect.gen(function* () {
           const correctionCount = params.workspace.sys.correctionCounts.get(intent.targetPath) ?? 0
 
@@ -178,6 +169,7 @@ function buildWriteIntentTool(params: {
           yield* params.stageWrite(intent)
           return { status: 'accepted', message: `Write to "${intent.targetPath}" validated and staged.` }
         }),
+        abortSignal ? { signal: abortSignal } : undefined,
       ),
   })
 }
@@ -201,9 +193,15 @@ function createExecuteNode(deps: AgentExecutorDeps) {
       return yield* new BadRequestError({ message: `Agent executor "${agentId}" is not registered.` })
     }
 
-    const thread = yield* tryAgentExecutorPromise('load-thread', 'Failed to load dispatched thread.', () =>
-      deps.db.findOne(TABLES.THREAD, { id: ensureRecordId(params.context.threadId, TABLES.THREAD) }, ThreadSchema),
-    ).pipe(Effect.withSpan('AgentExecutor.loadThread'))
+    const thread = yield* deps.db
+      .findOne(TABLES.THREAD, { id: ensureRecordId(params.context.threadId, TABLES.THREAD) }, ThreadSchema)
+      .pipe(
+        Effect.mapError(
+          (cause) =>
+            new AgentExecutorError({ operation: 'load-thread', message: 'Failed to load dispatched thread.', cause }),
+        ),
+        Effect.withSpan('AgentExecutor.loadThread'),
+      )
     if (!thread) {
       return yield* new NotFoundError({
         resource: TABLES.THREAD,
@@ -223,22 +221,34 @@ function createExecuteNode(deps: AgentExecutorDeps) {
     const userName = params.context.userName ?? 'User'
     const runtimeAdapters = yield* RuntimeAdaptersServiceTag
     const { linearInstallation, githubInstallation, indexedRepoContext } = yield* Effect.all({
-      linearInstallation: tryAgentExecutorTask(
-        'get-linear-installation',
-        `Failed to load Linear installation for org ${params.context.organizationId}.`,
-        () => getLinearInstallationByOrgId(deps.agentFactoryConfig.pluginRuntime, organizationRef),
-      ),
-      githubInstallation: tryAgentExecutorTask(
-        'get-github-installation',
-        `Failed to load GitHub installation for org ${params.context.organizationId}.`,
-        () =>
+      linearInstallation: Effect.tryPromise({
+        try: () => getLinearInstallationByOrgId(deps.agentFactoryConfig.pluginRuntime, organizationRef),
+        catch: (cause) =>
+          new AgentExecutorError({
+            operation: 'get-linear-installation',
+            message: `Failed to load Linear installation for org ${params.context.organizationId}.`,
+            cause,
+          }),
+      }),
+      githubInstallation: Effect.tryPromise({
+        try: () =>
           getGithubInstallationForOrganization(deps.agentFactoryConfig.pluginRuntime, params.context.organizationId),
-      ),
-      indexedRepoContext: tryAgentExecutorTask(
-        'build-indexed-repositories-context',
-        `Failed to build indexed repository context for org ${params.context.organizationId}.`,
-        () => buildIndexedRepositoriesContext(runtimeAdapters, params.context.organizationId),
-      ),
+        catch: (cause) =>
+          new AgentExecutorError({
+            operation: 'get-github-installation',
+            message: `Failed to load GitHub installation for org ${params.context.organizationId}.`,
+            cause,
+          }),
+      }),
+      indexedRepoContext: Effect.tryPromise({
+        try: () => buildIndexedRepositoriesContext(runtimeAdapters, params.context.organizationId),
+        catch: (cause) =>
+          new AgentExecutorError({
+            operation: 'build-indexed-repositories-context',
+            message: `Failed to build indexed repository context for org ${params.context.organizationId}.`,
+            cause,
+          }),
+      }),
     }).pipe(Effect.withSpan('AgentExecutor.loadInstallations'))
 
     const mode = params.executionMode ?? 'linear'
@@ -253,12 +263,12 @@ function createExecuteNode(deps: AgentExecutorDeps) {
     ]
 
     const turnHooks = yield* TurnHooksServiceTag
-    const agentResolutionValue = yield* tryAgentExecutorPromise(
-      'resolve-agent',
-      `Failed to resolve agent "${agentId}" for dispatched execution.`,
-      () =>
-        Promise.resolve(
-          turnHooks.resolveAgent?.({
+    const resolveAgent = turnHooks.resolveAgent
+    let agentResolutionValue: unknown = undefined
+    if (resolveAgent) {
+      agentResolutionValue = yield* Effect.tryPromise({
+        try: () =>
+          resolveAgent({
             agentId,
             mode: dispatchMode,
             thread,
@@ -272,8 +282,14 @@ function createExecuteNode(deps: AgentExecutorDeps) {
             additionalInstructionSections: dispatchInstructionSections,
             context: null,
           }),
-        ),
-    )
+        catch: (cause) =>
+          new AgentExecutorError({
+            operation: 'resolve-agent',
+            message: `Failed to resolve agent "${agentId}" for dispatched execution.`,
+            cause,
+          }),
+      })
+    }
     const agentResolution = isRecord(agentResolutionValue) ? agentResolutionValue : null
     const resolvedAgentId = readOptionalString(agentResolution?.agentId) ?? agentId
     const runtimeConfig = deps.agentFactoryConfig.getAgentRuntimeConfig({
@@ -290,32 +306,34 @@ function createExecuteNode(deps: AgentExecutorDeps) {
     })
 
     const agentFactoryConfig = deps.agentFactoryConfig
-    const rawTools = yield* tryAgentExecutorPromise(
-      'build-agent-tools',
-      `Failed to build agent tools for "${resolvedAgentId}".`,
-      () =>
-        Promise.resolve(
-          agentFactoryConfig.buildAgentTools({
-            agentId: resolvedAgentId,
-            orgId: organizationRef,
-            userId: userRef,
-            userName,
-            threadId: threadRef,
-            orgIdString: params.context.organizationId,
-            threadType: thread.type,
-            mode: dispatchMode,
-            linearInstalled: Boolean(linearInstallation),
-            onboardingActive: false,
-            githubInstalled: Boolean(githubInstallation),
-            provideRepoTool: indexedRepoContext.provideRepoTool,
-            defaultRepoSections: indexedRepoContext.defaultSectionsByAgent[resolvedAgentId],
-            memoryBlock: '',
-            onAppendMemoryBlock: () => undefined,
-            availableUploads: [],
-            includeExecutionPlanTools: false,
-          }),
-        ),
-    )
+    const rawTools = yield* Effect.tryPromise({
+      try: () =>
+        agentFactoryConfig.buildAgentTools({
+          agentId: resolvedAgentId,
+          orgId: organizationRef,
+          userId: userRef,
+          userName,
+          threadId: threadRef,
+          orgIdString: params.context.organizationId,
+          threadType: thread.type,
+          mode: dispatchMode,
+          linearInstalled: Boolean(linearInstallation),
+          onboardingActive: false,
+          githubInstalled: Boolean(githubInstallation),
+          provideRepoTool: indexedRepoContext.provideRepoTool,
+          defaultRepoSections: indexedRepoContext.defaultSectionsByAgent[resolvedAgentId],
+          memoryBlock: '',
+          onAppendMemoryBlock: () => undefined,
+          availableUploads: [],
+          includeExecutionPlanTools: false,
+        }),
+      catch: (cause) =>
+        new AgentExecutorError({
+          operation: 'build-agent-tools',
+          message: `Failed to build agent tools for "${resolvedAgentId}".`,
+          cause,
+        }),
+    })
     const tools = applyToolPolicy(rawTools, params.nodeSpec)
 
     const agentFactory = agentFactoryConfig.createAgent[resolvedAgentId]
@@ -335,11 +353,15 @@ function createExecuteNode(deps: AgentExecutorDeps) {
         stopWhen: [stepCountIs(maxSteps)],
       })
 
-      const result = yield* tryAgentExecutorPromise(
-        'generate-agent-dispatch',
-        `Agent "${resolvedAgentId}" failed to generate a dispatch result.`,
-        () => agent.generate({ prompt: buildDispatchPrompt(params.nodeSpec) }),
-      ).pipe(Effect.withSpan('AgentExecutor.generateLinearDispatch'))
+      const result = yield* Effect.tryPromise({
+        try: () => agent.generate({ prompt: buildDispatchPrompt(params.nodeSpec) }),
+        catch: (cause) =>
+          new AgentExecutorError({
+            operation: 'generate-agent-dispatch',
+            message: `Agent "${resolvedAgentId}" failed to generate a dispatch result.`,
+            cause,
+          }),
+      }).pipe(Effect.withSpan('AgentExecutor.generateLinearDispatch'))
       const outputCandidate = PlanNodeResultSubmissionSchema.safeParse(result.output)
       if (outputCandidate.success) {
         return outputCandidate.data
@@ -365,6 +387,7 @@ function createExecuteNode(deps: AgentExecutorDeps) {
         workspace,
         writeIntentValidatorService: deps.writeIntentValidatorService,
         stageWrite: (intent) => deps.nodeWorkspaceService.stageWrite(workspace, intent, 'validated'),
+        runPromise: deps.runPromise,
       }),
     }
     const agent = agentFactory({
@@ -376,11 +399,15 @@ function createExecuteNode(deps: AgentExecutorDeps) {
       stopWhen: [stepCountIs(maxSteps)],
     })
 
-    yield* tryAgentExecutorPromise(
-      'generate-write-intent-dispatch',
-      `Agent "${resolvedAgentId}" failed to generate a write-intent dispatch result.`,
-      () => agent.generate({ prompt: buildWriteIntentDispatchPrompt(params.nodeSpec) }),
-    ).pipe(Effect.withSpan('AgentExecutor.generateWriteIntentDispatch'))
+    yield* Effect.tryPromise({
+      try: () => agent.generate({ prompt: buildWriteIntentDispatchPrompt(params.nodeSpec) }),
+      catch: (cause) =>
+        new AgentExecutorError({
+          operation: 'generate-write-intent-dispatch',
+          message: `Agent "${resolvedAgentId}" failed to generate a write-intent dispatch result.`,
+          cause,
+        }),
+    }).pipe(Effect.withSpan('AgentExecutor.generateWriteIntentDispatch'))
 
     const finalResult = yield* deps.nodeWorkspaceService.finalize(workspace)
 
@@ -411,6 +438,7 @@ interface AgentExecutorDeps {
   db: SurrealDBService
   nodeWorkspaceService: ReturnType<typeof makeNodeWorkspaceService>
   writeIntentValidatorService: ReturnType<typeof makeWriteIntentValidatorService>
+  runPromise: <A, E>(effect: Effect.Effect<A, E, never>, options?: { signal?: AbortSignal }) => Promise<A>
 }
 
 export function makeAgentExecutorService(deps: AgentExecutorDeps) {
@@ -433,12 +461,14 @@ export const AgentExecutorServiceLive = Layer.effect(
     const db = yield* DatabaseServiceTag
     const nodeWorkspaceService = yield* NodeWorkspaceServiceTag
     const writeIntentValidatorService = yield* WriteIntentValidatorServiceTag
+    const runtimeBridge = yield* RuntimeBridgeTag
     return makeAgentExecutorService({
       agentConfig,
       agentFactoryConfig,
       db,
       nodeWorkspaceService,
       writeIntentValidatorService,
+      runPromise: runtimeBridge.runPromise,
     })
   }),
 )