ummaya 0.2.4 → 0.2.6

package/tui/src/tools/WebSearchTool/WebSearchTool.ts

@@ -1,158 +1,29 @@
-import type {
-  BetaContentBlock,
-  BetaWebSearchTool20250305,
-} from 'src/sdk-compat.js'
 import { getAPIProvider } from 'src/utils/model/providers.js'
 import type { PermissionResult } from 'src/utils/permissions/PermissionResult.js'
-import { z } from 'zod/v4'
-import { getFeatureValue_CACHED_MAY_BE_STALE } from '../../services/analytics/growthbook.js'
-// UMMAYA-original: CC queryModelWithStreaming replaced by FriendliAI IPC client.
-const queryModelWithStreaming = async (_opts?: unknown): Promise<never> => { throw new Error('not implemented') }
 import { buildTool, type ToolDef } from '../../Tool.js'
-import { lazySchema } from '../../utils/lazySchema.js'
-import { logError } from '../../utils/log.js'
-import { createUserMessage } from '../../utils/userMessageFactories.js'
-import { getMainLoopModel, getSmallFastModel } from '../../utils/model/model.js'
-import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
-import { asSystemPrompt } from '../../utils/systemPromptType.js'
+import { getMainLoopModel } from '../../utils/model/model.js'
+import { callWebSearch } from './call.js'
 import { getWebSearchPrompt, WEB_SEARCH_TOOL_NAME } from './prompt.js'
+import { mapWebSearchResultToToolResultBlockParam } from './resultBlock.js'
 import {
   getToolUseSummary,
   renderToolResultMessage,
   renderToolUseMessage,
   renderToolUseProgressMessage,
 } from './UI.js'
+import type { InputSchema, Output, OutputSchema, SearchResult } from './schemas.js'
+import { inputSchema, outputSchema } from './schemas.js'
 
-const inputSchema = lazySchema(() =>
-  z.strictObject({
-    query: z.string().min(2).describe('The search query to use'),
-    allowed_domains: z
-      .array(z.string())
-      .optional()
-      .describe('Only include search results from these domains'),
-    blocked_domains: z
-      .array(z.string())
-      .optional()
-      .describe('Never include search results from these domains'),
-  }),
-)
-type InputSchema = ReturnType<typeof inputSchema>
-
-type Input = z.infer<InputSchema>
-
-const searchResultSchema = lazySchema(() => {
-  const searchHitSchema = z.object({
-    title: z.string().describe('The title of the search result'),
-    url: z.string().describe('The URL of the search result'),
-  })
-
-  return z.object({
-    tool_use_id: z.string().describe('ID of the tool use'),
-    content: z.array(searchHitSchema).describe('Array of search hits'),
-  })
-})
-
-export type SearchResult = z.infer<ReturnType<typeof searchResultSchema>>
-
-const outputSchema = lazySchema(() =>
-  z.object({
-    query: z.string().describe('The search query that was executed'),
-    results: z
-      .array(z.union([searchResultSchema(), z.string()]))
-      .describe('Search results and/or text commentary from the model'),
-    durationSeconds: z
-      .number()
-      .describe('Time taken to complete the search operation'),
-  }),
-)
-type OutputSchema = ReturnType<typeof outputSchema>
-
-export type Output = z.infer<OutputSchema>
+export type { Output, SearchResult }
 
 // Re-export WebSearchProgress from centralized types to break import cycles
 export type { WebSearchProgress } from '../../types/tools.js'
 
 import type { WebSearchProgress } from '../../types/tools.js'
 
