ummaya 0.2.3 → 0.2.5

package/tui/src/tools/AgentTool/asyncAgentLifecycle.ts

@@ -0,0 +1,208 @@
+import { feature } from 'bun:bundle'
+import { clearInvokedSkillsForAgent } from '../../bootstrap/state.js'
+import { startAgentSummarization } from '../../services/AgentSummary/agentSummary.js'
+import {
+  type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+  logEvent,
+} from '../../services/analytics/index.js'
+import { clearDumpState } from '../../services/api/dumpPrompts.js'
+import type { AppState } from '../../state/AppState.js'
+import type { ToolUseContext } from '../../Tool.js'
+import {
+  completeAgentTask as completeAsyncAgent,
+  createActivityDescriptionResolver,
+  createProgressTracker,
+  enqueueAgentNotification,
+  failAgentTask as failAsyncAgent,
+  getProgressUpdate,
+  getTokenCountFromTracker,
+  isLocalAgentTask,
+  killAsyncAgent,
+  updateAgentProgress as updateAsyncAgentProgress,
+  updateProgressFromMessage,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import { asAgentId } from '../../types/ids.js'
+import type { Message as MessageType } from '../../types/message.js'
+import { AbortError, errorMessage } from '../../utils/errors.js'
+import type { CacheSafeParams } from '../../utils/forkedAgent.js'
+import { extractTextContent } from '../../utils/messages.js'
+import { classifyHandoffIfNeeded } from './agentToolHandoff.js'
+import { extractPartialResult } from './agentToolPartialResult.js'
+import { emitTaskProgress, getLastToolUseName } from './agentToolProgress.js'
+import {
+  finalizeAgentTool,
+  type FinalizeAgentToolMetadata,
+} from './agentToolResult.js'
+
+type SetAppState = (f: (prev: AppState) => AppState) => void
+
+export async function runAsyncAgentLifecycle({
+  taskId,
+  abortController,
+  makeStream,
+  metadata,
+  description,
+  toolUseContext,
+  rootSetAppState,
+  agentIdForCleanup,
+  enableSummarization,
+  getWorktreeResult,
+}: {
+  taskId: string
+  abortController: AbortController
+  makeStream: (
+    onCacheSafeParams: ((p: CacheSafeParams) => void) | undefined,
+  ) => AsyncGenerator<MessageType, void>
+  metadata: FinalizeAgentToolMetadata
+  description: string
+  toolUseContext: ToolUseContext
+  rootSetAppState: SetAppState
+  agentIdForCleanup: string
+  enableSummarization: boolean
+  getWorktreeResult: () => Promise<{
+    worktreePath?: string
+    worktreeBranch?: string
+  }>
+}): Promise<void> {
+  let stopSummarization: (() => void) | undefined
+  const agentMessages: MessageType[] = []
+  try {
+    const tracker = createProgressTracker()
+    const resolveActivity = createActivityDescriptionResolver(
+      toolUseContext.options.tools,
+    )
+    const onCacheSafeParams = enableSummarization
+      ? (params: CacheSafeParams) => {
+          const { stop } = startAgentSummarization(
+            taskId,
+            asAgentId(taskId),
+            params,
+            rootSetAppState,
+          )
+          stopSummarization = stop
+        }
+      : undefined
+    for await (const message of makeStream(onCacheSafeParams)) {
+      agentMessages.push(message)
+      rootSetAppState(prev => {
+        const task = prev.tasks[taskId]
+        if (!isLocalAgentTask(task) || !task.retain) return prev
+        const base = task.messages ?? []
+        return {
+          ...prev,
+          tasks: {
+            ...prev.tasks,
+            [taskId]: { ...task, messages: [...base, message] },
+          },
+        }
+      })
+      updateProgressFromMessage(
+        tracker,
+        message,
+        resolveActivity,
+        toolUseContext.options.tools,
+      )
+      updateAsyncAgentProgress(
+        taskId,
+        getProgressUpdate(tracker),
+        rootSetAppState,
+      )
+      const lastToolName = getLastToolUseName(message)
+      if (lastToolName) {
+        emitTaskProgress(
+          tracker,
+          taskId,
+          toolUseContext.toolUseId,
+          description,
+          metadata.startTime,
+          lastToolName,
+        )
+      }
+    }
+
+    stopSummarization?.()
+
+    const agentResult = finalizeAgentTool(agentMessages, taskId, metadata)
+
+    completeAsyncAgent(agentResult, rootSetAppState)
+
+    let finalMessage = extractTextContent(agentResult.content, '\n')
+
+    if (feature('TRANSCRIPT_CLASSIFIER')) {
+      const handoffWarning = await classifyHandoffIfNeeded({
+        agentMessages,
+        tools: toolUseContext.options.tools,
+        toolPermissionContext:
+          toolUseContext.getAppState().toolPermissionContext,
+        abortSignal: abortController.signal,
+        subagentType: metadata.agentType,
+        totalToolUseCount: agentResult.totalToolUseCount,
+      })
+      if (handoffWarning) {
+        finalMessage = `${handoffWarning}\n\n${finalMessage}`
+      }
+    }
+
+    const worktreeResult = await getWorktreeResult()
+
+    enqueueAgentNotification({
+      taskId,
+      description,
+      status: 'completed',
+      setAppState: rootSetAppState,
+      finalMessage,
+      usage: {
+        totalTokens: getTokenCountFromTracker(tracker),
+        toolUses: agentResult.totalToolUseCount,
+        durationMs: agentResult.totalDurationMs,
+      },
+      toolUseId: toolUseContext.toolUseId,
+      ...worktreeResult,
+    })
+  } catch (error) {
+    stopSummarization?.()
+    const asyncError =
+      error instanceof Error ? error : new Error(errorMessage(error))
+    if (asyncError instanceof AbortError) {
+      killAsyncAgent(taskId, rootSetAppState)
+      logEvent('tengu_agent_tool_terminated', {
+        agent_type:
+          metadata.agentType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+        model:
+          metadata.resolvedAgentModel as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+        duration_ms: Date.now() - metadata.startTime,
+        is_async: true,
+        is_built_in_agent: metadata.isBuiltInAgent,
+        reason:
+          'user_kill_async' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      })
+      const worktreeResult = await getWorktreeResult()
+      const partialResult = extractPartialResult(agentMessages)
+      enqueueAgentNotification({
+        taskId,
+        description,
+        status: 'killed',
+        setAppState: rootSetAppState,
+        toolUseId: toolUseContext.toolUseId,
+        finalMessage: partialResult,
+        ...worktreeResult,
+      })
+      return
+    }
+    const msg = errorMessage(asyncError)
+    failAsyncAgent(taskId, msg, rootSetAppState)
+    const worktreeResult = await getWorktreeResult()
+    enqueueAgentNotification({
+      taskId,
+      description,
+      status: 'failed',
+      error: msg,
+      setAppState: rootSetAppState,
+      toolUseId: toolUseContext.toolUseId,
+      ...worktreeResult,
+    })
+  } finally {
+    clearInvokedSkillsForAgent(agentIdForCleanup)
+    clearDumpState(agentIdForCleanup)
+  }
+}

package/tui/src/tools/AgentTool/asyncLifecycle.ts

@@ -0,0 +1,153 @@
+import type { AppState } from '../../state/AppState.js'
+import type { ToolUseContext, Tools } from '../../Tool.js'
+import { getSdkAgentProgressSummariesEnabled } from '../../bootstrap/state.js'
+import {
+  registerAsyncAgent,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import { asAgentId } from '../../types/ids.js'
+import type { AssistantMessage } from '../../types/message.js'
+import { runWithAgentContext } from '../../utils/agentContext.js'
+import { runWithCwdOverride } from '../../utils/cwd.js'
+import { getParentSessionId } from '../../utils/teammate.js'
+import { getTaskOutputPath } from '../../utils/task/diskOutput.js'
+import { toolMatchesName } from '../../Tool.js'
+import { BASH_TOOL_NAME } from '../BashTool/toolName.js'
+import { FILE_READ_TOOL_NAME } from '../FileReadTool/prompt.js'
+import { isForkSubagentEnabled } from './forkSubagent.js'
+import { buildAgentSupportMetadata } from './orchestrationSupport.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import { isBuiltInAgent } from './loadAgentsDir.js'
+import { runAgent } from './runAgent.js'
+import type {
+  AgentLifecycleMetadata,
+  AsyncLaunchedOutput,
+  WorktreeResult,
+} from './schemas.js'
+
+type SetAppState = (f: (prev: AppState) => AppState) => void
+type RunAgentParams = Parameters<typeof runAgent>[0]
+
+export function canReadAgentOutputFile(tools: Tools): boolean {
+  return tools.some(
+    tool => toolMatchesName(tool, FILE_READ_TOOL_NAME) || toolMatchesName(tool, BASH_TOOL_NAME),
+  )
+}
+
+export function registerAgentName({
+  name,
+  agentId,
+  rootSetAppState,
+}: {
+  readonly name?: string
+  readonly agentId: string
+  readonly rootSetAppState: SetAppState
+}): void {
+  if (!name) return
+  rootSetAppState(prev => {
+    const next = new Map(prev.agentNameRegistry)
+    next.set(name, asAgentId(agentId))
+    return { ...prev, agentNameRegistry: next }
+  })
+}
+
+export function launchAsyncAgent({
+  asyncAgentId,
+  description,
+  prompt,
+  selectedAgent,
+  rootSetAppState,
+  toolUseContext,
+  runAgentParams,
+  metadata,
+  assistantMessage,
+  cwdOverridePath,
+  isCoordinator,
+  getWorktreeResult,
+}: {
+  readonly asyncAgentId: string
+  readonly description: string
+  readonly prompt: string
+  readonly selectedAgent: AgentDefinition
+  readonly rootSetAppState: SetAppState
+  readonly toolUseContext: ToolUseContext
+  readonly runAgentParams: RunAgentParams
+  readonly metadata: AgentLifecycleMetadata
+  readonly assistantMessage?: AssistantMessage
+  readonly cwdOverridePath?: string
+  readonly isCoordinator: boolean
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+}): AsyncLaunchedOutput {
+  const agentBackgroundTask = registerAsyncAgent({
+    agentId: asyncAgentId,
+    description,
+    prompt,
+    selectedAgent,
+    setAppState: rootSetAppState,
+    toolUseId: toolUseContext.toolUseId,
+  })
+  const abortController = agentBackgroundTask.abortController
+  if (!abortController) {
+    throw new Error(
+      `Cannot launch agent ${agentBackgroundTask.agentId}: missing abort controller`,
+    )
+  }
+  void runWithAgentContext(
+    {
+      agentId: asyncAgentId,
+      parentSessionId: getParentSessionId(),
+      agentType: 'subagent',
+      subagentName: selectedAgent.agentType,
+      isBuiltIn: isBuiltInAgent(selectedAgent),
+      invokingRequestId: assistantMessage?.requestId,
+      invocationKind: 'spawn',
+      invocationEmitted: false,
+    },
+    async () =>
+      withCwd(cwdOverridePath, () =>
+        import('./asyncAgentLifecycle.js').then(
+          ({ runAsyncAgentLifecycle }) =>
+            runAsyncAgentLifecycle({
+              taskId: agentBackgroundTask.agentId,
+              abortController,
+              makeStream: onCacheSafeParams =>
+                runAgent({
+                  ...runAgentParams,
+                  override: {
+                    ...runAgentParams.override,
+                    agentId: asAgentId(agentBackgroundTask.agentId),
+                    abortController,
+                  },
+                  onCacheSafeParams,
+                }),
+              metadata,
+              description,
+              toolUseContext,
+              rootSetAppState,
+              agentIdForCleanup: asyncAgentId,
+              enableSummarization:
+                isCoordinator ||
+                isForkSubagentEnabled() ||
+                getSdkAgentProgressSummariesEnabled(),
+              getWorktreeResult,
+            }),
+        ),
+      ),
+  )
+  return {
+    isAsync: true,
+    status: 'async_launched',
+    agentId: agentBackgroundTask.agentId,
+    description,
+    prompt,
+    outputFile: getTaskOutputPath(agentBackgroundTask.agentId),
+    canReadOutputFile: canReadAgentOutputFile(toolUseContext.options.tools),
+    ...buildAgentSupportMetadata({
+      agentId: agentBackgroundTask.agentId,
+      parentToolUseId: toolUseContext.toolUseId,
+    }),
+  }
+}
+
+function withCwd<T>(cwdOverridePath: string | undefined, fn: () => T): T {
+  return cwdOverridePath ? runWithCwdOverride(cwdOverridePath, fn) : fn()
+}

package/tui/src/tools/AgentTool/backgroundedCompletion.ts

@@ -0,0 +1,126 @@
+import { feature } from 'bun:bundle'
+import { type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, logEvent } from '../../services/analytics/index.js'
+import {
+  completeAgentTask as completeAsyncAgent,
+  createProgressTracker,
+  enqueueAgentNotification,
+  failAgentTask as failAsyncAgent,
+  getTokenCountFromTracker,
+  killAsyncAgent,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import type { Message as MessageType } from '../../types/message.js'
+import { AbortError, errorMessage } from '../../utils/errors.js'
+import { extractTextContent } from '../../utils/messages.js'
+import { classifyHandoffIfNeeded } from './agentToolHandoff.js'
+import { extractPartialResult } from './agentToolPartialResult.js'
+import { finalizeAgentTool } from './agentToolResult.js'
+import { runAgent } from './runAgent.js'
+import type { AgentLifecycleMetadata, WorktreeResult } from './schemas.js'
+
+type SetAppState = Parameters<typeof completeAsyncAgent>[1]
+type RunAgentParams = Parameters<typeof runAgent>[0]
+
+export async function finishBackgroundedAgent({
+  backgroundedTaskId,
+  agentMessages,
+  metadata,
+  rootSetAppState,
+  tracker,
+  toolUseContext,
+  abortController,
+  description,
+  getWorktreeResult,
+}: {
+  readonly backgroundedTaskId: string
+  readonly agentMessages: MessageType[]
+  readonly metadata: AgentLifecycleMetadata
+  readonly rootSetAppState: SetAppState
+  readonly tracker: ReturnType<typeof createProgressTracker>
+  readonly toolUseContext: RunAgentParams['toolUseContext']
+  readonly abortController: AbortController
+  readonly description: string
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+}): Promise<void> {
+  const agentResult = finalizeAgentTool(agentMessages, backgroundedTaskId, metadata)
+  completeAsyncAgent(agentResult, rootSetAppState)
+  let finalMessage = extractTextContent(agentResult.content, '\n')
+  if (feature('TRANSCRIPT_CLASSIFIER')) {
+    const handoffWarning = await classifyHandoffIfNeeded({
+      agentMessages,
+      tools: toolUseContext.options.tools,
+      toolPermissionContext: toolUseContext.getAppState().toolPermissionContext,
+      abortSignal: abortController.signal,
+      subagentType: metadata.agentType,
+      totalToolUseCount: agentResult.totalToolUseCount,
+    })
+    if (handoffWarning) finalMessage = `${handoffWarning}\n\n${finalMessage}`
+  }
+  enqueueAgentNotification({
+    taskId: backgroundedTaskId,
+    description,
+    status: 'completed',
+    setAppState: rootSetAppState,
+    finalMessage,
+    usage: {
+      totalTokens: getTokenCountFromTracker(tracker),
+      toolUses: agentResult.totalToolUseCount,
+      durationMs: agentResult.totalDurationMs,
+    },
+    toolUseId: toolUseContext.toolUseId,
+    ...(await getWorktreeResult()),
+  })
+}
+
+export async function failBackgroundedAgent({
+  error,
+  backgroundedTaskId,
+  rootSetAppState,
+  metadata,
+  description,
+  agentMessages,
+  getWorktreeResult,
+  toolUseContext,
+}: {
+  readonly error: unknown
+  readonly backgroundedTaskId: string
+  readonly rootSetAppState: SetAppState
+  readonly metadata: AgentLifecycleMetadata
+  readonly description: string
+  readonly agentMessages: MessageType[]
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+  readonly toolUseContext: RunAgentParams['toolUseContext']
+}): Promise<void> {
+  const worktreeResult = await getWorktreeResult()
+  if (error instanceof AbortError) {
+    killAsyncAgent(backgroundedTaskId, rootSetAppState)
+    logEvent('tengu_agent_tool_terminated', {
+      agent_type: metadata.agentType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      model: metadata.resolvedAgentModel as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      duration_ms: Date.now() - metadata.startTime,
+      is_async: true,
+      is_built_in_agent: metadata.isBuiltInAgent,
+      reason: 'user_cancel_background' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+    })
+    enqueueAgentNotification({
+      taskId: backgroundedTaskId,
+      description,
+      status: 'killed',
+      setAppState: rootSetAppState,
+      toolUseId: toolUseContext.toolUseId,
+      finalMessage: extractPartialResult(agentMessages),
+      ...worktreeResult,
+    })
+    return
+  }
+  const errMsg = errorMessage(error)
+  failAsyncAgent(backgroundedTaskId, errMsg, rootSetAppState)
+  enqueueAgentNotification({
+    taskId: backgroundedTaskId,
+    description,
+    status: 'failed',
+    error: errMsg,
+    setAppState: rootSetAppState,
+    toolUseId: toolUseContext.toolUseId,
+    ...worktreeResult,
+  })
+}

package/tui/src/tools/AgentTool/backgroundedLifecycle.ts

@@ -0,0 +1,174 @@
+import { clearInvokedSkillsForAgent, getSdkAgentProgressSummariesEnabled } from '../../bootstrap/state.js'
+import { startAgentSummarization } from '../../services/AgentSummary/agentSummary.js'
+import { clearDumpState } from '../../services/api/dumpPrompts.js'
+import {
+  completeAgentTask as completeAsyncAgent,
+  createActivityDescriptionResolver,
+  createProgressTracker,
+  getProgressUpdate,
+  type LocalAgentTaskState,
+  updateAgentProgress as updateAsyncAgentProgress,
+  updateProgressFromMessage,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import { asAgentId } from '../../types/ids.js'
+import type { Message as MessageType } from '../../types/message.js'
+import { runWithAgentContext } from '../../utils/agentContext.js'
+import { errorMessage } from '../../utils/errors.js'
+import type { CacheSafeParams } from '../../utils/forkedAgent.js'
+import { logForDebugging } from '../../utils/debug.js'
+import { sleep } from '../../utils/sleep.js'
+import {
+  emitTaskProgress,
+  getLastToolUseName,
+} from './agentToolProgress.js'
+import type { AgentLifecycleMetadata, WorktreeResult } from './schemas.js'
+import { runAgent } from './runAgent.js'
+import {
+  failBackgroundedAgent,
+  finishBackgroundedAgent,
+} from './backgroundedCompletion.js'
+
+type SetAppState = Parameters<typeof completeAsyncAgent>[1]
+type RunAgentParams = Parameters<typeof runAgent>[0]
+type AgentContextCallback = Parameters<typeof runWithAgentContext>[0]
+
+export function continueForegroundAgentInBackground({
+  backgroundedTaskId,
+  task,
+  syncAgentContext,
+  agentIterator,
+  agentMessages,
+  runAgentParams,
+  rootSetAppState,
+  toolUseContext,
+  metadata,
+  description,
+  syncAgentId,
+  getWorktreeResult,
+}: {
+  readonly backgroundedTaskId: string
+  readonly task: LocalAgentTaskState
+  readonly syncAgentContext: AgentContextCallback
+  readonly agentIterator: AsyncIterator<MessageType>
+  readonly agentMessages: MessageType[]
+  readonly runAgentParams: RunAgentParams
+  readonly rootSetAppState: SetAppState
+  readonly toolUseContext: RunAgentParams['toolUseContext']
+  readonly metadata: AgentLifecycleMetadata
+  readonly description: string
+  readonly syncAgentId: string
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+}): void {
+  const abortController = task.abortController
+  if (!abortController) {
+    throw new Error(
+      `Cannot background agent ${backgroundedTaskId}: missing abort controller`,
+    )
+  }
+
+  void runWithAgentContext(syncAgentContext, async () => {
+    let stopBackgroundedSummarization: (() => void) | undefined
+    try {
+      await closeIterator(agentIterator)
+      const tracker = createProgressTracker()
+      const resolveActivity = createActivityDescriptionResolver(
+        toolUseContext.options.tools,
+      )
+      for (const existingMsg of agentMessages) {
+        updateProgressFromMessage(
+          tracker,
+          existingMsg,
+          resolveActivity,
+          toolUseContext.options.tools,
+        )
+      }
+      for await (const msg of runAgent({
+        ...runAgentParams,
+        isAsync: true,
+        override: {
+          ...runAgentParams.override,
+          agentId: asAgentId(backgroundedTaskId),
+          abortController,
+        },
+        onCacheSafeParams: getSdkAgentProgressSummariesEnabled()
+          ? (params: CacheSafeParams) => {
+              const { stop } = startAgentSummarization(
+                backgroundedTaskId,
+                asAgentId(backgroundedTaskId),
+                params,
+                rootSetAppState,
+              )
+              stopBackgroundedSummarization = stop
+            }
+          : undefined,
+      })) {
+        agentMessages.push(msg)
+        updateProgressFromMessage(
+          tracker,
+          msg,
+          resolveActivity,
+          toolUseContext.options.tools,
+        )
+        updateAsyncAgentProgress(
+          backgroundedTaskId,
+          getProgressUpdate(tracker),
+          rootSetAppState,
+        )
+        const lastToolName = getLastToolUseName(msg)
+        if (lastToolName) {
+          emitTaskProgress(
+            tracker,
+            backgroundedTaskId,
+            toolUseContext.toolUseId,
+            description,
+            metadata.startTime,
+            lastToolName,
+          )
+        }
+      }
+      await finishBackgroundedAgent({
+        backgroundedTaskId,
+        agentMessages,
+        metadata,
+        rootSetAppState,
+        tracker,
+        toolUseContext,
+        abortController,
+        description,
+        getWorktreeResult,
+      })
+    } catch (error) {
+      const handledError =
+        error instanceof Error ? error : new Error(errorMessage(error))
+      await failBackgroundedAgent({
+        error: handledError,
+        backgroundedTaskId,
+        rootSetAppState,
+        metadata,
+        description,
+        agentMessages,
+        getWorktreeResult,
+        toolUseContext,
+      })
+    } finally {
+      stopBackgroundedSummarization?.()
+      clearInvokedSkillsForAgent(syncAgentId)
+      clearDumpState(syncAgentId)
+    }
+  })
+}
+
+async function closeIterator(agentIterator: AsyncIterator<MessageType>): Promise<void> {
+  const close = agentIterator.return
+    ? agentIterator.return(undefined).then(
+        () => undefined,
+        error => {
+          logForDebugging(
+            `Agent iterator cleanup failed: ${errorMessage(error)}`,
+            { level: 'warn' },
+          )
+        },
+      )
+    : Promise.resolve()
+  await Promise.race([close, sleep(1000)])
+}