ummaya 0.2.4 → 0.2.6

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

@@ -0,0 +1,196 @@
+import { feature } from 'bun:bundle'
+import { z } from 'zod/v4'
+import type { AgentToolProgress, ShellProgress } from '../../types/tools.js'
+import type { PermissionMode } from '../../types/permissions.js'
+import { lazySchema } from '../../utils/lazySchema.js'
+import { permissionModeSchema } from '../../utils/permissions/PermissionMode.js'
+import { isForkSubagentEnabled } from './forkSubagentGate.js'
+import type { AgentSupportMetadata } from './orchestrationSupport.js'
+import { agentToolResultSchema, type AgentToolResult } from './agentToolResult.js'
+import { isAntBuild, isBackgroundTasksDisabled } from './runtimeConfig.js'
+
+const baseInputSchema = lazySchema(() =>
+  z.object({
+    description: z.string().describe('A short (3-5 word) description of the task'),
+    prompt: z.string().describe('The task for the agent to perform'),
+    subagent_type: z
+      .string()
+      .optional()
+      .describe('The type of specialized agent to use for this task'),
+    model: z
+      .enum(['sonnet', 'opus', 'haiku'])
+      .optional()
+      .describe(
+        "Optional model override for this agent. Takes precedence over the agent definition's model frontmatter. If omitted, uses the agent definition's model, or inherits from the parent.",
+      ),
+    run_in_background: z
+      .boolean()
+      .optional()
+      .describe(
+        'Set to true to run this agent in the background. You will be notified when it completes.',
+      ),
+  }),
+)
+
+const fullInputSchema = lazySchema(() => {
+  const multiAgentInputSchema = z.object({
+    name: z
+      .string()
+      .optional()
+      .describe(
+        'Name for the spawned agent. Makes it addressable via SendMessage({to: name}) while running.',
+      ),
+    team_name: z
+      .string()
+      .optional()
+      .describe('Team name for spawning. Uses current team context if omitted.'),
+    mode: permissionModeSchema()
+      .optional()
+      .describe('Permission mode for spawned teammate (e.g., "plan" to require plan approval).'),
+  })
+  return baseInputSchema()
+    .merge(multiAgentInputSchema)
+    .extend({
+      isolation: (isAntBuild()
+        ? z.enum(['worktree', 'remote'])
+        : z.enum(['worktree'])
+      )
+        .optional()
+        .describe(
+          isAntBuild()
+            ? 'Isolation mode. "worktree" creates a temporary git worktree so the agent works on an isolated copy of the repo. "remote" launches the agent in a remote CCR environment (always runs in background).'
+            : 'Isolation mode. "worktree" creates a temporary git worktree so the agent works on an isolated copy of the repo.',
+        ),
+      cwd: z
+        .string()
+        .optional()
+        .describe(
+          'Absolute path to run the agent in. Overrides the working directory for all filesystem and shell operations within this agent. Mutually exclusive with isolation: "worktree".',
+        ),
+    })
+})
+
+export const inputSchema = () => {
+  const schema = feature('KAIROS')
+    ? fullInputSchema()
+    : fullInputSchema().omit({ cwd: true })
+
+  // Keep the final feature-gated view uncached. Fork availability depends on
+  // session state, and memoizing this wrapper lets the first schema access
+  // leak field visibility into later accesses in the same process.
+  return isBackgroundTasksDisabled || isForkSubagentEnabled()
+    ? schema.omit({ run_in_background: true })
+    : schema
+}
+
+export type InputSchema = ReturnType<typeof inputSchema>
+
+export type AgentToolInput = z.infer<ReturnType<typeof baseInputSchema>> & {
+  readonly name?: string
+  readonly team_name?: string
+  readonly mode?: PermissionMode
+  readonly isolation?: 'worktree' | 'remote'
+  readonly cwd?: string
+}
+
+export const outputSchema = lazySchema(() => {
+  const syncOutputSchema = agentToolResultSchema().extend({
+    status: z.literal('completed'),
+    prompt: z.string(),
+    evidenceJoinKey: z.string(),
+    parentToolUseId: z.string(),
+    resumeToken: z.string(),
+    permissionFlow: z.literal('coordinator_parent_round_trip'),
+  })
+  const asyncOutputSchema = z.object({
+    status: z.literal('async_launched'),
+    agentId: z.string().describe('The ID of the async agent'),
+    description: z.string().describe('The description of the task'),
+    prompt: z.string().describe('The prompt for the agent'),
+    outputFile: z.string().describe('Path to the output file for checking agent progress'),
+    canReadOutputFile: z
+      .boolean()
+      .optional()
+      .describe('Whether the calling agent has Read/Bash tools to check progress'),
+    evidenceJoinKey: z.string(),
+    parentToolUseId: z.string(),
+    resumeToken: z.string(),
+    permissionFlow: z.literal('coordinator_parent_round_trip'),
+  })
+  return z.union([syncOutputSchema, asyncOutputSchema])
+})
+
+export type OutputSchema = ReturnType<typeof outputSchema>
+export type Progress = AgentToolProgress | ShellProgress
+
+export type AgentToolProgressCallback = (event: {
+  readonly toolUseID: string
+  readonly data: Progress
+}) => void
+
+export type WorktreeResult = {
+  readonly worktreePath?: string
+  readonly worktreeBranch?: string
+}
+
+export type CompletedAgentOutput = AgentToolResult &
+  AgentSupportMetadata &
+  WorktreeResult & {
+    readonly status: 'completed'
+    readonly prompt: string
+  }
+
+export type AsyncLaunchedOutput = AgentSupportMetadata & {
+  readonly isAsync?: true
+  readonly status: 'async_launched'
+  readonly agentId: string
+  readonly description: string
+  readonly prompt: string
+  readonly outputFile: string
+  readonly canReadOutputFile?: boolean
+}
+
+export type TeammateSpawnedOutput = AgentSupportMetadata & {
+  readonly status: 'teammate_spawned'
+  readonly prompt: string
+  readonly teammate_id: string
+  readonly agent_id: string
+  readonly agent_type?: string
+  readonly model?: string
+  readonly name: string
+  readonly color?: string
+  readonly tmux_session_name: string
+  readonly tmux_window_name: string
+  readonly tmux_pane_id: string
+  readonly team_name?: string
+  readonly is_splitpane?: boolean
+  readonly plan_mode_required?: boolean
+}
+
+export type RemoteLaunchedOutput = {
+  readonly status: 'remote_launched'
+  readonly taskId: string
+  readonly sessionUrl: string
+  readonly description: string
+  readonly prompt: string
+  readonly outputFile: string
+}
+
+export type AgentToolOutput =
+  | CompletedAgentOutput
+  | AsyncLaunchedOutput
+  | TeammateSpawnedOutput
+  | RemoteLaunchedOutput
+
+export type AgentToolCallResult = {
+  readonly data: AgentToolOutput
+}
+
+export type AgentLifecycleMetadata = {
+  readonly prompt: string
+  readonly resolvedAgentModel: string
+  readonly isBuiltInAgent: boolean
+  readonly startTime: number
+  readonly agentType: string
+  readonly isAsync: boolean
+}

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

