ummaya 0.2.3 → 0.2.5

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

@@ -0,0 +1,114 @@
+import { feature } from 'bun:bundle'
+import {
+  type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+  logEvent,
+} from '../../services/analytics/index.js'
+import type { AppState } from '../../state/AppState.js'
+import type { ToolPermissionContext, Tools } from '../../Tool.js'
+import type {
+  AssistantMessage,
+  Message as MessageType,
+} from '../../types/message.js'
+import { logForDebugging } from '../../utils/debug.js'
+import { isInProtectedNamespace } from '../../utils/envUtils.js'
+import {
+  buildTranscriptForClassifier,
+  classifyYoloAction,
+} from '../../utils/permissions/yoloClassifier.js'
+import { LEGACY_AGENT_TOOL_NAME } from './constants.js'
+
+function getLastHandoffAssistantMessage(
+  messages: readonly MessageType[],
+): AssistantMessage | undefined {
+  return messages.findLast(
+    (message): message is AssistantMessage => message.type === 'assistant',
+  )
+}
+
+export async function classifyHandoffIfNeeded({
+  agentMessages,
+  tools,
+  toolPermissionContext,
+  abortSignal,
+  subagentType,
+  totalToolUseCount,
+}: {
+  agentMessages: MessageType[]
+  tools: Tools
+  toolPermissionContext: AppState['toolPermissionContext']
+  abortSignal: AbortSignal
+  subagentType: string
+  totalToolUseCount: number
+}): Promise<string | null> {
+  if (feature('TRANSCRIPT_CLASSIFIER')) {
+    if (toolPermissionContext.mode !== 'auto') return null
+
+    const agentTranscript = buildTranscriptForClassifier(agentMessages, tools)
+    if (!agentTranscript) return null
+
+    const classifierResult = await classifyYoloAction(
+      agentMessages,
+      {
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: "Sub-agent has finished and is handing back control to the main agent. Review the sub-agent's work based on the block rules and let the main agent know if any file is dangerous (the main agent will see the reason).",
+          },
+        ],
+      },
+      tools,
+      toolPermissionContext as ToolPermissionContext,
+      abortSignal,
+    )
+
+    const handoffDecision = classifierResult.unavailable
+      ? 'unavailable'
+      : classifierResult.shouldBlock
+        ? 'blocked'
+        : 'allowed'
+    logEvent('tengu_auto_mode_decision', {
+      decision:
+        handoffDecision as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      toolName:
+        LEGACY_AGENT_TOOL_NAME as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      inProtectedNamespace: isInProtectedNamespace(),
+      classifierModel:
+        classifierResult.model as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      agentType:
+        subagentType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      toolUseCount: totalToolUseCount,
+      isHandoff: true,
+      agentMsgId: getLastHandoffAssistantMessage(agentMessages)?.message
+        .id as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      classifierStage:
+        classifierResult.stage as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      classifierStage1RequestId:
+        classifierResult.stage1RequestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      classifierStage1MsgId:
+        classifierResult.stage1MsgId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      classifierStage2RequestId:
+        classifierResult.stage2RequestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      classifierStage2MsgId:
+        classifierResult.stage2MsgId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+    })
+
+    if (classifierResult.shouldBlock) {
+      if (classifierResult.unavailable) {
+        logForDebugging(
+          'Handoff classifier unavailable, allowing sub-agent output with warning',
+          { level: 'warn' },
+        )
+        return `Note: The safety classifier was unavailable when reviewing this sub-agent's work. Please carefully verify the sub-agent's actions and output before acting on them.`
+      }
+
+      logForDebugging(
+        `Handoff classifier flagged sub-agent output: ${classifierResult.reason}`,
+        { level: 'warn' },
+      )
+      return `SECURITY WARNING: This sub-agent performed actions that may violate security policy. Reason: ${classifierResult.reason}. Review the sub-agent's actions carefully before acting on its output.`
+    }
+  }
+
+  return null
+}

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

