ummaya 0.2.4 → 0.2.5

package/tui/src/tools/TaskOutputTool/lookup.ts

@@ -0,0 +1,216 @@
+import type { ToolResult, ToolUseContext, ValidationResult } from '../../Tool.js'
+import type { TaskState } from '../../tasks/types.js'
+import { AbortError } from '../../utils/errors.js'
+import { extractTextContent } from '../../utils/messages.js'
+import { sleep } from '../../utils/sleep.js'
+import { getTaskOutput } from '../../utils/task/diskOutput.js'
+import { updateTaskState } from '../../utils/task/framework.js'
+import { buildAgentSupportMetadata } from '../AgentTool/orchestrationSupport.js'
+import type {
+  TaskOutputRecord,
+  TaskOutputToolInput,
+  TaskOutputToolOutput,
+} from './schemas.js'
+
+type TaskOutputProgressEvent = {
+  readonly toolUseID: string
+  readonly data: {
+    readonly type: 'waiting_for_task'
+    readonly taskDescription: string
+    readonly taskType: string
+    readonly evidenceJoinKey: string
+    readonly parentToolUseId: string
+    readonly resumeToken: string
+    readonly permissionFlow: 'coordinator_parent_round_trip'
+  }
+}
+
+type TaskOutputProgressCallback = (progress: TaskOutputProgressEvent) => void
+
+export class TaskOutputLookupError extends Error {
+  constructor(readonly taskId: string) {
+    super(`No task found with ID: ${taskId}`)
+    this.name = 'TaskOutputLookupError'
+  }
+}
+
+function assertNeverTask(task: never): never {
+  throw new TaskOutputLookupError(task)
+}
+
+async function readTaskOutput(task: TaskState): Promise<string> {
+  if (task.type !== 'local_bash') return getTaskOutput(task.id)
+  const taskOutput = task.shellCommand?.taskOutput
+  if (!taskOutput) return getTaskOutput(task.id)
+
+  const stdout = await taskOutput.getStdout()
+  const stderr = await taskOutput.getStderr()
+  return [stdout, stderr].filter(Boolean).join('\n')
+}
+
+export async function getTaskOutputData(
+  task: TaskState,
+): Promise<TaskOutputRecord> {
+  const output = await readTaskOutput(task)
+  const baseOutput = {
+    task_id: task.id,
+    task_type: task.type,
+    status: task.status,
+    description: task.description,
+    output,
+  }
+
+  switch (task.type) {
+    case 'local_bash':
+      return { ...baseOutput, exitCode: task.result?.code ?? null }
+    case 'local_agent': {
+      const cleanResult = task.result
+        ? extractTextContent(task.result.content, '\n')
+        : undefined
+      return {
+        ...baseOutput,
+        prompt: task.prompt,
+        result: cleanResult || output,
+        output: cleanResult || output,
+        error: task.error,
+      }
+    }
+    case 'remote_agent':
+      return { ...baseOutput, prompt: task.command }
+    case 'in_process_teammate':
+    case 'local_workflow':
+    case 'monitor_mcp':
+    case 'dream':
+      return baseOutput
+    default:
+      return assertNeverTask(task)
+  }
+}
+
+export async function validateTaskOutputInput(
+  { task_id }: TaskOutputToolInput,
+  { getAppState }: Pick<ToolUseContext, 'getAppState'>,
+): Promise<ValidationResult> {
+  if (!task_id) {
+    return { result: false, message: 'Task ID is required', errorCode: 1 }
+  }
+
+  const task = getAppState().tasks?.[task_id]
+  if (!task) {
+    return {
+      result: false,
+      message: `No task found with ID: ${task_id}`,
+      errorCode: 2,
+    }
+  }
+  return { result: true }
+}
+
+async function waitForTaskCompletion(
+  taskId: string,
+  getAppState: ToolUseContext['getAppState'],
+  timeoutMs: number,
+  abortController?: AbortController,
+): Promise<TaskState | null> {
+  const startTime = Date.now()
+  while (Date.now() - startTime < timeoutMs) {
+    if (abortController?.signal.aborted) throw new AbortError()
+
+    const task = getAppState().tasks?.[taskId]
+    if (!task) return null
+    if (task.status !== 'running' && task.status !== 'pending') return task
+
+    await sleep(100)
+  }
+
+  return getAppState().tasks?.[taskId] ?? null
+}
+
+function markTaskNotified(taskId: string, context: ToolUseContext): void {
+  updateTaskState(taskId, context.setAppState, task => ({
+    ...task,
+    notified: true,
+  }))
+}
+
+function isTaskStillActive(task: TaskState): boolean {
+  return task.status === 'running' || task.status === 'pending'
+}
+
+export async function callTaskOutputTool(
+  input: TaskOutputToolInput,
+  context: ToolUseContext,
+  onProgress?: TaskOutputProgressCallback,
+): Promise<ToolResult<TaskOutputToolOutput>> {
+  const task = context.getAppState().tasks?.[input.task_id]
+  if (!task) throw new TaskOutputLookupError(input.task_id)
+
+  const supportMetadata = buildAgentSupportMetadata({
+    taskId: input.task_id,
+    parentToolUseId: task.toolUseId ?? context.toolUseId,
+  })
+
+  if (!input.block) {
+    if (!isTaskStillActive(task)) {
+      markTaskNotified(input.task_id, context)
+      return {
+        data: {
+          retrieval_status: 'success',
+          task: await getTaskOutputData(task),
+          ...supportMetadata,
+        },
+      }
+    }
+    return {
+      data: {
+        retrieval_status: 'not_ready',
+        task: await getTaskOutputData(task),
+        ...supportMetadata,
+      },
+    }
+  }
+
+  onProgress?.({
+    toolUseID: `task-output-waiting-${Date.now()}`,
+    data: {
+      type: 'waiting_for_task',
+      taskDescription: task.description,
+      taskType: task.type,
+      ...supportMetadata,
+    },
+  })
+
+  const completedTask = await waitForTaskCompletion(
+    input.task_id,
+    context.getAppState,
+    input.timeout,
+    context.abortController,
+  )
+  if (!completedTask) {
+    return {
+      data: {
+        retrieval_status: 'timeout',
+        task: null,
+        ...supportMetadata,
+      },
+    }
+  }
+  if (isTaskStillActive(completedTask)) {
+    return {
+      data: {
+        retrieval_status: 'timeout',
+        task: await getTaskOutputData(completedTask),
+        ...supportMetadata,
+      },
+    }
+  }
+
+  markTaskNotified(input.task_id, context)
+  return {
+    data: {
+      retrieval_status: 'success',
+      task: await getTaskOutputData(completedTask),
+      ...supportMetadata,
+    },
+  }
+}

