ummaya 0.2.4 → 0.2.6

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

@@ -0,0 +1,83 @@
+import {
+  isLocalAgentTask,
+  updateAgentProgress as updateAsyncAgentProgress,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import type { ToolUseContext } from '../../Tool.js'
+import type { Message as MessageType } from '../../types/message.js'
+import { getTaskOutputPath } from '../../utils/task/diskOutput.js'
+import { buildAgentSupportMetadata } from './orchestrationSupport.js'
+import { continueForegroundAgentInBackground } from './backgroundedLifecycle.js'
+import { canReadAgentOutputFile } from './asyncLifecycle.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import { runAgent } from './runAgent.js'
+import type {
+  AgentLifecycleMetadata,
+  AsyncLaunchedOutput,
+  WorktreeResult,
+} from './schemas.js'
+
+type SetAppState = Parameters<typeof updateAsyncAgentProgress>[2]
+type RunAgentParams = Parameters<typeof runAgent>[0]
+type SyncAgentContext = Parameters<typeof continueForegroundAgentInBackground>[0]['syncAgentContext']
+
+export function maybeBackgroundAgent({
+  foregroundTaskId,
+  agentIterator,
+  agentMessages,
+  toolUseContext,
+  rootSetAppState,
+  description,
+  prompt,
+  syncAgentId,
+  runAgentParams,
+  metadata,
+  getWorktreeResult,
+  syncAgentContext,
+  stopForegroundSummarization,
+}: {
+  readonly foregroundTaskId: string
+  readonly agentIterator: AsyncIterator<MessageType>
+  readonly agentMessages: MessageType[]
+  readonly toolUseContext: ToolUseContext
+  readonly rootSetAppState: SetAppState
+  readonly description: string
+  readonly prompt: string
+  readonly syncAgentId: string
+  readonly selectedAgent: AgentDefinition
+  readonly runAgentParams: RunAgentParams
+  readonly metadata: AgentLifecycleMetadata
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+  readonly syncAgentContext: SyncAgentContext
+  readonly stopForegroundSummarization: () => void
+}): AsyncLaunchedOutput | undefined {
+  const task = toolUseContext.getAppState().tasks[foregroundTaskId]
+  if (!isLocalAgentTask(task) || !task.isBackgrounded) return undefined
+  stopForegroundSummarization()
+  continueForegroundAgentInBackground({
+    backgroundedTaskId: foregroundTaskId,
+    task,
+    syncAgentContext,
+    agentIterator,
+    agentMessages,
+    runAgentParams,
+    rootSetAppState,
+    toolUseContext,
+    metadata,
+    description,
+    syncAgentId,
+    getWorktreeResult,
+  })
+  return {
+    isAsync: true,
+    status: 'async_launched',
+    agentId: foregroundTaskId,
+    description,
+    prompt,
+    outputFile: getTaskOutputPath(foregroundTaskId),
+    canReadOutputFile: canReadAgentOutputFile(toolUseContext.options.tools),
+    ...buildAgentSupportMetadata({
+      agentId: foregroundTaskId,
+      parentToolUseId: toolUseContext.toolUseId,
+    }),
+  }
+}

package/tui/src/tools/AgentTool/foregroundDrain.tsx

@@ -0,0 +1,133 @@
+import {
+  createActivityDescriptionResolver,
+  createProgressTracker,
+  updateAgentProgress as updateAsyncAgentProgress,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import type { ToolUseContext } from '../../Tool.js'
+import type { Message as MessageType } from '../../types/message.js'
+import type { ForegroundTaskRuntime } from './foregroundTask.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import { runAgent } from './runAgent.js'
+import { maybeBackgroundAgent } from './foregroundBackground.js'
+import {
+  emitMessageProgress,
+  showBackgroundHintIfNeeded,
+  type AssistantProgressSource,
+} from './foregroundProgress.js'
+import type {
+  AgentLifecycleMetadata,
+  AgentToolProgressCallback,
+  AsyncLaunchedOutput,
+  WorktreeResult,
+} from './schemas.js'
+
+type SetAppState = Parameters<typeof updateAsyncAgentProgress>[2]
+type RunAgentParams = Parameters<typeof runAgent>[0]
+type SyncAgentContext = Parameters<typeof maybeBackgroundAgent>[0]['syncAgentContext']
+
+export async function drainForegroundMessages({
+  agentIterator,
+  agentMessages,
+  agentStartTime,
+  syncTracker,
+  syncResolveActivity,
+  foreground,
+  toolUseContext,
+  rootSetAppState,
+  description,
+  prompt,
+  syncAgentId,
+  selectedAgent,
+  runAgentParams,
+  metadata,
+  assistantMessage,
+  onProgress,
+  getWorktreeResult,
+  syncAgentContext,
+  stopForegroundSummarization,
+}: {
+  readonly agentIterator: AsyncIterator<MessageType>
+  readonly agentMessages: MessageType[]
+  readonly agentStartTime: number
+  readonly syncTracker: ReturnType<typeof createProgressTracker>
+  readonly syncResolveActivity: ReturnType<typeof createActivityDescriptionResolver>
+  readonly foreground: ForegroundTaskRuntime
+  readonly toolUseContext: ToolUseContext
+  readonly rootSetAppState: SetAppState
+  readonly description: string
+  readonly prompt: string
+  readonly syncAgentId: string
+  readonly selectedAgent: AgentDefinition
+  readonly runAgentParams: RunAgentParams
+  readonly metadata: AgentLifecycleMetadata
+  readonly assistantMessage: AssistantProgressSource | undefined
+  readonly onProgress: AgentToolProgressCallback | undefined
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+  readonly syncAgentContext: SyncAgentContext
+  readonly stopForegroundSummarization: () => void
+}): Promise<AsyncLaunchedOutput | undefined> {
+  let backgroundHintShown = false
+  while (true) {
+    backgroundHintShown = showBackgroundHintIfNeeded(
+      backgroundHintShown,
+      agentStartTime,
+      toolUseContext,
+    )
+    const raceResult = await nextForegroundRace(agentIterator, foreground)
+    if (raceResult.type === 'background' && foreground.taskId) {
+      const backgrounded = maybeBackgroundAgent({
+        foregroundTaskId: foreground.taskId,
+        agentIterator,
+        agentMessages,
+        toolUseContext,
+        rootSetAppState,
+        description,
+        prompt,
+        syncAgentId,
+        selectedAgent,
+        runAgentParams,
+        metadata,
+        getWorktreeResult,
+        syncAgentContext,
+        stopForegroundSummarization,
+      })
+      if (backgrounded) return backgrounded
+      continue
+    }
+    if (raceResult.type !== 'message') continue
+    if (raceResult.result.done) break
+    const message = raceResult.result.value
+    agentMessages.push(message)
+    emitMessageProgress({
+      message,
+      syncTracker,
+      syncResolveActivity,
+      foregroundTaskId: foreground.taskId,
+      toolUseContext,
+      rootSetAppState,
+      description,
+      agentStartTime,
+      syncAgentId,
+      assistantMessage,
+      onProgress,
+    })
+  }
+  return undefined
+}
+
+async function nextForegroundRace(
+  agentIterator: AsyncIterator<MessageType>,
+  foreground: ForegroundTaskRuntime,
+): Promise<
+  | { readonly type: 'background' }
+  | { readonly type: 'message'; readonly result: IteratorResult<MessageType, void> }
+> {
+  const nextMessagePromise = agentIterator.next()
+  if (!foreground.backgroundPromise) {
+    return { type: 'message', result: await nextMessagePromise }
+  }
+  return Promise.race([
+    nextMessagePromise.then(result => ({ type: 'message' as const, result })),
+    foreground.backgroundPromise,
+  ])
+}

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

@@ -0,0 +1,98 @@
+import { feature } from 'bun:bundle'
+import { type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS, logEvent } from '../../services/analytics/index.js'
+import type { ToolUseContext } from '../../Tool.js'
+import type { Message as MessageType } from '../../types/message.js'
+import { AbortError } from '../../utils/errors.js'
+import { logForDebugging } from '../../utils/debug.js'
+import { isSyntheticMessage } from '../../utils/messages.js'
+import { classifyHandoffIfNeeded } from './agentToolHandoff.js'
+import { finalizeAgentTool } from './agentToolResult.js'
+import { buildAgentSupportMetadata } from './orchestrationSupport.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import type {
+  AgentLifecycleMetadata,
+  AgentToolCallResult,
+  WorktreeResult,
+} from './schemas.js'
+
+export function logTermination(
+  metadata: AgentLifecycleMetadata,
+  isAsync: boolean,
+  reason: 'user_cancel_sync' | 'user_cancel_background',
+): void {
+  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: isAsync,
+    is_built_in_agent: metadata.isBuiltInAgent,
+    reason: reason as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+  })
+}
+
+export async function finalizeSyncAgent({
+  agentMessages,
+  syncAgentError,
+  metadata,
+  selectedAgent,
+  toolUseContext,
+  prompt,
+  syncAgentId,
+  worktreeResult,
+}: {
+  readonly agentMessages: MessageType[]
+  readonly syncAgentError?: Error
+  readonly metadata: AgentLifecycleMetadata
+  readonly selectedAgent: AgentDefinition
+  readonly toolUseContext: ToolUseContext
+  readonly prompt: string
+  readonly syncAgentId: string
+  readonly worktreeResult: WorktreeResult
+}): Promise<AgentToolCallResult> {
+  const lastMessage = agentMessages.findLast(
+    message => message.type !== 'system' && message.type !== 'progress',
+  )
+  if (lastMessage && isSyntheticMessage(lastMessage)) {
+    logTermination(metadata, false, 'user_cancel_sync')
+    throw new AbortError()
+  }
+
+  if (syncAgentError) {
+    const hasAssistantMessages = agentMessages.some(
+      message => message.type === 'assistant',
+    )
+    if (!hasAssistantMessages) throw syncAgentError
+    logForDebugging(
+      `Sync agent recovering from error with ${agentMessages.length} messages`,
+    )
+  }
+  const agentResult = finalizeAgentTool(agentMessages, syncAgentId, metadata)
+  if (feature('TRANSCRIPT_CLASSIFIER')) {
+    const handoffWarning = await classifyHandoffIfNeeded({
+      agentMessages,
+      tools: toolUseContext.options.tools,
+      toolPermissionContext: toolUseContext.getAppState().toolPermissionContext,
+      abortSignal: toolUseContext.abortController.signal,
+      subagentType: selectedAgent.agentType,
+      totalToolUseCount: agentResult.totalToolUseCount,
+    })
+    if (handoffWarning) {
+      agentResult.content = [
+        { type: 'text', text: handoffWarning },
+        ...agentResult.content,
+      ]
+    }
+  }
+  return {
+    data: {
+      status: 'completed',
+      prompt,
+      ...agentResult,
+      ...worktreeResult,
+      ...buildAgentSupportMetadata({
+        agentId: syncAgentId,
+        parentToolUseId: toolUseContext.toolUseId,
+      }),
+    },
+  }
+}