-function makeToolSchema(input: Input): BetaWebSearchTool20250305 {
-  return {
-    type: 'web_search_20250305',
-    name: 'web_search',
-    allowed_domains: input.allowed_domains,
-    blocked_domains: input.blocked_domains,
-    max_uses: 8, // Hardcoded to 8 searches maximum
-  }
-}
-
-function makeOutputFromSearchResponse(
-  result: BetaContentBlock[],
-  query: string,
-  durationSeconds: number,
-): Output {
-  // The result is a sequence of these blocks:
-  // - text to start -- always?
-  // [
-  //    - server_tool_use
-  //    - web_search_tool_result
-  //    - text and citation blocks intermingled
-  //  ]+  (this block repeated for each search)
-
-  const results: (SearchResult | string)[] = []
-  let textAcc = ''
-  let inText = true
-
-  for (const block of result) {
-    if (block.type === 'server_tool_use') {
-      if (inText) {
-        inText = false
-        if (textAcc.trim().length > 0) {
-          results.push(textAcc.trim())
-        }
-        textAcc = ''
-      }
-      continue
-    }
-
-    if (block.type === 'web_search_tool_result') {
-      // Handle error case - content is a WebSearchToolResultError
-      if (!Array.isArray(block.content)) {
-        const errorMessage = `Web search error: ${block.content.error_code}`
-        logError(new Error(errorMessage))
-        results.push(errorMessage)
-        continue
-      }
-      // Success case - add results to our collection
-      const hits = block.content.map(r => ({ title: r.title, url: r.url }))
-      results.push({
-        tool_use_id: block.tool_use_id,
-        content: hits,
-      })
-    }
-
-    if (block.type === 'text') {
-      if (inText) {
-        textAcc += block.text
-      } else {
-        inText = true
-        textAcc = block.text
-      }
-    }
-  }
-
-  if (textAcc.length) {
-    results.push(textAcc.trim())
-  }
-
-  return {
-    query,
-    results,
-    durationSeconds,
-  }
-}
-
 export const WebSearchTool = buildTool({
   name: WEB_SEARCH_TOOL_NAME,
-  searchHint: 'search the web for current information',
+  searchHint: 'search the web for current information and source evidence',
   maxResultSizeChars: 100_000,
   shouldDefer: true,
   async description(input) {
@@ -253,184 +124,9 @@ export const WebSearchTool = buildTool({
     return { result: true }
   },
   async call(input, context, _canUseTool, _parentMessage, onProgress) {
-    const startTime = performance.now()
-    const { query } = input
-    const userMessage = createUserMessage({
-      content: 'Perform a web search for the query: ' + query,
-    })
-    const toolSchema = makeToolSchema(input)
-
-    const useHaiku = getFeatureValue_CACHED_MAY_BE_STALE(
-      'tengu_plum_vx3',
-      false,
-    )
-
-    const appState = context.getAppState()
-    const queryStream = queryModelWithStreaming({
-      messages: [userMessage],
-      systemPrompt: asSystemPrompt([
-        'You are an assistant for performing a web search tool use',
-      ]),
-      thinkingConfig: useHaiku
-        ? { type: 'disabled' as const }
-        : context.options.thinkingConfig,
-      tools: [],
-      signal: context.abortController.signal,
-      options: {
-        getToolPermissionContext: async () => appState.toolPermissionContext,
-        model: useHaiku ? getSmallFastModel() : context.options.mainLoopModel,
-        toolChoice: useHaiku ? { type: 'tool', name: 'web_search' } : undefined,
-        isNonInteractiveSession: context.options.isNonInteractiveSession,
-        hasAppendSystemPrompt: !!context.options.appendSystemPrompt,
-        extraToolSchemas: [toolSchema],
-        querySource: 'web_search_tool',
-        agents: context.options.agentDefinitions.activeAgents,
-        mcpTools: [],
-        agentId: context.agentId,
-        effortValue: appState.effortValue,
-      },
-    })
-
-    const allContentBlocks: BetaContentBlock[] = []
-    let currentToolUseId = null
-    let currentToolUseJson = ''
-    let progressCounter = 0
-    const toolUseQueries = new Map() // Map of tool_use_id to query
-
-    for await (const event of queryStream) {
-      if (event.type === 'assistant') {
-        allContentBlocks.push(...event.message.content)
-        continue
-      }
-
-      // Track tool use ID when server_tool_use starts
-      if (
-        event.type === 'stream_event' &&
-        event.event?.type === 'content_block_start'
-      ) {
-        const contentBlock = event.event.content_block
-        if (contentBlock && contentBlock.type === 'server_tool_use') {
-          currentToolUseId = contentBlock.id
-          currentToolUseJson = ''
-          // Note: The ServerToolUseBlock doesn't contain input.query
-          // The actual query comes through input_json_delta events
-          continue
-        }
-      }
-
-      // Accumulate JSON for current tool use
-      if (
-        currentToolUseId &&
-        event.type === 'stream_event' &&
-        event.event?.type === 'content_block_delta'
-      ) {
-        const delta = event.event.delta
-        if (delta?.type === 'input_json_delta' && delta.partial_json) {
-          currentToolUseJson += delta.partial_json
-
-          // Try to extract query from partial JSON for progress updates
-          try {
-            // Look for a complete query field
-            const queryMatch = currentToolUseJson.match(
-              /"query"\s*:\s*"((?:[^"\\]|\\.)*)"/,
-            )
-            if (queryMatch && queryMatch[1]) {
-              // The regex properly handles escaped characters
-              const query = jsonParse('"' + queryMatch[1] + '"')
-
-              if (
-                !toolUseQueries.has(currentToolUseId) ||
-                toolUseQueries.get(currentToolUseId) !== query
-              ) {
-                toolUseQueries.set(currentToolUseId, query)
-                progressCounter++
-                if (onProgress) {
-                  onProgress({
-                    toolUseID: `search-progress-${progressCounter}`,
-                    data: {
-                      type: 'query_update',
-                      query,
-                    },
-                  })
-                }
-              }
-            }
-          } catch {
-            // Ignore parsing errors for partial JSON
-          }
-        }
-      }
-
-      // Yield progress when search results come in
-      if (
-        event.type === 'stream_event' &&
-        event.event?.type === 'content_block_start'
-      ) {
-        const contentBlock = event.event.content_block
-        if (contentBlock && contentBlock.type === 'web_search_tool_result') {
-          // Get the actual query that was used for this search
-          const toolUseId = contentBlock.tool_use_id
-          const actualQuery = toolUseQueries.get(toolUseId) || query
-          const content = contentBlock.content
-
-          progressCounter++
-          if (onProgress) {
-            onProgress({
-              toolUseID: toolUseId || `search-progress-${progressCounter}`,
-              data: {
-                type: 'search_results_received',
-                resultCount: Array.isArray(content) ? content.length : 0,
-                query: actualQuery,
-              },
-            })
-          }
-        }
-      }
-    }
-
-    // Process the final result
-    const endTime = performance.now()
-    const durationSeconds = (endTime - startTime) / 1000
-
-    const data = makeOutputFromSearchResponse(
-      allContentBlocks,
-      query,
-      durationSeconds,
-    )
-    return { data }
+    return callWebSearch(input, context, onProgress)
   },
   mapToolResultToToolResultBlockParam(output, toolUseID) {
-    const { query, results } = output
-
-    let formattedOutput = `Web search results for query: "${query}"\n\n`
-
-    // Process the results array - it can contain both string summaries and search result objects.
-    // Guard against null/undefined entries that can appear after JSON round-tripping
-    // (e.g., from compaction or transcript deserialization).
-    ;(results ?? []).forEach(result => {
-      if (result == null) {
-        return
-      }
-      if (typeof result === 'string') {
-        // Text summary
-        formattedOutput += result + '\n\n'
-      } else {
-        // Search result with links
-        if (result.content?.length > 0) {
-          formattedOutput += `Links: ${jsonStringify(result.content)}\n\n`
-        } else {
-          formattedOutput += 'No links found.\n\n'
-        }
-      }
-    })
-
-    formattedOutput +=
-      '\nREMINDER: You MUST include the sources above in your response to the user using markdown hyperlinks.'
-
-    return {
-      tool_use_id: toolUseID,
-      type: 'tool_result',
-      content: formattedOutput.trim(),
-    }
+    return mapWebSearchResultToToolResultBlockParam(output, toolUseID)
   },
 } satisfies ToolDef<InputSchema, Output, WebSearchProgress>)