@@ -0,0 +1,16 @@
+import type { Message as MessageType } from '../../types/message.js'
+import { extractTextContent } from '../../utils/messages.js'
+
+export function extractPartialResult(
+  messages: MessageType[],
+): string | undefined {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const message = messages.at(i)
+    if (message?.type !== 'assistant') continue
+    const text = extractTextContent(message.message.content, '\n')
+    if (text) {
+      return text
+    }
+  }
+  return undefined
+}

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

@@ -0,0 +1,32 @@
+import type { Message as MessageType } from '../../types/message.js'
+import {
+  getProgressUpdate,
+  type ProgressTracker,
+} from '../../tasks/LocalAgentTask/LocalAgentTask.js'
+import { emitTaskProgress as emitTaskProgressEvent } from '../../utils/task/sdkProgress.js'
+
+export function getLastToolUseName(message: MessageType): string | undefined {
+  if (message.type !== 'assistant') return undefined
+  const block = message.message.content.findLast(b => b.type === 'tool_use')
+  return block?.type === 'tool_use' ? block.name : undefined
+}
+
+export function emitTaskProgress(
+  tracker: ProgressTracker,
+  taskId: string,
+  toolUseId: string | undefined,
+  description: string,
+  startTime: number,
+  lastToolName: string,
+): void {
+  const progress = getProgressUpdate(tracker)
+  emitTaskProgressEvent({
+    taskId,
+    toolUseId,
+    description: progress.lastActivity?.activityDescription ?? description,
+    startTime,
+    totalTokens: progress.tokenCount,
+    toolUses: progress.toolUseCount,
+    lastToolName,
+  })
+}

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