@@ -0,0 +1,263 @@
+import type { Message as MessageType } from '../../types/message.js'
+import {
+  buildSourceVerification,
+  sourceVerificationSchema,
+  type SourceVerification,
+  type SourceVerificationEvidence,
+} from '../WebFetchTool/sourceVerification.js'
+import { WEB_FETCH_TOOL_NAME } from '../WebFetchTool/prompt.js'
+import { WEB_SEARCH_TOOL_NAME } from '../WebSearchTool/prompt.js'
+
+const MAX_PROPAGATED_SOURCE_EVIDENCE = 20
+
+type TrustedSourceToolName =
+  | typeof WEB_FETCH_TOOL_NAME
+  | typeof WEB_SEARCH_TOOL_NAME
+
+type TrustedToolResultText = {
+  readonly text: string
+  readonly toolName: TrustedSourceToolName
+}
+
+type SourceEvidenceDraft = {
+  toolId?: string
+  sourceUrl?: string | null
+  title?: string | null
+  observedAt?: string
+  citationHandle?: string
+  blockedOrUsed?: SourceVerificationEvidence['blockedOrUsed']
+  trust?: SourceVerificationEvidence['trust']
+  promptInjection?: SourceVerificationEvidence['promptInjection']
+  redacted?: boolean
+}
+
+function parseNullableField(value: string): string | null {
+  return value === 'none' ? null : value
+}
+
+function parseBlockedOrUsed(
+  value: string,
+): SourceVerificationEvidence['blockedOrUsed'] | undefined {
+  switch (value) {
+    case 'blocked':
+    case 'needs_input':
+      return value
+    default:
+      return undefined
+  }
+}
+
+function parsePromptInjection(
+  value: string,
+): SourceVerificationEvidence['promptInjection'] | undefined {
+  switch (value) {
+    case 'detected':
+    case 'not_detected':
+      return value
+    default:
+      return undefined
+  }
+}
+
+function parseRedacted(value: string): boolean | undefined {
+  switch (value) {
+    case 'true':
+      return true
+    case 'false':
+      return false
+    default:
+      return undefined
+  }
+}
+
+function appendSourceEvidenceDraft(
+  draft: SourceEvidenceDraft,
+  evidence: SourceVerificationEvidence[],
+  expectedToolId: TrustedSourceToolName,
+): void {
+  if (
+    draft.toolId === undefined ||
+    draft.toolId !== expectedToolId ||
+    draft.sourceUrl === undefined ||
+    draft.title === undefined ||
+    draft.observedAt === undefined ||
+    draft.citationHandle === undefined ||
+    draft.blockedOrUsed === undefined ||
+    draft.trust === undefined ||
+    draft.promptInjection === undefined ||
+    draft.redacted === undefined
+  ) {
+    return
+  }
+  evidence.push({
+    toolId: draft.toolId,
+    sourceUrl: draft.sourceUrl,
+    title: draft.title,
+    observedAt: draft.observedAt,
+    citationHandle: draft.citationHandle,
+    blockedOrUsed: draft.blockedOrUsed,
+    trust: draft.trust,
+    promptInjection: draft.promptInjection,
+    redacted: draft.redacted,
+  })
+}
+
+function parseSourceVerificationSegment(
+  segment: string,
+  evidence: SourceVerificationEvidence[],
+  expectedToolId: TrustedSourceToolName,
+): void {
+  let draft: SourceEvidenceDraft = {}
+  for (const line of segment.split('\n')) {
+    const separatorIndex = line.indexOf(':')
+    if (separatorIndex < 0) continue
+    const key = line.slice(0, separatorIndex).trim()
+    const value = line.slice(separatorIndex + 1).trim()
+    if (key === 'tool_id') {
+      appendSourceEvidenceDraft(draft, evidence, expectedToolId)
+      if (evidence.length >= MAX_PROPAGATED_SOURCE_EVIDENCE) return
+      draft = { toolId: value }
+      continue
+    }
+    switch (key) {
+      case 'source_url':
+        draft.sourceUrl = parseNullableField(value)
+        break
+      case 'title':
+        draft.title = parseNullableField(value)
+        break
+      case 'timestamp':
+        draft.observedAt = value
+        break
+      case 'citation_handle':
+        draft.citationHandle = value
+        break
+      case 'blocked_or_used':
+        draft.blockedOrUsed = parseBlockedOrUsed(value)
+        break
+      case 'trust':
+        draft.trust = value === 'untrusted_source' ? value : undefined
+        break
+      case 'prompt_injection':
+        draft.promptInjection = parsePromptInjection(value)
+        break
+      case 'redacted':
+        draft.redacted = parseRedacted(value)
+        break
+      default:
+        break
+    }
+  }
+  appendSourceEvidenceDraft(draft, evidence, expectedToolId)
+}
+
+function collectSourceVerificationSegments(text: string): string[] {
+  const segments: string[] = []
+  const openTag = '<source_verification>'
+  const closeTag = '</source_verification>'
+  let searchStart = 0
+  while (searchStart < text.length) {
+    const openIndex = text.indexOf(openTag, searchStart)
+    if (openIndex < 0) return segments
+    const contentStart = openIndex + openTag.length
+    const closeIndex = text.indexOf(closeTag, contentStart)
+    if (closeIndex < 0) return segments
+    segments.push(text.slice(contentStart, closeIndex))
+    searchStart = closeIndex + closeTag.length
+  }
+  return segments
+}
+
+function trustedSourceToolName(value: string | undefined): TrustedSourceToolName | undefined {
+  switch (value) {
+    case WEB_FETCH_TOOL_NAME:
+      return WEB_FETCH_TOOL_NAME
+    case WEB_SEARCH_TOOL_NAME:
+      return WEB_SEARCH_TOOL_NAME
+    default:
+      return undefined
+  }
+}
+
+function collectToolUseNameById(
+  agentMessages: readonly MessageType[],
+): ReadonlyMap<string, string> {
+  const toolUseNameById = new Map<string, string>()
+  for (const message of agentMessages) {
+    if (message.type !== 'assistant') continue
+    const { content } = message.message
+    if (!Array.isArray(content)) continue
+    for (const block of content) {
+      if (
+        block.type === 'tool_use' &&
+        typeof block.id === 'string' &&
+        typeof block.name === 'string'
+      ) {
+        toolUseNameById.set(block.id, block.name)
+      }
+    }
+  }
+  return toolUseNameById
+}
+
+function collectToolResultText(
+  message: MessageType,
+  toolUseNameById: ReadonlyMap<string, string>,
+): TrustedToolResultText[] {
+  if (message.type !== 'user') return []
+  const { content } = message.message
+  if (!Array.isArray(content)) return []
+  const texts: TrustedToolResultText[] = []
+  for (const block of content) {
+    if (block.type !== 'tool_result' || typeof block.tool_use_id !== 'string') {
+      continue
+    }
+    const toolName = trustedSourceToolName(toolUseNameById.get(block.tool_use_id))
+    if (toolName === undefined) continue
+    if (typeof block.content === 'string') {
+      texts.push({ text: block.content, toolName })
+      continue
+    }
+    if (!Array.isArray(block.content)) continue
+    for (const contentBlock of block.content) {
+      if (
+        typeof contentBlock === 'object' &&
+        contentBlock !== null &&
+        'text' in contentBlock &&
+        typeof contentBlock.text === 'string'
+      ) {
+        texts.push({ text: contentBlock.text, toolName })
+      }
+    }
+  }
+  return texts
+}
+
+function validatedSourceVerification(
+  evidence: readonly SourceVerificationEvidence[],
+): SourceVerification | undefined {
+  const verification = buildSourceVerification(evidence)
+  return sourceVerificationSchema.safeParse(verification).success
+    ? verification
+    : undefined
+}
+
+export function extractSourceVerificationFromMessages(
+  agentMessages: readonly MessageType[],
+): SourceVerification | undefined {
+  const evidence: SourceVerificationEvidence[] = []
+  const toolUseNameById = collectToolUseNameById(agentMessages)
+  for (const message of agentMessages) {
+    for (const { text, toolName } of collectToolResultText(message, toolUseNameById)) {
+      for (const segment of collectSourceVerificationSegments(text)) {
+        parseSourceVerificationSegment(segment, evidence, toolName)
+        if (evidence.length >= MAX_PROPAGATED_SOURCE_EVIDENCE) {
+          return validatedSourceVerification(
+            evidence.slice(0, MAX_PROPAGATED_SOURCE_EVIDENCE),
+          )
+        }
+      }
+    }
+  }
+  return evidence.length > 0 ? validatedSourceVerification(evidence) : undefined
+}

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