package/tui/src/tools/WebSearchTool/call.ts

@@ -0,0 +1,33 @@
+import type {
+  ToolCallProgress,
+  ToolResult,
+  ToolUseContext,
+} from '../../Tool.js'
+import type { WebSearchProgress } from '../../types/tools.js'
+import type { Input, Output } from './schemas.js'
+import { makeBlockedOutputFromProviderError } from './responseMapping.js'
+
+class WebSearchProviderUnavailableError extends Error {
+  readonly name = 'WebSearchProviderUnavailableError'
+
+  constructor() {
+    super(
+      'WebSearch is blocked: UMMAYA does not currently expose a policy-approved server-side web search provider for the FriendliAI execution path. Use WebFetch with an explicit public URL or a registered live adapter instead.',
+    )
+  }
+}
+
+export async function callWebSearch(
+  input: Input,
+  _context: ToolUseContext,
+  _onProgress?: ToolCallProgress<WebSearchProgress>,
+): Promise<ToolResult<Output>> {
+  const startTime = performance.now()
+  return {
+    data: makeBlockedOutputFromProviderError(
+      input.query,
+      startTime,
+      new WebSearchProviderUnavailableError(),
+    ),
+  }
+}

package/tui/src/tools/WebSearchTool/responseMapping.ts

@@ -0,0 +1,190 @@
+import type { BetaContentBlock } from 'src/sdk-compat.js'
+import { logError } from '../../utils/log.js'
+import { errorMessage } from '../../utils/errors.js'
+import {
+  buildSourceEvidence,
+  buildSourceVerification,
+  redactSourceVerificationText,
+  type SourceVerificationEvidence,
+} from '../WebFetchTool/sourceVerification.js'
+import { WEB_SEARCH_TOOL_NAME } from './prompt.js'
+import type { Output, SearchResult } from './schemas.js'
+
+type WebSearchResultBlock = {
+  readonly type: 'web_search_tool_result'
+  readonly tool_use_id: string
+  readonly content: unknown
+}
+
+type TextBlock = {
+  readonly type: 'text'
+  readonly text: string
+}
+
+type ServerToolUseBlock = {
+  readonly type: 'server_tool_use'
+}
+
+function isWebSearchResultBlock(
+  block: BetaContentBlock,
+): block is WebSearchResultBlock {
+  return block.type === 'web_search_tool_result'
+}
+
+function isTextBlock(block: BetaContentBlock): block is TextBlock {
+  return block.type === 'text' && 'text' in block && typeof block.text === 'string'
+}
+
+function isServerToolUseBlock(
+  block: BetaContentBlock,
+): block is ServerToolUseBlock {
+  return block.type === 'server_tool_use'
+}
+
+function isSearchHit(value: unknown): value is { readonly title: string; readonly url: string } {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    'title' in value &&
+    typeof value.title === 'string' &&
+    'url' in value &&
+    typeof value.url === 'string'
+  )
+}
+
+function providerErrorCode(content: unknown): string {
+  if (
+    typeof content === 'object' &&
+    content !== null &&
+    'error_code' in content &&
+    typeof content.error_code === 'string'
+  ) {
+    return content.error_code
+  }
+  return 'provider_error'
+}
+
+function appendSearchResultEvidence(
+  block: WebSearchResultBlock,
+  results: (SearchResult | string)[],
+  evidence: SourceVerificationEvidence[],
+): void {
+  if (!Array.isArray(block.content)) {
+    const message = `Web search error: ${providerErrorCode(block.content)}`
+    logError(new Error(message))
+    results.push(message)
+    evidence.push(
+      buildSourceEvidence({
+        toolId: WEB_SEARCH_TOOL_NAME,
+        sourceUrl: null,
+        title: 'Provider error',
+        blockedOrUsed: 'blocked',
+        rawText: message,
+      }),
+    )
+    return
+  }
+
+  const hits = block.content.filter(isSearchHit).map(hit => ({
+    title: hit.title,
+    url: hit.url,
+  }))
+  for (const hit of hits) {
+    evidence.push(
+      buildSourceEvidence({
+        toolId: WEB_SEARCH_TOOL_NAME,
+        sourceUrl: hit.url,
+        title: hit.title,
+        blockedOrUsed: 'needs_input',
+        rawText: `${hit.title}\n${hit.url}`,
+      }),
+    )
+  }
+  results.push({
+    tool_use_id: block.tool_use_id,
+    content: hits,
+  })
+}
+
+export function makeOutputFromSearchResponse(
+  result: readonly BetaContentBlock[],
+  query: string,
+  durationSeconds: number,
+): Output {
+  const results: (SearchResult | string)[] = []
+  const evidence: SourceVerificationEvidence[] = []
+  let textAcc = ''
+  let inText = true
+
+  for (const block of result) {
+    if (isServerToolUseBlock(block)) {
+      if (inText) {
+        inText = false
+        if (textAcc.trim().length > 0) {
+          results.push(textAcc.trim())
+        }
+        textAcc = ''
+      }
+      continue
+    }
+
+    if (isWebSearchResultBlock(block)) {
+      appendSearchResultEvidence(block, results, evidence)
+    }
+
+    if (isTextBlock(block)) {
+      if (inText) {
+        textAcc += block.text
+      } else {
+        inText = true
+        textAcc = block.text
+      }
+    }
+  }
+
+  if (textAcc.length) {
+    results.push(redactSourceVerificationText(textAcc.trim()))
+  }
+  if (evidence.length === 0) {
+    evidence.push(
+      buildSourceEvidence({
+        toolId: WEB_SEARCH_TOOL_NAME,
+        sourceUrl: null,
+        title: 'No source URL returned',
+        blockedOrUsed: 'needs_input',
+        rawText: query,
+      }),
+    )
+  }
+
+  return {
+    query,
+    results,
+    durationSeconds,
+    sourceVerification: buildSourceVerification(evidence),
+  }
+}
+
+export function makeBlockedOutputFromProviderError(
+  query: string,
+  startTime: number,
+  error: unknown,
+): Output {
+  const message = `Source verification blocked: ${errorMessage(error)}`
+  const safeMessage = redactSourceVerificationText(message)
+  logError(new Error(message))
+  return {
+    query,
+    results: [safeMessage],
+    durationSeconds: (performance.now() - startTime) / 1000,
+    sourceVerification: buildSourceVerification([
+      buildSourceEvidence({
+        toolId: WEB_SEARCH_TOOL_NAME,
+        sourceUrl: null,
+        title: 'Provider error',
+        blockedOrUsed: 'blocked',
+        rawText: message,
+      }),
+    ]),
+  }
+}