package/tui/src/tools/TaskOutputTool/render.tsx

@@ -0,0 +1,257 @@
+import React from 'react'
+import { FallbackToolUseErrorMessage } from '../../components/FallbackToolUseErrorMessage.js'
+import { FallbackToolUseRejectedMessage } from '../../components/FallbackToolUseRejectedMessage.js'
+import { MessageResponse } from '../../components/MessageResponse.js'
+import { Box, Text } from '../../ink.js'
+import { useShortcutDisplay } from '../../keybindings/useShortcutDisplay.js'
+import type { ProgressMessage } from '../../types/message.js'
+import { jsonParse } from '../../utils/slowOperations.js'
+import { countCharInString } from '../../utils/stringUtils.js'
+import type { ThemeName } from '../../utils/theme.js'
+import { AgentPromptDisplay, AgentResponseDisplay } from '../AgentTool/UI.js'
+import BashToolResultMessage from '../BashTool/BashToolResultMessage.js'
+import type { TaskOutputToolInput, TaskOutputToolOutput } from './schemas.js'
+import { taskOutputToolOutputSchema } from './schemas.js'
+
+type TaskOutputResultDisplayProps = {
+  readonly content: TaskOutputToolOutput | string
+  readonly verbose?: boolean
+  readonly theme: ThemeName
+}
+
+function parseTaskOutputResult(
+  content: TaskOutputToolOutput | string,
+): TaskOutputToolOutput {
+  if (typeof content !== 'string') return content
+  return taskOutputToolOutputSchema.parse(jsonParse(content))
+}
+
+function readTaskDescription(
+  message: ProgressMessage | undefined,
+): string | undefined {
+  const data = message?.data
+  if (typeof data !== 'object' || data === null) return undefined
+  if (!('taskDescription' in data)) return undefined
+  return typeof data.taskDescription === 'string'
+    ? data.taskDescription
+    : undefined
+}
+
+function StillRunningMessage(): React.ReactNode {
+  return (
+    <MessageResponse>
+      <Text dimColor={true}>Task is still running…</Text>
+    </MessageResponse>
+  )
+}
+
+export function renderTaskOutputUseMessage(
+  input: Partial<TaskOutputToolInput>,
+): React.ReactNode {
+  const { block = true } = input
+  return block ? '' : 'non-blocking'
+}
+
+export function renderTaskOutputUseTag(
+  input: Partial<TaskOutputToolInput>,
+): React.ReactNode {
+  if (!input.task_id) return null
+  return <Text dimColor> {input.task_id}</Text>
+}
+
+export function renderTaskOutputUseProgressMessage(
+  progressMessages: readonly ProgressMessage[],
+): React.ReactNode {
+  const taskDescription = readTaskDescription(
+    progressMessages[progressMessages.length - 1],
+  )
+  return (
+    <Box flexDirection="column">
+      {taskDescription && (
+        <Text>
+          {'  '}
+          {taskDescription}
+        </Text>
+      )}
+      <Text>
+        {'     '}Waiting for task{' '}
+        <Text dimColor>(esc to give additional instructions)</Text>
+      </Text>
+    </Box>
+  )
+}
+
+export function TaskOutputResultDisplay({
+  content,
+  verbose = false,
+  theme,
+}: TaskOutputResultDisplayProps): React.ReactNode {
+  const result = parseTaskOutputResult(content)
+  if (!result.task) {
+    return (
+      <MessageResponse>
+        <Text dimColor={true}>No task output available</Text>
+      </MessageResponse>
+    )
+  }
+
+  const { task } = result
+  const expandShortcut = useShortcutDisplay(
+    'app:toggleTranscript',
+    'Global',
+    'ctrl+o',
+  )
+
+  switch (task.task_type) {
+    case 'local_bash':
+      return (
+        <BashToolResultMessage
+          content={{
+            stdout: task.output,
+            stderr: '',
+            isImage: false,
+            dangerouslyDisableSandbox: true,
+            returnCodeInterpretation: task.error,
+          }}
+          verbose={verbose}
+        />
+      )
+    case 'local_agent':
+      return renderLocalAgentOutput(result, expandShortcut, verbose, theme)
+    case 'remote_agent':
+      return renderRemoteAgentOutput(task, expandShortcut, verbose)
+    case 'in_process_teammate':
+    case 'local_workflow':
+    case 'monitor_mcp':
+    case 'dream':
+      return renderGenericTaskOutput(task)
+  }
+}
+
+function renderLocalAgentOutput(
+  result: TaskOutputToolOutput,
+  expandShortcut: string,
+  verbose: boolean,
+  theme: ThemeName,
+): React.ReactNode {
+  const task = result.task
+  if (!task) return null
+
+  if (result.retrieval_status === 'success') {
+    if (!verbose) {
+      return (
+        <MessageResponse>
+          <Text dimColor={true}>Read output ({expandShortcut} to expand)</Text>
+        </MessageResponse>
+      )
+    }
+
+    const lineCount = task.result
+      ? countCharInString(task.result, '\n') + 1
+      : 0
+    return (
+      <Box flexDirection="column">
+        <Text>
+          {task.description} ({lineCount} lines)
+        </Text>
+        <Box flexDirection="column" paddingLeft={2} marginTop={1}>
+          {task.prompt && (
+            <AgentPromptDisplay prompt={task.prompt} theme={theme} dim={true} />
+          )}
+          {task.result && (
+            <Box marginTop={1}>
+              <AgentResponseDisplay
+                content={[{ type: 'text', text: task.result }]}
+                theme={theme}
+              />
+            </Box>
+          )}
+          {task.error && (
+            <Box flexDirection="column" marginTop={1}>
+              <Text color="error" bold={true}>
+                Error:
+              </Text>
+              <Box paddingLeft={2}>
+                <Text color="error">{task.error}</Text>
+              </Box>
+            </Box>
+          )}
+        </Box>
+      </Box>
+    )
+  }
+
+  if (result.retrieval_status === 'timeout' || task.status === 'running') {
+    return <StillRunningMessage />
+  }
+  if (result.retrieval_status === 'not_ready') return <StillRunningMessage />
+  return (
+    <MessageResponse>
+      <Text dimColor={true}>Task not ready</Text>
+    </MessageResponse>
+  )
+}
+
+function renderRemoteAgentOutput(
+  task: NonNullable<TaskOutputToolOutput['task']>,
+  expandShortcut: string,
+  verbose: boolean,
+): React.ReactNode {
+  return (
+    <Box flexDirection="column">
+      <Text>
+        {'  '}
+        {task.description} [{task.status}]
+      </Text>
+      {task.output && verbose && (
+        <Box paddingLeft={4} marginTop={1}>
+          <Text>{task.output}</Text>
+        </Box>
+      )}
+      {!verbose && task.output && (
+        <Text dimColor={true}>
+          {'     '}({expandShortcut} to expand)
+        </Text>
+      )}
+    </Box>
+  )
+}
+
+function renderGenericTaskOutput(
+  task: NonNullable<TaskOutputToolOutput['task']>,
+): React.ReactNode {
+  return (
+    <Box flexDirection="column">
+      <Text>
+        {'  '}
+        {task.description} [{task.status}]
+      </Text>
+      {task.output && (
+        <Box paddingLeft={4}>
+          <Text>{task.output.slice(0, 500)}</Text>
+        </Box>
+      )}
+    </Box>
+  )
+}
+
+export function renderTaskOutputResultMessage(
+  content: TaskOutputToolOutput | string,
+  verbose: boolean,
+  theme: ThemeName,
+): React.ReactNode {
+  return (
+    <TaskOutputResultDisplay content={content} verbose={verbose} theme={theme} />
+  )
+}
+
+export function renderTaskOutputRejectedMessage(): React.ReactNode {
+  return <FallbackToolUseRejectedMessage />
+}
+
+export function renderTaskOutputErrorMessage(
+  result: unknown,
+  verbose: boolean,
+): React.ReactNode {
+  return <FallbackToolUseErrorMessage result={result} verbose={verbose} />
+}