@@ -0,0 +1,161 @@
+import {
+  ALL_AGENT_DISALLOWED_TOOLS,
+  ASYNC_AGENT_ALLOWED_TOOLS,
+  CUSTOM_AGENT_DISALLOWED_TOOLS,
+  IN_PROCESS_TEAMMATE_ALLOWED_TOOLS,
+} from '../../constants/tools.js'
+import type { Tool, Tools } from '../../Tool.js'
+import { toolMatchesName } from '../../Tool.js'
+import { isAgentSwarmsEnabled } from '../../utils/agentSwarmsEnabled.js'
+import type { PermissionMode } from '../../utils/permissions/PermissionMode.js'
+import { permissionRuleValueFromString } from '../../utils/permissions/permissionRuleParser.js'
+import { isInProcessTeammate } from '../../utils/teammateContext.js'
+import { EXIT_PLAN_MODE_V2_TOOL_NAME } from '../ExitPlanModeTool/constants.js'
+import { AGENT_TOOL_NAME } from './constants.js'
+import type { AgentDefinition } from './loadAgentsDir.js'
+import { isProtectedWorkerTool } from './orchestrationSupport.js'
+
+export type ResolvedAgentTools = {
+  hasWildcard: boolean
+  validTools: string[]
+  invalidTools: string[]
+  resolvedTools: Tools
+  allowedAgentTypes?: string[]
+}
+
+export function filterToolsForAgent({
+  tools,
+  isBuiltIn,
+  isAsync = false,
+  permissionMode,
+}: {
+  tools: Tools
+  isBuiltIn: boolean
+  isAsync?: boolean
+  permissionMode?: PermissionMode
+}): Tools {
+  return tools.filter(tool => {
+    if (tool.name.startsWith('mcp__')) {
+      return true
+    }
+    if (isProtectedWorkerTool(tool.name)) {
+      return false
+    }
+    if (
+      toolMatchesName(tool, EXIT_PLAN_MODE_V2_TOOL_NAME) &&
+      permissionMode === 'plan'
+    ) {
+      return true
+    }
+    if (ALL_AGENT_DISALLOWED_TOOLS.has(tool.name)) {
+      return false
+    }
+    if (!isBuiltIn && CUSTOM_AGENT_DISALLOWED_TOOLS.has(tool.name)) {
+      return false
+    }
+    if (isAsync && !ASYNC_AGENT_ALLOWED_TOOLS.has(tool.name)) {
+      if (isAgentSwarmsEnabled() && isInProcessTeammate()) {
+        if (toolMatchesName(tool, AGENT_TOOL_NAME)) {
+          return true
+        }
+        if (IN_PROCESS_TEAMMATE_ALLOWED_TOOLS.has(tool.name)) {
+          return true
+        }
+      }
+      return false
+    }
+    return true
+  })
+}
+
+export function resolveAgentTools(
+  agentDefinition: Pick<
+    AgentDefinition,
+    'tools' | 'disallowedTools' | 'source' | 'permissionMode'
+  >,
+  availableTools: Tools,
+  isAsync = false,
+  isMainThread = false,
+): ResolvedAgentTools {
+  const {
+    tools: agentTools,
+    disallowedTools,
+    source,
+    permissionMode,
+  } = agentDefinition
+  const filteredAvailableTools = isMainThread
+    ? availableTools
+    : filterToolsForAgent({
+        tools: availableTools,
+        isBuiltIn: source === 'built-in',
+        isAsync,
+        permissionMode,
+      })
+
+  const disallowedToolSet = new Set(
+    disallowedTools?.map(toolSpec => {
+      const { toolName } = permissionRuleValueFromString(toolSpec)
+      return toolName
+    }) ?? [],
+  )
+
+  const allowedAvailableTools = filteredAvailableTools.filter(
+    tool => !disallowedToolSet.has(tool.name),
+  )
+
+  const hasWildcard =
+    agentTools === undefined ||
+    (agentTools.length === 1 && agentTools[0] === '*')
+  if (hasWildcard) {
+    return {
+      hasWildcard: true,
+      validTools: [],
+      invalidTools: [],
+      resolvedTools: allowedAvailableTools,
+    }
+  }
+
+  const availableToolMap = new Map<string, Tool>()
+  for (const tool of allowedAvailableTools) {
+    availableToolMap.set(tool.name, tool)
+  }
+
+  const validTools: string[] = []
+  const invalidTools: string[] = []
+  const resolved: Tool[] = []
+  const resolvedToolsSet = new Set<Tool>()
+  let allowedAgentTypes: string[] | undefined
+
+  for (const toolSpec of agentTools) {
+    const { toolName, ruleContent } = permissionRuleValueFromString(toolSpec)
+
+    if (toolName === AGENT_TOOL_NAME) {
+      if (ruleContent) {
+        allowedAgentTypes = ruleContent.split(',').map(s => s.trim())
+      }
+      if (!isMainThread) {
+        validTools.push(toolSpec)
+        continue
+      }
+    }
+
+    const tool = availableToolMap.get(toolName)
+    if (tool) {
+      validTools.push(toolSpec)
+      if (!resolvedToolsSet.has(tool)) {
+        resolved.push(tool)
+        resolvedToolsSet.add(tool)
+      }
+    } else {
+      invalidTools.push(toolSpec)
+    }
+  }
+
+  return {
+    hasWildcard: false,
+    validTools,
+    invalidTools,
+    resolvedTools: resolved,
+    allowedAgentTypes,
+  }
+}

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