package/tui/src/tools/WebSearchTool/resultBlock.ts

@@ -0,0 +1,47 @@
+import { jsonStringify } from '../../utils/slowOperations.js'
+import {
+  formatSourceVerifiedToolResult,
+  redactSourceVerificationText,
+  redactSourceVerificationUrl,
+} from '../WebFetchTool/sourceVerification.js'
+import type { Output } from './schemas.js'
+
+export function mapWebSearchResultToToolResultBlockParam(
+  output: Output,
+  toolUseID: string,
+): {
+  readonly tool_use_id: string
+  readonly type: 'tool_result'
+  readonly content: string
+} {
+  const { query, results, sourceVerification } = output
+  let formattedOutput = `Web search results for query: "${redactSourceVerificationText(query)}"\n\n`
+
+  for (const result of results) {
+    if (typeof result === 'string') {
+      formattedOutput += `${redactSourceVerificationText(result)}\n\n`
+      continue
+    }
+    if (result.content.length > 0) {
+      const safeContent = result.content.map(hit => ({
+        title: redactSourceVerificationText(hit.title),
+        url: redactSourceVerificationUrl(hit.url) ?? 'none',
+      }))
+      formattedOutput += `Links: ${jsonStringify(safeContent)}\n\n`
+    } else {
+      formattedOutput += 'No links found.\n\n'
+    }
+  }
+
+  formattedOutput +=
+    '\nREMINDER: You MUST include the sources above in your response to the user using markdown hyperlinks.'
+
+  return {
+    tool_use_id: toolUseID,
+    type: 'tool_result',
+    content: formatSourceVerifiedToolResult({
+      result: formattedOutput.trim(),
+      sourceVerification,
+    }),
+  }
+}

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