package/tui/src/tools/AgentTool/foregroundLifecycle.tsx

@@ -0,0 +1,237 @@
+import { clearInvokedSkillsForAgent, getSdkAgentProgressSummariesEnabled } from '../../bootstrap/state.js'
+import { startAgentSummarization } from '../../services/AgentSummary/agentSummary.js'
+import { clearDumpState } from '../../services/api/dumpPrompts.js'
+import type { AppState } from '../../state/AppState.js'
+import {
+  createActivityDescriptionResolver,
+  createProgressTracker,
+  unregisterAgentForeground,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import type { ToolUseContext } from '../../Tool.js'
+import { runWithAgentContext } from '../../utils/agentContext.js'
+import { runWithCwdOverride } from '../../utils/cwd.js'
+import { logForDebugging } from '../../utils/debug.js'
+import { AbortError, errorMessage, toError } from '../../utils/errors.js'
+import type { CacheSafeParams } from '../../utils/forkedAgent.js'
+import { getParentSessionId } from '../../utils/teammate.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import { isBuiltInAgent } from './loadAgentsDir.js'
+import { runAgent } from './runAgent.js'
+import { drainForegroundMessages } from './foregroundDrain.js'
+import { finalizeSyncAgent, logTermination } from './foregroundFinalize.js'
+import { emitInitialProgress } from './foregroundProgress.js'
+import {
+  enqueueCompletionSdkEvent,
+  setupForegroundTask,
+} from './foregroundTask.js'
+import type {
+  AgentLifecycleMetadata,
+  AgentToolCallResult,
+  AgentToolProgressCallback,
+  WorktreeResult,
+} from './schemas.js'
+
+type SetAppState = (f: (prev: AppState) => AppState) => void
+type RunAgentParams = Parameters<typeof runAgent>[0]
+type AssistantProgressSource = {
+  readonly requestId?: string
+  readonly message?: { readonly id?: string }
+}
+
+export async function launchSyncAgent({
+  syncAgentId,
+  selectedAgent,
+  prompt,
+  promptMessages,
+  description,
+  runAgentParams,
+  toolUseContext,
+  rootSetAppState,
+  metadata,
+  assistantMessage,
+  cwdOverridePath,
+  onProgress,
+  getWorktreeResult,
+}: {
+  readonly syncAgentId: string
+  readonly selectedAgent: AgentDefinition
+  readonly prompt: string
+  readonly promptMessages: RunAgentParams['promptMessages']
+  readonly description: string
+  readonly runAgentParams: RunAgentParams
+  readonly toolUseContext: ToolUseContext
+  readonly rootSetAppState: SetAppState
+  readonly metadata: AgentLifecycleMetadata
+  readonly assistantMessage?: AssistantProgressSource
+  readonly cwdOverridePath?: string
+  readonly onProgress?: AgentToolProgressCallback
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+}): Promise<AgentToolCallResult> {
+  const syncAgentContext = {
+    agentId: syncAgentId,
+    parentSessionId: getParentSessionId(),
+    agentType: 'subagent' as const,
+    subagentName: selectedAgent.agentType,
+    isBuiltIn: isBuiltInAgent(selectedAgent),
+    invokingRequestId: assistantMessage?.requestId,
+    invocationKind: 'spawn' as const,
+    invocationEmitted: false,
+  }
+  return runWithAgentContext(syncAgentContext, () =>
+    withCwd(cwdOverridePath, () =>
+      runForegroundLoop({
+        syncAgentId,
+        selectedAgent,
+        prompt,
+        promptMessages,
+        description,
+        runAgentParams,
+        toolUseContext,
+        rootSetAppState,
+        metadata,
+        assistantMessage,
+        onProgress,
+        getWorktreeResult,
+        syncAgentContext,
+      }),
+    ),
+  )
+}
+
+async function runForegroundLoop({
+  syncAgentId,
+  selectedAgent,
+  prompt,
+  promptMessages,
+  description,
+  runAgentParams,
+  toolUseContext,
+  rootSetAppState,
+  metadata,
+  assistantMessage,
+  onProgress,
+  getWorktreeResult,
+  syncAgentContext,
+}: {
+  readonly syncAgentId: string
+  readonly selectedAgent: AgentDefinition
+  readonly prompt: string
+  readonly promptMessages: RunAgentParams['promptMessages']
+  readonly description: string
+  readonly runAgentParams: RunAgentParams
+  readonly toolUseContext: ToolUseContext
+  readonly rootSetAppState: SetAppState
+  readonly metadata: AgentLifecycleMetadata
+  readonly assistantMessage?: { readonly message?: { readonly id?: string } }
+  readonly onProgress?: AgentToolProgressCallback
+  readonly getWorktreeResult: () => Promise<WorktreeResult>
+  readonly syncAgentContext: Parameters<typeof runWithAgentContext>[0]
+}): Promise<AgentToolCallResult> {
+  const agentMessages: RunAgentParams['promptMessages'] = []
+  const agentStartTime = Date.now()
+  const syncTracker = createProgressTracker()
+  const syncResolveActivity = createActivityDescriptionResolver(toolUseContext.options.tools)
+  const foreground = setupForegroundTask({
+    syncAgentId,
+    selectedAgent,
+    description,
+    prompt,
+    rootSetAppState,
+    toolUseContext,
+  })
+  let wasBackgrounded = false
+  let wasAborted = false
+  let syncAgentError: Error | undefined
+  let worktreeResult: WorktreeResult = {}
+  let stopForegroundSummarization: (() => void) | undefined
+  emitInitialProgress(promptMessages, prompt, syncAgentId, assistantMessage, onProgress)
+  const agentIterator = runAgent({
+    ...runAgentParams,
+    override: { ...runAgentParams.override, agentId: syncAgentId },
+    onCacheSafeParams:
+      foreground.taskId && getSdkAgentProgressSummariesEnabled()
+        ? (params: CacheSafeParams) => {
+            const { stop } = startAgentSummarization(
+              foreground.taskId,
+              syncAgentId,
+              params,
+              rootSetAppState,
+            )
+            stopForegroundSummarization = stop
+          }
+        : undefined,
+  })[Symbol.asyncIterator]()
+
+  try {
+    const backgrounded = await drainForegroundMessages({
+      agentIterator,
+      agentMessages,
+      agentStartTime,
+      syncTracker,
+      syncResolveActivity,
+      foreground,
+      toolUseContext,
+      rootSetAppState,
+      description,
+      prompt,
+      syncAgentId,
+      selectedAgent,
+      runAgentParams,
+      metadata,
+      assistantMessage,
+      onProgress,
+      getWorktreeResult,
+      syncAgentContext,
+      stopForegroundSummarization: () => stopForegroundSummarization?.(),
+    })
+    if (backgrounded) {
+      wasBackgrounded = true
+      return { data: backgrounded }
+    }
+  } catch (error) {
+    const syncError = error instanceof Error ? error : toError(error)
+    if (syncError instanceof AbortError) {
+      wasAborted = true
+      logTermination(metadata, false, 'user_cancel_sync')
+      throw syncError
+    }
+    logForDebugging(`Sync agent error: ${errorMessage(syncError)}`, { level: 'error' })
+    syncAgentError = syncError
+  } finally {
+    toolUseContext.setToolJSX?.(null)
+    stopForegroundSummarization?.()
+    if (foreground.taskId) {
+      unregisterAgentForeground(foreground.taskId, rootSetAppState)
+      if (!wasBackgrounded) {
+        enqueueCompletionSdkEvent({
+          foregroundTaskId: foreground.taskId,
+          toolUseContext,
+          syncTracker,
+          syncAgentError,
+          wasAborted,
+          description,
+          agentStartTime,
+        })
+      }
+    }
+    clearInvokedSkillsForAgent(syncAgentId)
+    if (!wasBackgrounded) clearDumpState(syncAgentId)
+    foreground.cancelAutoBackground?.()
+    if (!wasBackgrounded) worktreeResult = await getWorktreeResult()
+  }
+
+  return finalizeSyncAgent({
+    agentMessages,
+    syncAgentError,
+    metadata,
+    selectedAgent,
+    toolUseContext,
+    prompt,
+    syncAgentId,
+    worktreeResult,
+  })
+}
+
+function withCwd<T>(cwdOverridePath: string | undefined, fn: () => T): T {
+  return cwdOverridePath ? runWithCwdOverride(cwdOverridePath, fn) : fn()
+}