package/tui/src/tools/TaskOutputTool/schemas.ts

@@ -0,0 +1,55 @@
+import { z } from 'zod/v4'
+import type { TaskType } from '../../Task.js'
+import { lazySchema } from '../../utils/lazySchema.js'
+import { semanticBoolean } from '../../utils/semanticBoolean.js'
+
+export const inputSchema = lazySchema(() =>
+  z.strictObject({
+    task_id: z.string().describe('The task ID to get output from'),
+    block: semanticBoolean(z.boolean().default(true)).describe(
+      'Whether to wait for completion',
+    ),
+    timeout: z
+      .number()
+      .min(0)
+      .max(600000)
+      .default(30000)
+      .describe('Max wait time in ms'),
+  }),
+)
+
+const taskTypeSchema = z.enum([
+  'local_bash',
+  'local_agent',
+  'remote_agent',
+  'in_process_teammate',
+  'local_workflow',
+  'monitor_mcp',
+  'dream',
+] satisfies readonly TaskType[])
+
+export const taskOutputSchema = z.object({
+  task_id: z.string(),
+  task_type: taskTypeSchema,
+  status: z.string(),
+  description: z.string(),
+  output: z.string(),
+  exitCode: z.number().nullable().optional(),
+  error: z.string().optional(),
+  prompt: z.string().optional(),
+  result: z.string().optional(),
+})
+
+export const taskOutputToolOutputSchema = z.object({
+  retrieval_status: z.enum(['success', 'timeout', 'not_ready']),
+  task: taskOutputSchema.nullable(),
+  evidenceJoinKey: z.string(),
+  parentToolUseId: z.string(),
+  resumeToken: z.string(),
+  permissionFlow: z.literal('coordinator_parent_round_trip'),
+})
+
+export type InputSchema = ReturnType<typeof inputSchema>
+export type TaskOutputToolInput = z.infer<InputSchema>
+export type TaskOutputRecord = z.infer<typeof taskOutputSchema>
+export type TaskOutputToolOutput = z.infer<typeof taskOutputToolOutputSchema>