@@ -0,0 +1,47 @@
+import { z } from 'zod/v4'
+import { lazySchema } from '../../utils/lazySchema.js'
+import { sourceVerificationSchema } from '../WebFetchTool/sourceVerification.js'
+
+export const inputSchema = lazySchema(() =>
+  z.strictObject({
+    query: z.string().min(2).describe('The search query to use'),
+    allowed_domains: z
+      .array(z.string())
+      .optional()
+      .describe('Only include search results from these domains'),
+    blocked_domains: z
+      .array(z.string())
+      .optional()
+      .describe('Never include search results from these domains'),
+  }),
+)
+export type InputSchema = ReturnType<typeof inputSchema>
+export type Input = z.infer<InputSchema>
+
+const searchHitSchema = z.object({
+  title: z.string().describe('The title of the search result'),
+  url: z.string().describe('The URL of the search result'),
+})
+
+export const searchResultSchema = lazySchema(() =>
+  z.object({
+    tool_use_id: z.string().describe('ID of the tool use'),
+    content: z.array(searchHitSchema).describe('Array of search hits'),
+  }),
+)
+export type SearchResult = z.infer<ReturnType<typeof searchResultSchema>>
+
+export const outputSchema = lazySchema(() =>
+  z.object({
+    query: z.string().describe('The search query that was executed'),
+    results: z
+      .array(z.union([searchResultSchema(), z.string()]))
+      .describe('Search results and/or text commentary from the model'),
+    durationSeconds: z
+      .number()
+      .describe('Time taken to complete the search operation'),
+    sourceVerification: sourceVerificationSchema.optional(),
+  }),
+)
+export type OutputSchema = ReturnType<typeof outputSchema>
+export type Output = z.infer<OutputSchema>

package/tui/src/tools/WebSearchTool/toolSchema.ts

@@ -0,0 +1,12 @@
+import type { BetaWebSearchTool20250305 } from 'src/sdk-compat.js'
+import type { Input } from './schemas.js'
+
+export function makeToolSchema(input: Input): BetaWebSearchTool20250305 {
+  return {
+    type: 'web_search_20250305',
+    name: 'web_search',
+    allowed_domains: input.allowed_domains,
+    blocked_domains: input.blocked_domains,
+    max_uses: 8,
+  }
+}