ummaya 0.2.4 → 0.2.6

package/tui/src/services/api/ummaya/toolSelection.ts

@@ -0,0 +1,217 @@
+import { getAllBaseTools } from '../../../tools.js'
+import type { Tool, Tools } from '../../../Tool.js'
+import {
+  getAdapterToolByName,
+  selectTopKAdapterToolNamesForQuery,
+} from '../../../tools/AdapterTool/AdapterTool.js'
+import {
+  appendRouteDiagnostic,
+  hashRouteDiagnosticText,
+} from '../../../tools/AdapterTool/routeDiagnostics.js'
+import { selectRecoveredSupportToolNamesForQuery } from '../../../tools/ToolSearchTool/supportIntentHints.js'
+import {
+  getAdapterManifestHash,
+  listAdapters,
+  type AdapterManifestEntry,
+} from '../adapterManifest.js'
+import type { ProviderTurnEvidenceContext } from './evidence.js'
+import {
+  DOCUMENT_TOOL_NAME,
+  isExactLocalReadOnlyDocumentPrompt,
+} from '../../../tools/_shared/toolChoiceRepair/documentCompletionPatterns.js'
+
+const MAIN_THREAD_QUERY_SOURCE = 'repl_main_thread'
+const ADAPTER_CANDIDATE_LIMIT = 5
+const WORKSPACE_SUPPORT_TOOL_NAMES = new Set([
+  'workspace_grep',
+  'workspace_read',
+])
+const FORCEABLE_WORKSPACE_SUPPORT_TOOL_NAMES = new Set([
+  'workspace_bash',
+])
+
+function pushUniqueTool(target: Tool[], tool: Tool | undefined): void {
+  if (!tool) return
+  if (!target.some(candidate => candidate.name === tool.name)) {
+    target.push(tool)
+  }
+}
+
+function baseToolByName(name: string): Tool | undefined {
+  return getAllBaseTools().find(tool => tool.name === name)
+}
+
+function hasToolNamed(tools: readonly Tool[], name: string): boolean {
+  return tools.some(tool => tool.name === name)
+}
+
+function toolByName(tools: readonly Tool[], name: string): Tool | undefined {
+  return tools.find(tool => tool.name === name)
+}
+
+function pushUniqueAdapterEntry(
+  target: AdapterManifestEntry[],
+  entry: AdapterManifestEntry | undefined,
+): void {
+  if (!entry) return
+  if (!target.some(candidate => candidate.tool_id === entry.tool_id)) {
+    target.push(entry)
+  }
+}
+
+function hasWorkspaceSupportToolName(names: readonly string[]): boolean {
+  return names.some(name => WORKSPACE_SUPPORT_TOOL_NAMES.has(name))
+}
+
+function isMainThreadQuerySource(querySource: string): boolean {
+  return (
+    querySource === MAIN_THREAD_QUERY_SOURCE ||
+    querySource.startsWith(`${MAIN_THREAD_QUERY_SOURCE}:`)
+  )
+}
+
+function filterToAdapterSurface(params: {
+  readonly selected: readonly Tool[]
+  readonly selectedAdapterEntries: readonly AdapterManifestEntry[]
+  readonly forcedToolName?: string
+}): Tool[] {
+  if (params.selectedAdapterEntries.length === 0) {
+    return params.selected.filter(tool => tool.name === params.forcedToolName)
+  }
+  const selectedAdapterToolIds = new Set(
+    params.selectedAdapterEntries.map(entry => entry.tool_id),
+  )
+  return params.selected.filter(tool => {
+    if (tool.name === params.forcedToolName) return true
+    return selectedAdapterToolIds.has(tool.name)
+  })
+}
+
+export function selectProviderToolChoiceName(params: {
+  readonly tools: Tools
+  readonly userText: string
+  readonly forcedToolName?: string
+}): string | undefined {
+  if (params.forcedToolName) return params.forcedToolName
+  if (
+    isExactLocalReadOnlyDocumentPrompt(params.userText) &&
+    hasToolNamed(params.tools, DOCUMENT_TOOL_NAME)
+  ) {
+    return DOCUMENT_TOOL_NAME
+  }
+  for (const name of selectRecoveredSupportToolNamesForQuery(params.userText)) {
+    if (
+      FORCEABLE_WORKSPACE_SUPPORT_TOOL_NAMES.has(name) &&
+      baseToolByName(name)
+    ) {
+      return name
+    }
+  }
+  return undefined
+}
+
+export function shouldWaitForAdapterManifestForProviderRequest(params: {
+  readonly querySource: string
+  readonly userText: string
+}): boolean {
+  if (selectRecoveredSupportToolNamesForQuery(params.userText).length > 0) {
+    return false
+  }
+  if (isExactLocalReadOnlyDocumentPrompt(params.userText)) {
+    return false
+  }
+  return (
+    params.userText.trim().length > 0 &&
+    isMainThreadQuerySource(params.querySource)
+  )
+}
+
+export function selectProviderTools(params: {
+  readonly tools: Tools
+  readonly userText: string
+  readonly forcedToolName?: string
+  readonly disabledToolNames?: readonly string[]
+  readonly querySource: string
+  readonly hasCurrentTurnLocationContext?: boolean
+  readonly evidenceContext?: ProviderTurnEvidenceContext
+}): readonly Tool[] {
+  const adapters = listAdapters()
+  const adapterEntriesByToolId = new Map(
+    adapters.map(entry => [entry.tool_id, entry]),
+  )
+  const adapterToolIds = new Set(adapterEntriesByToolId.keys())
+  const disabledToolNames = new Set(params.disabledToolNames ?? [])
+  let selected: Tool[] = []
+  const selectedAdapterEntries: AdapterManifestEntry[] = []
+  const recoveredSupportToolNames = isExactLocalReadOnlyDocumentPrompt(params.userText)
+    ? []
+    : selectRecoveredSupportToolNamesForQuery(params.userText)
+  const useWorkspaceOnlySurface = hasWorkspaceSupportToolName(recoveredSupportToolNames)
+
+  for (const name of recoveredSupportToolNames) {
+    pushUniqueTool(selected, baseToolByName(name))
+  }
+
+  if (!useWorkspaceOnlySurface) {
+    for (const name of selectTopKAdapterToolNamesForQuery(
+      params.userText,
+      ADAPTER_CANDIDATE_LIMIT,
+      { hasCurrentTurnLocationContext: params.hasCurrentTurnLocationContext },
+    )) {
+      if (disabledToolNames.has(name)) continue
+      pushUniqueTool(selected, getAdapterToolByName(name))
+      pushUniqueAdapterEntry(selectedAdapterEntries, adapterEntriesByToolId.get(name))
+    }
+  }
+
+  if (params.forcedToolName) {
+    pushUniqueTool(
+      selected,
+      toolByName(params.tools, params.forcedToolName) ??
+        getAdapterToolByName(params.forcedToolName) ??
+        baseToolByName(params.forcedToolName),
+    )
+    pushUniqueAdapterEntry(
+      selectedAdapterEntries,
+      adapterEntriesByToolId.get(params.forcedToolName),
+    )
+    selected = selected.filter(tool => tool.name === params.forcedToolName)
+  }
+
+  if (useWorkspaceOnlySurface) {
+    selected = selected.filter(tool => {
+      if (tool.name === params.forcedToolName) return true
+      return recoveredSupportToolNames.includes(tool.name)
+    })
+  } else if (
+    isMainThreadQuerySource(params.querySource) &&
+    selectedAdapterEntries.length > 0
+  ) {
+    selected = filterToAdapterSurface({
+      selected,
+      selectedAdapterEntries,
+      forcedToolName: params.forcedToolName,
+    })
+  }
+
+  const adapterNames = selected
+    .map(tool => tool.name)
+    .filter(name => adapterEntriesByToolId.has(name))
+  if (adapterNames.length > 0 || process.env.UMMAYA_TUI_ROUTE_DIAGNOSTIC_FILE) {
+    appendRouteDiagnostic('adapter_selection', {
+      manifest_hash: getAdapterManifestHash(),
+      query_source: params.querySource,
+      schema_projection_level: 'top_k_concrete_adapter_schemas',
+      selected_tools: adapterNames,
+      final_adapter_tools: adapterNames,
+      disabled_provider_tools: [...disabledToolNames],
+      has_current_turn_location_context:
+        params.hasCurrentTurnLocationContext === true,
+      query_hash: hashRouteDiagnosticText(params.userText),
+      session_id: params.evidenceContext?.session_id ?? null,
+      correlation_id: params.evidenceContext?.correlation_id ?? null,
+      frame_hash: params.evidenceContext?.frame_hash ?? null,
+    })
+  }
+  return selected
+}