package/tui/src/tools/TaskOutputTool/serialization.ts

@@ -0,0 +1,36 @@
+import type { ToolResultBlockParam } from '@anthropic-ai/sdk/resources/index.mjs'
+import { formatTaskOutput } from '../../utils/task/outputFormatting.js'
+import type { TaskOutputToolOutput } from './schemas.js'
+
+export function mapTaskOutputResultToToolResultBlockParam(
+  data: TaskOutputToolOutput,
+  toolUseID: string,
+): ToolResultBlockParam {
+  const parts: string[] = []
+  parts.push(`<retrieval_status>${data.retrieval_status}</retrieval_status>`)
+  parts.push(`<evidence_join_key>${data.evidenceJoinKey}</evidence_join_key>`)
+  parts.push(`<parent_tool_use_id>${data.parentToolUseId}</parent_tool_use_id>`)
+  parts.push(`<resume_token>${data.resumeToken}</resume_token>`)
+  parts.push(`<permission_flow>${data.permissionFlow}</permission_flow>`)
+  if (data.task) {
+    parts.push(`<task_id>${data.task.task_id}</task_id>`)
+    parts.push(`<task_type>${data.task.task_type}</task_type>`)
+    parts.push(`<status>${data.task.status}</status>`)
+    if (data.task.exitCode !== undefined && data.task.exitCode !== null) {
+      parts.push(`<exit_code>${data.task.exitCode}</exit_code>`)
+    }
+    if (data.task.output?.trim()) {
+      const { content } = formatTaskOutput(
+        data.task.output,
+        data.task.task_id,
+      )
+      parts.push(`<output>\n${content.trimEnd()}\n</output>`)
+    }
+    if (data.task.error) parts.push(`<error>${data.task.error}</error>`)
+  }
+  return {
+    tool_use_id: toolUseID,
+    type: 'tool_result',
+    content: parts.join('\n\n'),
+  }
+}