@@ -0,0 +1,163 @@
+import type { BetaUsage as Usage } from '@anthropic-ai/sdk/resources/beta/messages/messages.mjs'
+import { z } from 'zod/v4'
+import {
+  type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+  logEvent,
+} from '../../services/analytics/index.js'
+import type {
+  AssistantMessage,
+  Message as MessageType,
+} from '../../types/message.js'
+import { lazySchema } from '../../utils/lazySchema.js'
+import {
+  sourceVerificationSchema,
+} from '../WebFetchTool/sourceVerification.js'
+import { extractSourceVerificationFromMessages } from './sourceVerificationPropagation.js'
+
+export const agentToolResultSchema = lazySchema(() =>
+  z.object({
+    agentId: z.string(),
+    agentType: z.string().optional(),
+    content: z.array(z.object({ type: z.literal('text'), text: z.string() })),
+    totalToolUseCount: z.number(),
+    totalDurationMs: z.number(),
+    totalTokens: z.number(),
+    usage: z.object({
+      input_tokens: z.number(),
+      output_tokens: z.number(),
+      cache_creation_input_tokens: z.number().nullable(),
+      cache_read_input_tokens: z.number().nullable(),
+      server_tool_use: z
+        .object({
+          web_search_requests: z.number(),
+          web_fetch_requests: z.number(),
+        })
+        .nullable(),
+      service_tier: z.enum(['standard', 'priority', 'batch']).nullable(),
+      cache_creation: z
+        .object({
+          ephemeral_1h_input_tokens: z.number(),
+          ephemeral_5m_input_tokens: z.number(),
+        })
+        .nullable(),
+    }),
+    sourceVerification: sourceVerificationSchema.optional(),
+  }),
+)
+
+export type AgentToolResult = z.input<ReturnType<typeof agentToolResultSchema>>
+
+export type FinalizeAgentToolMetadata = {
+  prompt: string
+  resolvedAgentModel: string
+  isBuiltInAgent: boolean
+  startTime: number
+  agentType: string
+  isAsync: boolean
+}
+
+function getLastAgentAssistantMessage(
+  messages: readonly MessageType[],
+): AssistantMessage | undefined {
+  return messages.findLast(
+    (message): message is AssistantMessage => message.type === 'assistant',
+  )
+}
+
+function getAgentTokenCountFromUsage(usage: Usage): number {
+  return (
+    usage.input_tokens +
+    (usage.cache_creation_input_tokens ?? 0) +
+    (usage.cache_read_input_tokens ?? 0) +
+    usage.output_tokens
+  )
+}
+
+export function countToolUses(messages: MessageType[]): number {
+  let count = 0
+  for (const m of messages) {
+    if (m.type === 'assistant') {
+      for (const block of m.message.content) {
+        if (block.type === 'tool_use') {
+          count++
+        }
+      }
+    }
+  }
+  return count
+}
+
+export function finalizeAgentTool(
+  agentMessages: MessageType[],
+  agentId: string,
+  metadata: FinalizeAgentToolMetadata,
+): AgentToolResult {
+  const {
+    prompt,
+    resolvedAgentModel,
+    isBuiltInAgent,
+    startTime,
+    agentType,
+    isAsync,
+  } = metadata
+
+  const lastAssistantMessage = getLastAgentAssistantMessage(agentMessages)
+  if (lastAssistantMessage === undefined) {
+    throw new Error('No assistant messages found')
+  }
+  let content = lastAssistantMessage.message.content.filter(
+    _ => _.type === 'text',
+  )
+  if (content.length === 0) {
+    for (let i = agentMessages.length - 1; i >= 0; i--) {
+      const message = agentMessages.at(i)
+      if (message?.type !== 'assistant') continue
+      const textBlocks = message.message.content.filter(_ => _.type === 'text')
+      if (textBlocks.length > 0) {
+        content = textBlocks
+        break
+      }
+    }
+  }
+
+  const totalTokens = getAgentTokenCountFromUsage(
+    lastAssistantMessage.message.usage,
+  )
+  const totalToolUseCount = countToolUses(agentMessages)
+
+  logEvent('tengu_agent_tool_completed', {
+    agent_type:
+      agentType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+    model:
+      resolvedAgentModel as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+    prompt_char_count: prompt.length,
+    response_char_count: content.length,
+    assistant_message_count: agentMessages.length,
+    total_tool_uses: totalToolUseCount,
+    duration_ms: Date.now() - startTime,
+    total_tokens: totalTokens,
+    is_built_in_agent: isBuiltInAgent,
+    is_async: isAsync,
+  })
+
+  const lastRequestId = lastAssistantMessage.requestId
+  if (lastRequestId) {
+    logEvent('tengu_cache_eviction_hint', {
+      scope:
+        'subagent_end' as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      last_request_id:
+        lastRequestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+    })
+  }
+
+  return {
+    agentId,
+    agentType,
+    content,
+    totalDurationMs: Date.now() - startTime,
+    totalTokens,
+    totalToolUseCount,
+    usage: lastAssistantMessage.message.usage,
+    sourceVerification: extractSourceVerificationFromMessages(agentMessages),
+  }
+}