@@ -0,0 +1,105 @@
+import type { AgentId } from '../../types/ids.js'
+import { asAgentId } from '../../types/ids.js'
+import { getCwd } from '../../utils/cwd.js'
+import { logForDebugging } from '../../utils/debug.js'
+import { writeAgentMetadata } from '../../utils/sessionStorage.js'
+import {
+  createAgentWorktree,
+  hasWorktreeChanges,
+  removeAgentWorktree,
+} from '../../utils/worktree.js'
+import { buildWorktreeNotice } from './forkSubagent.js'
+import { createUserMessage } from '../../utils/messages.js'
+import type { Message as MessageType } from 'src/types/message.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import type { WorktreeResult } from './schemas.js'
+
+export type AgentWorktreeInfo = {
+  readonly worktreePath: string
+  readonly worktreeBranch?: string
+  readonly headCommit?: string
+  readonly gitRoot?: string
+  readonly hookBased?: boolean
+}
+
+/** Mutable cleanup guard; cleanup must become idempotent after the first call. */
+export type MutableWorktreeState = {
+  current: AgentWorktreeInfo | null
+}
+
+export async function createWorktreeState({
+  isolation,
+  earlyAgentId,
+}: {
+  readonly isolation?: 'worktree' | 'remote'
+  readonly earlyAgentId: string
+}): Promise<MutableWorktreeState> {
+  if (isolation !== 'worktree') return { current: null }
+  return {
+    current: await createAgentWorktree(`agent-${earlyAgentId.slice(0, 8)}`),
+  }
+}
+
+export function appendForkWorktreeNotice({
+  isForkPath,
+  worktreePath,
+  promptMessages,
+}: {
+  readonly isForkPath: boolean
+  readonly worktreePath?: string
+  readonly promptMessages: MessageType[]
+}): void {
+  if (!isForkPath || !worktreePath) return
+  promptMessages.push(
+    createUserMessage({ content: buildWorktreeNotice(getCwd(), worktreePath) }),
+  )
+}
+
+export function buildWorktreeCleanup({
+  state,
+  earlyAgentId,
+  selectedAgent,
+  description,
+}: {
+  readonly state: MutableWorktreeState
+  readonly earlyAgentId: string
+  readonly selectedAgent: AgentDefinition
+  readonly description: string
+}): () => Promise<WorktreeResult> {
+  return async () => {
+    const worktreeInfo = state.current
+    if (!worktreeInfo) return {}
+    state.current = null
+    const { worktreePath, worktreeBranch, headCommit, gitRoot, hookBased } =
+      worktreeInfo
+    if (hookBased) {
+      logForDebugging(`Hook-based agent worktree kept at: ${worktreePath}`)
+      return { worktreePath }
+    }
+    if (headCommit) {
+      const changed = await hasWorktreeChanges(worktreePath, headCommit)
+      if (!changed) {
+        await removeAgentWorktree(worktreePath, worktreeBranch, gitRoot)
+        clearWorktreeMetadata(asAgentId(earlyAgentId), selectedAgent, description)
+        return {}
+      }
+    }
+    logForDebugging(`Agent worktree has changes, keeping: ${worktreePath}`)
+    return { worktreePath, worktreeBranch }
+  }
+}
+
+function clearWorktreeMetadata(
+  agentId: AgentId,
+  selectedAgent: AgentDefinition,
+  description: string,
+): void {
+  void writeAgentMetadata(agentId, {
+    agentType: selectedAgent.agentType,
+    description,
+  }).then(
+    () => undefined,
+    error =>
+      logForDebugging(`Failed to clear worktree metadata: ${String(error)}`),
+  )
+}