package/tui/src/tools/TaskStopTool/TaskStopTool.ts

@@ -4,6 +4,7 @@ import { buildTool, type ToolDef } from '../../Tool.js'
 import { stopTask } from '../../tasks/stopTask.js'
 import { lazySchema } from '../../utils/lazySchema.js'
 import { jsonStringify } from '../../utils/slowOperations.js'
+import { buildAgentSupportMetadata } from '../AgentTool/orchestrationSupport.js'
 import { DESCRIPTION, TASK_STOP_TOOL_NAME } from './prompt.js'
 import { renderToolResultMessage, renderToolUseMessage } from './UI.js'
 
@@ -30,6 +31,10 @@ const outputSchema = lazySchema(() =>
       .string()
       .optional()
       .describe('The command or description of the stopped task'),
+    evidenceJoinKey: z.string(),
+    parentToolUseId: z.string(),
+    resumeToken: z.string(),
+    permissionFlow: z.literal('coordinator_parent_round_trip'),
   }),
 )
 type OutputSchema = ReturnType<typeof outputSchema>
@@ -118,6 +123,10 @@ export const TaskStopTool = buildTool({
       getAppState,
       setAppState,
     })
+    const supportMetadata = buildAgentSupportMetadata({
+      taskId: result.taskId,
+      parentToolUseId: getAppState().tasks?.[id]?.toolUseId,
+    })
 
     return {
       data: {
@@ -125,6 +134,7 @@ export const TaskStopTool = buildTool({
         task_id: result.taskId,
         task_type: result.taskType,
         command: result.command,
+        ...supportMetadata,
       },
     }
   },