package/tui/src/services/api/ummaya/types.ts

@@ -0,0 +1,110 @@
+import type { Tools } from '../../../Tool.js'
+import type { QuerySource } from '../../../constants/querySource.js'
+import type { Message } from '../../../types/message.js'
+import type { SystemPrompt } from '../../../utils/systemPromptType.js'
+
+export type Usage = {
+  readonly input_tokens: number
+  readonly output_tokens: number
+  readonly cache_creation_input_tokens: number
+  readonly cache_read_input_tokens: number
+  readonly server_tool_use?: {
+    readonly web_search_requests?: number
+    readonly web_fetch_requests?: number
+  }
+  readonly service_tier?: string | null
+  readonly cache_creation?: {
+    readonly ephemeral_1h_input_tokens?: number
+    readonly ephemeral_5m_input_tokens?: number
+  }
+  readonly inference_geo?: string | null
+  readonly iterations?: number | null
+  readonly speed?: number | null
+}
+
+export const EMPTY_USAGE: Usage = {
+  input_tokens: 0,
+  output_tokens: 0,
+  cache_creation_input_tokens: 0,
+  cache_read_input_tokens: 0,
+  server_tool_use: { web_search_requests: 0, web_fetch_requests: 0 },
+  service_tier: null,
+  cache_creation: {
+    ephemeral_1h_input_tokens: 0,
+    ephemeral_5m_input_tokens: 0,
+  },
+  inference_geo: null,
+  iterations: null,
+  speed: null,
+}
+
+export type OpenAIMessage = {
+  readonly role: 'system' | 'user' | 'assistant' | 'tool'
+  readonly content: string
+  readonly tool_call_id?: string
+  readonly name?: string
+  readonly tool_calls?: readonly OpenAIToolCall[]
+}
+
+export type OpenAIToolCall = {
+  readonly id: string
+  readonly type: 'function'
+  readonly function: {
+    readonly name: string
+    readonly arguments: string
+  }
+}
+
+export type OpenAITool = {
+  readonly type: 'function'
+  readonly function: {
+    readonly name: string
+    readonly description: string
+    readonly parameters: Record<string, unknown>
+  }
+}
+
+export type ProviderOptions = {
+  readonly getToolPermissionContext: () => Promise<unknown>
+  readonly model: string
+  readonly isNonInteractiveSession: boolean
+  readonly querySource: QuerySource
+  readonly agents: readonly unknown[]
+  readonly allowedAgentTypes: readonly string[]
+  readonly mcpTools: readonly unknown[]
+  readonly toolChoice?: { readonly type: 'tool'; readonly name: string }
+  readonly disabledProviderToolNames?: readonly string[]
+  readonly fetchOverride?: (
+    input: string | URL | Request,
+    init?: RequestInit,
+  ) => Promise<Response>
+  readonly enablePromptCaching?: boolean
+  readonly taskBudget?: { readonly total: number; readonly remaining?: number }
+  readonly skipCacheWrite?: boolean
+  readonly maxOutputTokensOverride?: number
+}
+
+export type QueryModelParams = {
+  readonly messages: readonly Message[]
+  readonly systemPrompt: SystemPrompt
+  readonly thinkingConfig?: unknown
+  readonly tools: Tools
+  readonly signal: AbortSignal
+  readonly options: ProviderOptions
+}
+
+export type ProviderRequest = {
+  readonly model: string
+  readonly stream: boolean
+  readonly messages: readonly OpenAIMessage[]
+  readonly chat_template_kwargs?: { readonly enable_thinking: boolean }
+  readonly parse_reasoning?: boolean
+  readonly include_reasoning?: boolean
+  readonly tools?: readonly OpenAITool[]
+  readonly tool_choice?: {
+    readonly type: 'function'
+    readonly function: { readonly name: string }
+  }
+  readonly metadata?: Record<string, string>
+  readonly output_config?: Record<string, unknown>
+}

package/tui/src/services/api/ummaya/usage.ts

@@ -0,0 +1,30 @@
+import type { Usage } from './types.js'
+
+export function cleanupStream(stream: ReadableStream<unknown> | undefined): void {
+  void stream?.cancel()
+}
+
+function addNumber(left: number | undefined, right: number | undefined): number {
+  return (left ?? 0) + (right ?? 0)
+}
+
+export function updateUsage(usage: Usage, partUsage: Partial<Usage> | undefined): Usage {
+  if (!partUsage) return usage
+  return {
+    ...usage,
+    input_tokens: addNumber(usage.input_tokens, partUsage.input_tokens),
+    output_tokens: addNumber(usage.output_tokens, partUsage.output_tokens),
+    cache_creation_input_tokens: addNumber(
+      usage.cache_creation_input_tokens,
+      partUsage.cache_creation_input_tokens,
+    ),
+    cache_read_input_tokens: addNumber(
+      usage.cache_read_input_tokens,
+      partUsage.cache_read_input_tokens,
+    ),
+  }
+}
+
+export function accumulateUsage(totalUsage: Usage, messageUsage: Usage): Usage {
+  return updateUsage(totalUsage, messageUsage)
+}