diff --git a/server/api/chat/agentPromptCreation.ts b/server/api/chat/agentPromptCreation.ts index f8424f62d..1dfd8ca2f 100644 --- a/server/api/chat/agentPromptCreation.ts +++ b/server/api/chat/agentPromptCreation.ts @@ -1,5 +1,6 @@ export type AgentPromptSections = { toolCoordination: string + knowledgeBaseWorkflow: string publicAgentDiscipline: string agentQueryCrafting: string generalExecution: string @@ -14,6 +15,23 @@ export const agentPromptSections: AgentPromptSections = { - Pair every plan update with concrete tool execution; never emit a planning-only turn. - Example: run \`searchGlobal\` to list escalation IDs, then next turn call \`getSlackRelatedMessages\` per ID. `.trim(), + knowledgeBaseWorkflow: ` +### Knowledge Base Workflow +- Decide first whether the ask actually needs knowledge-base evidence; skip both \`ls\` and \`searchKnowledgeBase\` when other tools or existing context already cover the answer. +- Treat \`ls\` as the structure and metadata tool for KB: use it when the user asks what exists, where something lives, which files match a constraint such as PDF, or when a quick browse will make the next search materially sharper. +- Use \`ls\` alone when the question is about inventory, hierarchy, paths, or metadata rather than document contents. + - Use \`searchKnowledgeBase\` directly when the relevant collection, folder, file, or path is already known from the user query, agent prompt, prior tool output, or previously discovered IDs. + - Use \`ls\` before \`searchKnowledgeBase\` when you need to discover accessible collections, confirm a canonical path, inspect folder/file layout, collect file or folder IDs, or narrow the search to a metadata-defined subset such as PDFs inside a folder. + - \`ls\` and \`searchKnowledgeBase\` are complementary, not a mandatory pair; chain them only when browsing will materially improve the next search. + - Feel free to call \`ls\` in between turns or anytime if you think it would help sharpen scope, confirm structure, or avoid wasted KB searching. + - Keep \`ls\` cheap by default: start with \`depth: 1\` and \`metadata: false\`; increase depth or enable metadata only when the task needs deeper traversal or row details such as \`mime_type\`, timestamps, descriptions, or collection metadata. + - Put structural scoping in \`filters.targets\`, not inside the free-text query. \`targets\` can union multiple relevant KB locations inside the current allowed scope, including exact file IDs discovered from \`ls\`. +- Examples: + - Structure-only ask: answer "what is inside \`/Policies\`?" with \`ls({ target: { type: "path", collectionId: "kb-1", path: "/Policies" }, depth: 1, metadata: false })\`; do not call \`searchKnowledgeBase\` if the user only needs the listing. + - Filtered content ask: for "answer only from PDF files in Security policies", first call \`ls({ target: { type: "path", collectionId: "kb-1", path: "/Policies/Security" }, depth: 2, metadata: true })\`, keep only rows whose \`mime_type\` is PDF, then call \`searchKnowledgeBase({ query: "exception approval workflow", filters: { targets: [{ type: "file", fileId: "file-pdf-1" }, { type: "file", fileId: "file-pdf-2" }] }, limit: 5 })\`. + - Known scope ask: if the ask already names the exact KB location, call \`searchKnowledgeBase({ query: "contractor onboarding steps", filters: { targets: [{ type: "path", collectionId: "kb-1", path: "/HR/Onboarding/Checklist.md" }] }, limit: 5 })\`; skip \`ls\`. + - \`ls\` not useful: if the ask is not about KB, or the exact KB scope is already known and browsing will not improve precision, skip \`ls\`. + `.trim(), publicAgentDiscipline: ` ### Public Agent Discipline - Call \`list_custom_agents\` before delegating; log the evaluation and expect it to return \`null\` when nobody qualifies. @@ -35,6 +53,7 @@ export const agentPromptSections: AgentPromptSections = { - Show proactive drive: after each subtask plan, immediately schedule and run the necessary tools. - Combine sequential and parallel tool usage when safe; default to running more than one independent tool per turn. - Track dependencies explicitly so you never trigger a child tool before its parent results are analyzed. +- Prefer the shortest correct tool path; do not insert browsing or discovery steps when a precise search can answer directly. `.trim(), chainOfThought: ` ### Chain-of-Thought Commitment @@ -50,6 +69,7 @@ export const agentPromptSections: AgentPromptSections = { export function buildAgentPromptAddendum(): string { return [ agentPromptSections.toolCoordination, + agentPromptSections.knowledgeBaseWorkflow, agentPromptSections.publicAgentDiscipline, agentPromptSections.agentQueryCrafting, agentPromptSections.generalExecution, diff --git a/server/api/chat/jaf-adapter.ts b/server/api/chat/jaf-adapter.ts index 70bc8d190..fa09f7d6c 100644 --- a/server/api/chat/jaf-adapter.ts +++ b/server/api/chat/jaf-adapter.ts @@ -1,16 +1,19 @@ -import { z, type ZodType } from "zod" -import type { Tool } from "@xynehq/jaf" -import { ToolResponse } from "@xynehq/jaf" -import type { MinimalAgentFragment, Citation } from "./types" -import type { AgentRunContext } from "./agent-schemas" import { answerContextMapFromFragments } from "@/ai/context" import { getLogger } from "@/logger" import { Subsystem } from "@/types" import { Apps } from "@xyne/vespa-ts/types" +import type { Tool } from "@xynehq/jaf" +import { ToolResponse } from "@xynehq/jaf" +import { type ZodType, z } from "zod" +import type { AgentRunContext } from "./agent-schemas" +import type { Citation, MinimalAgentFragment } from "./types" const Logger = getLogger(Subsystem.Chat).child({ module: "jaf-adapter" }) -type ToolSchemaParameters = Tool["schema"]["parameters"] +type ToolSchemaParameters = Tool< + unknown, + AgentRunContext +>["schema"]["parameters"] const toToolSchemaParameters = (schema: ZodType): ToolSchemaParameters => schema as unknown as ToolSchemaParameters @@ -84,7 +87,7 @@ export function buildMCPJAFTools( ) } } - Logger.info( + Logger.debug( { connectorId, toolName, descLen: (toolDescription || "").length }, "[MCP] Registering tool for JAF agent", ) diff --git a/server/api/chat/jaf-logging.ts b/server/api/chat/jaf-logging.ts new file mode 100644 index 000000000..d621503bc --- /dev/null +++ b/server/api/chat/jaf-logging.ts @@ -0,0 +1,297 @@ +import { getLoggerWithChild } from "@/logger" +import { Subsystem } from "@/types" +import { getErrorMessage } from "@/utils" +import { type TraceEvent, getTextContent } from "@xynehq/jaf" + +const loggerWithChild = getLoggerWithChild(Subsystem.Chat, { + module: "jaf-logging", +}) + +export type JAFTraceLoggingContext = { + chatId: string + email: string + flow: "MessageAgents" | "DelegatedAgenticRun" + runId: string +} + +function truncateValue(value: string, maxLength = 160): string { + if (value.length <= maxLength) return value + return `${value.slice(0, maxLength - 1)}…` +} + +function summarizeToolResultPayload(result: any): string { + if (!result) { + return "No result returned." + } + const summaryCandidates: Array = [ + result?.data?.summary, + result?.data?.result, + ] + for (const candidate of summaryCandidates) { + if (typeof candidate === "string" && candidate.trim().length > 0) { + return truncateValue(candidate.trim(), 200) + } + } + if (typeof result?.data === "string") { + return truncateValue(result.data, 200) + } + try { + return truncateValue(JSON.stringify(result?.data ?? result), 200) + } catch { + return "Result unavailable." + } +} + +function formatToolArgumentsForLogging(args: Record): string { + if (!args || typeof args !== "object") { + return "{}" + } + const entries = Object.entries(args) + if (entries.length === 0) { + return "{}" + } + const parts = entries.map(([key, value]) => { + let serialized: string + if (typeof value === "string") { + serialized = `"${truncateValue(value, 80)}"` + } else if ( + typeof value === "number" || + typeof value === "boolean" || + value === null + ) { + serialized = String(value) + } else { + try { + serialized = truncateValue(JSON.stringify(value), 80) + } catch { + serialized = "[unserializable]" + } + } + return `${key}: ${serialized}` + }) + const combined = parts.join(", ") + return truncateValue(combined, 400) +} + +export function logJAFTraceEvent( + context: JAFTraceLoggingContext, + event: TraceEvent, +): void { + if (event.type === "before_tool_execution") { + return + } + + const logger = loggerWithChild({ email: context.email }) + const baseLog = { + chatId: context.chatId, + eventType: event.type, + flow: context.flow, + runId: context.runId, + } + + switch (event.type) { + case "run_start": + logger.info(baseLog, "[JAF] Run started") + return + + case "run_end": { + const outcome = ( + event.data as { outcome?: { status?: string; error?: any } } + )?.outcome + const runEndLog = { + ...baseLog, + errorDetail: + outcome?.status === "error" + ? getErrorMessage(outcome?.error) + : undefined, + errorTag: + outcome?.status === "error" && + outcome?.error && + typeof outcome.error === "object" && + "_tag" in outcome.error + ? String((outcome.error as { _tag?: string })._tag) + : undefined, + outcomeStatus: outcome?.status ?? "unknown", + } + if (outcome?.status === "error") { + logger.error(runEndLog, "[JAF] Run ended with error") + } else { + logger.info(runEndLog, "[JAF] Run completed") + } + return + } + + case "guardrail_violation": + logger.warn( + { + ...baseLog, + reason: event.data.reason, + stage: event.data.stage, + }, + "[JAF] Guardrail violation", + ) + return + + case "handoff_denied": + logger.warn( + { + ...baseLog, + from: event.data.from, + reason: event.data.reason, + to: event.data.to, + }, + "[JAF] Handoff denied", + ) + return + + case "decode_error": + logger.error( + { + ...baseLog, + errors: event.data.errors, + }, + "[JAF] Decode error", + ) + return + + case "turn_start": + logger.debug( + { + ...baseLog, + agentName: event.data.agentName, + turn: event.data.turn, + }, + "[JAF] Turn started", + ) + return + + case "turn_end": + logger.debug( + { + ...baseLog, + turn: event.data.turn, + }, + "[JAF] Turn ended", + ) + return + + case "tool_requests": + logger.debug( + { + ...baseLog, + toolCount: event.data.toolCalls.length, + toolNames: event.data.toolCalls.map((toolCall) => toolCall.name), + }, + "[JAF] Tool requests planned", + ) + return + + case "tool_call_start": + logger.debug( + { + ...baseLog, + args: formatToolArgumentsForLogging( + (event.data.args ?? {}) as Record, + ), + toolName: event.data.toolName, + }, + "[JAF] Tool call started", + ) + return + + case "tool_call_end": + if (event.data.error) { + logger.error( + { + ...baseLog, + error: event.data.error, + executionTimeMs: event.data.executionTime, + resultPreview: summarizeToolResultPayload(event.data.result), + status: event.data.status ?? "error", + toolName: event.data.toolName, + }, + "[JAF] Tool call failed", + ) + } else { + logger.debug( + { + ...baseLog, + executionTimeMs: event.data.executionTime, + resultPreview: summarizeToolResultPayload(event.data.result), + status: event.data.status ?? "completed", + toolName: event.data.toolName, + }, + "[JAF] Tool call completed", + ) + } + return + + case "assistant_message": { + const content = getTextContent(event.data.message.content) || "" + logger.debug( + { + ...baseLog, + contentLength: content.length, + contentPreview: truncateValue(content, 200), + hasToolCalls: + Array.isArray(event.data.message?.tool_calls) && + (event.data.message.tool_calls?.length ?? 0) > 0, + }, + "[JAF] Assistant message received", + ) + return + } + + case "token_usage": + logger.debug( + { + ...baseLog, + completionTokens: event.data.completion ?? 0, + promptTokens: event.data.prompt ?? 0, + totalTokens: event.data.total ?? 0, + }, + "[JAF] Token usage recorded", + ) + return + + case "clarification_requested": + logger.debug( + { + ...baseLog, + clarificationId: event.data.clarificationId, + optionsCount: event.data.options.length, + question: truncateValue(event.data.question, 200), + }, + "[JAF] Clarification requested", + ) + return + + case "clarification_provided": + logger.debug( + { + ...baseLog, + clarificationId: event.data.clarificationId, + selectedId: event.data.selectedId, + }, + "[JAF] Clarification provided", + ) + return + + case "final_output": + logger.debug( + { + ...baseLog, + outputLength: + typeof event.data.output === "string" + ? event.data.output.length + : 0, + }, + "[JAF] Final output emitted", + ) + return + + default: + logger.debug(baseLog, "[JAF] Trace event received") + return + } +} diff --git a/server/api/chat/jaf-provider.ts b/server/api/chat/jaf-provider.ts index 8dc40377a..8d2a501de 100644 --- a/server/api/chat/jaf-provider.ts +++ b/server/api/chat/jaf-provider.ts @@ -1,13 +1,11 @@ -import { getAISDKProviderByModel } from "@/ai/provider" -import { MODEL_CONFIGURATIONS } from "@/ai/modelConfig" -import type { - ModelProvider as JAFModelProvider, - Message as JAFMessage, - Agent as JAFAgent, -} from "@xynehq/jaf" -import { getTextContent } from "@xynehq/jaf" -import { Models, AIProviders } from "@/ai/types" +import fs from "fs" +import path from "path" import { ModelToProviderMap } from "@/ai/mappers" +import { MODEL_CONFIGURATIONS } from "@/ai/modelConfig" +import { getAISDKProviderByModel } from "@/ai/provider" +import { findImageByName, regex } from "@/ai/provider/base" +import { AIProviders, Models } from "@/ai/types" +import config from "@/config" import { getLogger, getLoggerWithChild } from "@/logger" import { Subsystem } from "@/types" import type { @@ -15,26 +13,28 @@ import type { JSONValue, LanguageModelV2CallOptions, LanguageModelV2Content, + LanguageModelV2FilePart, LanguageModelV2FunctionTool, LanguageModelV2Message, + LanguageModelV2ReasoningPart, + LanguageModelV2TextPart, LanguageModelV2ToolCall, + LanguageModelV2ToolCallPart, LanguageModelV2ToolChoice, LanguageModelV2ToolResultOutput, - LanguageModelV2TextPart, - LanguageModelV2FilePart, - LanguageModelV2ReasoningPart, - LanguageModelV2ToolCallPart, LanguageModelV2ToolResultPart, } from "@ai-sdk/provider" +import type { + Agent as JAFAgent, + Message as JAFMessage, + ModelProvider as JAFModelProvider, +} from "@xynehq/jaf" +import { getTextContent } from "@xynehq/jaf" import OpenAI from "openai" -import config from "@/config" -import { zodSchemaToJsonSchema } from "./jaf-provider-utils" -import path from "path" -import fs from "fs" -import { regex, findImageByName } from "@/ai/provider/base" import type { AgentRunContext } from "./agent-schemas" -import { getRecentImagesFromContext } from "./runContextUtils" import { raceWithStop, throwIfStopRequested } from "./agent-stop" +import { zodSchemaToJsonSchema } from "./jaf-provider-utils" +import { getRecentImagesFromContext } from "./runContextUtils" const { IMAGE_CONTEXT_CONFIG } = config const IMAGE_BASE_DIR = path.resolve( process.env.IMAGE_DIR || "downloads/xyne_images_db", @@ -48,8 +48,10 @@ const MIME_TYPE_MAP: Record = { ".webp": "image/webp", } -const Logger = getLogger(Subsystem.Chat) -const loggerWithChild = getLoggerWithChild(Subsystem.Chat) +const Logger = getLogger(Subsystem.Chat).child({ module: "jaf-provider" }) +const loggerWithChild = getLoggerWithChild(Subsystem.Chat, { + module: "jaf-provider", +}) const MIN_TURN_NUMBER = 1 const normalizeTurnNumber = (turn?: number | null): number => typeof turn === "number" && turn >= MIN_TURN_NUMBER ? turn : MIN_TURN_NUMBER @@ -64,7 +66,6 @@ type ImagePromptPart = { filePart: LanguageModelV2FilePart } - const findLastUserMessageIndex = ( messages: LanguageModelV2Message[], ): number => { @@ -81,17 +82,27 @@ const parseImageFileName = ( ): { docIndex: string; docId: string; imageNumber: string } | null => { const match = imageName.match(regex) if (!match) { - console.warn(`[JAF Provider] Invalid image reference: ${imageName}`) + Logger.debug({ imageName }, "Invalid image reference") return null } const [, docIndex, docId, imageNumber] = match if (docId.includes("..") || docId.includes("/") || docId.includes("\\")) { - console.warn(`[JAF Provider] Suspicious docId detected: ${docId}`) + Logger.warn( + { docId, imageName }, + "Suspicious docId detected in image reference", + ) return null } - if (imageNumber.includes("..") || imageNumber.includes("/") || imageNumber.includes("\\")) { - console.warn(`[JAF Provider] Suspicious imageNumber detected: ${imageNumber}`) - return null + if ( + imageNumber.includes("..") || + imageNumber.includes("/") || + imageNumber.includes("\\") + ) { + Logger.warn( + { imageName, imageNumber }, + "Suspicious imageNumber detected in image reference", + ) + return null } return { docIndex, docId, imageNumber } @@ -101,7 +112,7 @@ const buildLanguageModelImageParts = async ( imageFileNames: string[], ): Promise => { const loadStats = { success: 0, failed: 0, totalBytes: 0 } - + const parts = await Promise.all( imageFileNames.map(async (imageName): Promise => { const parsed = parseImageFileName(imageName) @@ -114,8 +125,9 @@ const buildLanguageModelImageParts = async ( const imageDir = path.join(IMAGE_BASE_DIR, docId) const resolvedPath = path.resolve(imageDir) if (!resolvedPath.startsWith(IMAGE_BASE_DIR)) { - console.warn( - `[JAF Provider] Rejecting image path outside base dir: ${imageDir}`, + Logger.warn( + { imageDir, imageName, resolvedPath }, + "Rejecting image path outside base dir", ) loadStats.failed++ return null @@ -127,8 +139,13 @@ const buildLanguageModelImageParts = async ( const imageBytes = await fs.promises.readFile(absolutePath) if (imageBytes.length > MAX_IMAGE_BYTES) { - console.warn( - `[JAF Provider] Skipping image ${absolutePath} due to size ${(imageBytes.length / (1024 * 1024)).toFixed(2)}MB`, + Logger.debug( + { + absolutePath, + imageName, + sizeMb: (imageBytes.length / (1024 * 1024)).toFixed(2), + }, + "Skipping oversized image", ) loadStats.failed++ return null @@ -137,8 +154,9 @@ const buildLanguageModelImageParts = async ( const extension = path.extname(absolutePath).toLowerCase() const mediaType = MIME_TYPE_MAP[extension] if (!mediaType) { - console.warn( - `[JAF Provider] Unsupported image format ${extension} for ${absolutePath}`, + Logger.debug( + { absolutePath, extension, imageName }, + "Unsupported image format for prompt attachment", ) loadStats.failed++ return null @@ -157,10 +175,12 @@ const buildLanguageModelImageParts = async ( }, } } catch (error) { - console.warn( - `[JAF Provider] Failed to load image ${imageName}: ${ - error instanceof Error ? error.message : error - }`, + Logger.debug( + { + err: error, + imageName, + }, + "Failed to load image for prompt attachment", ) loadStats.failed++ return null @@ -169,16 +189,16 @@ const buildLanguageModelImageParts = async ( ) const filtered = parts.filter( - (part): part is ImagePromptPart => part !== null + (part): part is ImagePromptPart => part !== null, ) - + // console.debug('[IMAGE addition][JAF Provider] Image loading complete:', { // requested: imageFileNames.length, // loaded: loadStats.success, // failed: loadStats.failed, // totalMB: (loadStats.totalBytes / (1024 * 1024)).toFixed(2), // }) - + return filtered } @@ -194,13 +214,13 @@ export const makeXyneJAFProvider = ( const modelConfig = MODEL_CONFIGURATIONS[model as Models] const actualModelId = modelConfig?.actualName ?? model - + // Check if this is a LiteLLM model - use OpenAI client directly like JAF does const providerType = ModelToProviderMap[model as Models] const runContext = state.context as unknown as AgentRunContext const stopSignal = runContext?.stopSignal ?? runContext?.stopController?.signal - + if (providerType === AIProviders.LiteLLM) { // Use OpenAI client directly for LiteLLM (same as JAF's makeLiteLLMProvider) const { LiteLLMBaseUrl, LiteLLMApiKey } = config @@ -221,8 +241,9 @@ export const makeXyneJAFProvider = ( }) // Build messages in OpenAI format - const messages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = [] - + const messages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = + [] + // Add system message messages.push({ role: "system", @@ -254,10 +275,13 @@ export const makeXyneJAFProvider = ( messages.push({ role: "assistant", content: text || null, - ...(toolCalls && toolCalls.length > 0 ? { tool_calls: toolCalls } : {}), + ...(toolCalls && toolCalls.length > 0 + ? { tool_calls: toolCalls } + : {}), }) } else if (message.role === "tool") { - const toolCallId = (message as { tool_call_id?: string }).tool_call_id + const toolCallId = (message as { tool_call_id?: string }) + .tool_call_id const content = getTextContent(message.content) || "" if (toolCallId) { messages.push({ @@ -291,19 +315,20 @@ export const makeXyneJAFProvider = ( )?.advancedConfig?.run // Create completion request - const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = { - model: actualModelId, - messages: messages, - temperature: agent.modelConfig?.temperature, - max_tokens: agent.modelConfig?.maxTokens, - ...(tools.length > 0 ? { tools } : {}), - ...(tools.length > 0 && advRun?.toolChoice - ? { tool_choice: advRun.toolChoice } - : {}), - ...(tools.length > 0 && advRun?.parallelToolCalls !== undefined - ? { parallel_tool_calls: advRun.parallelToolCalls } - : {}), - } + const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = + { + model: actualModelId, + messages: messages, + temperature: agent.modelConfig?.temperature, + max_tokens: agent.modelConfig?.maxTokens, + ...(tools.length > 0 ? { tools } : {}), + ...(tools.length > 0 && advRun?.toolChoice + ? { tool_choice: advRun.toolChoice } + : {}), + ...(tools.length > 0 && advRun?.parallelToolCalls !== undefined + ? { parallel_tool_calls: advRun.parallelToolCalls } + : {}), + } if (agent.outputCodec) { params.response_format = { type: "json_object" } @@ -318,11 +343,14 @@ export const makeXyneJAFProvider = ( // Return in JAF format (same as JAF's makeLiteLLMProvider) const choice = resp.choices[0] const message = choice?.message - + const toolCalls = message?.tool_calls ?.filter((tc) => tc.type === "function" && "function" in tc) ?.map((tc) => { - const toolCall = tc as Extract + const toolCall = tc as Extract< + typeof tc, + { type: "function"; function: any } + > return { id: toolCall.id || "", type: "function" as const, @@ -335,11 +363,13 @@ export const makeXyneJAFProvider = ( }, } }) - + return { message: { content: message?.content || null, - ...(toolCalls && toolCalls.length > 0 ? { tool_calls: toolCalls } : {}), + ...(toolCalls && toolCalls.length > 0 + ? { tool_calls: toolCalls } + : {}), }, } } @@ -348,6 +378,9 @@ export const makeXyneJAFProvider = ( const provider = getAISDKProviderByModel(model as Models) const languageModel = provider.languageModel(actualModelId) const pendingReview = runContext?.review?.pendingReview + const userEmail = runContext?.user?.email || "unknown" + const turnNumber = normalizeTurnNumber(runContext?.turnCount) + const providerLogger = loggerWithChild({ email: userEmail }) if (pendingReview) { try { await pendingReview @@ -358,7 +391,7 @@ export const makeXyneJAFProvider = ( turn: normalizeTurnNumber(runContext?.turnCount), error: error instanceof Error ? error.message : String(error), }, - "[JAF Provider] Pending review promise rejected; continuing LLM call" + "[JAF Provider] Pending review promise rejected; continuing LLM call", ) } throwIfStopRequested(stopSignal) @@ -387,17 +420,19 @@ export const makeXyneJAFProvider = ( email: runContext?.user?.email, turn: normalizeTurnNumber(runContext?.turnCount), currentTurnImages: runContext?.currentTurnArtifacts?.images?.map( - (img) => img.fileName + (img) => img.fileName, + ), + recentWindowImages: runContext?.recentImages?.map( + (img) => img.fileName, ), - recentWindowImages: runContext?.recentImages?.map((img) => img.fileName), selectedImages, }, - "[JAF Provider] Prepared image attachments for agent call" + "[JAF Provider] Prepared image attachments for agent call", ) - + if (selectedImages.length > 0) { const lastUserIndex = findLastUserMessageIndex(prompt) - + if (lastUserIndex !== -1) { const imageParts = await buildLanguageModelImageParts(selectedImages) if (imageParts.length > 0) { @@ -406,7 +441,10 @@ export const makeXyneJAFProvider = ( const userContent = userMessage.content const contentBefore = userContent.length for (const { label, filePart } of imageParts) { - userContent.push({ type: "text", text: label } as LanguageModelV2TextPart) + userContent.push({ + type: "text", + text: label, + } as LanguageModelV2TextPart) userContent.push(filePart) } // console.debug('[IMAGE addition][JAF Provider] Attached images to prompt:', { @@ -418,16 +456,26 @@ export const makeXyneJAFProvider = ( // userEmail: runContext.user?.email, // }) } else { - console.warn( - "[JAF Provider] Expected last user index to point to a user message but found", - userMessage?.role ?? "unknown", + providerLogger.warn( + { + messageRole: userMessage?.role ?? "unknown", + selectedImagesCount: selectedImages.length, + turn: turnNumber, + }, + "Expected last user index to resolve to a user message", ) } } else { - console.warn('[JAF Provider] No valid image parts built despite', selectedImages.length, 'selected') + providerLogger.debug( + { selectedImagesCount: selectedImages.length, turn: turnNumber }, + "No valid image parts built for selected images", + ) } } else { - console.warn('[JAF Provider] No user message found to attach', selectedImages.length, 'images to') + providerLogger.debug( + { selectedImagesCount: selectedImages.length, turn: turnNumber }, + "No user message found to attach selected images to", + ) } } @@ -462,31 +510,34 @@ export const makeXyneJAFProvider = ( } // Log the complete prompt and call options being sent to the LLM - const userEmail = runContext?.user?.email || "unknown" - const turnNumber = normalizeTurnNumber(runContext?.turnCount) - - Logger.debug({ - email: userEmail, - turn: turnNumber, - model: actualModelId, - agentName: agent.name, - messagesCount: callOptions.prompt.length, - toolsCount: tools.length, - maxOutputTokens: callOptions.maxOutputTokens, - temperature: callOptions.temperature, - hasResponseFormat: !!callOptions.responseFormat, - toolChoice: callOptions.toolChoice, - }, "[JAF Provider] LLM call parameters") + Logger.debug( + { + email: userEmail, + turn: turnNumber, + model: actualModelId, + agentName: agent.name, + messagesCount: callOptions.prompt.length, + toolsCount: tools.length, + maxOutputTokens: callOptions.maxOutputTokens, + temperature: callOptions.temperature, + hasResponseFormat: !!callOptions.responseFormat, + toolChoice: callOptions.toolChoice, + }, + "[JAF Provider] LLM call parameters", + ) // Sanitize prompt to avoid logging large file data buffers const sanitizedPrompt = sanitizePromptForLogging(callOptions.prompt) - - Logger.debug({ - email: userEmail, - turn: turnNumber, - model: actualModelId, - prompt: sanitizedPrompt, - }, "[JAF Provider] FULL PROMPT/MESSAGES ARRAY SENT TO LLM") + + Logger.debug( + { + email: userEmail, + turn: turnNumber, + model: actualModelId, + prompt: sanitizedPrompt, + }, + "[JAF Provider] FULL PROMPT/MESSAGES ARRAY SENT TO LLM", + ) throwIfStopRequested(stopSignal) const result = await raceWithStop( @@ -514,7 +565,7 @@ export const makeXyneJAFProvider = ( usage: result.usage, contentSummary, }, - "[JAF Provider] Raw LLM response" + "[JAF Provider] Raw LLM response", ) const message = convertResultToJAFMessage(result.content) @@ -586,9 +637,12 @@ const buildFunctionTools = ( const convertedSchema = zodSchemaToJsonSchema(zodSchema) as JSONSchema7 inputSchema = ensureObjectSchema(convertedSchema) } catch (error) { - console.warn( - `Failed to convert Zod schema to JSON for tool ${tool.schema.name}:`, - error, + Logger.warn( + { + err: error, + toolName: tool.schema.name, + }, + "Failed to convert Zod schema to JSON for tool", ) // Fallback to empty object schema inputSchema = { type: "object", properties: {} } @@ -802,12 +856,12 @@ const sanitizePromptForLogging = ( const sanitizedContent = message.content.map((part: any) => { if (part.type === "file") { const filePart = part as LanguageModelV2FilePart - const dataSize = Buffer.isBuffer(filePart.data) - ? filePart.data.length - : filePart.data instanceof Uint8Array + const dataSize = Buffer.isBuffer(filePart.data) ? filePart.data.length - : 0 - + : filePart.data instanceof Uint8Array + ? filePart.data.length + : 0 + return { type: "file", filename: filePart.filename, diff --git a/server/api/chat/message-agents.ts b/server/api/chat/message-agents.ts index 8ddde8c57..34b9d81cc 100644 --- a/server/api/chat/message-agents.ts +++ b/server/api/chat/message-agents.ts @@ -1,6 +1,6 @@ /** * MessageAgents - JAF-Based Agentic Architecture Implementation - * + * * New agentic flow with: * - Single agent with toDoWrite for planning * - Automatic turn-end review @@ -9,102 +9,32 @@ * - Task-based sequential execution */ -import type { Context } from "hono" -import { streamSSE } from "hono/streaming" -import { HTTPException } from "hono/http-exception" -import type { - AgentRunContext, - PlanState, - ToolExecutionRecord, - ReviewResult, - AutoReviewInput, - ToolFailureInfo, - ToolExpectation, - ToolExpectationAssignment, - FinalSynthesisState, - SubTask, - MCPVirtualAgentRuntime, - MCPToolDefinition, - CurrentTurnArtifacts, - ToolExecutionRecordWithResult, -} from "./agent-schemas" -import { ToolExpectationSchema, ReviewResultSchema } from "./agent-schemas" -import { getTracer, type Span } from "@/tracer" -import { getLogger, getLoggerWithChild } from "@/logger" -import { MessageRole, Subsystem, type UserMetadataType } from "@/types" -import config from "@/config" -import { Models, type ModelParams } from "@/ai/types" -import { getModelValueFromLabel } from "@/ai/modelConfig" import { - runStream, - generateRunId, - generateTraceId, - getTextContent, - ToolResponse, - ToolErrorCodes, - type Agent as JAFAgent, - type Message as JAFMessage, - type RunConfig as JAFRunConfig, - type RunState as JAFRunState, - type Tool, - type ToolCall, - type ToolResult, - type TraceEvent, -} from "@xynehq/jaf" -import { makeXyneJAFProvider } from "./jaf-provider" -import { getErrorMessage } from "@/utils" -import { ChatSSEvents, AgentReasoningStepType } from "@/shared/types" -import type { - MinimalAgentFragment, - Citation, - ImageCitation, - FragmentImageReference, -} from "./types" + answerContextMap, + answerContextMapFromFragments, + userContext, +} from "@/ai/context" +import { getModelValueFromLabel } from "@/ai/modelConfig" +import { extractBestDocumentsPrompt } from "@/ai/prompts" import { extractBestDocumentIndexes, getProviderByModel, jsonParseLLMOutput, } from "@/ai/provider" -import { extractBestDocumentsPrompt } from "@/ai/prompts" -import { answerContextMap, answerContextMapFromFragments, userContext } from "@/ai/context" -import { ConversationRole } from "@aws-sdk/client-bedrock-runtime" -import type { Message } from "@aws-sdk/client-bedrock-runtime" -import { - TOOL_SCHEMAS, - generateToolDescriptions, - validateToolInput, - ListCustomAgentsOutputSchema, - type ListCustomAgentsOutput, - type ToolOutput, - type ResourceAccessSummary, -} from "./tool-schemas" -import { - searchToCitation, - extractImageFileNames, - processMessage, - checkAndYieldCitationsForAgent, - extractFileIdsFromMessage, - isMessageWithContext, - processThreadResults, - safeDecodeURIComponent, -} from "./utils" -import { searchCollectionRAG, SearchEmailThreads, searchVespaInFiles } from "@/search/vespa" +import { type ModelParams, Models } from "@/ai/types" +import { executeAgentForWorkflowWithRag } from "@/api/agent/workflowAgentUtils" +import config from "@/config" import { - Apps, - AttachmentEntity, - KnowledgeBaseEntity, - SearchModes, - type VespaSearchResult, - type VespaSearchResults, -} from "@xyne/vespa-ts/types" -import { expandSheetIds } from "@/search/utils" -import type { ZodTypeAny } from "zod" -import { parseAttachmentMetadata } from "@/utils/parseAttachment" -import { db } from "@/db/client" + type SelectAgent, + getAgentByExternalIdWithPermissionCheck, +} from "@/db/agent" +import { storeAttachmentMetadata } from "@/db/attachment" import { insertChat, updateChatByExternalIdWithAuth } from "@/db/chat" import { insertChatTrace } from "@/db/chatTrace" +import { db } from "@/db/client" +import { getConnectorById } from "@/db/connector" import { getChatMessagesWithAuth, insertMessage } from "@/db/message" -import { storeAttachmentMetadata } from "@/db/attachment" +import { getUserPersonalizationByEmail } from "@/db/personalization" import { ChatType, type InsertChat, @@ -113,51 +43,90 @@ import { type SelectChat, type SelectMessage, } from "@/db/schema" +import { getToolsByConnectorId } from "@/db/tool" import { getUserAndWorkspaceByEmail } from "@/db/user" import { getUserAccessibleAgents } from "@/db/userAgentPermission" -import { executeAgentForWorkflowWithRag } from "@/api/agent/workflowAgentUtils" -import { getDateForAI } from "@/utils/index" -import googleTools from "./tools/google" -import { searchGlobalTool, fallbackTool } from "./tools/global" -import { searchKnowledgeBaseTool } from "./tools/knowledgeBase" -import { getSlackRelatedMessagesTool } from "./tools/slack/getSlackMessages" -import type { AttachmentMetadata } from "@/shared/types" -import { - evaluateAgentResourceAccess, - getUserConnectorState, - createEmptyConnectorState, - type UserConnectorState, -} from "./resource-access" +import { getPrecomputedDbContextIfNeeded } from "@/lib/databaseContext" +import { getLogger, getLoggerWithChild } from "@/logger" +import { expandSheetIds } from "@/search/utils" import { - getAgentByExternalIdWithPermissionCheck, - type SelectAgent, -} from "@/db/agent" -import { isCuid } from "@paralleldrive/cuid2" + SearchEmailThreads, + searchCollectionRAG, + searchVespaInFiles, +} from "@/search/vespa" +import { AgentReasoningStepType, ChatSSEvents } from "@/shared/types" +import type { AttachmentMetadata } from "@/shared/types" import { DEFAULT_TEST_AGENT_ID } from "@/shared/types" -import { parseAgentAppIntegrations } from "./tools/utils" -import { buildAgentPromptAddendum } from "./agentPromptCreation" -import { getConnectorById } from "@/db/connector" -import { getToolsByConnectorId } from "@/db/tool" -import { - buildMCPJAFTools, - type FinalToolsList, -} from "./jaf-adapter" -import { activeStreams } from "./stream" +import { type Span, getTracer } from "@/tracer" +import { MessageRole, Subsystem, type UserMetadataType } from "@/types" +import { getErrorMessage } from "@/utils" +import { getDateForAI } from "@/utils/index" +import { parseAttachmentMetadata } from "@/utils/parseAttachment" +import { ConversationRole } from "@aws-sdk/client-bedrock-runtime" +import type { Message } from "@aws-sdk/client-bedrock-runtime" import { Client } from "@modelcontextprotocol/sdk/client/index.js" -import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" import { SSEClientTransport, type SSEClientTransportOptions, } from "@modelcontextprotocol/sdk/client/sse.js" +import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" import { StreamableHTTPClientTransport, type StreamableHTTPClientTransportOptions, } from "@modelcontextprotocol/sdk/client/streamableHttp.js" +import { isCuid } from "@paralleldrive/cuid2" +import { + Apps, + AttachmentEntity, + KnowledgeBaseEntity, + SearchModes, + type VespaSearchResult, + type VespaSearchResults, +} from "@xyne/vespa-ts/types" +import { + type Agent as JAFAgent, + type Message as JAFMessage, + type RunConfig as JAFRunConfig, + type RunState as JAFRunState, + type Tool, + type ToolCall, + ToolErrorCodes, + ToolResponse, + type ToolResult, + type TraceEvent, + generateRunId, + generateTraceId, + getTextContent, + runStream, +} from "@xynehq/jaf" +import type { Context } from "hono" +import { HTTPException } from "hono/http-exception" +import { streamSSE } from "hono/streaming" +import type { ZodTypeAny } from "zod" +import type { + AgentRunContext, + AutoReviewInput, + CurrentTurnArtifacts, + FinalSynthesisState, + MCPToolDefinition, + MCPVirtualAgentRuntime, + PlanState, + ReviewResult, + SubTask, + ToolExecutionRecord, + ToolExecutionRecordWithResult, + ToolExpectation, + ToolExpectationAssignment, + ToolFailureInfo, +} from "./agent-schemas" +import { ReviewResultSchema, ToolExpectationSchema } from "./agent-schemas" import { isMessageAgentStopError, throwIfStopRequested } from "./agent-stop" +import { buildAgentPromptAddendum } from "./agentPromptCreation" import { parseMessageText } from "./chat" -import { getUserPersonalizationByEmail } from "@/db/personalization" import { getChunkCountPerDoc } from "./chunk-selection" -import { getPrecomputedDbContextIfNeeded } from "@/lib/databaseContext" +import { type FinalToolsList, buildMCPJAFTools } from "./jaf-adapter" +import { logJAFTraceEvent } from "./jaf-logging" +import { makeXyneJAFProvider } from "./jaf-provider" import { buildAgentSystemPromptContextBlock, enforceMetadataConstraintsOnSelection, @@ -168,6 +137,46 @@ import { sanitizeAgentSystemPromptSnapshot, withAgentSystemPromptMessage, } from "./message-agents-metadata" +import { + type UserConnectorState, + createEmptyConnectorState, + evaluateAgentResourceAccess, + getUserConnectorState, +} from "./resource-access" +import { activeStreams } from "./stream" +import { + type ListCustomAgentsOutput, + ListCustomAgentsOutputSchema, + type ResourceAccessSummary, + TOOL_SCHEMAS, + type ToolOutput, + generateToolDescriptions, + validateToolInput, +} from "./tool-schemas" +import { fallbackTool, searchGlobalTool } from "./tools/global" +import googleTools from "./tools/google" +import { + lsKnowledgeBaseTool, + searchKnowledgeBaseTool, +} from "./tools/knowledgeBaseFlow" +import { getSlackRelatedMessagesTool } from "./tools/slack/getSlackMessages" +import { parseAgentAppIntegrations } from "./tools/utils" +import type { + Citation, + FragmentImageReference, + ImageCitation, + MinimalAgentFragment, +} from "./types" +import { + checkAndYieldCitationsForAgent, + extractFileIdsFromMessage, + extractImageFileNames, + isMessageWithContext, + processMessage, + processThreadResults, + safeDecodeURIComponent, + searchToCitation, +} from "./utils" export { __messageAgentsMetadataInternals } from "./message-agents-metadata" @@ -204,7 +213,7 @@ function normalizeReviewFrequency(value: unknown): number { } const mutableAgentContext = ( - context: Readonly + context: Readonly, ): AgentRunContext => context as AgentRunContext const createEmptyTurnArtifacts = (): CurrentTurnArtifacts => ({ @@ -214,15 +223,13 @@ const createEmptyTurnArtifacts = (): CurrentTurnArtifacts => ({ images: [], }) - const reviewsAllowed = (context: AgentRunContext): boolean => !context.review.lockedByFinalSynthesis -function resolveAgenticModelId( - requestedModelId?: string | Models -): Models { +function resolveAgenticModelId(requestedModelId?: string | Models): Models { const hasAgenticOverride = - defaultBestModelAgenticMode && defaultBestModelAgenticMode !== ("" as Models) + defaultBestModelAgenticMode && + defaultBestModelAgenticMode !== ("" as Models) const fallback = hasAgenticOverride ? (defaultBestModelAgenticMode as Models) : (defaultBestModel as Models) @@ -231,12 +238,12 @@ function resolveAgenticModelId( } const toToolParameters = ( - schema: ZodTypeAny + schema: ZodTypeAny, ): Tool["schema"]["parameters"] => schema as unknown as Tool["schema"]["parameters"] function fragmentsToToolContexts( - fragments: MinimalAgentFragment[] | undefined + fragments: MinimalAgentFragment[] | undefined, ): ToolOutput["contexts"] { if (!fragments?.length) { return undefined @@ -286,34 +293,40 @@ async function streamReasoningStep( status?: string detail?: string [key: string]: unknown - } + }, ): Promise { if (!emitter) return try { // Build the step object to match agents.ts structure const step: Record = {} - + // Set type - infer from context if not provided if (extra?.type) { step.type = extra.type - } else if (extra?.iteration !== undefined && text.toLowerCase().includes('turn') && text.toLowerCase().includes('started')) { + } else if ( + extra?.iteration !== undefined && + text.toLowerCase().includes("turn") && + text.toLowerCase().includes("started") + ) { step.type = AgentReasoningStepType.Iteration } else { step.type = AgentReasoningStepType.LogMessage } - + if (extra?.iteration !== undefined) step.iteration = extra.iteration if (extra?.toolName !== undefined) step.toolName = extra.toolName if (extra?.status !== undefined) step.status = extra.status if (extra?.detail !== undefined) step.detail = extra.detail - + // Include any other extra properties in step - Object.keys(extra || {}).forEach(key => { - if (!['type', 'iteration', 'toolName', 'status', 'detail'].includes(key)) { + Object.keys(extra || {}).forEach((key) => { + if ( + !["type", "iteration", "toolName", "status", "detail"].includes(key) + ) { step[key] = (extra as any)[key] } }) - + await emitter({ text, step: Object.keys(step).length > 0 ? step : undefined, @@ -322,7 +335,7 @@ async function streamReasoningStep( } catch (error) { Logger.warn( { err: error instanceof Error ? error.message : String(error) }, - "Failed to stream reasoning step" + "Failed to stream reasoning step", ) } } @@ -368,9 +381,7 @@ function normalizeUserMessageForHistory(message: SelectMessage): string { } } -function buildConversationHistoryForAgentRun( - history: SelectMessage[], -): { +function buildConversationHistoryForAgentRun(history: SelectMessage[]): { jafHistory: JAFMessage[] llmHistory: Message[] } { @@ -409,9 +420,71 @@ export const __messageAgentsHistoryInternals = { const RECENT_IMAGE_WINDOW = 2 +function normalizeExcludedIdsForLogging(excludedIds: unknown): string[] { + if (Array.isArray(excludedIds)) { + return excludedIds + .map((value) => + typeof value === "string" + ? value + : value === null || value === undefined + ? "" + : String(value), + ) + .filter(Boolean) + } + if (excludedIds === null || excludedIds === undefined) { + return [] + } + const normalized = + typeof excludedIds === "string" ? excludedIds : String(excludedIds) + return normalized ? [normalized] : [] +} + +function buildContextTraceSnapshot( + context: AgentRunContext, +): Record { + return { + chatId: context.chat.externalId, + turnCount: context.turnCount, + currentSubTask: context.currentSubTask, + seenDocumentsCount: context.seenDocuments.size, + seenDocumentsSample: Array.from(context.seenDocuments).slice(0, 10), + allFragmentsCount: context.allFragments.length, + allImagesCount: context.allImages.length, + recentImagesCount: context.recentImages.length, + currentTurnFragmentCount: context.currentTurnArtifacts.fragments.length, + currentTurnImageCount: context.currentTurnArtifacts.images.length, + currentTurnToolOutputCount: context.currentTurnArtifacts.toolOutputs.length, + currentTurnExpectationCount: + context.currentTurnArtifacts.expectations.length, + toolCallHistoryCount: context.toolCallHistory.length, + failedToolsCount: context.failedTools.size, + availableAgentsCount: context.availableAgents.length, + usedAgentsCount: context.usedAgents.length, + ambiguityResolved: context.ambiguityResolved, + finalSynthesisRequested: context.finalSynthesis.requested, + finalSynthesisCompleted: context.finalSynthesis.completed, + finalSynthesisAckReceived: context.finalSynthesis.ackReceived, + } +} + +function logContextMutation( + context: AgentRunContext, + message: string, + details: Record = {}, +): void { + loggerWithChild({ email: context.user.email }).info( + { + ...buildContextTraceSnapshot(context), + ...details, + }, + message, + ) +} + function mergeFragmentLists( target: MinimalAgentFragment[], - incoming: MinimalAgentFragment[] + incoming: MinimalAgentFragment[], ): MinimalAgentFragment[] { if (!incoming.length) { return target @@ -434,7 +507,7 @@ function mergeFragmentLists( function mergeImageReferences( target: FragmentImageReference[], - incoming: FragmentImageReference[] + incoming: FragmentImageReference[], ): FragmentImageReference[] { if (!incoming.length) { return target @@ -456,11 +529,12 @@ function mergeImageReferences( } function extractImagesFromFragments( - fragments: MinimalAgentFragment[] + fragments: MinimalAgentFragment[], ): FragmentImageReference[] { const references: FragmentImageReference[] = [] for (const fragment of fragments) { - if (!Array.isArray(fragment.images) || fragment.images.length === 0) continue + if (!Array.isArray(fragment.images) || fragment.images.length === 0) + continue for (const image of fragment.images) { if (image?.fileName) { references.push(image) @@ -473,52 +547,84 @@ function extractImagesFromFragments( function recordFragmentsForContext( context: AgentRunContext, fragments: MinimalAgentFragment[], - turnNumber: number + turnNumber: number, ): void { if (!fragments.length) return - fragments.forEach((fragment) => context.seenDocuments.add(fragment.id)) + const seenDocumentsBefore = context.seenDocuments.size + const addedSeenDocumentIds: string[] = [] + fragments.forEach((fragment) => { + const docId = fragment.source?.docId + if (docId) { + if (!context.seenDocuments.has(docId)) { + addedSeenDocumentIds.push(docId) + } + context.seenDocuments.add(docId) + } + }) context.currentTurnArtifacts.fragments = mergeFragmentLists( context.currentTurnArtifacts.fragments, - fragments + fragments, ) context.allFragments = mergeFragmentLists(context.allFragments, fragments) const existingForTurn = context.turnFragments.get(turnNumber) ?? [] context.turnFragments.set( turnNumber, - mergeFragmentLists(existingForTurn, fragments) + mergeFragmentLists(existingForTurn, fragments), + ) + logContextMutation( + context, + "[MessageAgents][Context] Recorded fragments and updated seenDocuments", + { + turnNumber, + fragmentCount: fragments.length, + fragmentIds: fragments.map((fragment) => fragment.id), + addedSeenDocumentIds, + seenDocumentsBefore, + seenDocumentsAfter: context.seenDocuments.size, + turnFragmentCount: (context.turnFragments.get(turnNumber) || []).length, + }, ) const fragmentImages = extractImagesFromFragments(fragments) if (fragmentImages.length > 0) { context.currentTurnArtifacts.images = mergeImageReferences( context.currentTurnArtifacts.images, - fragmentImages + fragmentImages, ) - loggerWithChild({ email: context.user.email }).debug( + logContextMutation( + context, + "[MessageAgents][Context] Updated current turn images from fragments", { - chatId: context.chat.externalId, turnNumber, fragmentCount: fragments.length, imageNames: fragmentImages.map((img) => img.fileName), + addedImageCount: fragmentImages.length, }, - "[MessageAgents] Recorded new fragment images" ) } } function resetCurrentTurnArtifacts(context: AgentRunContext): void { + const previousArtifacts = context.currentTurnArtifacts context.currentTurnArtifacts = createEmptyTurnArtifacts() + logContextMutation( + context, + "[MessageAgents][Context] Reset current turn artifacts", + { + clearedFragmentCount: previousArtifacts.fragments.length, + clearedImageCount: previousArtifacts.images.length, + clearedExpectationCount: previousArtifacts.expectations.length, + clearedToolOutputCount: previousArtifacts.toolOutputs.length, + }, + ) } function finalizeTurnImages( context: AgentRunContext, - turnNumber: number + turnNumber: number, ): void { const imagesForTurn = context.currentTurnArtifacts.images context.imagesByTurn.set(turnNumber, [...imagesForTurn]) - context.allImages = mergeImageReferences( - context.allImages, - imagesForTurn - ) + context.allImages = mergeImageReferences(context.allImages, imagesForTurn) const recentTurns = Array.from(context.imagesByTurn.keys()) .sort((a, b) => a - b) .slice(-RECENT_IMAGE_WINDOW) @@ -531,16 +637,16 @@ function finalizeTurnImages( } context.recentImages = mergeImageReferences([], flattened) context.currentTurnArtifacts.images = [] - loggerWithChild({ email: context.user.email }).debug( + logContextMutation( + context, + "[MessageAgents][Context] Finalized turn images", { - chatId: context.chat.externalId, turnNumber, committedImages: imagesForTurn.map((img) => img.fileName), recentWindow: recentTurns, recentImages: context.recentImages.map((img) => img.fileName), allImagesCount: context.allImages.length, }, - "[MessageAgents] Updated turn image buffers" ) } @@ -568,7 +674,7 @@ function summarizeToolResultPayload(result: any): string { } function formatToolArgumentsForReasoning( - args: Record + args: Record, ): string { if (!args || typeof args !== "object") { return "{}" @@ -602,7 +708,7 @@ function formatToolArgumentsForReasoning( function buildTurnToolReasoningSummary( turnNumber: number, - records: ToolExecutionRecord[] + records: ToolExecutionRecord[], ): string { const lines = records.map((record, idx) => { const argsSummary = formatToolArgumentsForReasoning(record.arguments) @@ -620,7 +726,7 @@ type FragmentImageOptions = { function attachImagesToFragments( fragments: MinimalAgentFragment[], imageNames: string[], - options: FragmentImageOptions + options: FragmentImageOptions, ): MinimalAgentFragment[] { if (!Array.isArray(imageNames) || imageNames.length === 0) { return fragments @@ -635,7 +741,7 @@ function attachImagesToFragments( const fragmentId = getFragmentIdFromImageName( imageName, fragmentIndexMap, - fragments[0]?.id || "" + fragments[0]?.id || "", ) if (!fragmentId) continue const ref: FragmentImageReference = { @@ -683,7 +789,7 @@ function attachImagesToFragments( function getFragmentIdFromImageName( imageName: string, fragmentIndexMap: Map, - fallback = "" + fallback = "", ): string { const separatorIdx = imageName.indexOf("_") if (separatorIdx <= 0) return fallback @@ -692,9 +798,7 @@ function getFragmentIdFromImageName( return fragmentIndexMap.get(docIndex) ?? fallback } -function getMetadataLayers( - result: any -): Record[] { +function getMetadataLayers(result: any): Record[] { const layers: Record[] = [] const metadata = result?.metadata if (metadata && typeof metadata === "object") { @@ -709,7 +813,7 @@ function getMetadataLayers( function getMetadataValue( result: any, - key: string + key: string, ): T | undefined { if (result?.data && typeof result.data === "object" && key in result.data) { return (result.data as Record)[key] as T @@ -743,7 +847,9 @@ function formatPlanForPrompt(plan: PlanState | null): string { detailParts.push(`Tools: ${task.toolsRequired.join(", ")}`) } lines.push( - detailParts.length > 0 ? `${baseLine}\n ${detailParts.join(" | ")}` : baseLine + detailParts.length > 0 + ? `${baseLine}\n ${detailParts.join(" | ")}` + : baseLine, ) }) return lines.join("\n") @@ -760,7 +866,7 @@ function selectActiveSubTaskId(plan: PlanState | null): string | null { ] for (const status of priority) { const match = plan.subTasks.find( - (task) => task.status === status && task.id + (task) => task.status === status && task.id, ) if (match?.id) { return match.id @@ -803,10 +909,7 @@ function initializePlanState(plan: PlanState): string | null { activeId = selectActiveSubTaskId(plan) continue } - if ( - activeTask.status === "pending" || - activeTask.status === "blocked" - ) { + if (activeTask.status === "pending" || activeTask.status === "blocked") { activeTask.status = "in_progress" activeTask.error = undefined } @@ -819,11 +922,11 @@ function advancePlanAfterTool( context: AgentRunContext, toolName: string, wasSuccessful: boolean, - detail?: string + detail?: string, ): void { if (!context.plan || !context.currentSubTask) return const task = context.plan.subTasks.find( - (entry) => entry.id === context.currentSubTask + (entry) => entry.id === context.currentSubTask, ) if (!task) return normalizeSubTask(task) @@ -837,14 +940,13 @@ function advancePlanAfterTool( if (requiredTools.length === 0 || requiredTools.includes(toolName)) { task.status = "completed" task.completedAt = Date.now() - task.result = - detail || task.result || `Completed using ${toolName}` + task.result = detail || task.result || `Completed using ${toolName}` const previousTaskId = task.id const nextId = selectActiveSubTaskId(context.plan) if (nextId && nextId !== previousTaskId) { context.currentSubTask = nextId const nextTask = context.plan.subTasks.find( - (entry) => entry.id === nextId + (entry) => entry.id === nextId, ) if (nextTask && nextTask.status === "pending") { nextTask.status = "in_progress" @@ -861,13 +963,13 @@ function advancePlanAfterTool( } function formatClarificationsForPrompt( - clarifications: AgentRunContext["clarifications"] + clarifications: AgentRunContext["clarifications"], ): string { if (!clarifications?.length) return "" const formatted = clarifications .map( (clarification, idx) => - `${idx + 1}. Q: ${clarification.question}\n A: ${clarification.answer}` + `${idx + 1}. Q: ${clarification.question}\n A: ${clarification.answer}`, ) .join("\n") return formatted @@ -875,21 +977,26 @@ function formatClarificationsForPrompt( export function buildFinalSynthesisPayload( context: AgentRunContext, - fragmentsLimit = Math.max(12, context.allFragments.length || 1) + fragmentsLimit = Math.max(12, context.allFragments.length || 1), ): { systemPrompt: string; userMessage: string } { const fragments = context.allFragments const agentSystemPromptBlock = buildAgentSystemPromptContextBlock( - context.dedicatedAgentSystemPrompt + context.dedicatedAgentSystemPrompt, ) const agentSystemPromptSection = agentSystemPromptBlock ? `Agent System Prompt Context:\n${agentSystemPromptBlock}` : "" - const formattedFragments = formatFragmentsWithMetadata(fragments, fragmentsLimit) + const formattedFragments = formatFragmentsWithMetadata( + fragments, + fragmentsLimit, + ) const fragmentsSection = formattedFragments ? `Context Fragments:\n${formattedFragments}` : "" const planSection = formatPlanForPrompt(context.plan) - const clarificationSection = formatClarificationsForPrompt(context.clarifications) + const clarificationSection = formatClarificationsForPrompt( + context.clarifications, + ) const workspaceSection = context.userContext?.trim() ? `Workspace Context:\n${context.userContext}` : "" @@ -898,7 +1005,9 @@ export function buildFinalSynthesisPayload( `User Question:\n${context.message.text}`, agentSystemPromptSection, planSection ? `Execution Plan Snapshot:\n${planSection}` : "", - clarificationSection ? `Clarifications Resolved:\n${clarificationSection}` : "", + clarificationSection + ? `Clarifications Resolved:\n${clarificationSection}` + : "", workspaceSection, fragmentsSection, ].filter(Boolean) @@ -983,9 +1092,7 @@ export function buildFinalSynthesisPayload( return { systemPrompt, userMessage } } -function selectImagesForFinalSynthesis( - context: AgentRunContext -): { +function selectImagesForFinalSynthesis(context: AgentRunContext): { selected: string[] total: number dropped: string[] @@ -1036,7 +1143,7 @@ function selectImagesForFinalSynthesis( function buildAttachmentToolMessage( fragments: MinimalAgentFragment[], - summary: string + summary: string, ): JAFMessage { const resultPayload = ToolResponse.success({ summary, fragments }) const envelope = { @@ -1070,7 +1177,7 @@ function initializeAgentContext( stopController?: AbortController stopSignal?: AbortSignal modelId?: string - } + }, ): AgentRunContext { const finalSynthesis: FinalSynthesisState = { requested: false, @@ -1080,8 +1187,7 @@ function initializeAgentContext( ackReceived: false, } const currentTurnArtifacts = createEmptyTurnArtifacts() - - return { + const context: AgentRunContext = { user: { email: userEmail, workspaceId, @@ -1151,6 +1257,14 @@ function initializeAgentContext( options?.stopSignal?.aborted ?? false, } + logContextMutation(context, "[MessageAgents][Context] Initialized agent context", { + attachmentCount: attachments.length, + attachmentIds: attachments.map((attachment) => attachment.fileId), + hasAgentPrompt: !!options?.agentPrompt, + hasDedicatedAgentSystemPrompt: !!options?.dedicatedAgentSystemPrompt, + modelId: options?.modelId, + }) + return context } /** @@ -1159,7 +1273,7 @@ function initializeAgentContext( */ async function performAutomaticReview( input: AutoReviewInput, - fullContext: AgentRunContext + fullContext: AgentRunContext, ): Promise { const reviewContext: AgentRunContext = { ...fullContext, @@ -1171,7 +1285,7 @@ async function performAutomaticReview( tripReviewSpan.setAttribute("turn_number", input.turnNumber ?? -1) tripReviewSpan.setAttribute( "expected_results_count", - input.expectedResults?.length ?? 0 + input.expectedResults?.length ?? 0, ) let reviewResult: ReviewResult try { @@ -1183,12 +1297,15 @@ async function performAutomaticReview( expectedResults: input.expectedResults, delegationEnabled: fullContext.delegationEnabled, }, - fullContext.modelId + fullContext.modelId, ) } catch (error) { tripReviewSpan.setAttribute("error", true) tripReviewSpan.setAttribute("error_message", getErrorMessage(error)) - Logger.error(error, "Automatic review failed, falling back to default response") + Logger.error( + error, + "Automatic review failed, falling back to default response", + ) reviewResult = { status: "needs_attention", notes: `Automatic review fallback for turn ${input.turnNumber}: ${getErrorMessage(error)}`, @@ -1206,11 +1323,11 @@ async function performAutomaticReview( tripReviewSpan.setAttribute("review_status", reviewResult.status) tripReviewSpan.setAttribute( "recommendation", - reviewResult.recommendation ?? "unknown" + reviewResult.recommendation ?? "unknown", ) tripReviewSpan.setAttribute( "anomalies_detected", - reviewResult.anomaliesDetected ?? false + reviewResult.anomaliesDetected ?? false, ) tripReviewSpan.end() @@ -1222,17 +1339,28 @@ async function handleReviewOutcome( reviewResult: ReviewResult, iteration: number, focus: AutoReviewInput["focus"], - reasoningEmitter?: ReasoningEmitter + reasoningEmitter?: ReasoningEmitter, ): Promise { context.review.lastReviewResult = reviewResult context.review.lastReviewTurn = iteration context.ambiguityResolved = reviewResult.ambiguityResolved - context.review.outstandingAnomalies = - reviewResult.anomalies?.length ? reviewResult.anomalies : [] - context.review.clarificationQuestions = - reviewResult.clarificationQuestions?.length - ? reviewResult.clarificationQuestions - : [] + context.review.outstandingAnomalies = reviewResult.anomalies?.length + ? reviewResult.anomalies + : [] + context.review.clarificationQuestions = reviewResult.clarificationQuestions + ?.length + ? reviewResult.clarificationQuestions + : [] + logContextMutation(context, "[MessageAgents][Context] Review outcome applied", { + iteration, + focus, + recommendation: reviewResult.recommendation, + reviewStatus: reviewResult.status, + ambiguityResolved: reviewResult.ambiguityResolved, + anomaliesDetected: reviewResult.anomaliesDetected, + anomalies: context.review.outstandingAnomalies, + clarificationQuestions: context.review.clarificationQuestions, + }) await streamReasoningStep( reasoningEmitter, @@ -1246,21 +1374,24 @@ async function handleReviewOutcome( ambiguityResolved: reviewResult.ambiguityResolved, anomalies: reviewResult.anomalies, focus, - } + }, ) if ( reviewResult.anomaliesDetected || (reviewResult.anomalies?.length ?? 0) > 0 ) { - Logger.debug({ - turn: iteration, - anomalies: reviewResult.anomalies, - recommendation: reviewResult.recommendation, - planChangeNeeded: reviewResult.planChangeNeeded, - chatId: context.chat.externalId, - focus, - }, "[MessageAgents][Anomalies]") + Logger.debug( + { + turn: iteration, + anomalies: reviewResult.anomalies, + recommendation: reviewResult.recommendation, + planChangeNeeded: reviewResult.planChangeNeeded, + chatId: context.chat.externalId, + focus, + }, + "[MessageAgents][Anomalies]", + ) await streamReasoningStep( reasoningEmitter, `Anomalies detected: ${ @@ -1273,7 +1404,7 @@ async function handleReviewOutcome( ambiguityResolved: reviewResult.ambiguityResolved, anomalies: reviewResult.anomalies, focus, - } + }, ) } } @@ -1284,7 +1415,7 @@ type AttachmentPhaseMetadata = { } function getAttachmentPhaseMetadata( - context: AgentRunContext + context: AgentRunContext, ): AttachmentPhaseMetadata { return (context.chat.metadata as AttachmentPhaseMetadata) || {} } @@ -1309,7 +1440,7 @@ type ChatBootstrapResult = { } async function ensureChatAndPersistUserMessage( - params: ChatBootstrapParams + params: ChatBootstrapParams, ): Promise { const workspaceId = Number(params.workspace.id) const workspaceExternalId = String(params.workspace.externalId) @@ -1350,7 +1481,7 @@ async function ensureChatAndPersistUserMessage( tx, userEmail, String(userMessage.externalId), - params.attachmentMetadata + params.attachmentMetadata, ) if (storageErr) { attachmentError = storageErr @@ -1369,7 +1500,7 @@ async function ensureChatAndPersistUserMessage( tx, String(incomingChatId), String(params.email), - {} + {}, ) const conversationHistory = await getChatMessagesWithAuth( tx, @@ -1396,7 +1527,7 @@ async function ensureChatAndPersistUserMessage( tx, userEmail, String(userMessage.externalId), - params.attachmentMetadata + params.attachmentMetadata, ) if (storageErr) { attachmentError = storageErr @@ -1416,7 +1547,7 @@ async function storeAttachmentSafely( tx: Parameters[0], email: string, messageExternalId: string, - attachments: AttachmentMetadata[] + attachments: AttachmentMetadata[], ): Promise { try { await storeAttachmentMetadata(tx, messageExternalId, attachments, email) @@ -1424,7 +1555,7 @@ async function storeAttachmentSafely( } catch (error) { loggerWithChild({ email }).error( error, - `Failed to store attachment metadata for message ${messageExternalId}` + `Failed to store attachment metadata for message ${messageExternalId}`, ) return error as Error } @@ -1465,7 +1596,7 @@ async function prepareInitialAttachmentContext( userMetadata: UserMetadataType, query: string, email: string, - allowChunkCitations?: boolean + allowChunkCitations?: boolean, ): Promise<{ fragments: MinimalAgentFragment[]; summary: string } | null> { if (!fileIds?.length) { return null @@ -1490,112 +1621,112 @@ async function prepareInitialAttachmentContext( const span = tracer.startSpan("prepare_initial_attachment_context") try { - const combinedSearchResponse: VespaSearchResult[] = [] - let chunksPerDocument: number[] = [] - const targetChunks = 120 - - if (fileIds && fileIds.length > 0) { - const fileSearchSpan = span.startSpan("file_search") - let results - // Split into 3 groups - // Search each group - // Push results to combinedSearchResponse - const collectionFileIds = fileIds.filter( - (fid) => fid.startsWith("clf-") || fid.startsWith("att_"), - ) - const nonCollectionFileIds = fileIds.filter( - (fid) => !fid.startsWith("clf-") && !fid.startsWith("att"), + const combinedSearchResponse: VespaSearchResult[] = [] + let chunksPerDocument: number[] = [] + const targetChunks = 120 + + if (fileIds && fileIds.length > 0) { + const fileSearchSpan = span.startSpan("file_search") + let results + // Split into 3 groups + // Search each group + // Push results to combinedSearchResponse + const collectionFileIds = fileIds.filter( + (fid) => fid.startsWith("clf-") || fid.startsWith("att_"), + ) + const nonCollectionFileIds = fileIds.filter( + (fid) => !fid.startsWith("clf-") && !fid.startsWith("att"), + ) + const attachmentFileIds = fileIds.filter((fid) => fid.startsWith("attf_")) + if (nonCollectionFileIds && nonCollectionFileIds.length > 0) { + results = await searchVespaInFiles( + queryText, + email, + nonCollectionFileIds, + { + limit: fileIds?.length, + alpha: userAlpha, + }, ) - const attachmentFileIds = fileIds.filter((fid) => fid.startsWith("attf_")) - if (nonCollectionFileIds && nonCollectionFileIds.length > 0) { - results = await searchVespaInFiles( - queryText, - email, - nonCollectionFileIds, - { - limit: fileIds?.length, - alpha: userAlpha, - }, - ) - if (results.root.children) { - combinedSearchResponse.push(...results.root.children) - } + if (results.root.children) { + combinedSearchResponse.push(...results.root.children) } - if (collectionFileIds && collectionFileIds.length > 0) { - allowChunkCitations = true // for the case where kb files are in @ - results = await searchCollectionRAG( - queryText, - collectionFileIds, - undefined, - ) - if (results.root.children) { - combinedSearchResponse.push(...results.root.children) - } - } - if (attachmentFileIds && attachmentFileIds.length > 0) { - results = await searchVespaInFiles( - queryText, - email, - attachmentFileIds, - { - limit: fileIds?.length, - alpha: userAlpha, - rankProfile: SearchModes.AttachmentRank, - }, - ) - if (results.root.children) { - combinedSearchResponse.push(...results.root.children) - } + } + if (collectionFileIds && collectionFileIds.length > 0) { + allowChunkCitations = true // for the case where kb files are in @ + results = await searchCollectionRAG( + queryText, + collectionFileIds, + undefined, + ) + if (results.root.children) { + combinedSearchResponse.push(...results.root.children) } - - // Apply intelligent chunk selection based on document relevance and chunk scores - chunksPerDocument = await getChunkCountPerDoc( - combinedSearchResponse, - targetChunks, + } + if (attachmentFileIds && attachmentFileIds.length > 0) { + results = await searchVespaInFiles( + queryText, email, - fileSearchSpan, + attachmentFileIds, + { + limit: fileIds?.length, + alpha: userAlpha, + rankProfile: SearchModes.AttachmentRank, + }, ) - fileSearchSpan?.end() + if (results.root.children) { + combinedSearchResponse.push(...results.root.children) + } } - - if (threadIds && threadIds.length > 0) { - const threadSpan = span.startSpan("fetch_email_threads") - threadSpan.setAttribute("threadIds", JSON.stringify(threadIds)) - - try { - const threadResults = await SearchEmailThreads(threadIds, email) - if ( - threadResults.root.children && - threadResults.root.children.length > 0 - ) { - const existingDocIds = new Set( - combinedSearchResponse.map((child: any) => child.fields.docId), - ) - - // Use the helper function to process thread results - const { addedCount, threadInfo } = processThreadResults( - threadResults.root.children, - existingDocIds, - combinedSearchResponse, - ) - threadSpan.setAttribute("added_email_count", addedCount) - threadSpan.setAttribute( - "total_thread_emails_found", - threadResults.root.children.length, - ) - threadSpan.setAttribute("thread_info", JSON.stringify(threadInfo)) - } - } catch (error) { - loggerWithChild({ email: email }).error( - error, - `Error fetching email threads: ${getErrorMessage(error)}`, + + // Apply intelligent chunk selection based on document relevance and chunk scores + chunksPerDocument = await getChunkCountPerDoc( + combinedSearchResponse, + targetChunks, + email, + fileSearchSpan, + ) + fileSearchSpan?.end() + } + + if (threadIds && threadIds.length > 0) { + const threadSpan = span.startSpan("fetch_email_threads") + threadSpan.setAttribute("threadIds", JSON.stringify(threadIds)) + + try { + const threadResults = await SearchEmailThreads(threadIds, email) + if ( + threadResults.root.children && + threadResults.root.children.length > 0 + ) { + const existingDocIds = new Set( + combinedSearchResponse.map((child: any) => child.fields.docId), ) - threadSpan?.setAttribute("error", getErrorMessage(error)) + + // Use the helper function to process thread results + const { addedCount, threadInfo } = processThreadResults( + threadResults.root.children, + existingDocIds, + combinedSearchResponse, + ) + threadSpan.setAttribute("added_email_count", addedCount) + threadSpan.setAttribute( + "total_thread_emails_found", + threadResults.root.children.length, + ) + threadSpan.setAttribute("thread_info", JSON.stringify(threadInfo)) } - - threadSpan?.end() + } catch (error) { + loggerWithChild({ email: email }).error( + error, + `Error fetching email threads: ${getErrorMessage(error)}`, + ) + threadSpan?.setAttribute("error", getErrorMessage(error)) } - + + threadSpan?.end() + } + const precomputedDbContext = await getPrecomputedDbContextIfNeeded( combinedSearchResponse as VespaSearchResults[], query, @@ -1612,8 +1743,8 @@ async function prepareInitialAttachmentContext( allowChunkCitations, idx < chunksPerDocument.length ? chunksPerDocument[idx] : 0, precomputedDbContext, - ) - ) + ), + ), ) const summary = `User provided ${fragments.length} attachment fragment${ @@ -1643,18 +1774,29 @@ export async function beforeToolExecutionHook( toolName: string, args: any, context: AgentRunContext, - reasoningEmitter?: ReasoningEmitter + reasoningEmitter?: ReasoningEmitter, ): Promise { + const incomingExcludedIds = normalizeExcludedIdsForLogging(args?.excludedIds) + logContextMutation( + context, + "[beforeToolExecutionHook] Received tool args", + { + toolName, + args, + incomingExcludedIds, + incomingExcludedIdsCount: incomingExcludedIds.length, + }, + ) // 0. Validate input against schema const validation = validateToolInput(toolName, args) if (!validation.success) { await streamReasoningStep( reasoningEmitter, `⚠️ Invalid input for ${toolName}: ${validation.error.message}`, - { toolName } + { toolName }, ) Logger.warn( - `Tool input validation failed for ${toolName}: ${validation.error.message}` + `Tool input validation failed for ${toolName}: ${validation.error.message}`, ) // Don't block - let tool handle invalid input, but log it } @@ -1665,14 +1807,22 @@ export async function beforeToolExecutionHook( record.toolName === toolName && JSON.stringify(record.arguments) === JSON.stringify(args) && record.status === "success" && - Date.now() - record.startedAt.getTime() < 60000 // 1 minute + Date.now() - record.startedAt.getTime() < 60000, // 1 minute ) if (isDuplicate) { + logContextMutation( + context, + "[beforeToolExecutionHook] Skipping duplicate tool execution", + { + toolName, + args, + }, + ) await streamReasoningStep( reasoningEmitter, `Skipping redundant tool call to '${toolName}'.`, - { toolName, status: "skipped" } + { toolName, status: "skipped" }, ) return null // Skip execution } @@ -1680,23 +1830,60 @@ export async function beforeToolExecutionHook( // 2. Failed tool budget check const failureInfo = context.failedTools.get(toolName) if (failureInfo && failureInfo.count >= 3) { + logContextMutation( + context, + "[beforeToolExecutionHook] Blocking tool after repeated failures", + { + toolName, + args, + failureInfo, + }, + ) await streamReasoningStep( reasoningEmitter, `Tool '${toolName}' has failed ${failureInfo.count} times and is now blocked.`, - { toolName, status: "blocked" } + { toolName, status: "blocked" }, ) return null // Skip execution } // 3. Add excludedIds to prevent re-fetching seen documents - if (args.excludedIds !== undefined) { - const seenIds = Array.from(context.seenDocuments) + if (args?.excludedIds !== undefined) { + const providedExcludedIds = normalizeExcludedIdsForLogging(args.excludedIds) + const seenDocIds = Array.from(context.seenDocuments) + const mergedExcludedIds = [...providedExcludedIds, ...seenDocIds] + + logContextMutation( + context, + "[beforeToolExecutionHook] Merged excludedIds with seenDocuments", + { + toolName, + args, + providedExcludedIds, + providedExcludedIdsCount: providedExcludedIds.length, + seenDocumentIds: seenDocIds, + seenDocumentCount: seenDocIds.length, + mergedExcludedIds, + mergedExcludedIdsCount: mergedExcludedIds.length, + }, + ) + return { ...args, - excludedIds: [...(args.excludedIds || []), ...seenIds], + excludedIds: mergedExcludedIds, } } + logContextMutation( + context, + "[beforeToolExecutionHook] excludedIds not provided on tool args", + { + toolName, + args, + seenDocumentIds: Array.from(context.seenDocuments), + }, + ) + return args } @@ -1725,23 +1912,25 @@ export async function afterToolExecutionHook( gatheredFragmentsKeys: Set, expectedResult: ToolExpectation | undefined, turnNumber: number, - reasoningEmitter?: ReasoningEmitter + reasoningEmitter?: ReasoningEmitter, ): Promise { const { state, executionTime, status, args } = hookContext const context = state.context as AgentRunContext - // LOG: Hook entry point - loggerWithChild({ email: context.user.email }).debug( - { - toolName, - turnNumber, - status, - executionTime, - hasResult: !!result, - resultType: result ? typeof result : "null", - }, - "[afterToolExecutionHook] ENTRY - Tool execution completed" - ) + logContextMutation(context, "[afterToolExecutionHook] Processing tool result", { + toolName, + turnNumber, + status, + executionTime, + args, + hasResult: !!result, + resultType: result ? typeof result : "null", + resultStatus: result?.status, + resultError: result?.error, + resultMetadata: result?.metadata, + resultExcludedIds: normalizeExcludedIdsForLogging(result?.data?.excludedIds), + seenDocumentIds: Array.from(context.seenDocuments), + }) // 1. Create execution record const fallbackTurn = context.turnCount ?? MIN_TURN_NUMBER @@ -1754,7 +1943,7 @@ export async function afterToolExecutionHook( providedTurnNumber: turnNumber, fallbackTurnNumber: fallbackTurn, }, - "Tool turnNumber below minimum; normalizing to MIN_TURN_NUMBER" + "Tool turnNumber below minimum; normalizing to MIN_TURN_NUMBER", ) effectiveTurnNumber = MIN_TURN_NUMBER } @@ -1779,8 +1968,19 @@ export async function afterToolExecutionHook( : undefined, } - // 2. Add to history + // 2. Add tool call to history context.toolCallHistory.push(record) + logContextMutation( + context, + "[afterToolExecutionHook] Added tool execution record to history", + { + toolName, + turnNumber: effectiveTurnNumber, + recordStatus: record.status, + recordError: record.error, + historyLength: context.toolCallHistory.length, + }, + ) const toolFragments: MinimalAgentFragment[] = [] const addToolFragments = (fragments: MinimalAgentFragment[]) => { @@ -1832,13 +2032,13 @@ export async function afterToolExecutionHook( lastError: record.error!.message, lastAttempt: Date.now(), }) + logContextMutation(context, "[afterToolExecutionHook] Updated failed tool state", { + toolName, + failureInfo: context.failedTools.get(toolName), + }) } - advancePlanAfterTool( - context, - toolName, - status === "success" - ) + advancePlanAfterTool(context, toolName, status === "success") // 5. Extract and filter contexts const resultData = result?.data @@ -1852,7 +2052,10 @@ export async function afterToolExecutionHook( } } if (!Array.isArray(contexts) || contexts.length === 0) { - const legacyContexts = getMetadataValue(result, "contexts") + const legacyContexts = getMetadataValue( + result, + "contexts", + ) if (Array.isArray(legacyContexts)) { contexts = legacyContexts } @@ -1865,12 +2068,12 @@ export async function afterToolExecutionHook( totalContextsExtracted: contexts.length, gatheredFragmentsKeysSize: gatheredFragmentsKeys.size, }, - "[afterToolExecutionHook] Context extraction completed" + "[afterToolExecutionHook] Context extraction completed", ) if (Array.isArray(contexts) && contexts.length > 0) { const filteredContexts = contexts.filter( - (c: MinimalAgentFragment) => !gatheredFragmentsKeys.has(c.id) + (c: MinimalAgentFragment) => !gatheredFragmentsKeys.has(c.id), ) // LOG: Filtering results @@ -1881,7 +2084,7 @@ export async function afterToolExecutionHook( filteredContextsCount: filteredContexts.length, duplicatesFiltered: contexts.length - filteredContexts.length, }, - "[afterToolExecutionHook] Filtered out duplicate contexts" + "[afterToolExecutionHook] Filtered out duplicate contexts", ) if (filteredContexts.length > 0) { @@ -1890,16 +2093,22 @@ export async function afterToolExecutionHook( `Received ${filteredContexts.length} document${ filteredContexts.length === 1 ? "" : "s" }. Now filtering and ranking the most relevant ones.`, - { toolName } + { toolName }, ) - const metadataConstraints = extractMetadataConstraintsFromUserMessage(userMessage) + const metadataConstraints = + extractMetadataConstraintsFromUserMessage(userMessage) const { rankedCandidates, hasConstraints: hasMetadataConstraints, hasCompliantCandidates, - } = rankFragmentsByMetadataConstraints(filteredContexts, metadataConstraints) - const rankingCandidates = rankedCandidates.map((candidate) => candidate.fragment) + } = rankFragmentsByMetadataConstraints( + filteredContexts, + metadataConstraints, + ) + const rankingCandidates = rankedCandidates.map( + (candidate) => candidate.fragment, + ) const strictNoCompliantCandidates = hasMetadataConstraints && metadataConstraints.strict && @@ -1907,7 +2116,7 @@ export async function afterToolExecutionHook( const contextStrings = rankingCandidates.map( (fragment: MinimalAgentFragment, index: number) => - formatFragmentWithMetadata(fragment, index) + formatFragmentWithMetadata(fragment, index), ) // LOG: Prepared context strings for ranking @@ -1920,10 +2129,13 @@ export async function afterToolExecutionHook( metadataIncludeTerms: metadataConstraints.includeTerms, metadataExcludeTerms: metadataConstraints.excludeTerms, metadataStrict: metadataConstraints.strict, - compliantCandidateCount: rankedCandidates.filter((v) => v.compliant).length, - contextStringsSample: contextStrings.slice(0, 2).map(s => s.slice(0, 150)), + compliantCandidateCount: rankedCandidates.filter((v) => v.compliant) + .length, + contextStringsSample: contextStrings + .slice(0, 2) + .map((s) => s.slice(0, 150)), }, - "[afterToolExecutionHook] Prepared context strings for select_best_documents" + "[afterToolExecutionHook] Prepared context strings for select_best_documents", ) if (hasMetadataConstraints) { @@ -1932,8 +2144,8 @@ export async function afterToolExecutionHook( strictNoCompliantCandidates ? "Strict metadata constraints detected and no compliant documents were found." : hasCompliantCandidates - ? "Applied metadata constraints from the user request before ranking." - : "Detected metadata constraints in the user request but found no clearly compliant metadata matches.", + ? "Applied metadata constraints from the user request before ranking." + : "Detected metadata constraints in the user request but found no clearly compliant metadata matches.", { toolName, detail: [ @@ -1946,13 +2158,14 @@ export async function afterToolExecutionHook( ] .filter(Boolean) .join(" "), - } + }, ) } try { // LOG: Calling extractBestDocumentIndexes - const rankingModelId = (context.modelId as Models) || config.defaultBestModel + const rankingModelId = + (context.modelId as Models) || config.defaultBestModel loggerWithChild({ email: context.user.email }).debug( { toolName, @@ -1960,22 +2173,24 @@ export async function afterToolExecutionHook( contextStringsCount: contextStrings.length, modelId: rankingModelId, }, - "[afterToolExecutionHook][select_best_documents] CALLING extractBestDocumentIndexes" + "[afterToolExecutionHook][select_best_documents] CALLING extractBestDocumentIndexes", ) // Use extractBestDocumentIndexes to filter to best documents - const selectionSpan = getTracer("chat").startSpan("select_best_documents") + const selectionSpan = getTracer("chat").startSpan( + "select_best_documents", + ) selectionSpan.setAttribute("tool_name", toolName) selectionSpan.setAttribute("context_count", contextStrings.length) let bestDocIndexes: number[] = [] try { const rankingMessages = withAgentSystemPromptMessage( messagesWithNoErrResponse, - context.dedicatedAgentSystemPrompt + context.dedicatedAgentSystemPrompt, ) const rankingSystemPrompt = extractBestDocumentsPrompt( userMessage, - contextStrings + contextStrings, ) loggerWithChild({ email: context.user.email }).debug( { @@ -1984,13 +2199,13 @@ export async function afterToolExecutionHook( rankingSystemPrompt, rankingMessages, }, - "[afterToolExecutionHook][select_best_documents] FINAL LLM PAYLOAD" + "[afterToolExecutionHook][select_best_documents] FINAL LLM PAYLOAD", ) selectionSpan.setAttribute( "has_agent_system_prompt_snapshot", !!sanitizeAgentSystemPromptSnapshot( - context.dedicatedAgentSystemPrompt - ) + context.dedicatedAgentSystemPrompt, + ), ) bestDocIndexes = await extractBestDocumentIndexes( userMessage, @@ -2000,7 +2215,7 @@ export async function afterToolExecutionHook( json: false, stream: false, }, - rankingMessages + rankingMessages, ) selectionSpan.setAttribute("selected_count", bestDocIndexes.length) } catch (error) { @@ -2019,7 +2234,7 @@ export async function afterToolExecutionHook( bestDocIndexesType: typeof bestDocIndexes, isArray: Array.isArray(bestDocIndexes), }, - "[afterToolExecutionHook][select_best_documents] RESPONSE from extractBestDocumentIndexes" + "[afterToolExecutionHook][select_best_documents] RESPONSE from extractBestDocumentIndexes", ) if (bestDocIndexes.length > 0) { @@ -2035,16 +2250,15 @@ export async function afterToolExecutionHook( selectedDocs = enforceMetadataConstraintsOnSelection( selectedDocs, rankedCandidates, - metadataConstraints + metadataConstraints, ) if (selectedDocs.length === 0) { - const fallbackWhenSelectionEmpty = - strictNoCompliantCandidates - ? [] - : hasMetadataConstraints && - metadataConstraints.strict && - hasCompliantCandidates + const fallbackWhenSelectionEmpty = strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && + metadataConstraints.strict && + hasCompliantCandidates ? rankedCandidates .filter((candidate) => candidate.compliant) .map((candidate) => candidate.fragment) @@ -2054,7 +2268,7 @@ export async function afterToolExecutionHook( toolName, fallbackContextsCount: fallbackWhenSelectionEmpty.length, }, - "[afterToolExecutionHook] No contexts survived selection enforcement; using metadata-aware fallback" + "[afterToolExecutionHook] No contexts survived selection enforcement; using metadata-aware fallback", ) addToolFragments(fallbackWhenSelectionEmpty) } else { @@ -2065,10 +2279,10 @@ export async function afterToolExecutionHook( selectedDocsCount: selectedDocs.length, selectedDocIds: selectedDocs.map((d) => d.id), selectedDocTitles: selectedDocs.map( - (d) => d.source.title || "untitled" + (d) => d.source.title || "untitled", ), }, - "[afterToolExecutionHook][select_best_documents] Selected documents after filtering" + "[afterToolExecutionHook][select_best_documents] Selected documents after filtering", ) await streamReasoningStep( @@ -2081,7 +2295,7 @@ export async function afterToolExecutionHook( detail: selectedDocs .map((doc) => doc.source.title || doc.id) .join(", "), - } + }, ) let fragmentsForResult = selectedDocs @@ -2099,23 +2313,26 @@ export async function afterToolExecutionHook( const { imageFileNames: extractedImages } = extractImageFileNames( combinedContext, - vespaLikeResults + vespaLikeResults, ) if (extractedImages.length > 0) { const turnForImages = Math.max( context.turnCount ?? MIN_TURN_NUMBER, - MIN_TURN_NUMBER + MIN_TURN_NUMBER, ) const imageSpan = getTracer("chat").startSpan( - "tool_image_extraction" + "tool_image_extraction", ) imageSpan.setAttribute("tool_name", toolName) imageSpan.setAttribute("turn_number", turnForImages) - imageSpan.setAttribute("extracted_count", extractedImages.length) + imageSpan.setAttribute( + "extracted_count", + extractedImages.length, + ) imageSpan.setAttribute( "image_names_preview", - extractedImages.slice(0, 5).join(",") + extractedImages.slice(0, 5).join(","), ) // LOG: Before attaching images @@ -2127,7 +2344,7 @@ export async function afterToolExecutionHook( turnForImages, selectedDocsCount: selectedDocs.length, }, - "[afterToolExecutionHook] Extracted images from fragments - preparing to attach" + "[afterToolExecutionHook] Extracted images from fragments - preparing to attach", ) fragmentsForResult = attachImagesToFragments( @@ -2137,13 +2354,13 @@ export async function afterToolExecutionHook( turnNumber: turnForImages, sourceToolName: toolName, isUserAttachment: false, - } + }, ) imageSpan.setAttribute( "fragments_with_images", fragmentsForResult.filter( - (fragment) => fragment.images && fragment.images.length > 0 - ).length + (fragment) => fragment.images && fragment.images.length > 0, + ).length, ) imageSpan.end() @@ -2153,12 +2370,12 @@ export async function afterToolExecutionHook( toolName, attachedImagesCount: extractedImages.length, fragmentsWithImages: fragmentsForResult.filter( - (f) => f.images && f.images.length > 0 + (f) => f.images && f.images.length > 0, ).length, }, `[afterToolExecutionHook] Tracked ${extractedImages.length} image reference${ extractedImages.length === 1 ? "" : "s" - } from ${toolName} on turn ${turnForImages}` + } from ${toolName} on turn ${turnForImages}`, ) } } @@ -2166,10 +2383,11 @@ export async function afterToolExecutionHook( addToolFragments(fragmentsForResult) } } else { - const metadataFilteredFallback = - strictNoCompliantCandidates - ? [] - : hasMetadataConstraints && metadataConstraints.strict && hasCompliantCandidates + const metadataFilteredFallback = strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && + metadataConstraints.strict && + hasCompliantCandidates ? rankedCandidates .filter((candidate) => candidate.compliant) .map((candidate) => candidate.fragment) @@ -2182,7 +2400,7 @@ export async function afterToolExecutionHook( }, strictNoCompliantCandidates ? "[afterToolExecutionHook] Document ranking returned no results and strict metadata constraints had no compliant contexts" - : "[afterToolExecutionHook] Document ranking returned no results; retaining all filtered contexts" + : "[afterToolExecutionHook] Document ranking returned no results; retaining all filtered contexts", ) addToolFragments(metadataFilteredFallback) } @@ -2194,19 +2412,20 @@ export async function afterToolExecutionHook( error: error instanceof Error ? error.message : String(error), stack: error instanceof Error ? error.stack : undefined, }, - "[afterToolExecutionHook][select_best_documents] ERROR in document ranking" + "[afterToolExecutionHook][select_best_documents] ERROR in document ranking", ) await streamReasoningStep( reasoningEmitter, strictNoCompliantCandidates ? "Context ranking failed and strict metadata constraints had no compliant documents; returning no contexts." : "Context ranking failed, retaining all retrieved documents.", - { toolName } + { toolName }, ) - const fallbackContexts = - strictNoCompliantCandidates - ? [] - : hasMetadataConstraints && metadataConstraints.strict && hasCompliantCandidates + const fallbackContexts = strictNoCompliantCandidates + ? [] + : hasMetadataConstraints && + metadataConstraints.strict && + hasCompliantCandidates ? rankedCandidates .filter((candidate) => candidate.compliant) .map((candidate) => candidate.fragment) @@ -2234,11 +2453,12 @@ export async function afterToolExecutionHook( ) { const plan = result.data?.plan as PlanState | undefined if (plan) { - const nextStep = plan.subTasks?.[0]?.description || "No sub-tasks defined." + const nextStep = + plan.subTasks?.[0]?.description || "No sub-tasks defined." await streamReasoningStep( reasoningEmitter, `Plan created: ${plan.goal || "Goal not specified"}. Next step: ${nextStep}`, - { toolName } + { toolName }, ) } } @@ -2249,14 +2469,21 @@ export async function afterToolExecutionHook( typeof result === "object" && result.status === "success" ) { - const agents = (result?.data as { agents?: ListCustomAgentsOutput["agents"] })?.agents + const agents = ( + result?.data as { agents?: ListCustomAgentsOutput["agents"] } + )?.agents const agentCount = Array.isArray(agents) ? agents.length : 0 await streamReasoningStep( reasoningEmitter, agentCount ? `Found ${agentCount} suitable agent${agentCount === 1 ? "" : "s"}. Evaluating options...` : "No suitable custom agents found. Continuing with built-in tools.", - { toolName, detail: agentCount ? agents?.map((a) => a.agentName).join(", ") : undefined } + { + toolName, + detail: agentCount + ? agents?.map((a) => a.agentName).join(", ") + : undefined, + }, ) } @@ -2271,7 +2498,9 @@ export async function afterToolExecutionHook( (hookContext?.args as { agentId?: string })?.agentId const agentName = context.availableAgents.find((agent) => agent.agentId === agentId) - ?.agentName || agentId || "unknown agent" + ?.agentName || + agentId || + "unknown agent" const delegationFragments = buildDelegatedAgentFragments({ result, gatheredFragmentsKeys, @@ -2286,7 +2515,7 @@ export async function afterToolExecutionHook( await streamReasoningStep( reasoningEmitter, `Received response from '${agentName}' agent.`, - { toolName, detail: agentName } + { toolName, detail: agentName }, ) } @@ -2300,7 +2529,7 @@ export async function afterToolExecutionHook( await streamReasoningStep( reasoningEmitter, "Fallback analysis completed.", - { toolName } + { toolName }, ) } @@ -2311,19 +2540,26 @@ export async function afterToolExecutionHook( resultSummary: summarizeToolResultPayload(result), fragments: toolFragments, }) - - // LOG: Hook exit point - loggerWithChild({ email: context.user.email }).debug( + logContextMutation( + context, + "[afterToolExecutionHook] Recorded tool output for current turn", { toolName, turnNumber: effectiveTurnNumber, toolFragmentsCount: toolFragments.length, - totalCost: context.totalCost, - totalLatency: context.totalLatency, + toolFragmentIds: toolFragments.map((fragment) => fragment.id), + resultSummary: summarizeToolResultPayload(result), }, - "[afterToolExecutionHook] EXIT - Tool processing completed" ) + logContextMutation(context, "[afterToolExecutionHook] Completed tool result processing", { + toolName, + turnNumber: effectiveTurnNumber, + toolFragmentsCount: toolFragments.length, + totalCost: context.totalCost, + totalLatency: context.totalLatency, + }) + if (toolFragments.length > 0) { return ToolResponse.success(toolFragments) } @@ -2331,7 +2567,6 @@ export async function afterToolExecutionHook( return null } - export function buildDelegatedAgentFragments(opts: { result: any gatheredFragmentsKeys: Set @@ -2340,13 +2575,23 @@ export function buildDelegatedAgentFragments(opts: { turnNumber: number sourceToolName: string }): MinimalAgentFragment[] { - const { result, gatheredFragmentsKeys, agentId, agentName, turnNumber, sourceToolName } = opts + const { + result, + gatheredFragmentsKeys, + agentId, + agentName, + turnNumber, + sourceToolName, + } = opts const resultData = (result?.data as Record) || {} const citations = resultData.citations as Citation[] | undefined - const imageCitations = resultData.imageCitations as ImageCitation[] | undefined + const imageCitations = resultData.imageCitations as + | ImageCitation[] + | undefined const agentFragments: MinimalAgentFragment[] = [] const fragmentTurn = Math.max(turnNumber, MIN_TURN_NUMBER) - const normalizedAgentName = agentName || agentId || sourceToolName || "delegated_agent" + const normalizedAgentName = + agentName || agentId || sourceToolName || "delegated_agent" const normalizedAgentId = agentId || `agent:${sourceToolName}` const baseSource: Citation = { docId: normalizedAgentId, @@ -2361,7 +2606,8 @@ export function buildDelegatedAgentFragments(opts: { const textResult = typeof resultData.result === "string" ? (resultData.result as string) - : typeof (resultData as { agentResult?: string })?.agentResult === "string" + : typeof (resultData as { agentResult?: string })?.agentResult === + "string" ? ((resultData as { agentResult?: string }).agentResult as string) : typeof result?.result === "string" ? result.result @@ -2419,7 +2665,7 @@ export function buildDelegatedAgentFragments(opts: { const fragmentByDoc = new Map( agentFragments .filter((fragment) => fragment.source?.docId) - .map((fragment) => [fragment.source.docId!, fragment]) + .map((fragment) => [fragment.source.docId!, fragment]), ) for (const imageCitation of imageCitations) { @@ -2465,7 +2711,8 @@ export function extractExpectedResults(text: string): PendingExpectation[] { for (const entry of entries) { if (!entry || typeof entry !== "object") continue - const toolName = typeof entry.toolName === "string" ? entry.toolName.trim() : "" + const toolName = + typeof entry.toolName === "string" ? entry.toolName.trim() : "" if (!toolName) continue const expectationCandidate = { @@ -2480,7 +2727,7 @@ export function extractExpectedResults(text: string): PendingExpectation[] { if (!validation.success) { Logger.warn( { toolName, error: validation.error.format() }, - "Invalid expected_results entry emitted by agent" + "Invalid expected_results entry emitted by agent", ) continue } @@ -2494,11 +2741,11 @@ export function extractExpectedResults(text: string): PendingExpectation[] { function consumePendingExpectation( queue: PendingExpectation[], - toolName: string + toolName: string, ): PendingExpectation | undefined { if (!toolName) return undefined const idx = queue.findIndex( - (entry) => entry.toolName.toLowerCase() === toolName.toLowerCase() + (entry) => entry.toolName.toLowerCase() === toolName.toLowerCase(), ) if (idx === -1) { return undefined @@ -2517,10 +2764,7 @@ function safeJsonParse(text: string): unknown { function summarizePlan(plan: PlanState | null): string { Logger.debug({ plan }, "summarizePlan input") if (!plan) { - Logger.debug( - { summary: "No plan available." }, - "summarizePlan output" - ) + Logger.debug({ summary: "No plan available." }, "summarizePlan output") return "No plan available." } const steps = plan.subTasks @@ -2530,26 +2774,23 @@ function summarizePlan(plan: PlanState | null): string { task.toolsRequired?.length ? ` (tools: ${task.toolsRequired.join(", ")})` : "" - }` + }`, ) .join("\n") const summary = `Goal: ${plan.goal}\n${steps}` Logger.debug( { summary, subTaskCount: plan.subTasks.length }, - "summarizePlan output" + "summarizePlan output", ) return summary } function formatExpectationsForReview( - expectations?: ToolExpectationAssignment[] + expectations?: ToolExpectationAssignment[], ): string { Logger.debug({ expectations }, "formatExpectationsForReview input") if (!expectations || expectations.length === 0) { - Logger.debug( - { serialized: "[]" }, - "formatExpectationsForReview output" - ) + Logger.debug({ serialized: "[]" }, "formatExpectationsForReview output") return "[]" } const serialized = JSON.stringify(expectations, null, 2) @@ -2558,20 +2799,22 @@ function formatExpectationsForReview( expectationCount: expectations.length, serializedLength: serialized.length, }, - "formatExpectationsForReview output" + "formatExpectationsForReview output", ) return serialized } function formatToolOutputsForReview( - outputs: ToolExecutionRecordWithResult[] + outputs: ToolExecutionRecordWithResult[], ): string { if (!outputs || outputs.length === 0) { return "No tools executed this turn." } return outputs .map((output, idx) => { - const argsSummary = formatToolArgumentsForReasoning(output.arguments || {}) + const argsSummary = formatToolArgumentsForReasoning( + output.arguments || {}, + ) const fragmentSummary = output.fragments?.length ? `${output.fragments.length} fragment${output.fragments.length === 1 ? "" : "s"}` : "0 fragments" @@ -2581,9 +2824,7 @@ function formatToolOutputsForReview( } /** Format ToolExecutionRecord[] for review prompt, grouped by turn (for last-N-turns context). */ -function formatToolCallHistoryByTurn( - records: ToolExecutionRecord[] -): string { +function formatToolCallHistoryByTurn(records: ToolExecutionRecord[]): string { if (!records || records.length === 0) { return "No tool calls in this window." } @@ -2629,15 +2870,15 @@ export function buildReviewPromptFromContext( focus?: string turnNumber?: number }, - fallbackExpectations?: ToolExpectationAssignment[] + fallbackExpectations?: ToolExpectationAssignment[], ): { prompt: string; imageFileNames: string[] } { const turnExpectations = context.currentTurnArtifacts.expectations.length > 0 ? context.currentTurnArtifacts.expectations - : fallbackExpectations ?? [] + : (fallbackExpectations ?? []) const planSection = formatPlanForPrompt(context.plan) const clarificationsSection = formatClarificationsForPrompt( - context.clarifications + context.clarifications, ) const workspaceSection = context.userContext?.trim() ? `Workspace Context:\n${context.userContext.trim()}` @@ -2662,14 +2903,14 @@ export function buildReviewPromptFromContext( const expectationsSection = formatExpectationsForReview(turnExpectations) const fragmentsSection = answerContextMapFromFragments( context.allFragments, - Math.max(12, context.allFragments.length || 1) + Math.max(12, context.allFragments.length || 1), ) const currentImages = context.currentTurnArtifacts.images.map( - (image) => image.fileName + (image) => image.fileName, ) const additionalImages = Math.max( context.allImages.length - currentImages.length, - 0 + 0, ) const imageSection = `Current turn attachments: ${currentImages.length}\nAdditional images available from prior turns: ${additionalImages}` const reviewFocus = `Review Focus: ${options?.focus ?? "turn_end"} (Turn ${ @@ -2679,9 +2920,7 @@ export function buildReviewPromptFromContext( const userPromptSections = [ `User Question:\n${context.message.text}`, planSection ? `Execution Plan Snapshot:\n${planSection}` : "", - clarificationsSection - ? `Clarifications:\n${clarificationsSection}` - : "", + clarificationsSection ? `Clarifications:\n${clarificationsSection}` : "", workspaceSection, `Current Turn Tool Outputs:\n${toolOutputsSection}`, `Expectations:\n${expectationsSection}`, @@ -2705,7 +2944,7 @@ async function runReviewLLM( expectedResults?: ToolExpectationAssignment[] delegationEnabled?: boolean }, - modelOverride?: string + modelOverride?: string, ): Promise { const tracer = getTracer("chat") const reviewSpan = tracer.startSpan("review_llm_call") @@ -2713,7 +2952,7 @@ async function runReviewLLM( reviewSpan.setAttribute("turn_number", options?.turnNumber ?? -1) reviewSpan.setAttribute( "expected_results_count", - options?.expectedResults?.length ?? 0 + options?.expectedResults?.length ?? 0, ) Logger.debug( { @@ -2726,7 +2965,7 @@ async function runReviewLLM( email: context.user.email, chatId: context.chat.externalId, }, - "[MessageAgents][runReviewLLM] invoked - FULL expectedResults" + "[MessageAgents][runReviewLLM] invoked - FULL expectedResults", ) const modelId = (modelOverride as Models) || @@ -2737,14 +2976,8 @@ async function runReviewLLM( ? "- Delegation tools (list_custom_agents/run_public_agent) were disabled for this run; do not flag their absence." : "- If delegation tools are available, ensure list_custom_agents precedes run_public_agent when delegation is appropriate." - const { - prompt: userPrompt, - imageFileNames: currentImages, - } = buildReviewPromptFromContext( - context, - options, - options?.expectedResults - ) + const { prompt: userPrompt, imageFileNames: currentImages } = + buildReviewPromptFromContext(context, options, options?.expectedResults) Logger.debug( { email: context.user.email, @@ -2754,12 +2987,12 @@ async function runReviewLLM( reviewImages: currentImages, additionalImages: Math.max( context.allImages.length - currentImages.length, - 0 + 0, ), fragmentsCount: context.allFragments.length, toolOutputsCount: context.currentTurnArtifacts.toolOutputs.length, }, - "[MessageAgents][runReviewLLM] Context summary for review model" + "[MessageAgents][runReviewLLM] Context summary for review model", ) Logger.debug( @@ -2770,7 +3003,7 @@ async function runReviewLLM( turnNumber: options?.turnNumber, modelId, }, - "[MessageAgents][runReviewLLM] Preparing review LLM call" + "[MessageAgents][runReviewLLM] Preparing review LLM call", ) const params: ModelParams = { @@ -2805,7 +3038,7 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ anomaliesDetected: false, anomalies: ["Description of anomalies or ambiguities"], recommendation: "proceed", - ambiguityResolved: true + ambiguityResolved: true, })} - Use native JSON booleans (true/false) for every yes/no field. - Only emit keys defined in the schema; do not add prose outside the JSON object.`, @@ -2814,17 +3047,17 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ params.imageFileNames = currentImages } Logger.debug( - { + { email: context.user.email, chatId: context.chat.externalId, - modelId, + modelId, params, temperature: params.temperature, maxTokens: params.max_new_tokens, json: params.json, stream: params.stream, }, - "[MessageAgents][runReviewLLM] LLM params prepared" + "[MessageAgents][runReviewLLM] LLM params prepared", ) Logger.debug( { @@ -2832,16 +3065,16 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ chatId: context.chat.externalId, userPrompt, }, - "[MessageAgents][runReviewLLM] Review user prompt" + "[MessageAgents][runReviewLLM] Review user prompt", ) Logger.debug( - { + { email: context.user.email, chatId: context.chat.externalId, systemPrompt: params.systemPrompt, }, - "[MessageAgents][runReviewLLM] System prompt" + "[MessageAgents][runReviewLLM] System prompt", ) const messages: Message[] = [ @@ -2851,27 +3084,30 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ }, ] - const { text } = await getProviderByModel(modelId).converse( - messages, - params + const { text } = await getProviderByModel(modelId).converse(messages, params) + + Logger.debug( + { + email: context.user.email, + chatId: context.chat.externalId, + text, + }, + "[MessageAgents][runReviewLLM] Raw LLM response", ) - - Logger.debug({ - email: context.user.email, - chatId: context.chat.externalId, - text, - }, "[MessageAgents][runReviewLLM] Raw LLM response") if (!text) { throw new Error("LLM returned empty review response") } const parsed = jsonParseLLMOutput(text) - Logger.debug({ - email: context.user.email, - chatId: context.chat.externalId, - parsed, - }, "[MessageAgents][runReviewLLM] Parsed LLM response") + Logger.debug( + { + email: context.user.email, + chatId: context.chat.externalId, + parsed, + }, + "[MessageAgents][runReviewLLM] Parsed LLM response", + ) if (!parsed || typeof parsed !== "object") { Logger.error( @@ -2880,36 +3116,36 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ chatId: context.chat.externalId, raw: parsed, }, - "[MessageAgents][runReviewLLM] Invalid review payload" + "[MessageAgents][runReviewLLM] Invalid review payload", ) return buildDefaultReviewPayload( - `Review model returned invalid payload for turn ${options?.turnNumber ?? "unknown"}` + `Review model returned invalid payload for turn ${options?.turnNumber ?? "unknown"}`, ) } const validation = ReviewResultSchema.safeParse(parsed) if (!validation.success) { Logger.error( - { + { email: context.user.email, chatId: context.chat.externalId, - error: validation.error.format(), + error: validation.error.format(), raw: parsed, }, - "[MessageAgents][runReviewLLM] Review result does not match schema" + "[MessageAgents][runReviewLLM] Review result does not match schema", ) return buildDefaultReviewPayload( - `Review model response failed validation for turn ${options?.turnNumber ?? "unknown"}` + `Review model response failed validation for turn ${options?.turnNumber ?? "unknown"}`, ) } Logger.debug( - { + { email: context.user.email, chatId: context.chat.externalId, reviewResult: validation.data, }, - "[MessageAgents][runReviewLLM] Returning review result" + "[MessageAgents][runReviewLLM] Returning review result", ) Logger.debug( { @@ -2920,18 +3156,18 @@ Respond strictly in JSON matching this schema: ${JSON.stringify({ imageFileCount: currentImages.length, toolOutputsEvaluated: context.currentTurnArtifacts.toolOutputs.length, }, - "[MessageAgents][runReviewLLM] Review LLM call completed" + "[MessageAgents][runReviewLLM] Review LLM call completed", ) reviewSpan.setAttribute("model_id", modelId) reviewSpan.setAttribute("review_status", validation.data.status) reviewSpan.setAttribute("recommendation", validation.data.recommendation) reviewSpan.setAttribute( "anomalies_detected", - validation.data.anomaliesDetected ?? false + validation.data.anomaliesDetected ?? false, ) reviewSpan.setAttribute( "tool_feedback_count", - validation.data.toolFeedback.length + validation.data.toolFeedback.length, ) reviewSpan.end() return validation.data @@ -2941,6 +3177,7 @@ function buildInternalToolAdapters(): Tool[] { const baseTools = [ createToDoWriteTool(), searchGlobalTool, + lsKnowledgeBaseTool, searchKnowledgeBaseTool, ...googleTools, getSlackRelatedMessagesTool, @@ -2966,6 +3203,7 @@ const TOOL_ACCESS_REQUIREMENTS: Record = { requiredApp: Apps.GoogleCalendar, connectorFlag: "googleCalendarSynced", }, + ls: { requiredApp: Apps.KnowledgeBase }, searchKnowledgeBase: { requiredApp: Apps.KnowledgeBase }, searchGoogleContacts: { requiredApp: Apps.GoogleWorkspace, @@ -2999,10 +3237,7 @@ function filterToolsByAvailability( const rule = TOOL_ACCESS_REQUIREMENTS[tool.schema.name] if (!rule) return true - if ( - rule.connectorFlag && - !params.connectorState[rule.connectorFlag] - ) { + if (rule.connectorFlag && !params.connectorState[rule.connectorFlag]) { loggerWithChild({ email: params.email, agentId: params.agentId }).info( `Disabling tool ${tool.schema.name}: connector '${rule.connectorFlag}' unavailable.`, ) @@ -3039,14 +3274,11 @@ function createToDoWriteTool(): Tool { email: context.user.email, args, }, - "[toDoWrite] Execution started" + "[toDoWrite] Execution started", ) const validation = validateToolInput("toDoWrite", args) if (!validation.success) { - return ToolResponse.error( - "INVALID_INPUT", - validation.error.message - ) + return ToolResponse.error("INVALID_INPUT", validation.error.message) } const plan: PlanState = { @@ -3064,7 +3296,7 @@ function createToDoWriteTool(): Tool { subTaskCount: plan.subTasks.length, activeSubTaskId, }, - "[toDoWrite] Plan created" + "[toDoWrite] Plan created", ) return ToolResponse.success({ plan }) @@ -3072,8 +3304,10 @@ function createToDoWriteTool(): Tool { } } - -function buildDelegatedAgentQuery(baseQuery: string, context: AgentRunContext): string { +function buildDelegatedAgentQuery( + baseQuery: string, + context: AgentRunContext, +): string { const parts = [baseQuery.trim()] if (context.currentSubTask) { parts.push(`Active sub-task: ${context.currentSubTask}`) @@ -3088,13 +3322,9 @@ function buildDelegatedAgentQuery(baseQuery: string, context: AgentRunContext): } function buildCustomAgentTools(): Array> { - return [ - createListCustomAgentsTool(), - createRunPublicAgentTool(), - ] + return [createListCustomAgentsTool(), createRunPublicAgentTool()] } - function createListCustomAgentsTool(): Tool { return { schema: { @@ -3111,10 +3341,7 @@ function createListCustomAgentsTool(): Tool { }>("list_custom_agents", args) if (!validation.success) { - return ToolResponse.error( - "INVALID_INPUT", - validation.error.message - ) + return ToolResponse.error("INVALID_INPUT", validation.error.message) } const result = await listCustomAgentsSuitable({ @@ -3129,15 +3356,25 @@ function createListCustomAgentsTool(): Tool { }) Logger.debug( { params: validation.data, email: context.user.email }, - "[list_custom_agents] input params" + "[list_custom_agents] input params", ) Logger.debug( { selection: result, email: context.user.email }, - "[list_custom_agents] selection result" + "[list_custom_agents] selection result", ) const normalizedAgents = Array.isArray(result.agents) ? result.agents : [] mutableContext.availableAgents = normalizedAgents + logContextMutation( + mutableContext, + "[list_custom_agents] Updated availableAgents in context", + { + query: validation.data.query, + requiredCapabilities: validation.data.requiredCapabilities, + availableAgentIds: normalizedAgents.map((agent) => agent.agentId), + availableAgentNames: normalizedAgents.map((agent) => agent.agentName), + }, + ) return ToolResponse.success({ agents: normalizedAgents.length ? normalizedAgents : null, totalEvaluated: result.totalEvaluated, @@ -3162,10 +3399,7 @@ function createRunPublicAgentTool(): Tool { }>("run_public_agent", args) if (!validation.success) { - return ToolResponse.error( - "INVALID_INPUT", - validation.error.message - ) + return ToolResponse.error("INVALID_INPUT", validation.error.message) } if (!context.ambiguityResolved) { @@ -3173,18 +3407,16 @@ function createRunPublicAgentTool(): Tool { ToolErrorCodes.INVALID_INPUT, `Resolve ambiguity before running a custom agent. Unresolved: ${ context.clarifications.length - ? context.clarifications - .map((c) => c.question) - .join("; ") + ? context.clarifications.map((c) => c.question).join("; ") : "not specified" - }` + }`, ) } if (!context.availableAgents.length) { return ToolResponse.error( ToolErrorCodes.RESOURCE_UNAVAILABLE, - "No agents available. Run list_custom_agents this turn and select an agentId from its results." + "No agents available. Run list_custom_agents this turn and select an agentId from its results.", ) } @@ -3196,18 +3428,18 @@ function createRunPublicAgentTool(): Tool { agentName: a.agentName, })), }, - "[run_public_agent] Agent selection details" + "[run_public_agent] Agent selection details", ) const agentCapability = context.availableAgents.find( - (agent) => agent.agentId === validation.data.agentId + (agent) => agent.agentId === validation.data.agentId, ) if (!agentCapability) { return ToolResponse.error( ToolErrorCodes.NOT_FOUND, `Agent '${validation.data.agentId}' not found in availableAgents. Call list_custom_agents and use one of: ${context.availableAgents .map((a) => `${a.agentName} (${a.agentId})`) - .join("; ")}` + .join("; ")}`, ) } @@ -3219,24 +3451,33 @@ function createRunPublicAgentTool(): Tool { userEmail: context.user.email, workspaceExternalId: context.user.workspaceId, mcpAgents: context.mcpAgents, - parentTurn: Math.max(context.turnCount ?? MIN_TURN_NUMBER, MIN_TURN_NUMBER), + parentTurn: Math.max( + context.turnCount ?? MIN_TURN_NUMBER, + MIN_TURN_NUMBER, + ), stopSignal: context.stopSignal, }) Logger.debug( { params: validation.data, email: context.user.email }, - "[run_public_agent] input params" + "[run_public_agent] input params", ) Logger.debug( { toolOutput, email: context.user.email }, - "[run_public_agent] tool output" + "[run_public_agent] tool output", ) context.usedAgents.push(agentCapability.agentId) + logContextMutation( + context, + "[run_public_agent] Added agent to usedAgents", + { + selectedAgentId: agentCapability.agentId, + selectedAgentName: agentCapability.agentName, + query: validation.data.query, + }, + ) if (toolOutput.error) { - return ToolResponse.error( - "EXECUTION_FAILED", - toolOutput.error - ) + return ToolResponse.error("EXECUTION_FAILED", toolOutput.error) } const metadata = toolOutput.metadata || {} @@ -3282,7 +3523,9 @@ function createFinalSynthesisTool(): Tool { schema: { name: "synthesize_final_answer", description: TOOL_SCHEMAS.synthesize_final_answer.description, - parameters: toToolParameters(TOOL_SCHEMAS.synthesize_final_answer.inputSchema), + parameters: toToolParameters( + TOOL_SCHEMAS.synthesize_final_answer.inputSchema, + ), }, async execute(_args, context) { const mutableContext = mutableAgentContext(context) @@ -3290,12 +3533,19 @@ function createFinalSynthesisTool(): Tool { mutableContext.review.lockedByFinalSynthesis = true mutableContext.review.lockedAtTurn = mutableContext.turnCount ?? MIN_TURN_NUMBER + logContextMutation( + mutableContext, + "[MessageAgents][FinalSynthesis] Locked review state for synthesis", + { + lockedAtTurn: mutableContext.review.lockedAtTurn, + }, + ) loggerWithChild({ email: context.user.email }).info( { chatId: context.chat.externalId, turn: mutableContext.review.lockedAtTurn, }, - "[MessageAgents][FinalSynthesis] Review lock activated after synthesis tool call." + "[MessageAgents][FinalSynthesis] Review lock activated after synthesis tool call.", ) } if ( @@ -3304,7 +3554,7 @@ function createFinalSynthesisTool(): Tool { ) { return ToolResponse.error( "EXECUTION_FAILED", - "Final synthesis already completed for this run." + "Final synthesis already completed for this run.", ) } @@ -3312,7 +3562,7 @@ function createFinalSynthesisTool(): Tool { if (!streamAnswer) { return ToolResponse.error( "EXECUTION_FAILED", - "Streaming channel unavailable. Cannot deliver final answer." + "Streaming channel unavailable. Cannot deliver final answer.", ) } @@ -3326,7 +3576,7 @@ function createFinalSynthesisTool(): Tool { droppedImages: dropped, userAttachmentCount, }, - "[MessageAgents][FinalSynthesis] Image payload" + "[MessageAgents][FinalSynthesis] Image payload", ) const { systemPrompt, userMessage } = buildFinalSynthesisPayload(context) @@ -3337,13 +3587,23 @@ function createFinalSynthesisTool(): Tool { finalSynthesisSystemPrompt: systemPrompt, finalSynthesisUserMessage: userMessage, }, - "[MessageAgents][FinalSynthesis] Full context payload" + "[MessageAgents][FinalSynthesis] Full context payload", ) mutableContext.finalSynthesis.requested = true mutableContext.finalSynthesis.suppressAssistantStreaming = true mutableContext.finalSynthesis.completed = false mutableContext.finalSynthesis.streamedText = "" + logContextMutation( + mutableContext, + "[MessageAgents][FinalSynthesis] Updated final synthesis state to requested", + { + fragmentsCount, + selectedImages: selected, + totalImages: total, + droppedImages: dropped, + }, + ) await mutableContext.runtime?.emitReasoning?.({ text: `Initiating final synthesis with ${fragmentsCount} context fragments and ${selected.length}/${total} images (${userAttachmentCount} user attachments).`, @@ -3358,12 +3618,14 @@ function createFinalSynthesisTool(): Tool { limit: IMAGE_CONTEXT_CONFIG.maxImagesPerCall, totalImages: total, }, - "Final synthesis image limit enforced; dropped oldest references." + "Final synthesis image limit enforced; dropped oldest references.", ) } const modelId = - (context.modelId as Models) || (defaultBestModel as Models) || Models.Gpt_4o + (context.modelId as Models) || + (defaultBestModel as Models) || + Models.Gpt_4o const modelParams: ModelParams = { modelId, systemPrompt, @@ -3390,18 +3652,21 @@ function createFinalSynthesisTool(): Tool { toolOutputsThisTurn: context.currentTurnArtifacts.toolOutputs.length, imageNames: selected, }, - "[MessageAgents][FinalSynthesis] Context summary for synthesis call" + "[MessageAgents][FinalSynthesis] Context summary for synthesis call", + ) + + Logger.debug( + { + email: context.user.email, + chatId: context.chat.externalId, + modelId, + systemPrompt, + messagesCount: messages.length, + imagesProvided: selected.length, + }, + "[MessageAgents][FinalSynthesis] LLM call parameters", ) - Logger.debug({ - email: context.user.email, - chatId: context.chat.externalId, - modelId, - systemPrompt, - messagesCount: messages.length, - imagesProvided: selected.length, - }, "[MessageAgents][FinalSynthesis] LLM call parameters") - const provider = getProviderByModel(modelId) let streamedCharacters = 0 let estimatedCostUsd = 0 @@ -3421,6 +3686,15 @@ function createFinalSynthesisTool(): Tool { } context.finalSynthesis.completed = true + logContextMutation( + context, + "[MessageAgents][FinalSynthesis] Marked final synthesis as completed", + { + streamedCharacters, + estimatedCostUsd, + imagesProvided: selected, + }, + ) loggerWithChild({ email: context.user.email }).debug( { chatId: context.chat.externalId, @@ -3428,7 +3702,7 @@ function createFinalSynthesisTool(): Tool { estimatedCostUsd, imagesProvided: selected, }, - "[MessageAgents][FinalSynthesis] LLM call completed" + "[MessageAgents][FinalSynthesis] LLM call completed", ) await context.runtime?.emitReasoning?.({ @@ -3448,28 +3722,34 @@ function createFinalSynthesisTool(): Tool { }, { estimatedCostUsd, - } + }, ) } catch (error) { context.finalSynthesis.suppressAssistantStreaming = false context.finalSynthesis.requested = false context.finalSynthesis.completed = false + logContextMutation( + context, + "[MessageAgents][FinalSynthesis] Reset final synthesis state after failure", + { + error: error instanceof Error ? error.message : String(error), + }, + ) logger.error( { err: error instanceof Error ? error.message : String(error) }, - "Final synthesis tool failed." + "Final synthesis tool failed.", ) return ToolResponse.error( "EXECUTION_FAILED", `Failed to synthesize final answer: ${ error instanceof Error ? error.message : String(error) - }` + }`, ) } }, } } - /** * Build dynamic agent instructions including plan state and tool descriptions */ @@ -3485,9 +3765,13 @@ function buildAttachmentDirective(context: AgentRunContext): string { "User provided attachment context for this opening turn." // Include the actual fragment content, not just metadata - const fragmentsContent = context.allFragments.length > 0 - ? answerContextMapFromFragments(context.allFragments, context.allFragments.length) - : "No attachment content available." + const fragmentsContent = + context.allFragments.length > 0 + ? answerContextMapFromFragments( + context.allFragments, + context.allFragments.length, + ) + : "No attachment content available." return ` # ATTACHMENT-FIRST TURN @@ -3540,13 +3824,16 @@ function buildAgentInstructions( enabledToolNames: string[], dateForAI: string, agentPrompt?: string, - delegationEnabled = true + delegationEnabled = true, ): string { - const toolDescriptions = enabledToolNames.length > 0 - ? generateToolDescriptions(enabledToolNames) - : "No tools available yet. " + const toolDescriptions = + enabledToolNames.length > 0 + ? generateToolDescriptions(enabledToolNames) + : "No tools available yet. " - const agentSection = agentPrompt ? `\n\nAgent Constraints:\n${agentPrompt}` : "" + const agentSection = agentPrompt + ? `\n\nAgent Constraints:\n${agentPrompt}` + : "" const attachmentDirective = buildAttachmentDirective(context) const promptAddendum = buildAgentPromptAddendum() const reviewResultBlock = @@ -3565,9 +3852,13 @@ function buildAgentInstructions( planSection += "Steps:\n" context.plan.subTasks.forEach((task, i) => { const status = - task.status === "completed" ? "✓" : - task.status === "in_progress" ? "→" : - task.status === "failed" ? "✗" : "○" + task.status === "completed" + ? "✓" + : task.status === "in_progress" + ? "→" + : task.status === "failed" + ? "✗" + : "○" planSection += `${i + 1}. [${status}] ${task.description}\n` if (task.toolsRequired && task.toolsRequired.length > 0) { planSection += ` Tools: ${task.toolsRequired.join(", ")}\n` @@ -3575,7 +3866,8 @@ function buildAgentInstructions( }) planSection += "\n\n" } else { - planSection += "No plan exists yet. Use toDoWrite to create one.\n\n" + planSection += + "No plan exists yet. Use toDoWrite to create one.\n\n" } const delegationGuidance = delegationEnabled @@ -3631,7 +3923,7 @@ function buildAgentInstructions( "- Treat your self-review findings as mandatory: fix missed expectations, address anomalies, and resolve clarifications before progressing.", "- Log every required fix in the plan. Capture corrective sub-tasks and close them before moving forward.", "- Surface clarification questions to the user when intent is ambiguous.", - "" + "", ) } else { instructionLines.push( @@ -3641,11 +3933,10 @@ function buildAgentInstructions( "- Log every required fix directly in the plan so auditors can see alignment with the review.", "- When the review lists anomalies or ambiguity, capture each as a corrective sub-task (e.g., “Validate source for claim [2]”) and close it before moving forward.", "- Answer outstanding clarification questions immediately; if the user must respond, surface the exact question back to them.", - "" + "", ) } - instructionLines.push( "# PLANNING", USE_AGENT_SELF_REVIEW @@ -3689,7 +3980,7 @@ function buildAgentInstructions( "- After you decide which tools to call, emit a standalone expected-results block summarizing what each tool should achieve:", "", "[", - ' {', + " {", ' "toolName": "searchGlobal",', ' "goal": "Find Q4 ARR mentions",', ' "successCriteria": ["ARR keyword present", "Dated Q4"],', @@ -3706,12 +3997,11 @@ function buildAgentInstructions( "", "# FINAL SYNTHESIS", "- When research is complete and evidence is locked, CALL `synthesize_final_answer` (no arguments). This tool composes and streams the response.", - "- Never output the final answer directly—always go through the tool and then acknowledge completion." + "- Never output the final answer directly—always go through the tool and then acknowledge completion.", ) const finalInstructions = instructionLines.join("\n") - // Logger.debug({ // email: context.user.email, // chatId: context.chat.externalId, @@ -3743,7 +4033,7 @@ export async function MessageAgents(c: Context): Promise { const rootSpan = tracer.startSpan("MessageAgents") const { sub: email, workspaceId } = c.get(JwtPayloadKey) - + try { loggerWithChild({ email }).info("MessageAgents agentic flow starting") rootSpan.setAttribute("email", email) @@ -3765,11 +4055,11 @@ export async function MessageAgents(c: Context): Promise { toolsList?: Array<{ connectorId: string; tools: string[] }> selectedModelConfig?: string } = body - + if (!message) { throw new HTTPException(400, { message: "Message is required" }) } - + message = safeDecodeURIComponent(message) rootSpan.setAttribute("message", message) rootSpan.setAttribute("chatId", chatId || "new") @@ -3796,15 +4086,17 @@ export async function MessageAgents(c: Context): Promise { if (Array.isArray(modelConfig.capabilities)) { isReasoningEnabled = modelConfig.capabilities.includes("reasoning") enableWebSearch = modelConfig.capabilities.includes("websearch") - isDeepResearchEnabled = modelConfig.capabilities.includes("deepResearch") + isDeepResearchEnabled = + modelConfig.capabilities.includes("deepResearch") } else if (typeof modelConfig.capabilities === "object") { isReasoningEnabled = modelConfig.capabilities.reasoning === true enableWebSearch = modelConfig.capabilities.websearch === true - isDeepResearchEnabled = modelConfig.capabilities.deepResearch === true + isDeepResearchEnabled = + modelConfig.capabilities.deepResearch === true } } - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `Parsed model config for MessageAgents: model="${parsedModelId}", reasoning=${isReasoningEnabled}, websearch=${enableWebSearch}, deepResearch=${isDeepResearchEnabled}`, ) } catch (error) { @@ -3816,7 +4108,7 @@ export async function MessageAgents(c: Context): Promise { } } else { parsedModelId = config.defaultBestModel - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( "No model config provided to MessageAgents, using default", ) } @@ -3826,12 +4118,12 @@ export async function MessageAgents(c: Context): Promise { const convertedModelId = getModelValueFromLabel(parsedModelId) if (convertedModelId) { actualModelId = convertedModelId as string - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `Converted model label "${parsedModelId}" to value "${actualModelId}" for MessageAgents`, ) } else if (parsedModelId in Models) { actualModelId = parsedModelId - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `Using model ID "${parsedModelId}" directly for MessageAgents`, ) } else { @@ -3888,41 +4180,43 @@ export async function MessageAgents(c: Context): Promise { fileIds: [], threadIds: [], } - let attachmentsForContext = + let attachmentsForContext = extractedInfo?.fileIds.map((fileId) => ({ fileId, isImage: false, })) || [] const attachmentMetadata = parseAttachmentMetadata(c) - attachmentsForContext = attachmentsForContext.concat(attachmentMetadata.map((meta) => ({ - fileId: meta.fileId, - isImage: meta.isImage, - }))) + attachmentsForContext = attachmentsForContext.concat( + attachmentMetadata.map((meta) => ({ + fileId: meta.fileId, + isImage: meta.isImage, + })), + ) const threadIds = extractedInfo?.threadIds || [] const referencedFileIds = Array.from( new Set( attachmentsForContext .filter((meta) => !meta.isImage) - .flatMap((meta) => expandSheetIds(meta.fileId)) - ) + .flatMap((meta) => expandSheetIds(meta.fileId)), + ), ) const imageAttachmentFileIds = Array.from( - new Set(attachmentsForContext.filter((meta) => meta.isImage).map((meta) => meta.fileId)) + new Set( + attachmentsForContext + .filter((meta) => meta.isImage) + .map((meta) => meta.fileId), + ), ) const isMstWithAttachments = attachmentMetadata.length > 0 - const userAndWorkspace: InternalUserWorkspace = await getUserAndWorkspaceByEmail( - db, - workspaceId, - email - ) + const userAndWorkspace: InternalUserWorkspace = + await getUserAndWorkspaceByEmail(db, workspaceId, email) const rawUser = userAndWorkspace.user const rawWorkspace = userAndWorkspace.workspace const user = { id: Number(rawUser.id), email: String(rawUser.email), - timeZone: - typeof rawUser.timeZone === "string" ? rawUser.timeZone : "UTC", + timeZone: typeof rawUser.timeZone === "string" ? rawUser.timeZone : "UTC", } const workspace = { id: Number(rawWorkspace.id), @@ -3947,11 +4241,12 @@ export async function MessageAgents(c: Context): Promise { db, normalizedAgentId, workspace.id, - user.id + user.id, ) if (!agentRecord) { throw new HTTPException(403, { - message: "Access denied: You do not have permission to use this agent", + message: + "Access denied: You do not have permission to use this agent", }) } resolvedAgentId = String(agentRecord.externalId) @@ -4007,7 +4302,7 @@ export async function MessageAgents(c: Context): Promise { } catch (error) { loggerWithChild({ email }).error( error, - "Failed to persist user turn for MessageAgents" + "Failed to persist user turn for MessageAgents", ) const errMsg = error instanceof Error ? error.message : "Unknown persistence error" @@ -4033,7 +4328,7 @@ export async function MessageAgents(c: Context): Promise { db, resolvedAgentId, workspace.id, - user.id + user.id, ) if (!agentRecord) { throw new HTTPException(403, { @@ -4048,7 +4343,8 @@ export async function MessageAgents(c: Context): Promise { const hasExplicitAgent = Boolean(resolvedAgentId && agentPromptForLLM) const dedicatedAgentSystemPrompt = - typeof agentRecord?.prompt === "string" && agentRecord.prompt.trim().length > 0 + typeof agentRecord?.prompt === "string" && + agentRecord.prompt.trim().length > 0 ? agentRecord.prompt.trim() : undefined const delegationEnabled = !hasExplicitAgent @@ -4131,16 +4427,28 @@ export async function MessageAgents(c: Context): Promise { chatId: chatRecord.id as number, stopController, modelId: agenticModelId, - } + }, ) agentContextRef = agentContext agentContext.delegationEnabled = delegationEnabled + logContextMutation( + agentContext, + "[MessageAgents][Context] Updated delegationEnabled for primary run", + { + delegationEnabled, + hasExplicitAgent, + resolvedAgentId, + }, + ) // Build MCP connector tool map using the legacy agentic semantics const finalToolsMap: FinalToolsList = {} type FinalToolsEntry = FinalToolsList[string] type AdapterTool = FinalToolsEntry["tools"][number] - const connectorMetaById = new Map() + const connectorMetaById = new Map< + string, + { name?: string; description?: string } + >() if (toolsList && Array.isArray(toolsList) && toolsList.length > 0) { for (const item of toolsList) { @@ -4170,7 +4478,8 @@ export async function MessageAgents(c: Context): Promise { const client = new Client({ name: `connector-${connectorId}`, - version: (connector.config as { version?: string })?.version ?? "1.0", + version: + (connector.config as { version?: string })?.version ?? "1.0", }) const connectorNumericId = Number(connector.id) @@ -4188,27 +4497,34 @@ export async function MessageAgents(c: Context): Promise { const loadedMode = loadedConfig.mode || "sse" if (loadedUrl) { - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `Connecting to MCP client at ${loadedUrl} with mode: ${loadedMode}`, ) if (loadedMode === "streamable-http") { - const transportOptions: StreamableHTTPClientTransportOptions = { - requestInit: { headers: loadedHeaders }, - } + const transportOptions: StreamableHTTPClientTransportOptions = + { + requestInit: { headers: loadedHeaders }, + } await client.connect( - new StreamableHTTPClientTransport(new URL(loadedUrl), transportOptions), + new StreamableHTTPClientTransport( + new URL(loadedUrl), + transportOptions, + ), ) } else { const transportOptions: SSEClientTransportOptions = { requestInit: { headers: loadedHeaders }, } await client.connect( - new SSEClientTransport(new URL(loadedUrl), transportOptions), + new SSEClientTransport( + new URL(loadedUrl), + transportOptions, + ), ) } } else if (loadedConfig.command) { - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `Connecting to MCP Stdio client with command: ${loadedConfig.command}`, ) await client.connect( @@ -4233,7 +4549,11 @@ export async function MessageAgents(c: Context): Promise { mcpClients.push(client) let tools = [] try { - tools = await getToolsByConnectorId(db, workspace.id, connectorNumericId) + tools = await getToolsByConnectorId( + db, + workspace.id, + connectorNumericId, + ) } catch (error) { loggerWithChild({ email }).error( error, @@ -4243,11 +4563,13 @@ export async function MessageAgents(c: Context): Promise { } const filteredTools = tools.filter((tool) => { const toolExternalId = - typeof tool.externalId === "string" ? tool.externalId : undefined + typeof tool.externalId === "string" + ? tool.externalId + : undefined const isIncluded = !!toolExternalId && requestedToolIds.includes(toolExternalId) if (!isIncluded) { - loggerWithChild({ email }).info( + loggerWithChild({ email }).debug( `[MessageAgents][MCP] Tool ${toolExternalId}:${tool.toolName} not in requested toolExternalIds.`, ) } @@ -4371,22 +4693,39 @@ export async function MessageAgents(c: Context): Promise { ...customTools, ] agentContext.enabledTools = new Set( - allTools.map((tool) => tool.schema.name) + allTools.map((tool) => tool.schema.name), ) agentContext.mcpAgents = mcpAgentCandidates - loggerWithChild({ email }).info({ - totalToolBudget, - internalTools: internalTools.length, - directMcpTools: directMcpTools.length, - mcpAgents: mcpAgentCandidates.map((a) => a.agentId), - }, "[MessageAgents][MCP] Tool budget applied") - Logger.debug({ - enabledTools: Array.from(agentContext.enabledTools), - mcpAgentConnectors: Array.from(agentConnectorIds), - directMcpTools: directMcpTools.length, - email, - chatId: agentContext.chat.externalId, - }, "[MessageAgents] Tools exposed to LLM after filtering") + logContextMutation( + agentContext, + "[MessageAgents][Context] Updated enabled tools and MCP agents", + { + enabledTools: Array.from(agentContext.enabledTools), + mcpAgentIds: agentContext.mcpAgents.map((agent) => agent.agentId), + directMcpToolCount: directMcpTools.length, + internalToolCount: internalTools.length, + customToolCount: customTools.length, + }, + ) + loggerWithChild({ email }).debug( + { + totalToolBudget, + internalTools: internalTools.length, + directMcpTools: directMcpTools.length, + mcpAgents: mcpAgentCandidates.map((a) => a.agentId), + }, + "[MessageAgents][MCP] Tool budget applied", + ) + Logger.debug( + { + enabledTools: Array.from(agentContext.enabledTools), + mcpAgentConnectors: Array.from(agentConnectorIds), + directMcpTools: directMcpTools.length, + email, + chatId: agentContext.chat.externalId, + }, + "[MessageAgents] Tools exposed to LLM after filtering", + ) // Track gathered fragments const gatheredFragmentsKeys = new Set() @@ -4401,7 +4740,7 @@ export async function MessageAgents(c: Context): Promise { if (referencedFileIds.length > 0) { await streamReasoningStep( emitReasoningStep, - "Analyzing user-provided attachments..." + "Analyzing user-provided attachments...", ) initialAttachmentContext = await prepareInitialAttachmentContext( referencedFileIds, @@ -4416,7 +4755,7 @@ export async function MessageAgents(c: Context): Promise { emitReasoningStep, `Extracted ${initialAttachmentContext.fragments.length} context fragment${ initialAttachmentContext.fragments.length === 1 ? "" : "s" - } from attachments.` + } from attachments.`, ) } } @@ -4463,13 +4802,13 @@ export async function MessageAgents(c: Context): Promise { recordFragmentsForContext( agentContext, initialAttachmentContext.fragments, - MIN_TURN_NUMBER + MIN_TURN_NUMBER, ) initialSyntheticMessages.push( buildAttachmentToolMessage( initialAttachmentContext.fragments, - initialAttachmentContext.summary - ) + initialAttachmentContext.summary, + ), ) agentContext.chat.metadata = { ...agentContext.chat.metadata, @@ -4485,7 +4824,7 @@ export async function MessageAgents(c: Context): Promise { allTools.map((tool) => tool.schema.name), dateForAI, agentPromptForLLM, - delegationEnabled + delegationEnabled, ) } @@ -4501,15 +4840,16 @@ export async function MessageAgents(c: Context): Promise { const modelProvider = makeXyneJAFProvider() // Set up agent registry - const agentRegistry = new Map>([ - [jafAgent.name, jafAgent], - ]) + const agentRegistry = new Map< + string, + JAFAgent + >([[jafAgent.name, jafAgent]]) // Initialize run state const runId = generateRunId() const traceId = generateTraceId() const { jafHistory, llmHistory } = buildConversationHistoryForAgentRun( - previousConversationHistory + previousConversationHistory, ) const initialMessages: JAFMessage[] = [ ...jafHistory, @@ -4532,7 +4872,10 @@ export async function MessageAgents(c: Context): Promise { jafStreamingSpan.setAttribute("chat_external_id", chatRecord.externalId) jafStreamingSpan.setAttribute("run_id", runId) jafStreamingSpan.setAttribute("trace_id", traceId) - jafStreamingSpan.setAttribute("history_message_count", jafHistory.length) + jafStreamingSpan.setAttribute( + "history_message_count", + jafHistory.length, + ) jafStreamingSpan.setAttribute("history_seeded", jafHistory.length > 0) let turnSpan: Span | undefined const endTurnSpan = () => { @@ -4569,7 +4912,7 @@ export async function MessageAgents(c: Context): Promise { const recordExpectationsForTurn = ( turn: number, - expectations: PendingExpectation[] + expectations: PendingExpectation[], ) => { if (!expectations.length) { return @@ -4588,7 +4931,7 @@ export async function MessageAgents(c: Context): Promise { const ensureToolCallId = ( toolCall: ToolCallReference, turn: number, - index: number + index: number, ): string => { const mapKey = toolCall as object if (toolCall.id !== undefined && toolCall.id !== null) { @@ -4605,15 +4948,12 @@ export async function MessageAgents(c: Context): Promise { const buildTurnReviewInput = ( turn: number, - reviewFreq: number + reviewFreq: number, ): { reviewInput: AutoReviewInput } => { - const startTurn = Math.max( - MIN_TURN_NUMBER, - turn - reviewFreq + 1 - ) + const startTurn = Math.max(MIN_TURN_NUMBER, turn - reviewFreq + 1) const toolHistory = agentContext.toolCallHistory.filter( (record) => - record.turnNumber >= startTurn && record.turnNumber <= turn + record.turnNumber >= startTurn && record.turnNumber <= turn, ) const expectedResults: ToolExpectationAssignment[] = [] @@ -4633,14 +4973,17 @@ export async function MessageAgents(c: Context): Promise { } const runTurnEndReviewAndCleanup = async ( - turn: number + turn: number, ): Promise => { - Logger.debug({ - turn, - expectationHistoryKeys: Array.from(expectationHistory.keys()), - expectationsForThisTurn: expectationHistory.get(turn), - chatId: agentContext.chat.externalId, - }, "[DEBUG] Expectation history state at turn_end") + Logger.debug( + { + turn, + expectationHistoryKeys: Array.from(expectationHistory.keys()), + expectationsForThisTurn: expectationHistory.get(turn), + chatId: agentContext.chat.externalId, + }, + "[DEBUG] Expectation history state at turn_end", + ) try { if (!reviewsAllowed(agentContext)) { @@ -4650,13 +4993,13 @@ export async function MessageAgents(c: Context): Promise { chatId: agentContext.chat.externalId, lockedAtTurn: agentContext.review.lockedAtTurn, }, - "[MessageAgents] Review skipped because final synthesis was already requested." + "[MessageAgents] Review skipped because final synthesis was already requested.", ) return } if (!USE_AGENT_SELF_REVIEW) { const reviewFreq = normalizeReviewFrequency( - agentContext.review.reviewFrequency ?? DEFAULT_REVIEW_FREQUENCY + agentContext.review.reviewFrequency ?? DEFAULT_REVIEW_FREQUENCY, ) const runReviewThisTurn = turn === MIN_TURN_NUMBER || turn % reviewFreq === 0 @@ -4670,18 +5013,18 @@ export async function MessageAgents(c: Context): Promise { reviewFrequency: reviewFreq, chatId: agentContext.chat.externalId, }, - "[MessageAgents] Review skipped (runs every N turns)." + "[MessageAgents] Review skipped (runs every N turns).", ) } } else { await handleReviewOutcome( agentContext, buildDefaultReviewPayload( - "Self-review mode; no external review. State updated for continuity." + "Self-review mode; no external review. State updated for continuity.", ), turn, "turn_end", - emitReasoningStep + emitReasoningStep, ) } } catch (error) { @@ -4691,7 +5034,7 @@ export async function MessageAgents(c: Context): Promise { chatId: agentContext.chat.externalId, error: getErrorMessage(error), }, - "[MessageAgents] Turn-end review failed" + "[MessageAgents] Turn-end review failed", ) } finally { const attachmentState = getAttachmentPhaseMetadata(agentContext) @@ -4710,7 +5053,7 @@ export async function MessageAgents(c: Context): Promise { const runAndBroadcastReview = async ( context: AgentRunContext, reviewInput: AutoReviewInput, - iteration: number + iteration: number, ): Promise => { if (!reviewsAllowed(context)) { Logger.info( @@ -4720,7 +5063,7 @@ export async function MessageAgents(c: Context): Promise { lockedAtTurn: context.review.lockedAtTurn, focus: reviewInput.focus, }, - `[MessageAgents] Review skipped for focus '${reviewInput.focus}' due to final synthesis lock.` + `[MessageAgents] Review skipped for focus '${reviewInput.focus}' due to final synthesis lock.`, ) return null } @@ -4731,7 +5074,7 @@ export async function MessageAgents(c: Context): Promise { ) { Logger.warn( { turn: iteration, focus: reviewInput.focus }, - "[MessageAgents] No expected results recorded for review input." + "[MessageAgents] No expected results recorded for review input.", ) } @@ -4739,7 +5082,7 @@ export async function MessageAgents(c: Context): Promise { const pendingPromise = (async () => { const computedReview = await performAutomaticReview( reviewInput, - context + context, ) reviewResult = computedReview await handleReviewOutcome( @@ -4747,11 +5090,19 @@ export async function MessageAgents(c: Context): Promise { computedReview, iteration, reviewInput.focus, - emitReasoningStep + emitReasoningStep, ) })() context.review.pendingReview = pendingPromise + logContextMutation( + context, + "[MessageAgents][Context] Registered pending review promise", + { + iteration, + focus: reviewInput.focus, + }, + ) try { await pendingPromise if (!reviewResult) { @@ -4761,12 +5112,22 @@ export async function MessageAgents(c: Context): Promise { } finally { if (context.review.pendingReview === pendingPromise) { context.review.pendingReview = undefined + logContextMutation( + context, + "[MessageAgents][Context] Cleared pending review promise", + { + iteration, + focus: reviewInput.focus, + }, + ) } } } // Configure run with hooks - const runCfg: JAFRunConfig = { + const runCfg: JAFRunConfig & { + onEvent?: (event: TraceEvent) => void + } = { agentRegistry, modelProvider, maxTurns: 100, @@ -4778,23 +5139,21 @@ export async function MessageAgents(c: Context): Promise { onAfterToolExecution: async ( toolName: string, result: any, - hookContext: any + hookContext: any, ) => { const callIdRaw = hookContext?.toolCall?.id - const normalizedCallId = - hookContext?.toolCall - ? syntheticToolCallIds.get(hookContext.toolCall) ?? - (callIdRaw === undefined || callIdRaw === null - ? undefined - : String(callIdRaw)) - : undefined + const normalizedCallId = hookContext?.toolCall + ? (syntheticToolCallIds.get(hookContext.toolCall) ?? + (callIdRaw === undefined || callIdRaw === null + ? undefined + : String(callIdRaw))) + : undefined let expectationForCall: ToolExpectation | undefined if ( normalizedCallId && expectedResultsByCallId.has(normalizedCallId) ) { - expectationForCall = - expectedResultsByCallId.get(normalizedCallId) + expectationForCall = expectedResultsByCallId.get(normalizedCallId) expectedResultsByCallId.delete(normalizedCallId) } let turnForCall = normalizedCallId @@ -4803,14 +5162,9 @@ export async function MessageAgents(c: Context): Promise { if (normalizedCallId) { toolCallTurnMap.delete(normalizedCallId) } - if ( - turnForCall === undefined || - turnForCall < MIN_TURN_NUMBER - ) { + if (turnForCall === undefined || turnForCall < MIN_TURN_NUMBER) { turnForCall = - agentContext.turnCount ?? - currentTurn ?? - MIN_TURN_NUMBER + agentContext.turnCount ?? currentTurn ?? MIN_TURN_NUMBER } const content = await afterToolExecutionHook( toolName, @@ -4821,9 +5175,9 @@ export async function MessageAgents(c: Context): Promise { gatheredFragmentsKeys, expectationForCall, turnForCall, - emitReasoningStep + emitReasoningStep, ) - + return content }, } @@ -4889,7 +5243,6 @@ export async function MessageAgents(c: Context): Promise { fragmentsForCitations, yieldedImageCitations, email, - )) { if (citationEvent.citation) { const { index, item } = citationEvent.citation @@ -4918,17 +5271,36 @@ export async function MessageAgents(c: Context): Promise { emitReasoning: async (payload) => emitReasoningStep(payload as ReasoningPayload), } + logContextMutation( + agentContext, + "[MessageAgents][Context] Attached runtime callbacks", + { + hasStreamAnswerText: true, + hasEmitReasoning: true, + }, + ) const traceEventHandler = async (event: TraceEvent) => { if (event.type === "before_tool_execution") { return beforeToolExecutionHook( event.data.toolName, event.data.args, agentContext, - emitReasoningStep + emitReasoningStep, ) } return undefined } + runCfg.onEvent = (event) => { + logJAFTraceEvent( + { + chatId: agentContext.chat.externalId, + email, + flow: "MessageAgents", + runId, + }, + event, + ) + } Logger.debug( { @@ -4937,13 +5309,13 @@ export async function MessageAgents(c: Context): Promise { modelOverride: agenticModelId, email, }, - "[MessageAgents] Starting assistant call" + "[MessageAgents] Starting assistant call", ) for await (const evt of runStream( runState, runCfg, - traceEventHandler + traceEventHandler, )) { if (stream.closed) break @@ -4956,10 +5328,18 @@ export async function MessageAgents(c: Context): Promise { agentContext.turnCount = evt.data.turn currentTurn = evt.data.turn flushExpectationBufferToTurn(currentTurn) + logContextMutation( + agentContext, + "[MessageAgents][Context] Updated turnCount for turn start", + { + turnNumber: currentTurn, + agentName: evt.data.agentName, + }, + ) await streamReasoningStep( emitReasoningStep, `Turn ${currentTurn} started`, - { iteration: currentTurn } + { iteration: currentTurn }, ) break } @@ -4970,72 +5350,79 @@ export async function MessageAgents(c: Context): Promise { args: toolCall.args, })) const toolRequestsSpan = turnSpan?.startSpan("tool_requests") - toolRequestsSpan?.setAttribute("tool_calls_count", plannedTools.length) + toolRequestsSpan?.setAttribute( + "tool_calls_count", + plannedTools.length, + ) Logger.debug( { turn: currentTurn, plannedTools, chatId: agentContext.chat.externalId, }, - "[MessageAgents] Tool plan for turn" + "[MessageAgents] Tool plan for turn", ) for (const [idx, toolCall] of evt.data.toolCalls.entries()) { const normalizedCallId = ensureToolCallId( toolCall, currentTurn, - idx + idx, ) toolCallTurnMap.set(normalizedCallId, currentTurn) const assignedExpectation = consumePendingExpectation( pendingExpectations, - toolCall.name + toolCall.name, ) if (assignedExpectation) { expectedResultsByCallId.set( normalizedCallId, - assignedExpectation.expectation + assignedExpectation.expectation, ) } - const selectionSpan = toolRequestsSpan?.startSpan("tool_selection") + const selectionSpan = + toolRequestsSpan?.startSpan("tool_selection") selectionSpan?.setAttribute("tool_name", toolCall.name) selectionSpan?.setAttribute( "args", - JSON.stringify(toolCall.args ?? {}) + JSON.stringify(toolCall.args ?? {}), ) await streamReasoningStep( emitReasoningStep, `Tool selected: ${toolCall.name}`, - { toolName: toolCall.name } + { toolName: toolCall.name }, ) if (toolCall.name === "toDoWrite") { await streamReasoningStep( emitReasoningStep, "Formulating a step-by-step plan...", - { toolName: toolCall.name } + { toolName: toolCall.name }, ) } else if (toolCall.name === "list_custom_agents") { await streamReasoningStep( emitReasoningStep, "Searching for specialized agents that can help...", - { toolName: toolCall.name } + { toolName: toolCall.name }, ) } else if (toolCall.name === "run_public_agent") { - const agentId = (toolCall.args as { agentId?: string })?.agentId + const agentId = (toolCall.args as { agentId?: string }) + ?.agentId const agentName = agentContext.availableAgents.find( - (agent) => agent.agentId === agentId - )?.agentName || agentId || "selected agent" + (agent) => agent.agentId === agentId, + )?.agentName || + agentId || + "selected agent" await streamReasoningStep( emitReasoningStep, `Delegating sub-task to the '${agentName}' agent...`, - { toolName: toolCall.name, detail: agentName } + { toolName: toolCall.name, detail: agentName }, ) } else if (toolCall.name === "fall_back") { await streamReasoningStep( emitReasoningStep, "Initial strategy was unsuccessful. Activating fallback search to find an answer.", - { toolName: toolCall.name } + { toolName: toolCall.name }, ) } selectionSpan?.end() @@ -5049,22 +5436,25 @@ export async function MessageAgents(c: Context): Promise { toolStartSpan?.setAttribute("tool_name", evt.data.toolName) toolStartSpan?.setAttribute( "args", - JSON.stringify(evt.data.args ?? {}) + JSON.stringify(evt.data.args ?? {}), + ) + Logger.debug( + { + toolName: evt.data.toolName, + args: evt.data.args, + runId, + chatId: agentContext.chat.externalId, + turn: currentTurn, + }, + "[MessageAgents][Tool Start]", ) - Logger.debug({ - toolName: evt.data.toolName, - args: evt.data.args, - runId, - chatId: agentContext.chat.externalId, - turn: currentTurn, - }, "[MessageAgents][Tool Start]") await streamReasoningStep( emitReasoningStep, `Executing ${evt.data.toolName}...`, { toolName: evt.data.toolName, detail: JSON.stringify(evt.data.args ?? {}), - } + }, ) toolStartSpan?.end() break @@ -5075,22 +5465,25 @@ export async function MessageAgents(c: Context): Promise { toolEndSpan?.setAttribute("tool_name", evt.data.toolName) toolEndSpan?.setAttribute( "status", - evt.data.error ? "error" : evt.data.status ?? "completed" + evt.data.error ? "error" : (evt.data.status ?? "completed"), ) toolEndSpan?.setAttribute( "execution_time_ms", - evt.data.executionTime ?? 0 + evt.data.executionTime ?? 0, + ) + Logger.debug( + { + toolName: evt.data.toolName, + result: evt.data.result, + error: evt.data.error, + executionTime: evt.data.executionTime, + status: evt.data.error ? "error" : "success", + runId, + chatId: agentContext.chat.externalId, + turn: currentTurn, + }, + "[MessageAgents][Tool End]", ) - Logger.debug({ - toolName: evt.data.toolName, - result: evt.data.result, - error: evt.data.error, - executionTime: evt.data.executionTime, - status: evt.data.error ? "error" : "success", - runId, - chatId: agentContext.chat.externalId, - turn: currentTurn, - }, "[MessageAgents][Tool End]") await streamReasoningStep( emitReasoningStep, `Tool ${evt.data.toolName} completed`, @@ -5099,10 +5492,12 @@ export async function MessageAgents(c: Context): Promise { status: evt.data.error ? "error" : evt.data.status, detail: evt.data.error ? `Error: ${evt.data.error}` - : `Result: ${typeof evt.data.result === "string" - ? evt.data.result.slice(0, 800) - : JSON.stringify(evt.data.result).slice(0, 800)}`, - } + : `Result: ${ + typeof evt.data.result === "string" + ? evt.data.result.slice(0, 800) + : JSON.stringify(evt.data.result).slice(0, 800) + }`, + }, ) if (evt.data.error) { const newCount = @@ -5123,7 +5518,7 @@ export async function MessageAgents(c: Context): Promise { expectedResults: expectationHistory.get(currentTurn) || [], }, - currentTurn + currentTurn, ) } } @@ -5141,12 +5536,12 @@ export async function MessageAgents(c: Context): Promise { await streamReasoningStep( emitReasoningStep, "Turn complete. Reviewing progress and results...", - { iteration: evt.data.turn } + { iteration: evt.data.turn }, ) const currentTurnToolHistory = agentContext.toolCallHistory.filter( - (record) => record.turnNumber === evt.data.turn + (record) => record.turnNumber === evt.data.turn, ) if (currentTurnToolHistory.length > 0) { @@ -5154,15 +5549,15 @@ export async function MessageAgents(c: Context): Promise { emitReasoningStep, buildTurnToolReasoningSummary( evt.data.turn, - currentTurnToolHistory + currentTurnToolHistory, ), - { iteration: evt.data.turn } + { iteration: evt.data.turn }, ) } else { await streamReasoningStep( emitReasoningStep, `No tools were executed in turn ${evt.data.turn}.`, - { iteration: evt.data.turn } + { iteration: evt.data.turn }, ) } @@ -5186,7 +5581,7 @@ export async function MessageAgents(c: Context): Promise { expectedResults: expectationHistory.get(finalTurnNumber) || [], }, - finalTurnNumber + finalTurnNumber, ) } else { Logger.debug( @@ -5196,10 +5591,12 @@ export async function MessageAgents(c: Context): Promise { useSelfReview: USE_AGENT_SELF_REVIEW, chatId: agentContext.chat.externalId, }, - "[MessageAgents] Skipping run_end review (last turn = synthesis turn)." + "[MessageAgents] Skipping run_end review (last turn = synthesis turn).", ) } - loggerWithChild({ email }).info("Storing assistant response in database") + loggerWithChild({ email }).debug( + "Storing assistant response in database", + ) Logger.debug( { chatId: agentContext.chat.externalId, @@ -5209,7 +5606,7 @@ export async function MessageAgents(c: Context): Promise { imageCitationsCount: imageCitations.length, citationSample: citations.slice(0, 3), }, - "[MessageAgents][FinalSynthesis] LLM output preview" + "[MessageAgents][FinalSynthesis] LLM output preview", ) const totalCost = agentContext.totalCost const totalTokens = @@ -5238,15 +5635,18 @@ export async function MessageAgents(c: Context): Promise { } catch (error) { loggerWithChild({ email }).error( error, - "Failed to persist assistant response" + "Failed to persist assistant response", ) } - loggerWithChild({ email }).debug({ - answer, - citations: citations.length, - cost: totalCost, - tokens: totalTokens, - }, "Response generated successfully") + loggerWithChild({ email }).debug( + { + answer, + citations: citations.length, + cost: totalCost, + tokens: totalTokens, + }, + "Response generated successfully", + ) await stream.writeSSE({ event: ChatSSEvents.ResponseMetadata, data: JSON.stringify({ @@ -5271,11 +5671,12 @@ export async function MessageAgents(c: Context): Promise { hasToolCalls: Array.isArray(evt.data.message?.tool_calls) && (evt.data.message.tool_calls?.length ?? 0) > 0, - contentPreview: getTextContent(evt.data.message.content) - ?.slice(0, 200) || "", + contentPreview: + getTextContent(evt.data.message.content)?.slice(0, 200) || + "", chatId: agentContext.chat.externalId, }, - "[MessageAgents] Assistant output received" + "[MessageAgents] Assistant output received", ) const content = getTextContent(evt.data.message.content) || "" const hasToolCalls = @@ -5283,7 +5684,7 @@ export async function MessageAgents(c: Context): Promise { (evt.data.message.tool_calls?.length ?? 0) > 0 assistantSpan?.setAttribute("content_length", content.length) assistantSpan?.setAttribute("has_tool_calls", hasToolCalls) - + if (content) { const extractedExpectations = extractExpectedResults(content) Logger.debug({ @@ -5306,25 +5707,12 @@ export async function MessageAgents(c: Context): Promise { ) } pendingExpectations.push(...extractedExpectations) - agentContext.currentTurnArtifacts.expectations.push( - ...extractedExpectations - ) if (currentTurn > 0) { - Logger.debug({ - turn: currentTurn, - expectationsCount: extractedExpectations.length, - chatId: agentContext.chat.externalId, - }, "[DEBUG] Recording expectations for current turn") recordExpectationsForTurn( currentTurn, extractedExpectations ) } else { - Logger.debug({ - turn: currentTurn, - expectationsCount: extractedExpectations.length, - chatId: agentContext.chat.externalId, - }, "[DEBUG] Buffering expectations for future turn") expectationBuffer.push(...extractedExpectations) } } @@ -5333,7 +5721,7 @@ export async function MessageAgents(c: Context): Promise { if (hasToolCalls) { await streamReasoningStep( emitReasoningStep, - content || "Model planned tool usage." + content || "Model planned tool usage.", ) assistantSpan?.end() break @@ -5342,9 +5730,17 @@ export async function MessageAgents(c: Context): Promise { if (agentContext.finalSynthesis.suppressAssistantStreaming) { if (content?.trim()) { agentContext.finalSynthesis.ackReceived = true + logContextMutation( + agentContext, + "[MessageAgents][FinalSynthesis] Marked ackReceived from assistant message", + { + turnNumber: currentTurn, + contentPreview: truncateValue(content, 200), + }, + ) await streamReasoningStep( emitReasoningStep, - "Final synthesis acknowledged. Closing out the run." + "Final synthesis acknowledged. Closing out the run.", ) } assistantSpan?.end() @@ -5363,7 +5759,7 @@ export async function MessageAgents(c: Context): Promise { tokenUsageSpan.setAttribute("prompt_tokens", evt.data.prompt ?? 0) tokenUsageSpan.setAttribute( "completion_tokens", - evt.data.completion ?? 0 + evt.data.completion ?? 0, ) tokenUsageSpan.setAttribute("total_tokens", evt.data.total ?? 0) tokenUsageSpan.end() @@ -5371,7 +5767,9 @@ export async function MessageAgents(c: Context): Promise { } case "guardrail_violation": { - const guardrailSpan = jafStreamingSpan.startSpan("guardrail_violation") + const guardrailSpan = jafStreamingSpan.startSpan( + "guardrail_violation", + ) guardrailSpan.setAttribute("stage", evt.data.stage) guardrailSpan.setAttribute("reason", evt.data.reason) guardrailSpan.end() @@ -5382,7 +5780,7 @@ export async function MessageAgents(c: Context): Promise { const decodeSpan = jafStreamingSpan.startSpan("decode_error") decodeSpan.setAttribute( "errors", - JSON.stringify(evt.data.errors ?? []) + JSON.stringify(evt.data.errors ?? []), ) decodeSpan.end() break @@ -5399,16 +5797,16 @@ export async function MessageAgents(c: Context): Promise { case "clarification_requested": { const clarificationSpan = jafStreamingSpan.startSpan( - "clarification_requested" + "clarification_requested", ) clarificationSpan.setAttribute( "clarification_id", - evt.data.clarificationId + evt.data.clarificationId, ) clarificationSpan.setAttribute("question", evt.data.question) clarificationSpan.setAttribute( "options_count", - evt.data.options.length + evt.data.options.length, ) clarificationSpan.end() break @@ -5416,15 +5814,15 @@ export async function MessageAgents(c: Context): Promise { case "clarification_provided": { const clarificationProvidedSpan = jafStreamingSpan.startSpan( - "clarification_provided" + "clarification_provided", ) clarificationProvidedSpan.setAttribute( "clarification_id", - evt.data.clarificationId + evt.data.clarificationId, ) clarificationProvidedSpan.setAttribute( "selected_id", - evt.data.selectedId + evt.data.selectedId, ) clarificationProvidedSpan.end() break @@ -5445,7 +5843,7 @@ export async function MessageAgents(c: Context): Promise { } finalOutputSpan.setAttribute( "output_length", - typeof output === "string" ? output.length : 0 + typeof output === "string" ? output.length : 0, ) finalOutputSpan.end() break @@ -5456,25 +5854,27 @@ export async function MessageAgents(c: Context): Promise { const outcome = evt.data.outcome runEndSpan.setAttribute( "outcome_status", - outcome?.status ?? "unknown" + outcome?.status ?? "unknown", ) if (outcome?.status === "completed") { if (runCompletedAndPersisted) { runEndSpan.setAttribute("total_cost", agentContext.totalCost) runEndSpan.setAttribute( "total_tokens", - agentContext.tokenUsage.input + agentContext.tokenUsage.output + agentContext.tokenUsage.input + + agentContext.tokenUsage.output, ) runEndSpan.setAttribute("citations_count", citations.length) } else { const finalTurnNumber = Math.max( agentContext.turnCount ?? currentTurn ?? MIN_TURN_NUMBER, - MIN_TURN_NUMBER + MIN_TURN_NUMBER, ) const lastReviewTurn = agentContext.review.lastReviewTurn const skipRunEndReview = USE_AGENT_SELF_REVIEW || - (lastReviewTurn != null && lastReviewTurn === finalTurnNumber) + (lastReviewTurn != null && + lastReviewTurn === finalTurnNumber) if (!skipRunEndReview) { await runAndBroadcastReview( agentContext, @@ -5486,7 +5886,7 @@ export async function MessageAgents(c: Context): Promise { expectedResults: expectationHistory.get(finalTurnNumber) || [], }, - finalTurnNumber + finalTurnNumber, ) } else { Logger.debug( @@ -5496,10 +5896,12 @@ export async function MessageAgents(c: Context): Promise { useSelfReview: USE_AGENT_SELF_REVIEW, chatId: agentContext.chat.externalId, }, - "[MessageAgents] Skipping run_end review (self-review or already reviewed)." + "[MessageAgents] Skipping run_end review (self-review or already reviewed).", ) } - loggerWithChild({ email }).info("Storing assistant response in database") + loggerWithChild({ email }).debug( + "Storing assistant response in database", + ) Logger.debug( { chatId: agentContext.chat.externalId, @@ -5509,11 +5911,12 @@ export async function MessageAgents(c: Context): Promise { imageCitationsCount: imageCitations.length, citationSample: citations.slice(0, 3), }, - "[MessageAgents][FinalSynthesis] LLM output preview" + "[MessageAgents][FinalSynthesis] LLM output preview", ) const totalCost = agentContext.totalCost const totalTokens = - agentContext.tokenUsage.input + agentContext.tokenUsage.output + agentContext.tokenUsage.input + + agentContext.tokenUsage.output try { const assistantInsert = { chatId: chatRecord.id, @@ -5538,15 +5941,18 @@ export async function MessageAgents(c: Context): Promise { } catch (error) { loggerWithChild({ email }).error( error, - "Failed to persist assistant response" + "Failed to persist assistant response", ) } - loggerWithChild({ email }).debug({ - answer, - citations: citations.length, - cost: totalCost, - tokens: totalTokens, - }, "Response generated successfully") + loggerWithChild({ email }).debug( + { + answer, + citations: citations.length, + cost: totalCost, + tokens: totalTokens, + }, + "Response generated successfully", + ) runEndSpan.setAttribute("total_cost", totalCost) runEndSpan.setAttribute("total_tokens", totalTokens) runEndSpan.setAttribute("citations_count", citations.length) @@ -5670,7 +6076,7 @@ type ListAgentsParams = { } export async function listCustomAgentsSuitable( - params: ListAgentsParams + params: ListAgentsParams, ): Promise { const maxAgents = Math.min(Math.max(params.maxAgents ?? 5, 1), 10) let workspaceDbId = params.workspaceNumericId @@ -5678,11 +6084,12 @@ export async function listCustomAgentsSuitable( const mcpAgentsFromContext = params.mcpAgents ?? [] if (!workspaceDbId || !userDbId) { - const userAndWorkspace: InternalUserWorkspace = await getUserAndWorkspaceByEmail( - db, - params.workspaceExternalId, - params.userEmail - ) + const userAndWorkspace: InternalUserWorkspace = + await getUserAndWorkspaceByEmail( + db, + params.workspaceExternalId, + params.userEmail, + ) workspaceDbId = Number(userAndWorkspace.workspace.id) userDbId = Number(userAndWorkspace.user.id) } @@ -5692,7 +6099,7 @@ export async function listCustomAgentsSuitable( userDbId!, workspaceDbId!, 25, - 0 + 0, ) if (!accessibleAgents.length && mcpAgentsFromContext.length === 0) { @@ -5756,7 +6163,7 @@ export async function listCustomAgentsSuitable( `Select up to ${maxAgents} agents.`, "If no agent is unquestionably suitable, set agents to null.", "Only include an agent when you can cite concrete capability matches; otherwise leave it out.", - "You may return multiple agents when several are clearly relevant—rank the strongest ones first." + "You may return multiple agents when several are clearly relevant—rank the strongest ones first.", ].join(" ") const payload = [ @@ -5768,8 +6175,7 @@ export async function listCustomAgentsSuitable( formatAgentBriefsForPrompt(combinedBriefs), ].join("\n\n") - const modelId = - (defaultFastModel as Models) || (defaultBestModel as Models) + const modelId = (defaultFastModel as Models) || (defaultBestModel as Models) const modelParams: ModelParams = { modelId, json: true, @@ -5789,7 +6195,7 @@ export async function listCustomAgentsSuitable( const { text } = await getProviderByModel(modelId).converse( messages, - modelParams + modelParams, ) const parsed = jsonParseLLMOutput(text || "") @@ -5814,45 +6220,43 @@ export async function listCustomAgentsSuitable( loggerWithChild({ email: params.userEmail }).warn( { issue: validation.error.format() }, - "LLM agent selection output invalid, falling back to heuristic scoring" + "LLM agent selection output invalid, falling back to heuristic scoring", ) } catch (error) { loggerWithChild({ email: params.userEmail }).error( { err: error }, - "LLM agent selection failed, falling back to heuristic scoring" + "LLM agent selection failed, falling back to heuristic scoring", ) } loggerWithChild({ email: params.userEmail }).info( - { + { query: params.query, totalAgents: combinedBriefs.length, maxAgents, }, - "Using heuristic agent selection mechanism (LLM-based selection not available or failed)" + "Using heuristic agent selection mechanism (LLM-based selection not available or failed)", ) return buildHeuristicAgentSelection( combinedBriefs, params.query, maxAgents, - totalEvaluated + totalEvaluated, ) } -export async function executeCustomAgent( - params: { - agentId: string - query: string - userEmail: string - workspaceExternalId: string - contextSnippet?: string - maxTokens?: number - parentTurn?: number - mcpAgents?: MCPVirtualAgentRuntime[] - stopSignal?: AbortSignal - } -): Promise { +export async function executeCustomAgent(params: { + agentId: string + query: string + userEmail: string + workspaceExternalId: string + contextSnippet?: string + maxTokens?: number + parentTurn?: number + mcpAgents?: MCPVirtualAgentRuntime[] + stopSignal?: AbortSignal +}): Promise { const turnInfo = typeof params.parentTurn === "number" ? `\n\nTurn info: Parent turn number is ${params.parentTurn}. Continue numbering from here.` @@ -5924,10 +6328,7 @@ export async function executeCustomAgent( }, } } catch (error) { - Logger.error( - { err: error }, - "executeCustomAgent encountered an error" - ) + Logger.error({ err: error }, "executeCustomAgent encountered an error") return { result: "Agent execution threw an exception", error: getErrorMessage(error), @@ -5953,7 +6354,7 @@ type DelegatedAgentRunParams = { } async function runDelegatedAgentWithMessageAgents( - params: DelegatedAgentRunParams + params: DelegatedAgentRunParams, ): Promise { const logger = loggerWithChild({ email: params.userEmail }) const delegateModelId = resolveAgenticModelId(defaultBestModel) @@ -5962,7 +6363,7 @@ async function runDelegatedAgentWithMessageAgents( const userAndWorkspace = await getUserAndWorkspaceByEmail( db, params.workspaceExternalId, - params.userEmail + params.userEmail, ) const rawUser = userAndWorkspace.user const rawWorkspace = userAndWorkspace.workspace @@ -5970,7 +6371,9 @@ async function runDelegatedAgentWithMessageAgents( id: Number(rawUser.id), email: String(rawUser.email), timeZone: - typeof rawUser.timeZone === "string" ? rawUser.timeZone : "Asia/Kolkata", + typeof rawUser.timeZone === "string" + ? rawUser.timeZone + : "Asia/Kolkata", } const workspace = { id: Number(rawWorkspace.id), @@ -5980,7 +6383,7 @@ async function runDelegatedAgentWithMessageAgents( db, params.agentId, workspace.id, - user.id + user.id, ) if (!agentRecord) { @@ -5993,13 +6396,15 @@ async function runDelegatedAgentWithMessageAgents( const agentPromptForLLM = JSON.stringify(agentRecord) const dedicatedAgentSystemPrompt = - typeof agentRecord.prompt === "string" && agentRecord.prompt.trim().length > 0 + typeof agentRecord.prompt === "string" && + agentRecord.prompt.trim().length > 0 ? agentRecord.prompt.trim() : undefined const userCtxString = userContext(userAndWorkspace) const userTimezone = user.timeZone || "Asia/Kolkata" const dateForAI = getDateForAI({ userTimeZone: userTimezone }) - const attachmentsForContext: Array<{ fileId: string; isImage: boolean }> = [] + const attachmentsForContext: Array<{ fileId: string; isImage: boolean }> = + [] let connectorState = createEmptyConnectorState() try { @@ -6007,7 +6412,7 @@ async function runDelegatedAgentWithMessageAgents( } catch (error) { logger.warn( error, - "[DelegatedAgenticRun] Failed to load connector state; assuming no connectors" + "[DelegatedAgenticRun] Failed to load connector state; assuming no connectors", ) } @@ -6026,12 +6431,21 @@ async function runDelegatedAgentWithMessageAgents( workspaceNumericId: workspace.id, stopSignal: params.stopSignal, modelId: delegateModelId, - } + }, ) agentContext.delegationEnabled = false agentContext.ambiguityResolved = true agentContext.maxOutputTokens = params.maxTokens agentContext.mcpAgents = params.mcpAgents ?? [] + logContextMutation( + agentContext, + "[DelegatedAgenticRun][Context] Updated delegated agent context defaults", + { + delegationEnabled: agentContext.delegationEnabled, + maxOutputTokens: agentContext.maxOutputTokens, + mcpAgentCount: agentContext.mcpAgents.length, + }, + ) const allowedAgentApps = deriveAllowedAgentApps(agentPromptForLLM) const baseInternalTools = buildInternalToolAdapters() @@ -6064,7 +6478,16 @@ async function runDelegatedAgentWithMessageAgents( ...directMcpTools, ] agentContext.enabledTools = new Set( - allTools.map((tool) => tool.schema.name) + allTools.map((tool) => tool.schema.name), + ) + logContextMutation( + agentContext, + "[DelegatedAgenticRun][Context] Updated enabled tools", + { + enabledTools: Array.from(agentContext.enabledTools), + directMcpToolCount: directMcpTools.length, + internalToolCount: internalTools.length, + }, ) const gatheredFragmentsKeys = new Set() @@ -6075,7 +6498,7 @@ async function runDelegatedAgentWithMessageAgents( allTools.map((tool) => tool.schema.name), dateForAI, agentPromptForLLM, - false + false, ) } @@ -6129,7 +6552,7 @@ async function runDelegatedAgentWithMessageAgents( const recordExpectationsForTurn = ( turn: number, - expectations: PendingExpectation[] + expectations: PendingExpectation[], ) => { if (!expectations.length) return const existing = expectationHistory.get(turn) || [] @@ -6146,7 +6569,7 @@ async function runDelegatedAgentWithMessageAgents( const ensureToolCallId = ( toolCall: ToolCallReference, turn: number, - index: number + index: number, ): string => { const mapKey = toolCall as object if (toolCall.id !== undefined && toolCall.id !== null) { @@ -6163,12 +6586,11 @@ async function runDelegatedAgentWithMessageAgents( const buildTurnReviewInput = ( turn: number, - reviewFreq: number + reviewFreq: number, ): { reviewInput: AutoReviewInput } => { const startTurn = Math.max(MIN_TURN_NUMBER, turn - reviewFreq + 1) const toolHistory = agentContext.toolCallHistory.filter( - (record) => - record.turnNumber >= startTurn && record.turnNumber <= turn + (record) => record.turnNumber >= startTurn && record.turnNumber <= turn, ) const expectedResults: ToolExpectationAssignment[] = [] @@ -6187,15 +6609,16 @@ async function runDelegatedAgentWithMessageAgents( } } - const runTurnEndReviewAndCleanup = async ( - turn: number - ): Promise => { - Logger.debug({ - turn, - expectationHistoryKeys: Array.from(expectationHistory.keys()), - expectationsForThisTurn: expectationHistory.get(turn), - chatId: agentContext.chat.externalId, - }, "[DelegatedAgenticRun][DEBUG] Expectation history state at turn_end") + const runTurnEndReviewAndCleanup = async (turn: number): Promise => { + Logger.debug( + { + turn, + expectationHistoryKeys: Array.from(expectationHistory.keys()), + expectationsForThisTurn: expectationHistory.get(turn), + chatId: agentContext.chat.externalId, + }, + "[DelegatedAgenticRun][DEBUG] Expectation history state at turn_end", + ) try { if (!reviewsAllowed(agentContext)) { @@ -6205,13 +6628,13 @@ async function runDelegatedAgentWithMessageAgents( chatId: agentContext.chat.externalId, lockedAtTurn: agentContext.review.lockedAtTurn, }, - "[DelegatedAgenticRun] Review skipped because final synthesis was already requested." + "[DelegatedAgenticRun] Review skipped because final synthesis was already requested.", ) return } if (!USE_AGENT_SELF_REVIEW) { const reviewFreq = normalizeReviewFrequency( - agentContext.review.reviewFrequency ?? DEFAULT_REVIEW_FREQUENCY + agentContext.review.reviewFrequency ?? DEFAULT_REVIEW_FREQUENCY, ) const runReviewThisTurn = turn === MIN_TURN_NUMBER || turn % reviewFreq === 0 @@ -6225,26 +6648,29 @@ async function runDelegatedAgentWithMessageAgents( reviewFrequency: reviewFreq, chatId: agentContext.chat.externalId, }, - "[DelegatedAgenticRun] Review skipped (runs every N turns)." + "[DelegatedAgenticRun] Review skipped (runs every N turns).", ) } } else { await handleReviewOutcome( agentContext, buildDefaultReviewPayload( - "Self-review mode; no external review. State updated for continuity." + "Self-review mode; no external review. State updated for continuity.", ), turn, "turn_end", - emitReasoningStep + emitReasoningStep, ) } } catch (error) { - Logger.error({ - turn, - chatId: agentContext.chat.externalId, - error: getErrorMessage(error), - }, "[DelegatedAgenticRun] Turn-end review failed") + Logger.error( + { + turn, + chatId: agentContext.chat.externalId, + error: getErrorMessage(error), + }, + "[DelegatedAgenticRun] Turn-end review failed", + ) } finally { const attachmentState = getAttachmentPhaseMetadata(agentContext) if (attachmentState.initialAttachmentPhase) { @@ -6266,7 +6692,7 @@ async function runDelegatedAgentWithMessageAgents( const runAndBroadcastReview = async ( context: AgentRunContext, reviewInput: AutoReviewInput, - iteration: number + iteration: number, ): Promise => { if (!reviewsAllowed(context)) { Logger.info( @@ -6276,7 +6702,7 @@ async function runDelegatedAgentWithMessageAgents( lockedAtTurn: context.review.lockedAtTurn, focus: reviewInput.focus, }, - `[DelegatedAgenticRun] Review skipped for focus '${reviewInput.focus}' due to final synthesis lock.` + `[DelegatedAgenticRun] Review skipped for focus '${reviewInput.focus}' due to final synthesis lock.`, ) return null } @@ -6287,24 +6713,23 @@ async function runDelegatedAgentWithMessageAgents( ) { Logger.warn( { turn: iteration, focus: reviewInput.focus }, - "[DelegatedAgenticRun] No expected results recorded for review input." + "[DelegatedAgenticRun] No expected results recorded for review input.", ) } - const reviewResult = await performAutomaticReview( - reviewInput, - context - ) + const reviewResult = await performAutomaticReview(reviewInput, context) await handleReviewOutcome( context, reviewResult, iteration, reviewInput.focus, - emitReasoningStep + emitReasoningStep, ) return reviewResult } - const runCfg: JAFRunConfig = { + const runCfg: JAFRunConfig & { + onEvent?: (event: TraceEvent) => void + } = { agentRegistry, modelProvider, maxTurns: Math.min(DELEGATED_RUN_MAX_TURNS, 100), @@ -6315,21 +6740,17 @@ async function runDelegatedAgentWithMessageAgents( onAfterToolExecution: async ( toolName: string, result: any, - hookContext: any + hookContext: any, ) => { const callIdRaw = hookContext?.toolCall?.id - const normalizedCallId = - hookContext?.toolCall - ? syntheticToolCallIds.get(hookContext.toolCall) ?? - (callIdRaw === undefined || callIdRaw === null - ? undefined - : String(callIdRaw)) - : undefined + const normalizedCallId = hookContext?.toolCall + ? (syntheticToolCallIds.get(hookContext.toolCall) ?? + (callIdRaw === undefined || callIdRaw === null + ? undefined + : String(callIdRaw))) + : undefined let expectationForCall: ToolExpectation | undefined - if ( - normalizedCallId && - expectedResultsByCallId.has(normalizedCallId) - ) { + if (normalizedCallId && expectedResultsByCallId.has(normalizedCallId)) { expectationForCall = expectedResultsByCallId.get(normalizedCallId) expectedResultsByCallId.delete(normalizedCallId) } @@ -6339,14 +6760,8 @@ async function runDelegatedAgentWithMessageAgents( if (normalizedCallId) { toolCallTurnMap.delete(normalizedCallId) } - if ( - turnForCall === undefined || - turnForCall < MIN_TURN_NUMBER - ) { - turnForCall = - agentContext.turnCount ?? - currentTurn ?? - MIN_TURN_NUMBER + if (turnForCall === undefined || turnForCall < MIN_TURN_NUMBER) { + turnForCall = agentContext.turnCount ?? currentTurn ?? MIN_TURN_NUMBER } return afterToolExecutionHook( toolName, @@ -6357,7 +6772,7 @@ async function runDelegatedAgentWithMessageAgents( gatheredFragmentsKeys, expectationForCall, turnForCall, - emitReasoningStep + emitReasoningStep, ) }, } @@ -6368,7 +6783,18 @@ async function runDelegatedAgentWithMessageAgents( event.data.toolName, event.data.args, agentContext, - emitReasoningStep + emitReasoningStep, + ) + } + runCfg.onEvent = (event) => { + logJAFTraceEvent( + { + chatId: agentContext.chat.externalId, + email: params.userEmail, + flow: "DelegatedAgenticRun", + runId, + }, + event, ) } @@ -6383,6 +6809,10 @@ async function runDelegatedAgentWithMessageAgents( emitReasoning: async (payload) => emitReasoningStep(payload as ReasoningPayload), } + logContextMutation(agentContext, "[DelegatedAgenticRun][Context] Attached runtime callbacks", { + hasStreamAnswerText: true, + hasEmitReasoning: true, + }) let currentTurn = MIN_TURN_NUMBER let runCompleted = false @@ -6392,7 +6822,7 @@ async function runDelegatedAgentWithMessageAgents( for await (const evt of runStream( runState, runCfg, - traceEventHandler + traceEventHandler, )) { throwIfStopRequested(params.stopSignal) switch (evt.type) { @@ -6400,10 +6830,17 @@ async function runDelegatedAgentWithMessageAgents( agentContext.turnCount = evt.data.turn currentTurn = evt.data.turn flushExpectationBufferToTurn(currentTurn) + logContextMutation( + agentContext, + "[DelegatedAgenticRun][Context] Updated turnCount for turn start", + { + turnNumber: currentTurn, + }, + ) await streamReasoningStep( emitReasoningStep, `Turn ${currentTurn} started`, - { iteration: currentTurn } + { iteration: currentTurn }, ) break @@ -6412,23 +6849,23 @@ async function runDelegatedAgentWithMessageAgents( const normalizedCallId = ensureToolCallId( toolCall, currentTurn, - idx + idx, ) toolCallTurnMap.set(normalizedCallId, currentTurn) const assignedExpectation = consumePendingExpectation( pendingExpectations, - toolCall.name + toolCall.name, ) if (assignedExpectation) { expectedResultsByCallId.set( normalizedCallId, - assignedExpectation.expectation + assignedExpectation.expectation, ) } await streamReasoningStep( emitReasoningStep, `Tool selected: ${toolCall.name}`, - { toolName: toolCall.name } + { toolName: toolCall.name }, ) } break @@ -6440,7 +6877,7 @@ async function runDelegatedAgentWithMessageAgents( { toolName: evt.data.toolName, detail: JSON.stringify(evt.data.args ?? {}), - } + }, ) break @@ -6453,10 +6890,12 @@ async function runDelegatedAgentWithMessageAgents( status: evt.data.error ? "error" : evt.data.status, detail: evt.data.error ? `Error: ${evt.data.error}` - : `Result: ${typeof evt.data.result === "string" - ? evt.data.result.slice(0, 800) - : JSON.stringify(evt.data.result).slice(0, 800)}`, - } + : `Result: ${ + typeof evt.data.result === "string" + ? evt.data.result.slice(0, 800) + : JSON.stringify(evt.data.result).slice(0, 800) + }`, + }, ) if (evt.data.error) { const newCount = @@ -6477,7 +6916,7 @@ async function runDelegatedAgentWithMessageAgents( expectedResults: expectationHistory.get(currentTurn) || [], }, - currentTurn + currentTurn, ) } } @@ -6491,30 +6930,29 @@ async function runDelegatedAgentWithMessageAgents( await streamReasoningStep( emitReasoningStep, "Turn complete. Reviewing progress and results...", - { iteration: evt.data.turn } + { iteration: evt.data.turn }, ) - const currentTurnToolHistory = - agentContext.toolCallHistory.filter( - (record) => record.turnNumber === evt.data.turn - ) + const currentTurnToolHistory = agentContext.toolCallHistory.filter( + (record) => record.turnNumber === evt.data.turn, + ) if (currentTurnToolHistory.length > 0) { await streamReasoningStep( emitReasoningStep, buildTurnToolReasoningSummary( evt.data.turn, - currentTurnToolHistory + currentTurnToolHistory, ), - { iteration: evt.data.turn } + { iteration: evt.data.turn }, ) } else { - await streamReasoningStep( - emitReasoningStep, - `No tools were executed in turn ${evt.data.turn}.`, - { iteration: evt.data.turn } - ) - } + await streamReasoningStep( + emitReasoningStep, + `No tools were executed in turn ${evt.data.turn}.`, + { iteration: evt.data.turn }, + ) + } break } @@ -6554,7 +6992,7 @@ async function runDelegatedAgentWithMessageAgents( if (hasToolCalls) { await streamReasoningStep( emitReasoningStep, - content || "Model planned tool usage." + content || "Model planned tool usage.", ) break } @@ -6562,9 +7000,17 @@ async function runDelegatedAgentWithMessageAgents( if (agentContext.finalSynthesis.suppressAssistantStreaming) { if (content?.trim()) { agentContext.finalSynthesis.ackReceived = true + logContextMutation( + agentContext, + "[DelegatedAgenticRun][FinalSynthesis] Marked ackReceived from assistant message", + { + turnNumber: currentTurn, + contentPreview: truncateValue(content, 200), + }, + ) await streamReasoningStep( emitReasoningStep, - "Final synthesis acknowledged. Closing out the run." + "Final synthesis acknowledged. Closing out the run.", ) } break @@ -6595,7 +7041,7 @@ async function runDelegatedAgentWithMessageAgents( if (outcome?.status === "completed") { const finalTurnNumber = Math.max( agentContext.turnCount ?? currentTurn ?? MIN_TURN_NUMBER, - MIN_TURN_NUMBER + MIN_TURN_NUMBER, ) const lastReviewTurn = agentContext.review.lastReviewTurn const skipRunEndReview = @@ -6612,7 +7058,7 @@ async function runDelegatedAgentWithMessageAgents( expectedResults: expectationHistory.get(finalTurnNumber) || [], }, - finalTurnNumber + finalTurnNumber, ) } else { Logger.debug( @@ -6622,7 +7068,7 @@ async function runDelegatedAgentWithMessageAgents( useSelfReview: USE_AGENT_SELF_REVIEW, chatId: agentContext.chat.externalId, }, - "[DelegatedAgenticRun] Skipping run_end review." + "[DelegatedAgenticRun] Skipping run_end review.", ) } runCompleted = true @@ -6688,7 +7134,7 @@ async function runDelegatedAgentWithMessageAgents( yieldedCitations, fragmentsForCitations, yieldedImageCitations, - params.userEmail + params.userEmail, )) { if (event.citation) { citations.push(event.citation.item) @@ -6699,8 +7145,7 @@ async function runDelegatedAgentWithMessageAgents( } } - const finalAnswer = - answerForCitations || "Agent did not return any text." + const finalAnswer = answerForCitations || "Agent did not return any text." return { result: finalAnswer, @@ -6741,11 +7186,11 @@ type ExecuteMcpAgentOptions = { async function executeMcpAgent( agentId: string, query: string, - options: ExecuteMcpAgentOptions + options: ExecuteMcpAgentOptions, ): Promise { const connectorId = agentId.replace(/^mcp:/, "") const mcpAgent = options.mcpAgents?.find( - (agent) => agent.agentId === agentId || agent.connectorId === connectorId + (agent) => agent.agentId === agentId || agent.connectorId === connectorId, ) if (!mcpAgent) { return { @@ -6766,7 +7211,7 @@ async function executeMcpAgent( const toolList = mcpAgent.tools .map( (tool, idx) => - `${idx + 1}. ${tool.toolName} - ${tool.description ?? "No description provided"}` + `${idx + 1}. ${tool.toolName} - ${tool.description ?? "No description provided"}`, ) .join("\n") @@ -6790,9 +7235,11 @@ async function executeMcpAgent( let selectedToolName = mcpAgent.tools[0]?.toolName let selectedArgs: Record = {} let selectionRationale = "Heuristic default selection." - let selectedToolsArray: - | Array<{ toolName: string; arguments?: Record; rationale?: string }> - | null = null + let selectedToolsArray: Array<{ + toolName: string + arguments?: Record + rationale?: string + }> | null = null try { const provider = getProviderByModel(modelId) @@ -6844,33 +7291,33 @@ async function executeMcpAgent( (selectionResponse as { toolCalls?: typeof selectionResponse.tool_calls }) ?.toolCalls const calls = Array.isArray(responseToolCalls) - ? responseToolCalls.map((tc: any) => ({ - toolName: - tc.name ?? - tc.function?.name ?? - (typeof tc.toolName === "string" ? tc.toolName : undefined), - arguments: (() => { - const rawArgs = - tc.arguments ?? - tc.function?.arguments ?? - (typeof tc.args === "string" ? tc.args : "{}") - if (typeof rawArgs === "object" && rawArgs !== null) { - return rawArgs as Record - } - if (typeof rawArgs === "string") { - try { - return JSON.parse(rawArgs) - } catch { - return {} - } + ? responseToolCalls.map((tc: any) => ({ + toolName: + tc.name ?? + tc.function?.name ?? + (typeof tc.toolName === "string" ? tc.toolName : undefined), + arguments: (() => { + const rawArgs = + tc.arguments ?? + tc.function?.arguments ?? + (typeof tc.args === "string" ? tc.args : "{}") + if (typeof rawArgs === "object" && rawArgs !== null) { + return rawArgs as Record + } + if (typeof rawArgs === "string") { + try { + return JSON.parse(rawArgs) + } catch { + return {} } - return {} - })(), - rationale: - tc.rationale ?? - (typeof tc.reason === "string" ? tc.reason : undefined), - })) - : null + } + return {} + })(), + rationale: + tc.rationale ?? + (typeof tc.reason === "string" ? tc.reason : undefined), + })) + : null if (calls && calls.length > 0) { selectedToolsArray = calls @@ -7012,7 +7459,7 @@ type AgentBrief = { function buildAgentBrief( agent: any, - resourceAccess?: ResourceAccessSummary[] + resourceAccess?: ResourceAccessSummary[], ): AgentBrief { const integrations = extractIntegrationKeys(agent.appIntegrations) const domains = deriveDomainsFromIntegrations(integrations) @@ -7055,13 +7502,13 @@ Description: ${brief.description || "N/A"} Capabilities: ${brief.capabilities.join(", ") || "N/A"} Domains: ${brief.domains.join(", ")} Estimated cost: ${brief.estimatedCost} -Resource readiness: ${summarizeResourceAccess(brief.resourceAccess)}` +Resource readiness: ${summarizeResourceAccess(brief.resourceAccess)}`, ) .join("\n\n") } function summarizeResourceAccess( - resourceAccess?: ResourceAccessSummary[] + resourceAccess?: ResourceAccessSummary[], ): string { if (!resourceAccess || resourceAccess.length === 0) { return "unknown" @@ -7089,19 +7536,17 @@ function buildHeuristicAgentSelection( briefs: AgentBrief[], query: string, maxAgents: number, - totalEvaluated: number + totalEvaluated: number, ): ListCustomAgentsOutput { const tokens = query.toLowerCase().split(/\s+/) const scored = briefs.map((brief) => { const text = `${brief.agentName} ${brief.description} ${brief.capabilities.join(" ")}`.toLowerCase() const baseScore = - tokens.reduce( - (acc, token) => (text.includes(token) ? acc + 1 : acc), - 0 - ) / Math.max(tokens.length, 1) + tokens.reduce((acc, token) => (text.includes(token) ? acc + 1 : acc), 0) / + Math.max(tokens.length, 1) const penalty = brief.resourceAccess?.some( - (entry) => entry.status === "missing" + (entry) => entry.status === "missing", ) ? 0.3 : brief.resourceAccess?.some((entry) => entry.status === "partial") diff --git a/server/api/chat/tool-schemas.ts b/server/api/chat/tool-schemas.ts index a6cf8ab3b..e6e63a7ef 100644 --- a/server/api/chat/tool-schemas.ts +++ b/server/api/chat/tool-schemas.ts @@ -1,6 +1,6 @@ /** * Universal Tool Schema System for JAF Agentic Architecture - * + * * Defines strict input/output schemas for ALL tools to ensure: * - LLM outputs conform to structured formats * - Type safety and validation @@ -8,6 +8,7 @@ * - Easy tool discovery and documentation */ +import type { JSONSchema7, JSONSchema7Definition } from "json-schema" import { z } from "zod" import { Apps } from "@xyne/vespa-ts/types" import { @@ -17,9 +18,24 @@ import { SubTaskSchema, } from "./agent-schemas" import type { Entity, MailParticipant } from "@xyne/vespa-ts/types" +import { zodSchemaToJsonSchema } from "./jaf-provider-utils" import { timeRangeSchema } from "./tools/schemas" - -export type { ListCustomAgentsInput, RunPublicAgentInput } from "./agent-schemas" +import { + LsKnowledgeBaseInputSchema, + LS_KNOWLEDGE_BASE_TOOL_DESCRIPTION, + SEARCH_KNOWLEDGE_BASE_TOOL_DESCRIPTION, + SearchKnowledgeBaseInputSchema, +} from "./tools/knowledgeBaseFlow" + +export type { + ListCustomAgentsInput, + RunPublicAgentInput, +} from "./agent-schemas" +export type { + KnowledgeBaseTarget, + LsKnowledgeBaseToolParams, + SearchKnowledgeBaseToolParams, +} from "./tools/knowledgeBaseFlow" // ============================================================================ // UNIVERSAL TOOL SCHEMA STRUCTURE @@ -38,7 +54,7 @@ export interface ToolSchema { export enum ToolCategory { Planning = "planning", Search = "search", - Metadata = "metadata", + Metadata = "metadata", Agent = "agent", Clarification = "clarification", Review = "review", @@ -56,22 +72,72 @@ export interface ToolExample { // BASE SCHEMAS // ============================================================================ +const STANDARD_LIMIT_DESCRIPTION = + "Maximum number of results to return. Keep this small for precision-first retrieval and increase only when broader coverage is necessary." + +const STANDARD_OFFSET_DESCRIPTION = + "Pagination offset. Use it after reviewing the current page to continue from the next unseen results." + +const FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION = + "Previously seen result document `docId`s to suppress on follow-up searches. Prefer prior `fragment.source.docId` values. Do not pass collection, folder, file, path, or fragment IDs." + // Pagination schema export const PaginationSchema = z.object({ - limit: z.number().min(1).max(100).optional().default(20).describe("Maximum number of results to return (1-100)"), - offset: z.number().min(0).optional().default(0).describe("Number of results to skip for pagination"), + limit: z + .number() + .min(1) + .max(100) + .optional() + .default(20) + .describe(STANDARD_LIMIT_DESCRIPTION), + offset: z + .number() + .min(0) + .optional() + .default(0) + .describe(STANDARD_OFFSET_DESCRIPTION), }) // Sort direction schema -export const SortDirectionSchema = z.enum(["asc", "desc"]).optional() +export const SortDirectionSchema = z + .enum(["asc", "desc"]) + .describe( + "Sort direction. Use `desc` for newest-first or highest-priority-first ordering when supported, and `asc` for oldest-first ordering.", + ) + .optional() // Mail participant schema -export const MailParticipantSchema = z.object({ - from: z.array(z.string().email()).optional().describe("Sender email addresses"), - to: z.array(z.string().email()).optional().describe("Recipient email addresses"), - cc: z.array(z.string().email()).optional().describe("CC email addresses"), - bcc: z.array(z.string().email()).optional().describe("BCC email addresses"), -}).optional() +export const MailParticipantSchema = z + .object({ + from: z + .array(z.string()) + .optional() + .describe( + "Sender identifiers as strings. Email addresses are best; full names or organization names are also accepted when needed.", + ), + to: z + .array(z.string()) + .optional() + .describe( + "Primary recipient identifiers as strings. Email addresses are best; full names or organization names are also accepted when needed.", + ), + cc: z + .array(z.string()) + .optional() + .describe( + "CC recipient identifiers as strings. Email addresses are best; full names or organization names are also accepted when needed.", + ), + bcc: z + .array(z.string()) + .optional() + .describe( + "BCC recipient identifiers as strings. Email addresses are best; full names or organization names are also accepted when needed.", + ), + }) + .describe( + "Structured Gmail participant filter object with optional `from`, `to`, `cc`, and `bcc` string arrays. Use only the fields that are explicitly relevant to the query.", + ) + .optional() // ============================================================================ // TOOL OUTPUT SCHEMAS @@ -79,25 +145,34 @@ export const MailParticipantSchema = z.object({ // Standard tool output with contexts export const ToolOutputSchema = z.object({ - result: z.string().describe("Human-readable summary of the tool execution result"), - contexts: z.array(z.object({ - id: z.string(), - content: z.string(), - source: z.object({ - docId: z.string(), - title: z.string(), - url: z.string().default(""), - app: z.string(), - entity: z.any().optional(), - // Preserve rich citation metadata for downstream consumers (KB, Slack threads, etc.) - itemId: z.string().optional(), - clId: z.string().optional(), - page_title: z.string().optional(), - threadId: z.string().optional(), - parentThreadId: z.string().optional(), - }).catchall(z.any()), - confidence: z.number().min(0).max(1), - })).optional().describe("Retrieved context fragments"), + result: z + .string() + .describe("Human-readable summary of the tool execution result"), + contexts: z + .array( + z.object({ + id: z.string(), + content: z.string(), + source: z + .object({ + docId: z.string(), + title: z.string(), + url: z.string().default(""), + app: z.string(), + entity: z.any().optional(), + // Preserve rich citation metadata for downstream consumers (KB, Slack threads, etc.) + itemId: z.string().optional(), + clId: z.string().optional(), + page_title: z.string().optional(), + threadId: z.string().optional(), + parentThreadId: z.string().optional(), + }) + .catchall(z.any()), + confidence: z.number().min(0).max(1), + }), + ) + .optional() + .describe("Retrieved context fragments"), error: z.string().optional().describe("Error message if execution failed"), metadata: z .record(z.string(), z.any()) @@ -111,7 +186,6 @@ export type ToolOutput = z.infer // PLANNING TOOL SCHEMAS // ============================================================================ - // toDoWrite input schema export const ToDoWriteInputSchema = z.object({ goal: z.string().describe("The overarching goal to accomplish"), @@ -133,33 +207,47 @@ export const ToDoWriteOutputSchema = z.object({ export type ToDoWriteOutput = z.infer -// ============================================================================ +// ============================================================================ // FINAL SYNTHESIS TOOL SCHEMAS // ============================================================================ export const SynthesizeFinalAnswerInputSchema = z .object({}) - .describe("No arguments allowed. Invoke only when you are fully ready to deliver the final answer.") + .describe( + "No arguments allowed. Invoke only when you are fully ready to deliver the final answer.", + ) export const SynthesizeFinalAnswerOutputSchema = z.object({ result: z .string() - .describe("Confirmation that the final synthesis was executed (the actual answer is streamed to the user)."), + .describe( + "Confirmation that the final synthesis was executed (the actual answer is streamed to the user).", + ), streamed: z .boolean() - .describe("Indicates whether the answer was streamed to the user directly during tool execution."), + .describe( + "Indicates whether the answer was streamed to the user directly during tool execution.", + ), metadata: z .object({ textLength: z.number().describe("Characters streamed to the user."), - totalImagesAvailable: z.number().describe("Total tracked images across the run."), - imagesProvided: z.number().describe("Images forwarded to the final synthesis call, post limit."), + totalImagesAvailable: z + .number() + .describe("Total tracked images across the run."), + imagesProvided: z + .number() + .describe("Images forwarded to the final synthesis call, post limit."), }) .partial() .optional(), }) -export type SynthesizeFinalAnswerInput = z.infer -export type SynthesizeFinalAnswerOutput = z.infer +export type SynthesizeFinalAnswerInput = z.infer< + typeof SynthesizeFinalAnswerInputSchema +> +export type SynthesizeFinalAnswerOutput = z.infer< + typeof SynthesizeFinalAnswerOutputSchema +> // ============================================================================ // SEARCH TOOL SCHEMAS @@ -167,134 +255,184 @@ export type SynthesizeFinalAnswerOutput = z.infer // Gmail search input export const SearchGmailInputSchema = z.object({ - query: z.string().describe("Email search query"), + query: z + .string() + .optional() + .describe( + "Optional short email-content retrieval query string. Omit it when participant, label, or time filters already define the request well enough.", + ), participants: MailParticipantSchema, - labels: z.array(z.string()).optional().describe("Gmail labels to filter by"), + labels: z + .array(z.string()) + .optional() + .describe( + "Optional Gmail label strings to narrow the search, for example `INBOX`, `UNREAD`, `IMPORTANT`, `SENT`, or category labels.", + ), timeRange: timeRangeSchema, - limit: z.number().optional(), - offset: z.number().optional(), + limit: z.number().optional().describe(STANDARD_LIMIT_DESCRIPTION), + offset: z.number().optional().describe(STANDARD_OFFSET_DESCRIPTION), sortBy: SortDirectionSchema, - excludedIds: z.array(z.string()).optional(), + excludedIds: z + .array(z.string()) + .optional() + .describe(FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION), }) export type SearchGmailInput = z.infer // Drive search input export const SearchDriveInputSchema = z.object({ - query: z.string().describe("Drive file search query"), - owner: z.string().email().optional().describe("Filter by file owner email"), - filetype: z.array(z.string()).optional().describe("File entity types (e.g., 'document', 'spreadsheet')"), - timeRange: timeRangeSchema, - limit: z.number().optional(), - offset: z.number().optional(), - sortBy: SortDirectionSchema, - excludedIds: z.array(z.string()).optional(), -}) - -export type SearchDriveInput = z.infer - -export type SearchKnowledgeBaseToolParams = { - query: string - limit?: number - offset?: number - collectionId?: string - folderId?: string - fileId?: string - excludedIds?: string[] -} - -export const SearchKnowledgeBaseInputSchema: z.ZodType = z.object({ query: z - .string() - .min(1) - .describe("Keywords or phrases to search within the knowledge base"), - limit: z - .number() - .min(1) - .max(25) - .optional() - .describe("Maximum number of KB results to return"), - offset: z - .number() - .min(0) - .optional() - .describe("Pagination offset for KB results"), - collectionId: z .string() .optional() - .describe("Restrict search to a single knowledge base collection"), - folderId: z + .describe( + "Optional short content or title retrieval query string for Drive files. Omit it when owner, file-type, or time filters already define the request well enough.", + ), + owner: z .string() .optional() - .describe("Restrict search to a collection folder"), - fileId: z - .string() + .describe( + "Optional Drive owner identifier string. Email is preferred; owner display name can also work.", + ), + filetype: z + .array(z.string()) .optional() - .describe("Restrict search to a collection file"), + .describe( + "Optional Drive file-type strings. Valid values come from Drive entity types such as `docs`, `sheets`, `slides`, `presentation`, `pdf`, `folder`, `image`, `video`, `audio`, or `zip`.", + ), + timeRange: timeRangeSchema, + limit: z.number().optional().describe(STANDARD_LIMIT_DESCRIPTION), + offset: z.number().optional().describe(STANDARD_OFFSET_DESCRIPTION), + sortBy: SortDirectionSchema, excludedIds: z .array(z.string()) .optional() - .describe("Document IDs to exclude"), + .describe(FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION), }) +export type SearchDriveInput = z.infer +//NOTE : Knowledgebase scehma is defined in knowledgeBaseFlow file since it has some specific types related to KB targets and projections. export type SearchKnowledgeBaseInput = z.infer< typeof SearchKnowledgeBaseInputSchema > // Calendar search input export const SearchCalendarInputSchema = z.object({ - query: z.string().describe("Calendar event search query"), - attendees: z.array(z.string().email()).optional().describe("Filter by attendee emails"), - status: z.enum(["confirmed", "tentative", "cancelled"]).optional().describe("Event status"), + query: z.string().describe( + "Short meeting/topic retrieval query for calendar events. Put attendee, status, and time constraints in the dedicated filters.", + ), + attendees: z + .array(z.string()) + .optional() + .describe( + "Optional attendee identifiers as strings. Email addresses are preferred; attendee display names can also work.", + ), + status: z + .enum(["confirmed", "tentative", "cancelled"]) + .optional() + .describe("Optional event status filter."), timeRange: timeRangeSchema, - limit: z.number().optional(), - offset: z.number().optional(), + limit: z.number().optional().describe(STANDARD_LIMIT_DESCRIPTION), + offset: z.number().optional().describe(STANDARD_OFFSET_DESCRIPTION), sortBy: SortDirectionSchema, - excludedIds: z.array(z.string()).optional(), + excludedIds: z + .array(z.string()) + .optional() + .describe(FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION), }) export type SearchCalendarInput = z.infer // Google Contacts search input export const SearchGoogleContactsInputSchema = z.object({ - query: z.string().describe("Contact search query (name, email, phone, etc.)"), - limit: z.number().optional(), - offset: z.number().optional(), - excludedIds: z.array(z.string()).optional(), + query: z + .string() + .describe( + "Person or company identifier such as a name, email, phone number, or title. Use this to disambiguate people before searching other sources.", + ), + limit: z.number().optional().describe(STANDARD_LIMIT_DESCRIPTION), + offset: z.number().optional().describe(STANDARD_OFFSET_DESCRIPTION), + excludedIds: z + .array(z.string()) + .optional() + .describe(FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION), }) -export type SearchGoogleContactsInput = z.infer +export type SearchGoogleContactsInput = z.infer< + typeof SearchGoogleContactsInputSchema +> // Slack messages input -export const GetSlackMessagesInputSchema = z.object({ - filter_query: z.string().optional().describe("Keywords to search within messages"), - channel_name: z.string().optional().describe("Specific channel name"), - user_email: z.string().email().optional().describe("Filter by user email"), - date_from: z.string().optional().describe("Start date (ISO 8601)"), - date_to: z.string().optional().describe("End date (ISO 8601)"), - limit: z.number().optional(), - offset: z.number().optional(), - order_direction: SortDirectionSchema, +export const GetSlackRelatedMessagesInputSchema = z.object({ + query: z + .string() + .optional() + .describe( + "Optional short Slack message-content query string. Omit it when channel, author, mentions, or time filters already define the request well enough.", + ), + channelName: z + .string() + .optional() + .describe( + "Optional Slack channel name string, such as `eng-launches`. Pass the human-facing channel name, not a Slack channel ID.", + ), + user: z + .string() + .optional() + .describe( + "Optional Slack user identifier string to restrict messages by author. Email is preferred; display name can also work.", + ), + mentions: z + .array(z.string()) + .optional() + .describe( + "Optional list of mentioned-user identifier strings, usually emails or usernames, to find messages that mention specific people.", + ), + timeRange: timeRangeSchema, + limit: z.number().optional().describe(STANDARD_LIMIT_DESCRIPTION), + offset: z.number().optional().describe(STANDARD_OFFSET_DESCRIPTION), + sortBy: SortDirectionSchema, + excludedIds: z + .array(z.string()) + .optional() + .describe(FOLLOW_UP_EXCLUDED_IDS_DESCRIPTION), }) -export type GetSlackMessagesInput = z.infer +export const GetSlackMessagesInputSchema = GetSlackRelatedMessagesInputSchema + +export type GetSlackMessagesInput = z.infer< + typeof GetSlackRelatedMessagesInputSchema +> // Slack user profile input export const GetSlackUserProfileInputSchema = z.object({ - user_email: z.string().email().describe("Email of user whose Slack profile to retrieve"), + user_email: z + .string() + .email() + .describe("Email of user whose Slack profile to retrieve"), }) -export type GetSlackUserProfileInput = z.infer +export type GetSlackUserProfileInput = z.infer< + typeof GetSlackUserProfileInputSchema +> // ============================================================================ // AGENT TOOL SCHEMAS @@ -317,24 +455,30 @@ const ResourceAccessSummarySchema = z.object({ export const ListCustomAgentsOutputSchema = z.object({ agents: z - .array(z.object({ - agentId: z.string(), - agentName: z.string(), - description: z.string(), - capabilities: z.array(z.string()), - domains: z.array(z.string()), - suitabilityScore: z.number().min(0).max(1), - confidence: z.number().min(0).max(1), - estimatedCost: z.enum(["low", "medium", "high"]), - averageLatency: z.number(), - resourceAccess: z.array(ResourceAccessSummarySchema).optional(), - })) + .array( + z.object({ + agentId: z.string(), + agentName: z.string(), + description: z.string(), + capabilities: z.array(z.string()), + domains: z.array(z.string()), + suitabilityScore: z.number().min(0).max(1), + confidence: z.number().min(0).max(1), + estimatedCost: z.enum(["low", "medium", "high"]), + averageLatency: z.number(), + resourceAccess: z.array(ResourceAccessSummarySchema).optional(), + }), + ) .nullable() - .describe("Ordered list of best-fit agents. Return null when no agent is sufficiently certain."), + .describe( + "Ordered list of best-fit agents. Return null when no agent is sufficiently certain.", + ), totalEvaluated: z.number(), }) -export type ListCustomAgentsOutput = z.infer +export type ListCustomAgentsOutput = z.infer< + typeof ListCustomAgentsOutputSchema +> export type ResourceAccessItem = z.infer export type ResourceAccessSummary = z.infer @@ -348,7 +492,9 @@ export const ReviewAgentOutputSchema = z.object({ recommendation: z .enum(["proceed", "gather_more", "clarify_query", "replan"]) .describe("Next action recommendation"), - metExpectations: z.boolean().describe("Whether tool expectations were satisfied"), + metExpectations: z + .boolean() + .describe("Whether tool expectations were satisfied"), unmetExpectations: z .array(z.string()) .describe("List of expectation goals that remain unmet"), @@ -362,12 +508,8 @@ export const ReviewAgentOutputSchema = z.object({ toolFeedback: z .array(ToolReviewFindingSchema) .describe("Per-tool assessment for this turn"), - anomaliesDetected: z - .boolean() - .describe("Whether any anomalies were found"), - anomalies: z - .array(z.string()) - .describe("Descriptions of detected anomalies"), + anomaliesDetected: z.boolean().describe("Whether any anomalies were found"), + anomalies: z.array(z.string()).describe("Descriptions of detected anomalies"), }) export type ReviewAgentOutput = z.infer @@ -399,30 +541,14 @@ export const TOOL_SCHEMAS: Record = { // Planning Tools toDoWrite: { name: "toDoWrite", - description: "Create or update an execution plan with sequential tasks. MUST be called first before any other tool.", + description: + "Create or update an execution plan with sequential tasks. MUST be called first before any other tool.", category: ToolCategory.Planning, inputSchema: ToDoWriteInputSchema, outputSchema: ToDoWriteOutputSchema, - examples: [{ - input: { - goal: "Find what Alex says about Q4", - subTasks: [ - { - id: "task_1", - description: "Identify which Alex user is referring to", - status: "pending" as const, - toolsRequired: ["searchGoogleContacts"], - }, - { - id: "task_2", - description: "Search all of Alex's Q4 communications", - status: "pending" as const, - toolsRequired: ["searchGmail", "getSlackMessages", "searchDriveFiles"], - }, - ], - }, - output: { - plan: { + examples: [ + { + input: { goal: "Find what Alex says about Q4", subTasks: [ { @@ -431,17 +557,41 @@ export const TOOL_SCHEMAS: Record = { status: "pending" as const, toolsRequired: ["searchGoogleContacts"], }, + { + id: "task_2", + description: "Search all of Alex's Q4 communications", + status: "pending" as const, + toolsRequired: [ + "searchGmail", + "getSlackRelatedMessages", + "searchDriveFiles", + ], + }, ], }, + output: { + plan: { + goal: "Find what Alex says about Q4", + subTasks: [ + { + id: "task_1", + description: "Identify which Alex user is referring to", + status: "pending" as const, + toolsRequired: ["searchGoogleContacts"], + }, + ], + }, + }, + scenario: "Creating a task-based plan for ambiguous query", }, - scenario: "Creating a task-based plan for ambiguous query", - }], + ], }, // Search Tools searchGlobal: { name: "searchGlobal", - description: "Search across all accessible data sources. Use for broad searches when specific app is unknown.", + description: + "Search across all accessible data sources when the likely source is unclear. Prefer a more specific tool when the query already points clearly to Gmail, Drive, Slack, Calendar, Contacts, or a known KB location.", category: ToolCategory.Search, inputSchema: SearchGlobalInputSchema, outputSchema: ToolOutputSchema, @@ -449,15 +599,108 @@ export const TOOL_SCHEMAS: Record = { searchKnowledgeBase: { name: "searchKnowledgeBase", - description: "Search the user's knowledge base collections and return relevant document fragments with citations.", + description: SEARCH_KNOWLEDGE_BASE_TOOL_DESCRIPTION, category: ToolCategory.Search, inputSchema: SearchKnowledgeBaseInputSchema, outputSchema: ToolOutputSchema, + examples: [ + { + scenario: "Search a known KB folder directly without browsing first", + input: { + query: "security review exception policy", + filters: { + targets: [ + { + type: "path" as const, + collectionId: "kb-collection-123", + path: "/Policies/Security", + }, + ], + }, + limit: 5, + excludedIds: ["doc-prev-1"], + }, + output: { + result: "Found relevant policy fragments in the targeted KB folder.", + contexts: [], + }, + }, + { + scenario: + "Search only the PDF files identified from a prior ls call", + input: { + query: "vendor risk questionnaire requirements", + filters: { + targets: [ + { + type: "file" as const, + fileId: "kb-file-pdf-1", + }, + { + type: "file" as const, + fileId: "kb-file-pdf-2", + }, + ], + }, + limit: 5, + }, + output: { + result: "Searched only the selected PDF documents from the folder.", + contexts: [], + }, + }, + ], + }, + + ls: { + name: "ls", + description: LS_KNOWLEDGE_BASE_TOOL_DESCRIPTION, + category: ToolCategory.Metadata, + inputSchema: LsKnowledgeBaseInputSchema, + outputSchema: ToolOutputSchema, + examples: [ + { + scenario: "Inspect a known collection root before deciding whether to search inside it", + input: { + target: { + type: "collection" as const, + collectionId: "kb-collection-123", + }, + depth: 1, + limit: 20, + metadata: false, + }, + output: { + result: "Listed the top-level folders and files in the collection.", + contexts: [], + }, + }, + { + scenario: + "Inspect a folder with metadata so you can keep only PDF files for a later KB search", + input: { + target: { + type: "path" as const, + collectionId: "kb-collection-123", + path: "/Policies/Security", + }, + depth: 2, + limit: 50, + metadata: true, + }, + output: { + result: + "Listed files with mime types and timestamps so PDF rows can be selected for targeted search.", + contexts: [], + }, + }, + ], }, searchGmail: { name: "searchGmail", - description: "Search Gmail messages with filters for participants, labels, and time range.", + description: + "Search Gmail messages by content, participants, labels, and time range. Use participant filters for people and organizations instead of stuffing them into the query.", category: ToolCategory.Search, inputSchema: SearchGmailInputSchema, outputSchema: ToolOutputSchema, @@ -465,7 +708,8 @@ export const TOOL_SCHEMAS: Record = { searchDriveFiles: { name: "searchDriveFiles", - description: "Search Google Drive files with filters for owner, file type, and time range.", + description: + "Search Google Drive files by title/content with optional owner, file-type, and time filters.", category: ToolCategory.Search, inputSchema: SearchDriveInputSchema, outputSchema: ToolOutputSchema, @@ -473,7 +717,8 @@ export const TOOL_SCHEMAS: Record = { searchCalendarEvents: { name: "searchCalendarEvents", - description: "Search Google Calendar events with filters for attendees, status, and time range.", + description: + "Search Google Calendar events by topic with optional attendee, status, and time filters.", category: ToolCategory.Search, inputSchema: SearchCalendarInputSchema, outputSchema: ToolOutputSchema, @@ -481,24 +726,29 @@ export const TOOL_SCHEMAS: Record = { searchGoogleContacts: { name: "searchGoogleContacts", - description: "Search Google Contacts by name, email, or phone. Use for disambiguating person names.", + description: + "Search Google Contacts by name, email, phone, or title. Use this first when person identity is ambiguous.", category: ToolCategory.Search, inputSchema: SearchGoogleContactsInputSchema, outputSchema: ToolOutputSchema, - prerequisites: ["Must be used before contacting people with ambiguous names"], + prerequisites: [ + "Must be used before contacting people with ambiguous names", + ], }, - getSlackMessages: { - name: "getSlackMessages", - description: "Search Slack messages with flexible filters for channel, user, time range.", + getSlackRelatedMessages: { + name: "getSlackRelatedMessages", + description: + "Search Slack messages with flexible filters for content, channel, author, mentions, and time range. When neither query nor scope is provided, the live tool defaults to recent Slack history.", category: ToolCategory.Search, - inputSchema: GetSlackMessagesInputSchema, + inputSchema: GetSlackRelatedMessagesInputSchema, outputSchema: ToolOutputSchema, }, getSlackUserProfile: { name: "getSlackUserProfile", - description: "Get a user's Slack profile by email address.", + description: + "Get a user's Slack profile by email address. Use when you need identity, channel, or profile metadata before deeper Slack search.", category: ToolCategory.Metadata, inputSchema: GetSlackUserProfileInputSchema, outputSchema: ToolOutputSchema, @@ -530,7 +780,8 @@ export const TOOL_SCHEMAS: Record = { { agentId: "agent_renewal_nav", agentName: "Renewal Navigator", - description: "Summarizes customer renewals and risks across ACME accounts.", + description: + "Summarizes customer renewals and risks across ACME accounts.", capabilities: ["renewal_strategy", "deal_health"], domains: ["salesforce", "revops"], suitabilityScore: 0.93, @@ -548,7 +799,8 @@ export const TOOL_SCHEMAS: Record = { { agentId: "agent_revops_deepdive", agentName: "RevOps Deep Dive", - description: "Explains revenue risk drivers for enterprise deals.", + description: + "Explains revenue risk drivers for enterprise deals.", capabilities: ["renewal_strategy", "pipeline_insights"], domains: ["revops"], suitabilityScore: 0.81, @@ -588,16 +840,20 @@ export const TOOL_SCHEMAS: Record = { scenario: "Delegate Delta Airlines renewal recap to Renewal Navigator", input: { agentId: "agent_renewal_nav", - query: "Summarize Delta Airlines Q3 renewal blockers and list owners for each blocker.", - context: "Delta Airlines confirmed as DAL GTM account; timeframe Q3 FY25.", + query: + "Summarize Delta Airlines Q3 renewal blockers and list owners for each blocker.", + context: + "Delta Airlines confirmed as DAL GTM account; timeframe Q3 FY25.", maxTokens: 900, }, output: { - result: "Renewal Navigator summarized Delta Airlines blockers with owners.", + result: + "Renewal Navigator summarized Delta Airlines blockers with owners.", contexts: [ { id: "delta-renewal-summary", - content: "Risk stems from security review and pending legal redlines.", + content: + "Risk stems from security review and pending legal redlines.", source: { docId: "drive-doc-123", title: "Delta Renewal Brief", @@ -620,7 +876,8 @@ export const TOOL_SCHEMAS: Record = { // Fallback Tool fall_back: { name: "fall_back", - description: "Generate reasoning about why search failed when max iterations reached. Used automatically by system.", + description: + "Generate reasoning about why search failed when max iterations reached. Used automatically by system.", category: ToolCategory.Fallback, inputSchema: FallbackToolInputSchema, outputSchema: FallbackToolOutputSchema, @@ -662,7 +919,7 @@ export function getToolSchema(toolName: string): ToolSchema | undefined { */ export function validateToolInput( toolName: string, - input: unknown + input: unknown, ): { success: true; data: T } | { success: false; error: z.ZodError } { const schema = getToolSchema(toolName) if (!schema) { @@ -682,7 +939,7 @@ export function validateToolInput( */ export function validateToolOutput( toolName: string, - output: unknown + output: unknown, ): { success: true; data: T } | { success: false; error: z.ZodError } { const schema = getToolSchema(toolName) if (!schema) { @@ -697,6 +954,85 @@ export function validateToolOutput( } } +function isJsonSchemaObject( + value: JSONSchema7Definition | undefined, +): value is JSONSchema7 { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function formatJsonSchemaType(schema: JSONSchema7): string { + if (Array.isArray(schema.enum) && schema.enum.length > 0) { + return `enum(${schema.enum.map((value) => JSON.stringify(value)).join(", ")})` + } + + if (schema.type === "array") { + const itemType = Array.isArray(schema.items) + ? Array.from( + new Set( + schema.items + .filter(isJsonSchemaObject) + .map((item) => formatJsonSchemaType(item)), + ), + ).join(" | ") || "value" + : isJsonSchemaObject(schema.items) + ? formatJsonSchemaType(schema.items) + : "value" + return `array<${itemType}>` + } + + if (typeof schema.type === "string") { + return schema.type + } + + if (Array.isArray(schema.type)) { + return schema.type.join(" | ") + } + + if (Array.isArray(schema.anyOf) && schema.anyOf.length > 0) { + const parts = schema.anyOf + .filter(isJsonSchemaObject) + .map((entry) => formatJsonSchemaType(entry)) + return Array.from(new Set(parts)).join(" | ") || "value" + } + + if (schema.properties) return "object" + if (schema.additionalProperties) return "record" + + return "value" +} + +function formatParameterLines( + schema: JSONSchema7, + pathPrefix = "", + depth = 0, +): string[] { + const properties = schema.properties ?? {} + const required = new Set(schema.required ?? []) + const lines: string[] = [] + + for (const [name, definition] of Object.entries(properties)) { + if (!isJsonSchemaObject(definition)) continue + + const fullPath = pathPrefix ? `${pathPrefix}.${name}` : name + const requiredLabel = required.has(name) ? "required" : "optional" + const description = definition.description + ? `: ${definition.description}` + : "" + + lines.push( + `${" ".repeat(depth)}- \`${fullPath}\` (${formatJsonSchemaType(definition)}, ${requiredLabel})${description}`, + ) + + if (depth >= 1) continue + + if (definition.properties) { + lines.push(...formatParameterLines(definition, fullPath, depth + 1)) + } + } + + return lines +} + /** * Generate tool descriptions for LLM prompt */ @@ -712,21 +1048,30 @@ export function generateToolDescriptions(toolNames: string[]): string { desc += `**Description**: ${schema.description}\n\n` if (schema.prerequisites && schema.prerequisites.length > 0) { - desc += `**Prerequisites**:\n${schema.prerequisites.map(p => `- ${p}`).join('\n')}\n\n` + desc += `**Prerequisites**:\n${schema.prerequisites.map((p) => `- ${p}`).join("\n")}\n\n` } - desc += `**Input Schema**: Use the defined Zod schema for this tool\n\n` + const jsonSchema = zodSchemaToJsonSchema(schema.inputSchema) + const parameterLines = formatParameterLines(jsonSchema) + + if (parameterLines.length > 0) { + desc += `**Parameters**:\n${parameterLines.join("\n")}\n\n` + } else { + desc += `**Parameters**: No arguments.\n\n` + } if (schema.examples && schema.examples.length > 0) { - desc += `**Example**:\n` - desc += `Scenario: ${schema.examples[0].scenario}\n` - desc += `\`\`\`json\n${JSON.stringify(schema.examples[0].input, null, 2)}\n\`\`\`\n\n` + desc += `**Examples**:\n` + for (const example of schema.examples.slice(0, 2)) { + desc += `Scenario: ${example.scenario}\n` + desc += `\`\`\`json\n${JSON.stringify(example.input, null, 2)}\n\`\`\`\n\n` + } } descriptions.push(desc) } - return descriptions.join('\n---\n\n') + return descriptions.join("\n---\n\n") } /** diff --git a/server/api/chat/tools/global/index.ts b/server/api/chat/tools/global/index.ts index 788b50161..fcafec408 100644 --- a/server/api/chat/tools/global/index.ts +++ b/server/api/chat/tools/global/index.ts @@ -72,7 +72,7 @@ export const searchGlobalTool: Tool = { schema: { name: "searchGlobal", description: - "Search across all connected applications and data sources to find relevant information. This is the primary search tool that can look through emails, documents, messages, and other content.", + "Search across all connected applications and data sources when the likely source is unclear. Prefer a more specific tool when the query already points clearly to Gmail, Drive, Slack, Calendar, Contacts, or a known knowledge-base location.", parameters: toToolSchemaParameters(searchGlobalToolSchema), }, async execute(params: WithExcludedIds, context: Ctx) { diff --git a/server/api/chat/tools/google/calendar.ts b/server/api/chat/tools/google/calendar.ts index ca0d48f7e..87c5172d5 100644 --- a/server/api/chat/tools/google/calendar.ts +++ b/server/api/chat/tools/google/calendar.ts @@ -14,12 +14,14 @@ const calendarSearchToolSchema = z.object({ query: createQuerySchema(GoogleApps.Calendar, true), attendees: z .array(z.string()) - .describe("Filter events by attendee name or email addresses") + .describe( + "Optional attendee identifier strings. Email addresses are preferred; attendee display names can also work.", + ) .optional(), status: z .enum(["confirmed", "tentative", "cancelled"]) .describe( - "Filter events by status. Available statuses: 'confirmed', 'tentative', 'cancelled'", + "Optional event status enum. Valid values are `confirmed`, `tentative`, and `cancelled`.", ) .optional(), ...baseToolParams, @@ -38,7 +40,7 @@ export const searchCalendarEventsTool: Tool = { schema: { name: "searchCalendarEvents", description: - "Retrieve calendar events and meetings from Google Calendar. Search by event title, attendees, or time period. Ideal for scheduling analysis, meeting preparation, and availability checking.", + "Search Google Calendar events by meeting topic with optional attendee, status, and time filters. Use attendee and time fields for scheduling or meeting-history queries instead of overloading the query text.", parameters: toToolSchemaParameters(calendarSearchToolSchema), }, async execute( diff --git a/server/api/chat/tools/google/contacts.ts b/server/api/chat/tools/google/contacts.ts index b4c47af23..873e505c8 100644 --- a/server/api/chat/tools/google/contacts.ts +++ b/server/api/chat/tools/google/contacts.ts @@ -27,7 +27,7 @@ export const searchGoogleContactsTool: Tool = { schema: { name: "searchGoogleContacts", description: - "Find people and contact information from Google Contacts. Search by name, email or organization. Useful for contact lookup, networking, and communication planning.", + "Search Google Contacts for people or organizations by name, email, phone number, title, or company. Use this to disambiguate identity before searching other apps.", parameters: toToolSchemaParameters(contactsSearchToolSchema), }, async execute( diff --git a/server/api/chat/tools/google/drive.ts b/server/api/chat/tools/google/drive.ts index 0b78bc2a2..61481a018 100644 --- a/server/api/chat/tools/google/drive.ts +++ b/server/api/chat/tools/google/drive.ts @@ -12,12 +12,17 @@ import { baseToolParams, createQuerySchema } from "../schemas" const driveSearchToolSchema = z.object({ query: createQuerySchema(GoogleApps.Drive), - owner: z.string().describe("Filter files by owner").optional(), + owner: z + .string() + .describe( + "Optional Drive owner identifier string. Email is preferred; owner display name can also work.", + ) + .optional(), ...baseToolParams, filetype: z .array(z.nativeEnum(DriveEntity)) .describe( - `Filter files by type. Available types: ${Object.values(DriveEntity) + `Optional Drive file-type enum values. Valid values are ${Object.values(DriveEntity) .map((e) => `'${e}'`) .join(", ")}.`, ) @@ -37,7 +42,7 @@ export const searchDriveFilesTool: Tool = { schema: { name: "searchDriveFiles", description: - "Access and search files in Google Drive. Find documents, spreadsheets, presentations, PDFs, and folders by name, content, owner, or file type. Essential for document management and collaboration.", + "Search Google Drive files by title/content with optional owner, file-type, and time filters. Use file types when the ask is constrained to PDFs, folders, spreadsheets, or other specific Drive entities.", parameters: toToolSchemaParameters(driveSearchToolSchema), }, async execute( diff --git a/server/api/chat/tools/google/gmail.ts b/server/api/chat/tools/google/gmail.ts index c994640cc..bbd67e6b2 100644 --- a/server/api/chat/tools/google/gmail.ts +++ b/server/api/chat/tools/google/gmail.ts @@ -12,12 +12,38 @@ import { baseToolParams, createQuerySchema } from "../schemas" export const participantsSchema = z .object({ - from: z.array(z.string().describe("From email addresses")).optional(), - to: z.array(z.string().describe("To email addresses")).optional(), - cc: z.array(z.string().describe("CC email addresses")).optional(), - bcc: z.array(z.string().describe("BCC email addresses")).optional(), + from: z + .array( + z.string().describe( + "Sender identifier string. Email is preferred; full name or organization name can also work.", + ), + ) + .optional(), + to: z + .array( + z.string().describe( + "Primary recipient identifier string. Email is preferred; full name or organization name can also work.", + ), + ) + .optional(), + cc: z + .array( + z.string().describe( + "CC recipient identifier string. Email is preferred; full name or organization name can also work.", + ), + ) + .optional(), + bcc: z + .array( + z.string().describe( + "BCC recipient identifier string. Email is preferred; full name or organization name can also work.", + ), + ) + .optional(), }) - .describe("Email participants filter") + .describe( + "Structured Gmail participant filter object with optional `from`, `to`, `cc`, and `bcc` string arrays.", + ) const gmailSearchToolSchema = z.object({ query: createQuerySchema(GoogleApps.Gmail), @@ -25,7 +51,7 @@ const gmailSearchToolSchema = z.object({ labels: z .array(z.string().describe("Gmail label")) .describe( - "Filter emails by Gmail labels. labels are 'IMPORTANT', 'STARRED', 'UNREAD', 'CATEGORY_PERSONAL', 'CATEGORY_SOCIAL', 'CATEGORY_PROMOTIONS', 'CATEGORY_UPDATES', 'CATEGORY_FORUMS', 'DRAFT', 'SENT', 'INBOX', 'SPAM', 'TRASH'.", + "Optional Gmail label strings used to narrow the search. Common values include `IMPORTANT`, `STARRED`, `UNREAD`, `CATEGORY_PERSONAL`, `CATEGORY_SOCIAL`, `CATEGORY_PROMOTIONS`, `CATEGORY_UPDATES`, `CATEGORY_FORUMS`, `DRAFT`, `SENT`, `INBOX`, `SPAM`, and `TRASH`.", ) .optional(), participants: participantsSchema @@ -48,7 +74,7 @@ export const searchGmailTool: Tool = { schema: { name: "searchGmail", description: - "Find and retrieve emails. Can search by keywords, filter by sender/recipient, time period, labels, or simply fetch recent emails when no query is provided.", + "Search Gmail messages by content with optional participant, label, and time filters. Omit the query when the sender/recipient/time constraints already define the request well enough.", parameters: toToolSchemaParameters(gmailSearchToolSchema), }, async execute(params: WithExcludedIds, context: Ctx) { diff --git a/server/api/chat/tools/knowledgeBase.ts b/server/api/chat/tools/knowledgeBase.ts deleted file mode 100644 index 7dbfb2c92..000000000 --- a/server/api/chat/tools/knowledgeBase.ts +++ /dev/null @@ -1,139 +0,0 @@ -import type { Tool } from "@xynehq/jaf" -import { ToolErrorCodes, ToolResponse } from "@xynehq/jaf" -import type { ZodType } from "zod" -import { Apps } from "@xyne/vespa-ts/types" -import { getErrorMessage } from "@/utils" -import { getLogger } from "@/logger" -import { Subsystem } from "@/types" -import type { Ctx } from "./types" -import { - SearchKnowledgeBaseInputSchema, - type SearchKnowledgeBaseToolParams, -} from "../tool-schemas" -import { parseAgentAppIntegrations } from "./utils" -import { - buildKnowledgeBaseCollectionSelections, - KnowledgeBaseScope, - type KnowledgeBaseSelection, -} from "@/api/chat/knowledgeBaseSelections" -import { executeVespaSearch } from "./global" - -const Logger = getLogger(Subsystem.Chat) -type ToolSchemaParameters = Tool["schema"]["parameters"] -const toToolSchemaParameters = ( - schema: ZodType, -): ToolSchemaParameters => schema as unknown as ToolSchemaParameters - -const buildOverrideSelections = ( - params: SearchKnowledgeBaseToolParams, -): KnowledgeBaseSelection[] => { - const selection: KnowledgeBaseSelection = {} - if (params.collectionId) selection.collectionIds = [params.collectionId] - if (params.folderId) selection.collectionFolderIds = [params.folderId] - if (params.fileId) selection.collectionFileIds = [params.fileId] - - return Object.keys(selection).length ? [selection] : [] -} - -export const searchKnowledgeBaseTool: Tool< - SearchKnowledgeBaseToolParams, - Ctx -> = { - schema: { - name: "searchKnowledgeBase", - description: - "Search the user's knowledge base collections and return relevant document fragments with citations.", - parameters: toToolSchemaParameters( - SearchKnowledgeBaseInputSchema, - ), - }, - async execute(params, context) { - const email = context.user.email - if (!email) { - return ToolResponse.error( - ToolErrorCodes.MISSING_REQUIRED_FIELD, - "User email not found while executing searchKnowledgeBase", - { toolName: "searchKnowledgeBase" }, - ) - } - - const query = params.query?.trim() - if (!query) { - return ToolResponse.error( - ToolErrorCodes.MISSING_REQUIRED_FIELD, - "Query cannot be empty for knowledge base search", - { toolName: "searchKnowledgeBase" }, - ) - } - - try { - const agentPrompt = context.agentPrompt - const { selectedItems } = parseAgentAppIntegrations(agentPrompt) - const scope = agentPrompt - ? KnowledgeBaseScope.AgentScoped - : KnowledgeBaseScope.UserOwned - - const baseSelections = await buildKnowledgeBaseCollectionSelections({ - scope, - email, - selectedItems, - }) - - const overrides = buildOverrideSelections(params) - const collectionSelections = - overrides.length > 0 ? overrides : baseSelections - - Logger.info( - { - email, - scope, - baseSelectionCount: baseSelections.length, - overrideSelectionCount: overrides.length, - appliedSelectionCount: collectionSelections.length, - selectedItemKeys: Object.keys( - selectedItems as Record, - ).length, - }, - "[MessageAgents][searchKnowledgeBaseTool] Using KnowledgeBaseScope for KB search", - ) - - if (!collectionSelections.length) { - return ToolResponse.error( - ToolErrorCodes.EXECUTION_FAILED, - "No accessible knowledge base collections found for this user", - { toolName: "searchKnowledgeBase" }, - ) - } - - const fragments = await executeVespaSearch({ - email, - query, - app: Apps.KnowledgeBase, - agentAppEnums: [Apps.KnowledgeBase], - limit: params.limit, - offset: params.offset ?? 0, - excludedIds: params.excludedIds, - collectionSelections, - selectedItems, - userId: context.user.numericId ?? undefined, - workspaceId: context.user.workspaceNumericId ?? undefined, - }) - - if (!fragments.length) { - return ToolResponse.error( - ToolErrorCodes.EXECUTION_FAILED, - "No knowledge base results found for the query.", - { toolName: "searchKnowledgeBase" }, - ) - } - - return ToolResponse.success(fragments) - } catch (error) { - return ToolResponse.error( - ToolErrorCodes.EXECUTION_FAILED, - `Knowledge base search failed: ${getErrorMessage(error)}`, - { toolName: "searchKnowledgeBase" }, - ) - } - }, -} diff --git a/server/api/chat/tools/knowledgeBaseFlow.ts b/server/api/chat/tools/knowledgeBaseFlow.ts new file mode 100644 index 000000000..ee6e99117 --- /dev/null +++ b/server/api/chat/tools/knowledgeBaseFlow.ts @@ -0,0 +1,1469 @@ +import type { Tool } from "@xynehq/jaf" +import { ToolErrorCodes, ToolResponse } from "@xynehq/jaf" +import { Apps } from "@xyne/vespa-ts/types" +import { and, eq, inArray, isNull } from "drizzle-orm" +import { z, type ZodType } from "zod" +import { + buildKnowledgeBaseCollectionSelections, + KnowledgeBaseScope, + type KnowledgeBaseSelection, +} from "@/api/chat/knowledgeBaseSelections" +import { db } from "@/db/client" +import { + getCollectionById, + getCollectionItemById, + getCollectionLsProjection, + getCollectionsByOwner, + recordCollectionLsProjectionError, + upsertCollectionLsProjection, +} from "@/db/knowledgeBase" +import { + collectionItems, + type Collection, + type CollectionItem, + type CollectionLsProjection, +} from "@/db/schema" +import { getUserByEmail } from "@/db/user" +import { getLogger } from "@/logger" +import { Subsystem } from "@/types" +import { getErrorMessage } from "@/utils" +import { executeVespaSearch } from "./global" +import type { Ctx } from "./types" +import { parseAgentAppIntegrations } from "./utils" + +const Logger = getLogger(Subsystem.Chat) + +// Narrows JAF tool parameter typing from a Zod schema. +type ToolSchemaParameters = Tool["schema"]["parameters"] +// Bridges a Zod schema into the tool parameter type expected by JAF. +const toToolSchemaParameters = ( + schema: ZodType, +): ToolSchemaParameters => schema as unknown as ToolSchemaParameters + +const KNOWLEDGE_BASE_TARGET_DESCRIPTION = + "A discriminated knowledge-base target object for browse/search. Set `type` to one of `collection`, `folder`, `file`, or `path`, then provide only the matching ID/path fields for that variant." + +const KNOWLEDGE_BASE_OFFSET_DESCRIPTION = + "Pagination offset. Use it after reviewing the current page to continue from the next unseen rows or fragments." + +const KNOWLEDGE_BASE_EXCLUDED_IDS_DESCRIPTION = + "Previously seen result document `docId`s to suppress on follow-up KB searches. Prefer `fragment.source.docId` values from prior results. Do not pass collection, folder, file, path, or fragment IDs." + +export const LS_KNOWLEDGE_BASE_TOOL_DESCRIPTION = [ + "Browse the caller's accessible knowledge-base namespace.", + "Use it to discover collections, inspect folder/file layout, confirm canonical paths, answer inventory or metadata questions directly, or obtain IDs for a later `searchKnowledgeBase.filters.targets` call.", + "It is especially useful when the user wants answers constrained by structure or metadata such as a specific folder, collection, file set, or file type like PDFs.", + "Skip `ls` only when the exact KB scope is already known and browsing will not improve the answer.", + "Start shallow with `depth: 1` and `metadata: false` if unsure; but you are always free to enable metadata or deepen traversal only when the task truly needs row details or more hierarchy.", +].join(" ") + +export const SEARCH_KNOWLEDGE_BASE_TOOL_DESCRIPTION = [ + "Search document content inside the caller's accessible knowledge-base scope and return cited fragments.", + "Use it directly when the task is about document contents and the relevant KB scope is already known or broad KB search is acceptable.", + "Pair it with `ls` when you need structural scoping, canonical-path confirmation, or file preselection such as searching only .txt files from a folder.", + "If the collection, folder, file, or path is known, pass it in `filters.targets`; file targets can come from prior `ls` output.", + "`filters.targets` narrows search by location, while `excludedIds` should contain previously seen document/result IDs to avoid rereading the same hits.", +].join(" ") + +export const KnowledgeBaseTargetSchema = z.discriminatedUnion("type", [ + z.object({ + type: z.literal("collection"), + collectionId: z + .string() + .describe( + "Knowledge-base collection row ID as a string, typically a UUID. Reuse `ls` output directly here: for a collection row, pass `entries[i].id`; for a previously targeted `ls` response, pass `target.collection_id`. This stays a collection DB ID through KB search and is translated downstream into Vespa `clId` filtering. Do not pass a folder ID, file ID, or path here.", + ), + }).describe( + "Object shape: `{ type: \"collection\", collectionId: string }`. Targets an entire collection root. Best when the user names a known collection or you want to browse/search everything inside it.", + ), + z.object({ + type: z.literal("folder"), + folderId: z + .string() + .describe( + "Knowledge-base folder row ID as a string, typically a UUID. Reuse `ls` output directly here: when an `ls` entry has `type: \"folder\"`, pass that row's `id` as `folderId`. This is later translated into KB folder selections and then Vespa `clFd` filtering. Do not pass a collection ID, file ID, or path here.", + ), + }).describe( + "Object shape: `{ type: \"folder\", folderId: string }`. Targets a folder subtree inside a collection. Useful after `ls` returns a folder ID or the folder is already known.", + ), + z.object({ + type: z.literal("file"), + fileId: z + .string() + .describe( + "Knowledge-base file row ID as a string, typically a UUID. Reuse `ls` output directly here: when an `ls` entry has `type: \"file\"`, pass that row's `id` as `fileId`. This is later translated into the file's Vespa document `docId` filtering downstream. Do not pass a collection ID, folder ID, or path here.", + ), + }).describe( + "Object shape: `{ type: \"file\", fileId: string }`. Targets one exact file. Use for pinpointed browsing/search when the relevant document is already known.", + ), + z.object({ + type: z.literal("path"), + collectionId: z + .string() + .describe( + "Knowledge-base collection row ID as a string, typically a UUID. Required with `type: \"path\"` so the path is resolved inside the correct collection. Reuse `ls` output directly here with `entries[i].collection_id` or `target.collection_id` from a prior targeted `ls` response.", + ), + path: z + .string() + .describe( + "Collection-relative path string such as `/`, `/Policies`, `/Policies/Security`, or `/Policies/Security.md`. Reuse `ls` output directly here with `entries[i].path` or `target.path` from a prior targeted `ls` response. A missing leading slash is accepted and will be canonicalized. `path: \"/\"` means the collection root. `.` and `..` path segments are invalid. The resolved path is then translated into collection, folder, or file search scope before Vespa filtering.", + ), + }).describe( + "Object shape: `{ type: \"path\", collectionId: string, path: string }`. Targets a collection-relative path when the location is known or easier to express than raw folder/file IDs.", + ), +]).describe(KNOWLEDGE_BASE_TARGET_DESCRIPTION) + +export type KnowledgeBaseTarget = z.infer + +export const LsKnowledgeBaseInputSchema = z.object({ + target: KnowledgeBaseTargetSchema.optional().describe( + "Optional KB location to browse. Omit it to list accessible collections. Provide a collection, folder, file, or path target when you already know where to inspect or when the user asked about a specific location.", + ), + depth: z + .number() + .int() + .min(1) + .max(5) + .optional() + .default(1) + .describe( + "Traversal depth from the target. `1` lists immediate children only. Start shallow and increase depth only when the task truly needs more hierarchy.", + ), + limit: z + .number() + .int() + .min(1) + .max(100) + .optional() + .describe( + "Maximum number of browse rows to return from the flattened listing. Keep this small for discovery and page with `offset` when needed.", + ), + offset: z + .number() + .int() + .min(0) + .optional() + .default(0) + .describe(KNOWLEDGE_BASE_OFFSET_DESCRIPTION), + metadata: z + .boolean() + .optional() + .default(false) + .describe( + "Return persisted row metadata when true. Leave false for normal navigation; enable when you need details like description, mime type for filtering PDFs or other file types, timestamps, or collection metadata.", + ), +}).describe( + "Browse accessible knowledge-base collections, folders, and files. Use for navigation and scope discovery, not for full-text retrieval.", +) + +export type LsKnowledgeBaseToolParams = z.infer< + typeof LsKnowledgeBaseInputSchema +> +export const SearchKnowledgeBaseInputSchema = z.object({ + query: z + .string() + .min(1) + .describe( + "Short, content-focused KB retrieval query. Use the semantic terms you expect inside documents, not navigation instructions. If the scope is known, narrow with `filters.targets` instead of stuffing paths or folder names into the query.", + ), + filters: z + .object({ + targets: z + .array(KnowledgeBaseTargetSchema) + .min(1) + .optional() + .describe( + "Optional union of KB locations to search inside the current allowed scope. Each target may be a collection root, folder subtree, exact file, or collection-relative path. Use this when the user query or prior `ls` output tells you where to search; file targets are especially useful after `ls` identifies a subset such as PDFs.", + ), + }) + .optional() + .describe( + "Optional structural scope for KB search. Omit it when a broad search across the caller's allowed KB scope is appropriate.", + ), + limit: z + .number() + .int() + .min(1) + .max(25) + .optional() + .describe( + "Maximum number of KB fragments to return (up to 25). Keep this tight for precision-first retrieval; raise it only when the user needs broader coverage.", + ), + offset: z + .number() + .int() + .min(0) + .optional() + .describe(KNOWLEDGE_BASE_OFFSET_DESCRIPTION), + excludedIds: z + .array(z.string()) + .optional() + .describe(KNOWLEDGE_BASE_EXCLUDED_IDS_DESCRIPTION), +}).describe( + "Full-text search over document content in the caller's accessible knowledge-base scope.", +) + +export type SearchKnowledgeBaseToolParams = z.infer< + typeof SearchKnowledgeBaseInputSchema +> + +// Reuses the global Vespa search implementation shape for dependency injection in tests. +type SearchExecutor = typeof executeVespaSearch + +// Carries the resolved KB scope and selection state for the current tool execution. +type KnowledgeBaseScopeState = { + email: string + scope: KnowledgeBaseScope + selectedItems: Partial> + baseSelections: KnowledgeBaseSelection[] +} + +// Represents a navigable KB item inside an in-memory tree or projection snapshot. +type KnowledgeBaseNavigationNode = { + id: string + parent_id: string | null + collection_id: string + type: "folder" | "file" + name: string + path: string +} + +const CollectionLsProjectionPayloadSchema = z.object({ + rootIds: z.array(z.string()), + childrenByParentId: z.record(z.string(), z.array(z.string())), + nodesById: z.record( + z.string(), + z.object({ + id: z.string(), + parent_id: z.string().nullable(), + collection_id: z.string(), + type: z.enum(["folder", "file"]), + name: z.string(), + path: z.string(), + }), + ), + nodeIdByPath: z.record(z.string(), z.string()), +}) + +// Describes the persisted latest-only projection payload stored for ls traversal. +type CollectionLsProjectionPayload = z.infer< + typeof CollectionLsProjectionPayloadSchema +> + +// Holds the normalized navigation state needed to resolve and traverse a collection. +type KnowledgeBaseNavigationSnapshot = { + collection: Collection + rootIds: string[] + childrenByParentId: Map + nodesById: Map + nodeIdByPath: Map +} + +// Groups the DB and projection operations the KB flow depends on. +type KnowledgeBaseRepository = { + getUserByEmail: (email: string) => Promise<{ id: number } | null> + listUserScopedCollections: (userId: number) => Promise + getCollectionById: (collectionId: string) => Promise + getCollectionItemById: (itemId: string) => Promise + listCollectionItems: (collectionId: string) => Promise + listCollectionItemsByIds?: ( + collectionId: string, + itemIds: string[], + ) => Promise + getCollectionLsProjection?: ( + collectionId: string, + ) => Promise + upsertCollectionLsProjection?: (params: { + collectionId: string + projection: CollectionLsProjectionPayload + lsCollectionProjectionUpdatedAt: Date + lastError?: string | null + }) => Promise + recordCollectionLsProjectionError?: ( + collectionId: string, + lastError: string, + ) => Promise +} + +// Represents a fully resolved KB target that browse/search can safely operate on. +type ResolvedKnowledgeBaseTarget = + | { + // Used when the target resolves to the collection root itself. + kind: "collection" + targetType: KnowledgeBaseTarget["type"] + collection: Collection + path: "/" + snapshot: KnowledgeBaseNavigationSnapshot + } + | { + // Used when the target resolves to a concrete folder or file node. + kind: "node" + targetType: KnowledgeBaseTarget["type"] + collection: Collection + node: KnowledgeBaseNavigationNode + path: string + snapshot: KnowledgeBaseNavigationSnapshot + } + +// Returns the merged search selections along with the resolved targets that produced them. +type SearchSelectionBuildResult = { + selections: KnowledgeBaseSelection[] + resolvedTargets: ResolvedKnowledgeBaseTarget[] +} + +// Models one final ls row returned to the agent. +type LsEntry = { + id: string + type: "collection" | "folder" | "file" + name: string + path: string + collection_id?: string + parent_id?: string | null + depth: number + details?: Record +} + +// Captures a traversed node before optional live metadata hydration. +type LsEntrySeed = { + id: string + type: "folder" | "file" + name: string + path: string + collection_id: string + parent_id: string | null + depth: number +} + +// Separates direct-scan and projection-backed snapshots within one request. +type NavigationCache = { + direct: Map + projection: Map +} + +// Wires the production DB-backed repository used by the KB flow. +const defaultKnowledgeBaseRepository: KnowledgeBaseRepository = { + async getUserByEmail(email) { + const [user] = await getUserByEmail(db, email) + return user ? { id: user.id } : null + }, + async listUserScopedCollections(userId) { + return getCollectionsByOwner(db, userId) + }, + async getCollectionById(collectionId) { + return getCollectionById(db, collectionId) + }, + async getCollectionItemById(itemId) { + return getCollectionItemById(db, itemId) + }, + async listCollectionItems(collectionId) { + return db + .select() + .from(collectionItems) + .where( + and( + eq(collectionItems.collectionId, collectionId), + isNull(collectionItems.deletedAt), + ), + ) + }, + async listCollectionItemsByIds(collectionId, itemIds) { + if (!itemIds.length) return [] + return db + .select() + .from(collectionItems) + .where( + and( + eq(collectionItems.collectionId, collectionId), + inArray(collectionItems.id, itemIds), + isNull(collectionItems.deletedAt), + ), + ) + }, + async getCollectionLsProjection(collectionId) { + return getCollectionLsProjection(db, collectionId) + }, + async upsertCollectionLsProjection(params) { + return upsertCollectionLsProjection(db, params) + }, + async recordCollectionLsProjectionError(collectionId, lastError) { + await recordCollectionLsProjectionError(db, collectionId, lastError) + }, +} + +// Creates request-local caches so repeated resolutions do not rescan the same collection. +function createNavigationCache(): NavigationCache { + return { + direct: new Map(), + projection: new Map(), + } +} + +// Canonicalizes collection-relative KB paths and rejects unsupported path segments. +export function canonicalizeKnowledgeBasePath(path: string): string { + const rawPath = path.trim() + if (!rawPath) return "/" + + const normalized = rawPath.startsWith("/") ? rawPath : `/${rawPath}` + const collapsed = normalized.replace(/\/+/g, "/") + const segments = collapsed.split("/").filter(Boolean) + + if (segments.some((segment) => segment === "." || segment === "..")) { + throw new Error("Knowledge base paths cannot contain '.' or '..'") + } + + return segments.length ? `/${segments.join("/")}` : "/" +} + +// Builds the canonical absolute path for a persisted KB item row. +function buildItemCanonicalPath( + item: Pick, +): string { + return canonicalizeKnowledgeBasePath( + item.path === "/" ? `/${item.name}` : `${item.path}/${item.name}`, + ) +} + +// Converts a persisted KB row into the lightweight navigation node shape. +function buildNavigationNode( + item: Pick< + CollectionItem, + "id" | "parentId" | "collectionId" | "type" | "name" | "path" + >, +): KnowledgeBaseNavigationNode { + return { + id: item.id, + parent_id: item.parentId, + collection_id: item.collectionId, + type: item.type, + name: item.name, + path: buildItemCanonicalPath(item), + } +} + +// Keeps traversal order stable and folder-first for ls responses. +function sortItemsForTraversal(items: CollectionItem[]): CollectionItem[] { + return [...items].sort((left, right) => { + if (left.type !== right.type) { + return left.type === "folder" ? -1 : 1 + } + if (left.position !== right.position) { + return left.position - right.position + } + return left.name.localeCompare(right.name) + }) +} + +// Builds an in-memory collection tree directly from live item rows. +function buildCollectionTree( + collection: Collection, + items: CollectionItem[], +): KnowledgeBaseNavigationSnapshot { + const nodesById = new Map() + const nodeIdByPath = new Map() + const childItemsByParentId = new Map() + + for (const item of items) { + const siblings = childItemsByParentId.get(item.parentId ?? null) ?? [] + siblings.push(item) + childItemsByParentId.set(item.parentId ?? null, siblings) + } + + const childrenByParentId = new Map() + let rootIds: string[] = [] + + for (const item of items) { + const node = buildNavigationNode(item) + nodesById.set(node.id, node) + nodeIdByPath.set(node.path, node.id) + } + + rootIds = sortItemsForTraversal(childItemsByParentId.get(null) ?? []).map( + (item) => item.id, + ) + + for (const [parentId, children] of childItemsByParentId.entries()) { + if (!parentId) continue + childrenByParentId.set( + parentId, + sortItemsForTraversal(children).map((item) => item.id), + ) + } + + return { + collection, + rootIds, + childrenByParentId, + nodesById, + nodeIdByPath, + } +} + +// Produces the normalized projection payload persisted for target-scoped ls. +function buildCollectionLsProjection( + items: CollectionItem[], +): CollectionLsProjectionPayload { + const nodesById: Record = {} + const nodeIdByPath: Record = {} + const childItemsByParentId = new Map() + + for (const item of items) { + const siblings = childItemsByParentId.get(item.parentId ?? null) ?? [] + siblings.push(item) + childItemsByParentId.set(item.parentId ?? null, siblings) + + const node = buildNavigationNode(item) + nodesById[node.id] = node + nodeIdByPath[node.path] = node.id + } + + const rootIds = sortItemsForTraversal( + childItemsByParentId.get(null) ?? [], + ).map((item) => item.id) + + const childrenByParentId: Record = {} + for (const [parentId, children] of childItemsByParentId.entries()) { + if (!parentId) continue + childrenByParentId[parentId] = sortItemsForTraversal(children).map( + (item) => item.id, + ) + } + + return { + rootIds, + childrenByParentId, + nodesById, + nodeIdByPath, + } +} + +// Rehydrates an in-memory navigation snapshot from a stored projection row. +function buildSnapshotFromProjection( + collection: Collection, + projection: CollectionLsProjectionPayload, +): KnowledgeBaseNavigationSnapshot { + return { + collection, + rootIds: projection.rootIds, + childrenByParentId: new Map( + Object.entries(projection.childrenByParentId).map( + ([parentId, childIds]) => [parentId, childIds], + ), + ), + nodesById: new Map( + Object.entries(projection.nodesById).map(([nodeId, node]) => [ + nodeId, + node, + ]), + ), + nodeIdByPath: new Map( + Object.entries(projection.nodeIdByPath).map(([path, nodeId]) => [ + path, + nodeId, + ]), + ), + } +} + +// Validates and parses the stored projection payload before traversal. +function parseCollectionLsProjection( + projection: unknown, +): CollectionLsProjectionPayload { + return CollectionLsProjectionPayloadSchema.parse(projection) +} + +// Compares a stored projection against the collection staleness watermark. +function isProjectionStale( + collection: Collection, + projection: CollectionLsProjection | null, +): boolean { + if (!projection) return true + + return ( + projection.lsCollectionProjectionUpdatedAt.getTime() < + collection.collectionSourceUpdatedAt.getTime() + ) +} + +// Resolves the current caller's KB scope and base selections from agent context. +async function buildScopeState( + context: Pick, +): Promise { + const email = context.user.email + const { selectedItems } = parseAgentAppIntegrations(context.agentPrompt) + const scope = context.agentPrompt + ? KnowledgeBaseScope.AgentScoped + : KnowledgeBaseScope.UserOwned + const baseSelections = await buildKnowledgeBaseCollectionSelections({ + scope, + email, + selectedItems, + }) + + return { + email, + scope, + selectedItems, + baseSelections, + } +} + +// Lists only the collections available under the caller's effective KB scope. +async function listScopedCollections( + repo: KnowledgeBaseRepository, + scopeState: KnowledgeBaseScopeState, +): Promise { + if (scopeState.scope !== KnowledgeBaseScope.AgentScoped) { + const user = await repo.getUserByEmail(scopeState.email) + if (!user) { + throw new Error( + "User email not found while resolving knowledge base scope", + ) + } + return repo.listUserScopedCollections(user.id) + } + + const collectionIds = new Set() + for (const selection of scopeState.baseSelections) { + selection.collectionIds?.forEach((id) => collectionIds.add(id)) + + for (const folderId of selection.collectionFolderIds ?? []) { + const folder = await repo.getCollectionItemById(folderId) + if (folder) collectionIds.add(folder.collectionId) + } + + for (const fileId of selection.collectionFileIds ?? []) { + const file = await repo.getCollectionItemById(fileId) + if (file) collectionIds.add(file.collectionId) + } + } + + const collections = await Promise.all( + Array.from(collectionIds).map((collectionId) => + repo.getCollectionById(collectionId), + ), + ) + + return collections.filter( + (collection): collection is Collection => !!collection, + ) +} + +// Loads a live navigation snapshot by scanning all items in one collection. +async function loadDirectNavigationSnapshot( + collectionId: string, + repo: KnowledgeBaseRepository, + cache: NavigationCache, +): Promise { + const cached = cache.direct.get(collectionId) + if (cached) return cached + + const collection = await repo.getCollectionById(collectionId) + if (!collection) { + throw new Error(`Knowledge base collection '${collectionId}' was not found`) + } + + const items = await repo.listCollectionItems(collectionId) + const snapshot = buildCollectionTree(collection, items) + cache.direct.set(collectionId, snapshot) + return snapshot +} + +// Rebuilds and persists the latest projection snapshot for one collection. +async function rebuildCollectionProjection( + collection: Collection, + repo: KnowledgeBaseRepository, +): Promise { + const items = await repo.listCollectionItems(collection.id) + const projection = buildCollectionLsProjection(items) + + if (repo.upsertCollectionLsProjection) { + await repo.upsertCollectionLsProjection({ + collectionId: collection.id, + projection, + lsCollectionProjectionUpdatedAt: collection.collectionSourceUpdatedAt, + lastError: null, + }) + } + + return buildSnapshotFromProjection(collection, projection) +} + +// Loads a projection-backed snapshot and lazily rebuilds it when stale or missing. Main entry function +async function loadProjectionNavigationSnapshot( + collectionId: string, + repo: KnowledgeBaseRepository, + cache: NavigationCache, +): Promise { + const cached = cache.projection.get(collectionId) + if (cached) return cached + + const collection = await repo.getCollectionById(collectionId) + if (!collection) { + throw new Error(`Knowledge base collection '${collectionId}' was not found`) + } + + if (!repo.getCollectionLsProjection) { + return loadDirectNavigationSnapshot(collectionId, repo, cache) + } + + const existingProjection = await repo.getCollectionLsProjection(collectionId) + if (!isProjectionStale(collection, existingProjection)) { + const snapshot = buildSnapshotFromProjection( + collection, + parseCollectionLsProjection(existingProjection?.projection), + ) + cache.projection.set(collectionId, snapshot) + return snapshot + } + + try { + const snapshot = await rebuildCollectionProjection(collection, repo) + cache.projection.set(collectionId, snapshot) + return snapshot + } catch (error) { + const errorMessage = getErrorMessage(error) + try { + await repo.recordCollectionLsProjectionError?.(collectionId, errorMessage) + } catch {} + + return loadDirectNavigationSnapshot(collectionId, repo, cache) + } +} + +// Checks whether a node falls within a folder-scoped agent selection. +function isDescendantOfFolder( + snapshot: KnowledgeBaseNavigationSnapshot, + nodeId: string, + folderId: string, +): boolean { + let current = snapshot.nodesById.get(nodeId) + while (current) { + if (current.id === folderId) return true + current = current.parent_id + ? snapshot.nodesById.get(current.parent_id) + : undefined + } + return false +} + +// Enforces that a resolved target does not widen the current KB scope. +function isResolvedTargetAllowed( + resolvedTarget: ResolvedKnowledgeBaseTarget, + scopeState: KnowledgeBaseScopeState, + scopedCollectionIds: Set, +): boolean { + if (scopeState.scope !== KnowledgeBaseScope.AgentScoped) { + return scopedCollectionIds.has(resolvedTarget.collection.id) + } + + if (!scopeState.baseSelections.length) return false + + return scopeState.baseSelections.some((selection) => { + if (selection.collectionIds?.includes(resolvedTarget.collection.id)) { + return true + } + + if (resolvedTarget.kind !== "node") { + return false + } + + if ( + resolvedTarget.node.type === "file" && + selection.collectionFileIds?.includes(resolvedTarget.node.id) + ) { + return true + } + + return ( + selection.collectionFolderIds?.some((folderId) => + isDescendantOfFolder( + resolvedTarget.snapshot, + resolvedTarget.node.id, + folderId, + ), + ) ?? false + ) + }) +} + +// Rejects collection-root/path targets that point at collections outside the caller scope +// before we load or rebuild their navigation snapshot. +function isTargetCollectionPreAuthorized( + target: KnowledgeBaseTarget, + scopedCollectionIds: Set, +): boolean { + switch (target.type) { + case "collection": + case "path": + return scopedCollectionIds.has(target.collectionId) + default: + return true + } +} + +// Converts one resolved target into the search selection shape expected downstream. +function toSelectionForResolvedTarget( + resolvedTarget: ResolvedKnowledgeBaseTarget, +): KnowledgeBaseSelection { + if (resolvedTarget.kind === "collection") { + return { collectionIds: [resolvedTarget.collection.id] } + } + + if (resolvedTarget.node.type === "folder") { + return { collectionFolderIds: [resolvedTarget.node.id] } + } + + return { collectionFileIds: [resolvedTarget.node.id] } +} + +// Unions collection, folder, and file selections while removing duplicates. +function mergeSelections( + selections: KnowledgeBaseSelection[], +): KnowledgeBaseSelection[] { + if (!selections.length) return [] + + const collectionIds = new Set() + const collectionFolderIds = new Set() + const collectionFileIds = new Set() + + for (const selection of selections) { + selection.collectionIds?.forEach((id) => collectionIds.add(id)) + selection.collectionFolderIds?.forEach((id) => collectionFolderIds.add(id)) + selection.collectionFileIds?.forEach((id) => collectionFileIds.add(id)) + } + + const merged: KnowledgeBaseSelection = {} + if (collectionIds.size) merged.collectionIds = Array.from(collectionIds) + if (collectionFolderIds.size) { + merged.collectionFolderIds = Array.from(collectionFolderIds) + } + if (collectionFileIds.size) { + merged.collectionFileIds = Array.from(collectionFileIds) + } + + return Object.keys(merged).length ? [merged] : [] +} + +// Resolves a target contract into a collection root or concrete node snapshot. Translation +async function resolveKnowledgeBaseTarget( + target: KnowledgeBaseTarget, + repo: KnowledgeBaseRepository, + cache: NavigationCache, +): Promise { + switch (target.type) { + case "collection": { + const snapshot = await loadProjectionNavigationSnapshot( + target.collectionId, + repo, + cache, + ) + return { + kind: "collection", + targetType: target.type, + collection: snapshot.collection, + path: "/", + snapshot, + } + } + case "path": { + const snapshot = await loadProjectionNavigationSnapshot( + target.collectionId, + repo, + cache, + ) + const path = canonicalizeKnowledgeBasePath(target.path) + if (path === "/") { + return { + kind: "collection", + targetType: target.type, + collection: snapshot.collection, + path, + snapshot, + } + } + + const nodeId = snapshot.nodeIdByPath.get(path) + if (!nodeId) { + throw new Error( + `Knowledge base path '${path}' was not found in collection '${target.collectionId}'`, + ) + } + + const node = snapshot.nodesById.get(nodeId) + if (!node) { + throw new Error( + `Knowledge base path '${path}' could not be resolved in collection '${target.collectionId}'`, + ) + } + + return { + kind: "node", + targetType: target.type, + collection: snapshot.collection, + node, + path, + snapshot, + } + } + case "folder": + case "file": { + const itemId = target.type === "folder" ? target.folderId : target.fileId + const directItem = await repo.getCollectionItemById(itemId) + if (!directItem || directItem.type !== target.type) { + throw new Error( + `Knowledge base ${target.type} '${itemId}' was not found`, + ) + } + + const snapshot = await loadProjectionNavigationSnapshot( + directItem.collectionId, + repo, + cache, + ) + const node = snapshot.nodesById.get(itemId) + if (!node || node.type !== target.type) { + throw new Error( + `Knowledge base ${target.type} '${itemId}' was not found`, + ) + } + + return { + kind: "node", + targetType: target.type, + collection: snapshot.collection, + node, + path: node.path, + snapshot, + } + } + } +} + +// Resolves search targets and maps them into the scoped KB selection format. +async function buildSearchSelections( + params: SearchKnowledgeBaseToolParams, + scopeState: KnowledgeBaseScopeState, + repo: KnowledgeBaseRepository, +): Promise { + if (!params.filters?.targets?.length) { + return { + selections: scopeState.baseSelections, + resolvedTargets: [], + } + } + + const scopedCollections = await listScopedCollections(repo, scopeState) + const scopedCollectionIds = new Set( + scopedCollections.map((collection) => collection.id), + ) + const cache = createNavigationCache() + const resolvedTargets: ResolvedKnowledgeBaseTarget[] = [] + const selections: KnowledgeBaseSelection[] = [] + + for (const target of params.filters.targets) { + if (!isTargetCollectionPreAuthorized(target, scopedCollectionIds)) { + throw new Error( + "Requested knowledge base target is outside the current KB scope", + ) + } + + const resolvedTarget = await resolveKnowledgeBaseTarget(target, repo, cache) + if ( + !isResolvedTargetAllowed(resolvedTarget, scopeState, scopedCollectionIds) + ) { + throw new Error( + "Requested knowledge base target is outside the current KB scope", + ) + } + + resolvedTargets.push(resolvedTarget) + selections.push(toSelectionForResolvedTarget(resolvedTarget)) + } + + return { + selections: mergeSelections(selections), + resolvedTargets, + } +} + +// Shapes a collection row for no-target ls responses. +function createCollectionLsEntry( + collection: Collection, + includeMetadata: boolean, +): LsEntry { + const entry: LsEntry = { + id: collection.id, + type: "collection", + name: collection.name, + path: "/", + depth: 0, + } + + if (includeMetadata) { + entry.details = { + total_items: collection.totalItems, + name: collection.name, + last_updated_by_email: collection.lastUpdatedByEmail, + description: collection.description, + updated_at: collection.updatedAt, + created_at: collection.createdAt, + metadata: collection.metadata, + } + } + + return entry +} + +// Captures the minimal ls row state needed before metadata hydration. +function createSeedEntry( + node: KnowledgeBaseNavigationNode, + depth: number, +): LsEntrySeed { + return { + id: node.id, + type: node.type, + name: node.name, + path: node.path, + collection_id: node.collection_id, + parent_id: node.parent_id, + depth, + } +} + +// Combines traversal output with optional live item metadata for the response. +function createItemLsEntry( + seed: LsEntrySeed, + item: CollectionItem | null, + includeMetadata: boolean, +): LsEntry { + const entry: LsEntry = { + id: seed.id, + type: seed.type, + name: seed.name, + path: seed.path, + collection_id: seed.collection_id, + parent_id: seed.parent_id, + depth: seed.depth, + } + + if (includeMetadata && item) { + entry.details = { + type: item.type, + name: item.name, + collection_id: item.collectionId, + mime_type: item.mimeType, + updated_at: item.updatedAt, + created_at: item.createdAt, + metadata: item.metadata, + ...(item.type === "folder" + ? { total_file_count: item.totalFileCount } + : {}), + } + } + + return entry +} + +function summarizeResolvedTargetForLog( + resolvedTarget: ResolvedKnowledgeBaseTarget, +) { + return { + kind: resolvedTarget.kind, + targetType: resolvedTarget.targetType, + collectionId: resolvedTarget.collection.id, + collectionName: resolvedTarget.collection.name, + path: resolvedTarget.path, + ...(resolvedTarget.kind === "node" + ? { + node: { + id: resolvedTarget.node.id, + type: resolvedTarget.node.type, + name: resolvedTarget.node.name, + parentId: resolvedTarget.node.parent_id, + }, + } + : {}), + } +} + +// Traverses a resolved target into ordered ls seeds up to the requested depth. +function flattenLsEntries( + resolvedTarget: ResolvedKnowledgeBaseTarget, + depthLimit: number, +): LsEntrySeed[] { + if (resolvedTarget.kind === "node" && resolvedTarget.node.type === "file") { + return [createSeedEntry(resolvedTarget.node, 0)] + } + + const rootIds = + resolvedTarget.kind === "collection" + ? resolvedTarget.snapshot.rootIds + : (resolvedTarget.snapshot.childrenByParentId.get( + resolvedTarget.node.id, + ) ?? []) + + const entries: LsEntrySeed[] = [] + const stack = rootIds + .slice() + .reverse() + .map((nodeId) => ({ + nodeId, + depth: 1, + })) + + while (stack.length > 0) { + const current = stack.pop() + if (!current) continue + + const node = resolvedTarget.snapshot.nodesById.get(current.nodeId) + if (!node) continue + + entries.push(createSeedEntry(node, current.depth)) + + if (current.depth >= depthLimit || node.type !== "folder") { + continue + } + + const childIds = + resolvedTarget.snapshot.childrenByParentId.get(node.id) ?? [] + for (let index = childIds.length - 1; index >= 0; index -= 1) { + stack.push({ + nodeId: childIds[index]!, + depth: current.depth + 1, + }) + } + } + + return entries +} + +// Loads only the persisted rows needed to hydrate paginated ls entries. +async function hydrateCollectionItemsForEntries( + collectionId: string, + itemIds: string[], + repo: KnowledgeBaseRepository, +): Promise> { + if (!itemIds.length) return new Map() + + const items = repo.listCollectionItemsByIds + ? await repo.listCollectionItemsByIds(collectionId, itemIds) + : (await repo.listCollectionItems(collectionId)).filter((item) => + itemIds.includes(item.id), + ) + + return new Map(items.map((item) => [item.id, item])) +} + +// Builds the paginated ls payload for a resolved collection, folder, or file target. +async function buildLsEntries( + resolvedTarget: ResolvedKnowledgeBaseTarget, + params: Pick< + LsKnowledgeBaseToolParams, + "depth" | "limit" | "offset" | "metadata" + >, + repo: KnowledgeBaseRepository, +): Promise<{ entries: LsEntry[]; total: number }> { + const rawEntries = flattenLsEntries(resolvedTarget, params.depth ?? 1) + const offset = params.offset ?? 0 + const limit = params.limit + const paginatedEntries = + typeof limit === "number" + ? rawEntries.slice(offset, offset + limit) + : rawEntries.slice(offset) + + const hydratedItems = params.metadata + ? await hydrateCollectionItemsForEntries( + resolvedTarget.collection.id, + paginatedEntries.map((entry) => entry.id), + repo, + ) + : new Map() + + return { + entries: paginatedEntries.map((entry) => + createItemLsEntry( + entry, + params.metadata ? (hydratedItems.get(entry.id) ?? null) : null, + params.metadata ?? false, + ), + ), + total: rawEntries.length, + } +} + +// Executes the internal ls tool over the caller's scoped KB view. +export async function executeLsKnowledgeBase( + params: LsKnowledgeBaseToolParams, + context: Ctx, + repo: KnowledgeBaseRepository = defaultKnowledgeBaseRepository, +) { + const email = context.user.email + if (!email) { + return ToolResponse.error( + ToolErrorCodes.MISSING_REQUIRED_FIELD, + "User email not found while executing ls", + { toolName: "ls" }, + ) + } + + try { + const scopeState = await buildScopeState(context) + const collections = await listScopedCollections(repo, scopeState) + if (!collections.length) { + return ToolResponse.error( + ToolErrorCodes.EXECUTION_FAILED, + "No accessible knowledge base collections found for this user", + { toolName: "ls" }, + ) + } + + const offset = params.offset ?? 0 + const limit = params.limit + + if (!params.target) { + const entries = collections.map((collection) => + createCollectionLsEntry(collection, params.metadata ?? false), + ) + const paginatedEntries = + typeof limit === "number" + ? entries.slice(offset, offset + limit) + : entries.slice(offset) + const response = { + target: null, + entries: paginatedEntries, + total: entries.length, + offset, + limit: limit ?? paginatedEntries.length, + } + + Logger.info( + { + email, + scope: scopeState.scope, + target: null, + result: response, + }, + "[KnowledgeBase][ls] Returning top-level browse result", + ) + + return ToolResponse.success(response) + } + + const scopedCollectionIds = new Set( + collections.map((collection) => collection.id), + ) + if ( + !isTargetCollectionPreAuthorized(params.target, scopedCollectionIds) + ) { + return ToolResponse.error( + ToolErrorCodes.PERMISSION_DENIED, + "Requested knowledge base target is outside the current KB scope", + { toolName: "ls" }, + ) + } + + const cache = createNavigationCache() + const resolvedTarget = await resolveKnowledgeBaseTarget( + params.target, + repo, + cache, + ) + + if ( + !isResolvedTargetAllowed(resolvedTarget, scopeState, scopedCollectionIds) + ) { + return ToolResponse.error( + ToolErrorCodes.PERMISSION_DENIED, + "Requested knowledge base target is outside the current KB scope", + { toolName: "ls" }, + ) + } + + const { entries, total } = await buildLsEntries( + resolvedTarget, + params, + repo, + ) + const response = { + target: { + type: + resolvedTarget.kind === "collection" + ? "collection" + : resolvedTarget.node.type, + collection_id: resolvedTarget.collection.id, + id: + resolvedTarget.kind === "collection" + ? resolvedTarget.collection.id + : resolvedTarget.node.id, + path: resolvedTarget.path, + }, + entries, + total, + offset, + limit: limit ?? entries.length, + } + + Logger.info( + { + email, + scope: scopeState.scope, + requestedTarget: params.target, + resolvedTarget: summarizeResolvedTargetForLog(resolvedTarget), + result: response, + }, + "[KnowledgeBase][ls] Returning targeted browse result", + ) + + return ToolResponse.success(response) + } catch (error) { + return ToolResponse.error( + ToolErrorCodes.EXECUTION_FAILED, + `Knowledge base browse failed: ${getErrorMessage(error)}`, + { toolName: "ls" }, + ) + } +} + +// Executes the internal KB search tool after resolving scoped target filters. Main function entry point for serachKB tool. +export async function executeSearchKnowledgeBase( + params: SearchKnowledgeBaseToolParams, + context: Ctx, + options: { + repo?: KnowledgeBaseRepository + searchExecutor?: SearchExecutor + } = {}, +) { + const email = context.user.email + if (!email) { + return ToolResponse.error( + ToolErrorCodes.MISSING_REQUIRED_FIELD, + "User email not found while executing searchKnowledgeBase", + { toolName: "searchKnowledgeBase" }, + ) + } + + const query = params.query?.trim() + if (!query) { + return ToolResponse.error( + ToolErrorCodes.MISSING_REQUIRED_FIELD, + "Query cannot be empty for knowledge base search", + { toolName: "searchKnowledgeBase" }, + ) + } + + const repo = options.repo ?? defaultKnowledgeBaseRepository + const searchExecutor = options.searchExecutor ?? executeVespaSearch + + try { + const scopeState = await buildScopeState(context) + Logger.info( + { + email, + scope: scopeState.scope, + baseSelectionCount: scopeState.baseSelections.length, + selectedItemKeys: Object.keys( + scopeState.selectedItems as Record, + ).length, + requestedTargetCount: params.filters?.targets?.length ?? 0, + }, + "[MessageAgents][searchKnowledgeBaseTool] Using KnowledgeBaseScope for KB search", + ) + + const { selections, resolvedTargets } = await buildSearchSelections( + params, + scopeState, + repo, + ) + if (!selections.length) { + return ToolResponse.error( + ToolErrorCodes.EXECUTION_FAILED, + "No accessible knowledge base collections found for this user", + { toolName: "searchKnowledgeBase" }, + ) + } + + Logger.info( + { + email, + query, + requestedTargets: params.filters?.targets ?? [], + resolvedTargets: resolvedTargets.map(summarizeResolvedTargetForLog), + baseSelections: scopeState.baseSelections, + finalCollectionSelections: selections, + selectedKnowledgeItemIds: + scopeState.selectedItems[Apps.KnowledgeBase] ?? [], + excludedIds: params.excludedIds ?? [], + }, + "[KnowledgeBase][search] Resolved KB targets into collection selections for Vespa", + ) + + const fragments = await searchExecutor({ + email, + query, + app: Apps.KnowledgeBase, + agentAppEnums: [Apps.KnowledgeBase], + limit: params.limit, + offset: params.offset ?? 0, + excludedIds: params.excludedIds, + collectionSelections: selections, + selectedItems: scopeState.selectedItems, + userId: context.user.numericId ?? undefined, + workspaceId: context.user.workspaceNumericId ?? undefined, + }) + + if (!fragments.length) { + return ToolResponse.error( + ToolErrorCodes.EXECUTION_FAILED, + "No knowledge base results found for the query.", + { toolName: "searchKnowledgeBase" }, + ) + } + + return ToolResponse.success(fragments) + } catch (error) { + return ToolResponse.error( + ToolErrorCodes.EXECUTION_FAILED, + `Knowledge base search failed: ${getErrorMessage(error)}`, + { toolName: "searchKnowledgeBase" }, + ) + } +} +export const lsKnowledgeBaseTool: Tool = { + schema: { + name: "ls", + description: LS_KNOWLEDGE_BASE_TOOL_DESCRIPTION, + parameters: toToolSchemaParameters(LsKnowledgeBaseInputSchema), + }, + execute: executeLsKnowledgeBase, +} +export const searchKnowledgeBaseTool: Tool = + { + schema: { + name: "searchKnowledgeBase", + description: SEARCH_KNOWLEDGE_BASE_TOOL_DESCRIPTION, + parameters: toToolSchemaParameters( + SearchKnowledgeBaseInputSchema, + ), + }, + execute: executeSearchKnowledgeBase, + } + +export const __knowledgeBaseFlowInternals = { + buildCollectionLsProjection, + buildCollectionTree, + buildItemCanonicalPath, + buildLsEntries, + buildSearchSelections, + buildSnapshotFromProjection, + createCollectionLsEntry, + createItemLsEntry, + flattenLsEntries, + isDescendantOfFolder, + isProjectionStale, + isResolvedTargetAllowed, + listScopedCollections, + mergeSelections, + parseCollectionLsProjection, + rebuildCollectionProjection, + resolveKnowledgeBaseTarget, +} diff --git a/server/api/chat/tools/schemas.ts b/server/api/chat/tools/schemas.ts index 380452c18..9e71f623b 100644 --- a/server/api/chat/tools/schemas.ts +++ b/server/api/chat/tools/schemas.ts @@ -66,7 +66,7 @@ export const limitSchema = z .min(1) .max(100) .describe( - "Maximum number of results to return. Default behavior is to return 20 results.", + "Maximum number of results to return as an integer between 1 and 100. Default is 20. Keep this small for precision-first retrieval and page with `offset` when needed.", ) .default(20) @@ -74,13 +74,22 @@ export const offsetSchema = z .number() .min(0) .describe( - "Number of results to skip from the beginning, useful for pagination.", + "Pagination offset as a non-negative integer. Use it after reviewing the current page to continue from the next unseen results.", ) .optional() export const sortBySchema = z .enum(["asc", "desc"]) - .describe("Sort order of results. Accepts 'asc' or 'desc'.") + .describe( + "Sort direction. Valid values are `asc` and `desc`. Use `desc` for newest-first or most-recent-first ordering when supported.", + ) + .optional() + +export const excludedIdsSchema = z + .array(z.string()) + .describe( + "Previously seen result document `docId`s to suppress on follow-up searches. Prefer prior `fragment.source.docId` values. Do not pass collection, folder, file, path, or fragment IDs.", + ) .optional() // Common time range schema @@ -88,11 +97,17 @@ export const timeRangeSchema = z .object({ startTime: z .string() - .describe(`Start time in ${config.llmTimeFormat} format`), - endTime: z.string().describe(`End time in ${config.llmTimeFormat} format`), + .describe( + `Inclusive start time as a string in ${config.llmTimeFormat} format.`, + ), + endTime: z + .string() + .describe( + `Inclusive end time as a string in ${config.llmTimeFormat} format.`, + ), }) .describe( - `Filter within a specific time range. Example: { startTime: ${config.llmTimeFormat}, endTime: ${config.llmTimeFormat} }`, + `Optional time-range object with string fields \`{ startTime, endTime }\` in ${config.llmTimeFormat} format. Use it when the query is bounded by an explicit time window.`, ) .optional() @@ -121,4 +136,5 @@ export const baseToolParams = { offset: offsetSchema, sortBy: sortBySchema, timeRange: timeRangeSchema, + excludedIds: excludedIdsSchema, } diff --git a/server/api/chat/tools/slack/getSlackMessages.ts b/server/api/chat/tools/slack/getSlackMessages.ts index db5c78714..89747b024 100644 --- a/server/api/chat/tools/slack/getSlackMessages.ts +++ b/server/api/chat/tools/slack/getSlackMessages.ts @@ -41,15 +41,19 @@ const getSlackRelatedMessagesSchema = z.object({ channelName: z .string() .optional() - .describe("Name of specific channel to search within"), + .describe( + "Optional Slack channel name string, such as `eng-launches`. Pass the human-facing channel name, not a Slack channel ID.", + ), user: z .string() .optional() - .describe("Name or Email of specific user whose messages to retrieve"), + .describe( + "Optional Slack user identifier string to restrict messages by author. Email is preferred; display name can also work.", + ), mentions: z .array(z.string()) .describe( - "Filter messages that mention specific users. Provide usernames or email (e.g., '@john.doe' or john.d@domain.in)", + "Optional list of mentioned-user identifier strings, usually emails or usernames, to find messages that mention specific people.", ) .optional(), ...baseToolParams, @@ -66,7 +70,7 @@ export const getSlackRelatedMessagesTool: Tool< schema: { name: "getSlackRelatedMessages", description: - "Unified tool to retrieve Slack messages with flexible filtering options. Can search by channel, user, time range, thread, or any combination. Automatically includes thread messages when found. Use this single tool for all Slack message retrieval needs.", + "Search Slack messages with flexible filters for content, channel, author, mentions, and time range. Automatically includes thread replies when thread roots are found, and defaults to recent Slack history when no query or scope is provided.", parameters: toToolSchemaParameters( getSlackRelatedMessagesSchema, ), @@ -146,6 +150,7 @@ export const getSlackRelatedMessagesTool: Tool< offset: searchOptions.offset, timestampRange: normalizedTimestampRange, agentChannelIds: channelIds.length > 0 ? channelIds : undefined, + excludeDocIds: params.excludedIds || [], mentions: params.mentions && params.mentions.length > 0 ? params.mentions @@ -203,6 +208,16 @@ export const getSlackRelatedMessagesTool: Tool< } } + // Base-search exclusions are pushed down into Vespa. Keep this pass for + // appended thread items fetched separately from the main search. + const excludedDocIds = new Set(params.excludedIds || []) + if (excludedDocIds.size > 0) { + allItems = allItems.filter((item) => { + const citation = searchToCitation(item) + return !excludedDocIds.has(citation.docId) + }) + } + const fragments: MinimalAgentFragment[] = await Promise.all( allItems.map( async ( diff --git a/server/db/knowledgeBase.ts b/server/db/knowledgeBase.ts index 3646a7ff4..8c3933875 100644 --- a/server/db/knowledgeBase.ts +++ b/server/db/knowledgeBase.ts @@ -1,7 +1,9 @@ import { collections, + collectionLsProjections, collectionItems, type Collection, + type CollectionLsProjection, type NewCollection, type CollectionItem, type NewCollectionItem, @@ -15,6 +17,109 @@ import { and, asc, desc, eq, isNull, sql, or, inArray } from "drizzle-orm" import { UploadStatus } from "@/shared/types" import { getUserByEmail } from "./user" +export const touchCollectionLsStructure = async ( + trx: TxnOrClient, + collectionId: string, +): Promise => { + await trx + .update(collections) + .set({ + collectionSourceUpdatedAt: sql`NOW()`, + updatedAt: sql`NOW()`, + }) + .where(eq(collections.id, collectionId)) +} + +export const touchCollectionLsStructures = async ( + trx: TxnOrClient, + collectionIds: string[], +): Promise => { + const uniqueCollectionIds = [...new Set(collectionIds.filter(Boolean))] + if (!uniqueCollectionIds.length) return + + await trx + .update(collections) + .set({ + collectionSourceUpdatedAt: sql`NOW()`, + updatedAt: sql`NOW()`, + }) + .where(inArray(collections.id, uniqueCollectionIds)) +} + +export const getCollectionLsProjection = async ( + trx: TxnOrClient, + collectionId: string, +): Promise => { + const [result] = await trx + .select() + .from(collectionLsProjections) + .where(eq(collectionLsProjections.collectionId, collectionId)) + + return result || null +} + +export const upsertCollectionLsProjection = async ( + trx: TxnOrClient, + params: { + collectionId: string + projection: Record + lsCollectionProjectionUpdatedAt: Date + lastError?: string | null + }, +): Promise => { + const [result] = await trx + .insert(collectionLsProjections) + .values({ + collectionId: params.collectionId, + projection: params.projection, + lsCollectionProjectionUpdatedAt: params.lsCollectionProjectionUpdatedAt, + lastError: params.lastError ?? null, + }) + .onConflictDoUpdate({ + target: collectionLsProjections.collectionId, + set: { + projection: params.projection, + lsCollectionProjectionUpdatedAt: params.lsCollectionProjectionUpdatedAt, + lastError: params.lastError ?? null, + updatedAt: sql`NOW()`, + }, + }) + .returning() + + if (!result) { + throw new Error("Failed to upsert collection ls projection") + } + + return result +} + +export const recordCollectionLsProjectionError = async ( + trx: TxnOrClient, + collectionId: string, + lastError: string, +): Promise => { + await trx + .insert(collectionLsProjections) + .values({ + collectionId, + projection: { + rootIds: [], + childrenByParentId: {}, + nodesById: {}, + nodeIdByPath: {}, + }, + lsCollectionProjectionUpdatedAt: new Date(0), + lastError, + }) + .onConflictDoUpdate({ + target: collectionLsProjections.collectionId, + set: { + lastError, + updatedAt: sql`NOW()`, + }, + }) +} + // Collection CRUD operations export const createCollection = async ( trx: TxnOrClient, @@ -114,6 +219,8 @@ export const softDeleteCollection = async ( trx: TxnOrClient, collectionId: string, ): Promise => { + await touchCollectionLsStructure(trx, collectionId) + const [result] = await trx .update(collections) .set({ @@ -225,6 +332,8 @@ export const softDeleteCollectionItem = async ( throw new Error("Collection item not found") } + await touchCollectionLsStructure(trx, item.collectionId) + // If it's a folder, recursively delete all items inside it if (item.type === "folder") { const markDescendantsAsDeleted = async ( @@ -457,6 +566,8 @@ export const createFolder = async ( metadata, }) + await touchCollectionLsStructure(trx, collectionId) + // Update collection total count await updateCollectionTotalCount(trx, collectionId, 1) @@ -541,6 +652,8 @@ export const createFileItem = async ( metadata, }) + await touchCollectionLsStructure(trx, collectionId) + // Update collection total count await updateCollectionTotalCount(trx, collectionId, 1) diff --git a/server/db/schema/knowledgeBase.ts b/server/db/schema/knowledgeBase.ts index 7626ad9fb..b8df22654 100644 --- a/server/db/schema/knowledgeBase.ts +++ b/server/db/schema/knowledgeBase.ts @@ -42,6 +42,9 @@ export const collections = pgTable( retryCount: integer("retry_count").default(0).notNull(), // Track processing retry attempts metadata: jsonb("metadata").default({}).notNull(), permissions: jsonb("permissions").default(sql`'[]'::jsonb`), + collectionSourceUpdatedAt: timestamp("ls_projection_source_updated_at") + .defaultNow() + .notNull(), createdAt: timestamp("created_at").defaultNow().notNull(), updatedAt: timestamp("updated_at").defaultNow().notNull(), deletedAt: timestamp("deleted_at"), @@ -49,9 +52,7 @@ export const collections = pgTable( }, (table) => ({ // Ensure unique names per owner (excluding soft-deleted items) - uniqueOwnerName: uniqueIndex( - "unique_owner_collection_name_not_deleted", - ) + uniqueOwnerName: uniqueIndex("unique_owner_collection_name_not_deleted") .on(table.ownerId, table.name) .where(sql`${table.deletedAt} IS NULL`), // Index for finding collections by owner @@ -105,7 +106,10 @@ export const collectionItems = pgTable( processingInfo: jsonb("processing_info").default({}).notNull(), processedAt: timestamp("processed_at"), - uploadStatus: varchar("upload_status", { length: 20 }).default(UploadStatus.PENDING).notNull().$type(), + uploadStatus: varchar("upload_status", { length: 20 }) + .default(UploadStatus.PENDING) + .notNull() + .$type(), statusMessage: text("status_message"), // Stores error messages, processing details, or success info retryCount: integer("retry_count").default(0).notNull(), // Track processing retry attempts metadata: jsonb("metadata").default({}).notNull(), @@ -143,9 +147,34 @@ export const collectionItems = pgTable( }), ) +export const collectionLsProjections = pgTable( + "collection_ls_projections", + { + collectionId: uuid("collection_id") + .primaryKey() + .references(() => collections.id, { onDelete: "cascade" }), + projection: jsonb("projection").notNull(), + lsCollectionProjectionUpdatedAt: timestamp( + "built_from_source_updated_at", + ).notNull(), + createdAt: timestamp("created_at").defaultNow().notNull(), + updatedAt: timestamp("updated_at").defaultNow().notNull(), + lastError: text("last_error"), + }, + (table) => ({ + idxBuiltFromSourceUpdatedAt: index( + "idx_collection_ls_projection_built_from_source_updated_at", + ).on(table.lsCollectionProjectionUpdatedAt), + }), +) + // Relations definitions using Drizzle ORM relations() function export const collectionsRelations = relations(collections, ({ many, one }) => ({ items: many(collectionItems), + lsProjection: one(collectionLsProjections, { + fields: [collections.id], + references: [collectionLsProjections.collectionId], + }), owner: one(users, { fields: [collections.ownerId], references: [users.id], @@ -194,11 +223,24 @@ export const collectionItemsRelations = relations( }), ) +export const collectionLsProjectionsRelations = relations( + collectionLsProjections, + ({ one }) => ({ + collection: one(collections, { + fields: [collectionLsProjections.collectionId], + references: [collections.id], + }), + }), +) + // Type definitions for use in the application export type Collection = typeof collections.$inferSelect export type NewCollection = typeof collections.$inferInsert export type CollectionItem = typeof collectionItems.$inferSelect export type NewCollectionItem = typeof collectionItems.$inferInsert +export type CollectionLsProjection = typeof collectionLsProjections.$inferSelect +export type NewCollectionLsProjection = + typeof collectionLsProjections.$inferInsert // Helper types export type Folder = CollectionItem & { type: "folder" } diff --git a/server/queue/fileProcessor.ts b/server/queue/fileProcessor.ts index da5685513..3f3c8733b 100644 --- a/server/queue/fileProcessor.ts +++ b/server/queue/fileProcessor.ts @@ -17,6 +17,21 @@ import { updateParentStatus } from "@/db/knowledgeBase" const Logger = getLogger(Subsystem.Queue) +export function mergeCollectionItemMetadata( + existingMetadata: unknown, + updates: Record, +): Record { + const baseMetadata = + typeof existingMetadata === "object" && existingMetadata !== null + ? { ...(existingMetadata as Record) } + : {} + + return { + ...baseMetadata, + ...updates, + } +} + function extractMarkdownTitle(content: string): string { const lines = content.split("\n") let inFrontmatter = false @@ -106,7 +121,11 @@ async function handleRetryFailure( .where(eq(collectionItems.id, entityId)) // If it's a file that failed, trigger parent status update - if (entityType === ProcessingJobType.FILE && parentId !== undefined && collectionId) { + if ( + entityType === ProcessingJobType.FILE && + parentId !== undefined && + collectionId + ) { if (parentId) { await updateParentStatus(db, parentId, false) } else { @@ -249,7 +268,7 @@ async function processFileJob(jobData: FileProcessingJob, startTime: number) { ) // Extract title for markdown files - let pageTitle:string="" + let pageTitle: string = "" if (getBaseMimeType(file.mimeType || "") === "text/markdown") { try { const fileContent = fileBuffer.toString("utf-8") @@ -272,7 +291,7 @@ async function processFileJob(jobData: FileProcessingJob, startTime: number) { // Handle multiple processing results (e.g., for spreadsheets with multiple sheets) let totalChunksCount = 0 let newVespaDocId = "" - if(processingResults.length > 0 && 'totalSheets' in processingResults[0]) { + if (processingResults.length > 0 && "totalSheets" in processingResults[0]) { newVespaDocId = `${file.vespaDocId}_sheet_${(processingResults[0] as SheetProcessingResult).totalSheets}` } else { newVespaDocId = file.vespaDocId @@ -280,24 +299,26 @@ async function processFileJob(jobData: FileProcessingJob, startTime: number) { for (const [resultIndex, processingResult] of processingResults.entries()) { // Create Vespa document with proper fileName (matching original logic) const targetPath = file.path - + // Reconstruct the original filePath (full path from collection root) - const reconstructedFilePath = targetPath === "/" - ? file.fileName - : targetPath.substring(1) + file.fileName // Remove leading "/" and add filename - + const reconstructedFilePath = + targetPath === "/" + ? file.fileName + : targetPath.substring(1) + file.fileName // Remove leading "/" and add filename + let vespaFileName = targetPath === "/" - ? file.collectionName + targetPath + reconstructedFilePath // Uses full path for root - : file.collectionName + targetPath + file.fileName // Uses filename for nested + ? file.collectionName + targetPath + reconstructedFilePath // Uses full path for root + : file.collectionName + targetPath + file.fileName // Uses filename for nested // For sheet processing results, append sheet information to fileName let docId = file.vespaDocId - if ('sheetName' in processingResult) { + if ("sheetName" in processingResult) { const sheetResult = processingResult as SheetProcessingResult - vespaFileName = processingResults.length > 1 - ? `${vespaFileName} / ${sheetResult.sheetName}` - : vespaFileName + vespaFileName = + processingResults.length > 1 + ? `${vespaFileName} / ${sheetResult.sheetName}` + : vespaFileName docId = sheetResult.docId } else if (processingResults.length > 1) { // For non-sheet files with multiple results, append index @@ -320,23 +341,30 @@ async function processFileJob(jobData: FileProcessingJob, startTime: number) { image_chunks_pos: processingResult.image_chunks_pos, chunks_map: processingResult.chunks_map, image_chunks_map: processingResult.image_chunks_map, - pageTitle : pageTitle, - metadata: JSON.stringify({ - originalFileName: file.originalName || file.fileName, - uploadedBy: file.uploadedByEmail || "system", - chunksCount: processingResult.chunks.length + processingResult.image_chunks.length, - imageChunksCount: processingResult.image_chunks.length, - processingMethod: getBaseMimeType(file.mimeType || "text/plain"), - ...(processingResult.processingMethod && { pdfProcessingMethod: processingResult.processingMethod }), - ...(pageTitle && { pageTitle }), - lastModified: Date.now(), - ...(('sheetName' in processingResult) && { - sheetName: (processingResult as SheetProcessingResult).sheetName, - sheetIndex: (processingResult as SheetProcessingResult).sheetIndex, - totalSheets: (processingResult as SheetProcessingResult).totalSheets, + pageTitle: pageTitle, + metadata: JSON.stringify( + mergeCollectionItemMetadata(file.metadata, { + originalFileName: file.originalName || file.fileName, + uploadedBy: file.uploadedByEmail || "system", + chunksCount: + processingResult.chunks.length + + processingResult.image_chunks.length, + imageChunksCount: processingResult.image_chunks.length, + processingMethod: getBaseMimeType(file.mimeType || "text/plain"), + ...(processingResult.processingMethod && { + pdfProcessingMethod: processingResult.processingMethod, + }), + ...(pageTitle && { pageTitle }), + lastModified: Date.now(), + ...("sheetName" in processingResult && { + sheetName: (processingResult as SheetProcessingResult).sheetName, + sheetIndex: (processingResult as SheetProcessingResult) + .sheetIndex, + totalSheets: (processingResult as SheetProcessingResult) + .totalSheets, + }), }), - ...(typeof file.metadata === "object" && file.metadata !== null && { ...file.metadata }), - }), + ), createdBy: file.uploadedByEmail || "system", duration: 0, mimeType: getBaseMimeType(file.mimeType || "text/plain"), @@ -349,19 +377,25 @@ async function processFileJob(jobData: FileProcessingJob, startTime: number) { // Insert into Vespa await insert(vespaDoc, KbItemsSchema) - totalChunksCount += processingResult.chunks.length + processingResult.image_chunks.length + totalChunksCount += + processingResult.chunks.length + processingResult.image_chunks.length } // Update status to completed with processing method metadata const chunksCount = totalChunksCount - + // Prepare metadata for database record - use last processing result for method info const lastResult = processingResults[processingResults.length - 1] - const dbMetadata = { + const dbMetadata = mergeCollectionItemMetadata(file.metadata, { chunksCount, - imageChunksCount: processingResults.reduce((sum, r) => sum + r.image_chunks.length, 0), - ...(lastResult.processingMethod && { pdfProcessingMethod: lastResult.processingMethod }), - } + imageChunksCount: processingResults.reduce( + (sum, r) => sum + r.image_chunks.length, + 0, + ), + ...(lastResult.processingMethod && { + pdfProcessingMethod: lastResult.processingMethod, + }), + }) await db .update(collectionItems) diff --git a/server/search/utils.ts b/server/search/utils.ts index 8418eaf9b..2c3a7331a 100644 --- a/server/search/utils.ts +++ b/server/search/utils.ts @@ -263,6 +263,18 @@ export async function extractCollectionVespaIds( } } + Logger.info( + { + collectionSelections: options.collectionSelections, + vespaFieldTranslation: { + clId: result.collectionIds ?? [], + clFd: result.collectionFolderIds ?? [], + docId: result.collectionFileIds ?? [], + }, + }, + "[extractCollectionVespaIds] Translated collection selections to Vespa field filters", + ) + return result } diff --git a/server/search/vespa.ts b/server/search/vespa.ts index 36731ea76..3a20d0dbc 100644 --- a/server/search/vespa.ts +++ b/server/search/vespa.ts @@ -110,17 +110,24 @@ export const searchVespaAgent = async ( AgentApps: Apps[] | null, options: Partial = {}, ) => { - - Logger.info( - `[searchVespaAgent] options.collectionSelections: ${JSON.stringify(options.collectionSelections)}`, - ) Logger.info( - `[searchVespaAgent] options.selectedItem: ${JSON.stringify(options.selectedItem)}`, + { + query, + email, + app, + entity, + agentApps: AgentApps, + collectionSelections: options.collectionSelections ?? [], + selectedItem: options.selectedItem ?? {}, + }, + "[searchVespaAgent] Incoming agent search options", ) const driveIds = await extractDriveIds(options, email) const processedCollectionSelections = await extractCollectionVespaIds(options) - Logger.debug( + Logger.info( { + driveIds, + processedCollectionSelections, hasCollectionIds: Boolean( processedCollectionSelections.collectionIds?.length, ), @@ -130,8 +137,11 @@ export const searchVespaAgent = async ( hasFileIds: Boolean( processedCollectionSelections.collectionFileIds?.length, ), + selectedItemApps: Object.keys( + (options.selectedItem ?? {}) as Record, + ), }, - "[searchVespaAgent] Processed selections summary", + "[searchVespaAgent] Prepared translated Vespa search inputs", ) // Send permissionId if available, otherwise send email diff --git a/server/tests/knowledgeBaseFlow.test.ts b/server/tests/knowledgeBaseFlow.test.ts new file mode 100644 index 000000000..9254ebc85 --- /dev/null +++ b/server/tests/knowledgeBaseFlow.test.ts @@ -0,0 +1,896 @@ +import { describe, expect, mock, test } from "bun:test" +import { Apps, KnowledgeBaseEntity } from "@xyne/vespa-ts/types" +import type { + Collection, + CollectionItem, + CollectionLsProjection, +} from "@/db/schema" +import type { MinimalAgentFragment } from "@/api/chat/types" + +process.env.ENCRYPTION_KEY ??= + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" +process.env.SERVICE_ACCOUNT_ENCRYPTION_KEY ??= + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" + +const { + __knowledgeBaseFlowInternals, + canonicalizeKnowledgeBasePath, + executeLsKnowledgeBase, + executeSearchKnowledgeBase, +} = await import("@/api/chat/tools/knowledgeBaseFlow") +const { mergeCollectionItemMetadata } = await import("@/queue/fileProcessor") + +const createCollection = (overrides: Partial): Collection => ({ + id: "collection-default", + workspaceId: 1, + ownerId: 1, + name: "Default Collection", + description: null, + vespaDocId: "cl-default", + isPrivate: true, + totalItems: 0, + lastUpdatedByEmail: "owner@example.com", + lastUpdatedById: 1, + uploadStatus: "completed" as any, + statusMessage: null, + retryCount: 0, + metadata: {}, + permissions: [], + collectionSourceUpdatedAt: new Date("2025-01-02T00:00:00.000Z"), + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-02T00:00:00.000Z"), + deletedAt: null, + via_apiKey: false, + ...overrides, +}) + +const createItem = (overrides: Partial): CollectionItem => ({ + id: "item-default", + collectionId: "collection-default", + parentId: null, + workspaceId: 1, + ownerId: 1, + name: "default", + type: "folder", + path: "/", + position: 0, + vespaDocId: "clfd-default", + totalFileCount: 0, + originalName: null, + storagePath: null, + storageKey: null, + mimeType: null, + fileSize: null, + checksum: null, + uploadedByEmail: "owner@example.com", + uploadedById: 1, + lastUpdatedByEmail: "owner@example.com", + lastUpdatedById: 1, + processingInfo: {}, + processedAt: null, + uploadStatus: "completed" as any, + statusMessage: null, + retryCount: 0, + metadata: {}, + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-02T00:00:00.000Z"), + deletedAt: null, + ...overrides, +}) + +const collectionAlpha = createCollection({ + id: "collection-alpha", + name: "Alpha", + description: "Alpha docs", + totalItems: 4, + metadata: { team: "platform" }, +}) + +const collectionBeta = createCollection({ + id: "collection-beta", + name: "Beta", + description: "Beta docs", + totalItems: 1, + metadata: { team: "sales" }, +}) + +const projectsFolder = createItem({ + id: "folder-projects", + collectionId: collectionAlpha.id, + name: "Projects", + type: "folder", + path: "/", + position: 0, + totalFileCount: 2, + metadata: { section: "root" }, +}) + +const apiFolder = createItem({ + id: "folder-api", + collectionId: collectionAlpha.id, + parentId: projectsFolder.id, + name: "API", + type: "folder", + path: "/Projects/", + position: 0, + totalFileCount: 1, + metadata: { section: "nested" }, +}) + +const specFile = createItem({ + id: "file-spec", + collectionId: collectionAlpha.id, + parentId: apiFolder.id, + name: "spec.md", + type: "file", + path: "/Projects/API/", + position: 0, + vespaDocId: "clf-spec", + originalName: "spec.md", + mimeType: "text/markdown", + metadata: { language: "md" }, +}) + +const readmeFile = createItem({ + id: "file-readme", + collectionId: collectionAlpha.id, + name: "README.txt", + type: "file", + path: "/", + position: 1, + vespaDocId: "clf-readme", + originalName: "README.txt", + mimeType: "text/plain", + metadata: { language: "txt" }, +}) + +const betaFile = createItem({ + id: "file-beta", + collectionId: collectionBeta.id, + name: "beta.txt", + type: "file", + path: "/", + position: 0, + vespaDocId: "clf-beta", + originalName: "beta.txt", + mimeType: "text/plain", +}) + +type LsKnowledgeBaseToolParams = Parameters[0] +type SearchKnowledgeBaseOptions = NonNullable< + Parameters[2] +> +type SearchExecutor = NonNullable + +function cloneMetadata( + metadata: unknown, +): Record | null | undefined { + if (!metadata || typeof metadata !== "object" || Array.isArray(metadata)) { + return undefined + } + return { ...(metadata as Record) } +} + +function withLsDefaults( + params: Partial, +): LsKnowledgeBaseToolParams { + return { + depth: 1, + offset: 0, + metadata: false, + ...params, + } +} + +function buildFixtures() { + const collections = [ + { + ...collectionAlpha, + metadata: cloneMetadata(collectionAlpha.metadata), + }, + { + ...collectionBeta, + metadata: cloneMetadata(collectionBeta.metadata), + }, + ] + const items = [ + { ...projectsFolder, metadata: cloneMetadata(projectsFolder.metadata) }, + { ...apiFolder, metadata: cloneMetadata(apiFolder.metadata) }, + { ...specFile, metadata: cloneMetadata(specFile.metadata) }, + { ...readmeFile, metadata: cloneMetadata(readmeFile.metadata) }, + { ...betaFile, metadata: cloneMetadata(betaFile.metadata) }, + ] + + return { collections, items } +} + +function createEmptyProjectionRow( + collectionId: string, + lsCollectionProjectionUpdatedAt = new Date(0), +): CollectionLsProjection { + return { + collectionId, + projection: { + rootIds: [], + childrenByParentId: {}, + nodesById: {}, + nodeIdByPath: {}, + }, + lsCollectionProjectionUpdatedAt, + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-01T00:00:00.000Z"), + lastError: null, + } +} + +function createRepo(options?: { + initialProjections?: Record + failProjectionUpsert?: boolean +}) { + const { collections, items } = buildFixtures() + const counters = { + listCollectionItems: 0, + listCollectionItemsByIds: 0, + getCollectionLsProjection: 0, + upsertCollectionLsProjection: 0, + recordCollectionLsProjectionError: 0, + } + const projectionRows = new Map() + Object.entries(options?.initialProjections ?? {}).forEach( + ([collectionId, row]) => { + if (row) projectionRows.set(collectionId, row) + }, + ) + + const repo = { + counters, + collections, + items, + projectionRows, + async getUserByEmail() { + return { id: 1 } + }, + async listUserScopedCollections() { + return collections + }, + async getCollectionById(collectionId: string) { + return ( + collections.find((collection) => collection.id === collectionId) ?? null + ) + }, + async getCollectionItemById(itemId: string) { + return items.find((item) => item.id === itemId) ?? null + }, + async listCollectionItems(collectionId: string) { + counters.listCollectionItems += 1 + return items.filter((item) => item.collectionId === collectionId) + }, + async listCollectionItemsByIds(collectionId: string, itemIds: string[]) { + counters.listCollectionItemsByIds += 1 + return items.filter( + (item) => + item.collectionId === collectionId && itemIds.includes(item.id), + ) + }, + async getCollectionLsProjection(collectionId: string) { + counters.getCollectionLsProjection += 1 + return projectionRows.get(collectionId) ?? null + }, + async upsertCollectionLsProjection(params: any) { + counters.upsertCollectionLsProjection += 1 + if (options?.failProjectionUpsert) { + throw new Error("projection upsert failed") + } + const row: CollectionLsProjection = { + collectionId: params.collectionId, + projection: params.projection, + lsCollectionProjectionUpdatedAt: params.lsCollectionProjectionUpdatedAt, + createdAt: + projectionRows.get(params.collectionId)?.createdAt ?? + new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-03T00:00:00.000Z"), + lastError: params.lastError ?? null, + } + projectionRows.set(params.collectionId, row) + return row + }, + async recordCollectionLsProjectionError( + collectionId: string, + lastError: string, + ) { + counters.recordCollectionLsProjectionError += 1 + const existing = + projectionRows.get(collectionId) ?? + createEmptyProjectionRow(collectionId) + projectionRows.set(collectionId, { + ...existing, + lastError, + updatedAt: new Date("2025-01-04T00:00:00.000Z"), + }) + }, + } + + return repo +} + +function createAgentPrompt(itemIds: string[]) { + return JSON.stringify({ + appIntegrations: { + knowledge_base: { + selectedAll: false, + itemIds, + }, + }, + }) +} + +function createContext(itemIds: string[]) { + return { + user: { + email: "tester@example.com", + numericId: 7, + workspaceNumericId: 9, + }, + agentPrompt: createAgentPrompt(itemIds), + } as any +} + +describe("canonicalizeKnowledgeBasePath", () => { + test("normalizes root, repeated separators, trailing separators, and preserves case", () => { + expect(canonicalizeKnowledgeBasePath("")).toBe("/") + expect(canonicalizeKnowledgeBasePath("////")).toBe("/") + expect(canonicalizeKnowledgeBasePath("Projects//API/Specs/")).toBe( + "/Projects/API/Specs", + ) + expect(canonicalizeKnowledgeBasePath("/TeamA/API")).toBe("/TeamA/API") + }) + + test("rejects dot segments", () => { + expect(() => canonicalizeKnowledgeBasePath("/Projects/./API")).toThrow() + expect(() => canonicalizeKnowledgeBasePath("/Projects/../API")).toThrow() + }) +}) + +describe("knowledge base target resolution", () => { + test("resolves collection, folder, file, path, and root path targets", async () => { + const repo = createRepo() + const cache = { direct: new Map(), projection: new Map() } + + const collectionTarget = + await __knowledgeBaseFlowInternals.resolveKnowledgeBaseTarget( + { type: "collection", collectionId: collectionAlpha.id }, + repo, + cache, + ) + expect(collectionTarget.kind).toBe("collection") + expect(collectionTarget.collection.id).toBe(collectionAlpha.id) + expect(collectionTarget.path).toBe("/") + + const folderTarget = + await __knowledgeBaseFlowInternals.resolveKnowledgeBaseTarget( + { type: "folder", folderId: projectsFolder.id }, + repo, + cache, + ) + expect(folderTarget.kind).toBe("node") + if (folderTarget.kind !== "node") { + throw new Error("Expected folder target to resolve to a node") + } + expect(folderTarget.node.id).toBe(projectsFolder.id) + expect(folderTarget.path).toBe("/Projects") + + const fileTarget = + await __knowledgeBaseFlowInternals.resolveKnowledgeBaseTarget( + { type: "file", fileId: specFile.id }, + repo, + cache, + ) + expect(fileTarget.kind).toBe("node") + if (fileTarget.kind !== "node") { + throw new Error("Expected file target to resolve to a node") + } + expect(fileTarget.node.id).toBe(specFile.id) + expect(fileTarget.path).toBe("/Projects/API/spec.md") + + const pathTarget = + await __knowledgeBaseFlowInternals.resolveKnowledgeBaseTarget( + { + type: "path", + collectionId: collectionAlpha.id, + path: "//Projects//API/", + }, + repo, + cache, + ) + expect(pathTarget.kind).toBe("node") + if (pathTarget.kind !== "node") { + throw new Error("Expected path target to resolve to a node") + } + expect(pathTarget.node.id).toBe(apiFolder.id) + expect(pathTarget.path).toBe("/Projects/API") + + const rootPathTarget = + await __knowledgeBaseFlowInternals.resolveKnowledgeBaseTarget( + { + type: "path", + collectionId: collectionAlpha.id, + path: "/", + }, + repo, + cache, + ) + expect(rootPathTarget.kind).toBe("collection") + expect(rootPathTarget.collection.id).toBe(collectionAlpha.id) + expect(rootPathTarget.path).toBe("/") + }) +}) + +describe("lsKnowledgeBase", () => { + test("lists accessible collections without scanning collection trees", async () => { + const repo = createRepo() + const result = await executeLsKnowledgeBase( + withLsDefaults({ + metadata: true, + }), + createContext([`cl-${collectionAlpha.id}`, `cl-${collectionBeta.id}`]), + repo, + ) + + expect(result.status).toBe("success") + expect(repo.counters.listCollectionItems).toBe(0) + expect( + (result as any).data.entries.map((entry: any) => entry.name), + ).toEqual(["Alpha", "Beta"]) + expect((result as any).data.entries[0].details).toEqual({ + total_items: 4, + name: "Alpha", + last_updated_by_email: "owner@example.com", + description: "Alpha docs", + updated_at: collectionAlpha.updatedAt, + created_at: collectionAlpha.createdAt, + metadata: { team: "platform" }, + }) + }) + + test("lists collection contents with depth, pagination, and metadata", async () => { + const repo = createRepo() + const result = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "collection", collectionId: collectionAlpha.id }, + depth: 2, + limit: 2, + offset: 1, + metadata: true, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + + expect(result.status).toBe("success") + expect(repo.counters.listCollectionItems).toBe(1) + expect(repo.counters.upsertCollectionLsProjection).toBe(1) + expect((result as any).data.total).toBe(3) + expect((result as any).data.entries.map((entry: any) => entry.id)).toEqual([ + apiFolder.id, + readmeFile.id, + ]) + expect((result as any).data.entries[0].details.total_file_count).toBe(1) + expect((result as any).data.entries[1].details.mime_type).toBe("text/plain") + }) + + test("lists folder, file, and path targets correctly", async () => { + const repo = createRepo() + + const folderResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "folder", folderId: projectsFolder.id }, + depth: 2, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + expect(folderResult.status).toBe("success") + expect( + (folderResult as any).data.entries.map((entry: any) => entry.id), + ).toEqual([apiFolder.id, specFile.id]) + + const fileResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "file", fileId: specFile.id }, + metadata: true, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + expect(fileResult.status).toBe("success") + expect((fileResult as any).data.entries).toHaveLength(1) + expect((fileResult as any).data.entries[0]).toMatchObject({ + id: specFile.id, + depth: 0, + path: "/Projects/API/spec.md", + }) + + const pathResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects/API", + }, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + expect(pathResult.status).toBe("success") + expect( + (pathResult as any).data.entries.map((entry: any) => entry.id), + ).toEqual([specFile.id]) + }) + + test("rejects out-of-scope collection and path targets before loading projections", async () => { + const collectionRepo = createRepo() + const collectionResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "collection", collectionId: collectionAlpha.id }, + depth: 2, + }), + createContext([`cl-${collectionBeta.id}`]), + collectionRepo, + ) + + expect(collectionResult.status).toBe("error") + expect((collectionResult as any).error.message).toContain( + "outside the current KB scope", + ) + expect(collectionRepo.counters.listCollectionItems).toBe(0) + expect(collectionRepo.counters.getCollectionLsProjection).toBe(0) + + const pathRepo = createRepo() + const pathResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects/API", + }, + }), + createContext([`cl-${collectionBeta.id}`]), + pathRepo, + ) + + expect(pathResult.status).toBe("error") + expect((pathResult as any).error.message).toContain( + "outside the current KB scope", + ) + expect(pathRepo.counters.listCollectionItems).toBe(0) + expect(pathRepo.counters.getCollectionLsProjection).toBe(0) + }) + + test("projection build correctness and traversal match stage 1 output", async () => { + const repo = createRepo() + const alphaItems = repo.items.filter( + (item) => item.collectionId === collectionAlpha.id, + ) + const projection = + __knowledgeBaseFlowInternals.buildCollectionLsProjection(alphaItems) + + expect(projection.rootIds).toEqual([projectsFolder.id, readmeFile.id]) + expect(projection.childrenByParentId[projectsFolder.id]).toEqual([ + apiFolder.id, + ]) + expect(projection.nodeIdByPath["/Projects/API/spec.md"]).toBe(specFile.id) + + const collection = repo.collections.find( + (item) => item.id === collectionAlpha.id, + )! + const directSnapshot = __knowledgeBaseFlowInternals.buildCollectionTree( + collection, + alphaItems, + ) + const projectionSnapshot = + __knowledgeBaseFlowInternals.buildSnapshotFromProjection( + collection, + __knowledgeBaseFlowInternals.buildCollectionLsProjection(alphaItems), + ) + + const directEntries = __knowledgeBaseFlowInternals.flattenLsEntries( + { + kind: "collection", + targetType: "collection", + collection, + path: "/", + snapshot: directSnapshot, + }, + 3, + ) + const projectionEntries = __knowledgeBaseFlowInternals.flattenLsEntries( + { + kind: "collection", + targetType: "collection", + collection, + path: "/", + snapshot: projectionSnapshot, + }, + 3, + ) + + expect(projectionEntries).toEqual(directEntries) + }) + + test("lazily rebuilds stale projection and then reuses persisted projection without scanning again", async () => { + const repo = createRepo({ + initialProjections: { + [collectionAlpha.id]: createEmptyProjectionRow( + collectionAlpha.id, + new Date("2024-12-31T00:00:00.000Z"), + ), + }, + }) + + const firstResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "collection", collectionId: collectionAlpha.id }, + depth: 2, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + + expect(firstResult.status).toBe("success") + expect(repo.counters.listCollectionItems).toBe(1) + expect(repo.counters.upsertCollectionLsProjection).toBe(1) + expect(repo.projectionRows.get(collectionAlpha.id)?.lastError).toBeNull() + + const secondResult = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "collection", collectionId: collectionAlpha.id }, + depth: 2, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + + expect(secondResult.status).toBe("success") + expect(repo.counters.listCollectionItems).toBe(1) + expect(repo.counters.getCollectionLsProjection).toBeGreaterThanOrEqual(2) + }) + + test("falls back to direct-db traversal and records last_error when projection rebuild fails", async () => { + const repo = createRepo({ + failProjectionUpsert: true, + }) + + const result = await executeLsKnowledgeBase( + withLsDefaults({ + target: { type: "collection", collectionId: collectionAlpha.id }, + depth: 2, + }), + createContext([`cl-${collectionAlpha.id}`]), + repo, + ) + + expect(result.status).toBe("success") + expect(repo.counters.listCollectionItems).toBe(2) + expect(repo.counters.recordCollectionLsProjectionError).toBe(1) + expect(repo.projectionRows.get(collectionAlpha.id)?.lastError).toContain( + "projection upsert failed", + ) + }) +}) + +describe("searchKnowledgeBase", () => { + test("agent-scoped KB access does not widen", async () => { + const repo = createRepo() + const searchExecutor = mock( + async (): Promise => [], + ) + const result = await executeSearchKnowledgeBase( + { + query: "api", + filters: { + targets: [ + { + type: "collection", + collectionId: collectionAlpha.id, + }, + ], + }, + }, + createContext([`clfd-${projectsFolder.id}`]), + { + repo, + searchExecutor, + }, + ) + + expect(result.status).toBe("error") + expect((result as any).error.message).toContain( + "outside the current KB scope", + ) + expect(searchExecutor).not.toHaveBeenCalled() + }) + + test("rejects out-of-scope collection and path search targets before resolving snapshots", async () => { + const collectionRepo = createRepo() + const collectionSearchExecutor = mock( + async (): Promise => [], + ) + const collectionResult = await executeSearchKnowledgeBase( + { + query: "api", + filters: { + targets: [ + { + type: "collection", + collectionId: collectionAlpha.id, + }, + ], + }, + }, + createContext([`cl-${collectionBeta.id}`]), + { + repo: collectionRepo, + searchExecutor: collectionSearchExecutor, + }, + ) + + expect(collectionResult.status).toBe("error") + expect((collectionResult as any).error.message).toContain( + "outside the current KB scope", + ) + expect(collectionRepo.counters.listCollectionItems).toBe(0) + expect(collectionRepo.counters.getCollectionLsProjection).toBe(0) + expect(collectionSearchExecutor).not.toHaveBeenCalled() + + const pathRepo = createRepo() + const pathSearchExecutor = mock( + async (): Promise => [], + ) + const pathResult = await executeSearchKnowledgeBase( + { + query: "api", + filters: { + targets: [ + { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects/API", + }, + ], + }, + }, + createContext([`cl-${collectionBeta.id}`]), + { + repo: pathRepo, + searchExecutor: pathSearchExecutor, + }, + ) + + expect(pathResult.status).toBe("error") + expect((pathResult as any).error.message).toContain( + "outside the current KB scope", + ) + expect(pathRepo.counters.listCollectionItems).toBe(0) + expect(pathRepo.counters.getCollectionLsProjection).toBe(0) + expect(pathSearchExecutor).not.toHaveBeenCalled() + }) + + test("maps filters.targets into the current KB search path and preserves citations", async () => { + const repo = createRepo() + const fragments: MinimalAgentFragment[] = [ + { + id: "file-spec:0", + content: "API search hit", + source: { + docId: "file-spec", + title: "spec.md", + url: "https://example.com/spec", + app: Apps.KnowledgeBase, + entity: KnowledgeBaseEntity.File, + }, + confidence: 0.91, + }, + ] + + let capturedSelections: unknown = null + const searchExecutor = mock( + async (options: any): Promise => { + capturedSelections = options.collectionSelections + return fragments + }, + ) as SearchExecutor + + const result = await executeSearchKnowledgeBase( + { + query: "api contract", + filters: { + targets: [ + { + type: "path", + collectionId: collectionAlpha.id, + path: "//Projects//API/", + }, + ], + }, + limit: 5, + offset: 2, + excludedIds: ["skip-me"], + }, + createContext([`clfd-${projectsFolder.id}`]), + { + repo, + searchExecutor, + }, + ) + + expect(searchExecutor).toHaveBeenCalledTimes(1) + expect(capturedSelections).toEqual([ + { + collectionFolderIds: [apiFolder.id], + }, + ]) + expect(result).toEqual({ + status: "success", + data: fragments, + }) + }) +}) + +describe("stage 3 metadata integrity", () => { + test("merges file processing metadata instead of replacing upload-time metadata", () => { + expect( + mergeCollectionItemMetadata( + { + originalPath: "Projects/API/spec.md", + originalFileName: "spec.md", + wasOverwritten: true, + }, + { + chunksCount: 12, + imageChunksCount: 2, + pdfProcessingMethod: "ocr", + }, + ), + ).toEqual({ + originalPath: "Projects/API/spec.md", + originalFileName: "spec.md", + wasOverwritten: true, + chunksCount: 12, + imageChunksCount: 2, + pdfProcessingMethod: "ocr", + }) + }) + + test("prefers computed processing metadata over stale upload-time values", () => { + expect( + mergeCollectionItemMetadata( + { + chunksCount: 1, + imageChunksCount: 0, + processingMethod: "application/pdf", + pageTitle: "Old Title", + sheetName: "Sheet 0", + sheetIndex: 0, + totalSheets: 1, + }, + { + chunksCount: 12, + imageChunksCount: 2, + processingMethod: "text/plain", + pageTitle: "Fresh Title", + sheetName: "Sheet 1", + sheetIndex: 1, + totalSheets: 3, + }, + ), + ).toEqual({ + chunksCount: 12, + imageChunksCount: 2, + processingMethod: "text/plain", + pageTitle: "Fresh Title", + sheetName: "Sheet 1", + sheetIndex: 1, + totalSheets: 3, + }) + }) +}) diff --git a/server/tests/knowledgeBaseTargetTranslation.test.ts b/server/tests/knowledgeBaseTargetTranslation.test.ts new file mode 100644 index 000000000..90351819a --- /dev/null +++ b/server/tests/knowledgeBaseTargetTranslation.test.ts @@ -0,0 +1,672 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test" +import { Apps, KnowledgeBaseEntity, SearchModes } from "@xyne/vespa-ts/types" +import type { MinimalAgentFragment } from "@/api/chat/types" + +process.env.ENCRYPTION_KEY ??= + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" +process.env.SERVICE_ACCOUNT_ENCRYPTION_KEY ??= + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" + +const mockLogger = { + error: mock(() => {}), + info: mock(() => {}), + warn: mock(() => {}), + debug: mock(() => {}), + child() { + return mockLogger + }, +} + +const mockGetAllFolderIds = mock(async (folderIds: string[]) => { + if (folderIds.includes("folder-projects")) { + return ["folder-api"] + } + return [] +}) + +const mockGetCollectionFilesVespaIds = mock(async (fileIds: string[]) => + fileIds.map((fileId) => ({ + vespaDocId: + { + "file-spec": "clf-spec", + "file-beta": "clf-beta", + }[fileId] ?? null, + })), +) + +mock.module("@/logger", () => ({ + getLogger: () => mockLogger, + getLoggerWithChild: () => () => mockLogger, + Subsystem: { + Chat: "Chat", + Api: "Api", + Vespa: "Vespa", + }, +})) + +mock.module("@/db/client", () => ({ + db: {}, +})) + +mock.module("@/db/connector", () => ({ + db: {}, + insertConnector: mock(async () => null), + getConnectors: mock(async () => []), + getConnectorByApp: mock(async () => null), + getConnector: mock(async () => null), + getOAuthConnectorWithCredentials: mock(async () => null), + getMicrosoftAuthConnectorWithCredentials: mock(async () => null), + getConnectorByExternalId: mock(async () => null), + getConnectorById: mock(async () => null), + getConnectorByAppAndEmailId: mock(async () => null), + updateConnector: mock(async () => null), + deleteConnector: mock(async () => null), + deleteOauthConnector: mock(async () => null), + loadConnectorState: mock(async () => null), + saveConnectorState: mock(async () => undefined), + getDatabaseConnectorForUser: mock(async () => null), + getOrCreateDatabaseConnectorKbCollectionId: mock(async () => null), + getDatabaseConnectorExternalIdByKbCollectionId: mock(async () => null), + clearDatabaseConnectorKbCollectionId: mock(async () => undefined), +})) + +mock.module("@/db/user", () => ({ + getPublicUserAndWorkspaceByEmail: mock(async () => []), + getUserAndWorkspaceByEmail: mock(async () => []), + getUserAndWorkspaceByOnlyEmail: mock(async () => []), + getUserByEmail: mock(async () => []), + createUser: mock(async () => null), + saveRefreshTokenToDB: mock(async () => undefined), + deleteRefreshTokenFromDB: mock(async () => undefined), + getUserById: mock(async () => null), + getUserMetaData: mock(async () => null), + getUsersByWorkspace: mock(async () => []), + getAllLoggedInUsers: mock(async () => []), + getAllIngestedUsers: mock(async () => []), + updateUser: mock(async () => null), + updateUserTimezone: mock(async () => null), + getUserFromJWT: mock(async () => null), + createUserApiKey: mock(async () => null), +})) + +mock.module("@/db/knowledgeBase", () => ({ + touchCollectionLsStructure: mock(async () => undefined), + touchCollectionLsStructures: mock(async () => undefined), + getCollectionById: mock(async () => null), + getCollectionItemById: mock(async () => null), + getCollectionLsProjection: mock(async () => null), + getCollectionsByOwner: mock(async () => []), + getAccessibleCollections: mock(async () => []), + recordCollectionLsProjectionError: mock(async () => undefined), + upsertCollectionLsProjection: mock(async () => null), + createCollection: mock(async () => null), + updateCollection: mock(async () => null), + softDeleteCollection: mock(async () => null), + createCollectionItem: mock(async () => null), + getCollectionItemsByParent: mock(async () => []), + getCollectionItemByPath: mock(async () => null), + updateCollectionItem: mock(async () => null), + softDeleteCollectionItem: mock(async () => null), + updateCollectionTotalCount: mock(async () => undefined), + updateFolderTotalCount: mock(async () => undefined), + updateParentFolderCounts: mock(async () => undefined), + createFolder: mock(async () => null), + createFileItem: mock(async () => null), + getAllCollectionItems: mock(async () => []), + getParentItems: mock(async () => []), + getAllFolderIds: mockGetAllFolderIds, + getCollectionFilesVespaIds: mockGetCollectionFilesVespaIds, + getCollectionItemsStatusByCollections: mock(async () => []), + getAllCollectionAndFolderItems: mock(async () => ({ + fileIds: [], + folderIds: [], + })), + getAllFolderItems: mock(async () => []), + getCollectionFoldersItemIds: mock(async () => []), + getCollectionFileByItemId: mock(async () => null), + createCollectionFile: mock(async () => null), + updateCollectionFile: mock(async () => null), + softDeleteCollectionFile: mock(async () => null), + generateStorageKey: mock(() => "storage-key"), + generateFileVespaDocId: mock(() => "file-vespa-id"), + generateFolderVespaDocId: mock(() => "folder-vespa-id"), + generateCollectionVespaDocId: mock(() => "collection-vespa-id"), + markParentAsProcessing: mock(async () => undefined), + updateParentStatus: mock(async () => undefined), + getRecordBypath: mock(async () => null), +})) + +const { executeLsKnowledgeBase, executeSearchKnowledgeBase } = await import( + "@/api/chat/tools/knowledgeBaseFlow" +) +const { extractCollectionVespaIds } = await import("@/search/utils") +const { sharedVespaService } = await import("@/search/vespaService") + +const collectionAlpha = { + id: "collection-alpha", + workspaceId: 1, + ownerId: 1, + name: "Alpha", + description: "Alpha docs", + vespaDocId: "cl-alpha", + isPrivate: true, + totalItems: 4, + lastUpdatedByEmail: "owner@example.com", + lastUpdatedById: 1, + uploadStatus: "completed", + statusMessage: null, + retryCount: 0, + metadata: {}, + permissions: [], + collectionSourceUpdatedAt: new Date("2025-01-02T00:00:00.000Z"), + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-02T00:00:00.000Z"), + deletedAt: null, + via_apiKey: false, +} as any + +const collectionBeta = { + ...collectionAlpha, + id: "collection-beta", + name: "Beta", + description: "Beta docs", + vespaDocId: "cl-beta", + totalItems: 1, +} as any + +const projectsFolder = { + id: "folder-projects", + collectionId: collectionAlpha.id, + parentId: null, + workspaceId: 1, + ownerId: 1, + name: "Projects", + type: "folder", + path: "/", + position: 0, + vespaDocId: "clfd-projects", + totalFileCount: 2, + originalName: null, + storagePath: null, + storageKey: null, + mimeType: null, + fileSize: null, + checksum: null, + uploadedByEmail: "owner@example.com", + uploadedById: 1, + lastUpdatedByEmail: "owner@example.com", + lastUpdatedById: 1, + processingInfo: {}, + processedAt: null, + uploadStatus: "completed", + statusMessage: null, + retryCount: 0, + metadata: {}, + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-02T00:00:00.000Z"), + deletedAt: null, +} as any + +const apiFolder = { + ...projectsFolder, + id: "folder-api", + parentId: projectsFolder.id, + name: "API", + path: "/Projects/", + totalFileCount: 1, + vespaDocId: "clfd-api", +} as any + +const specFile = { + ...projectsFolder, + id: "file-spec", + parentId: apiFolder.id, + name: "spec.md", + type: "file", + path: "/Projects/API/", + vespaDocId: "clf-spec", + totalFileCount: 0, + originalName: "spec.md", + mimeType: "text/markdown", +} as any + +const readmeFile = { + ...specFile, + id: "file-readme", + parentId: null, + name: "README.txt", + path: "/", + position: 1, + vespaDocId: "clf-readme", + originalName: "README.txt", + mimeType: "text/plain", +} as any + +const betaFile = { + ...specFile, + id: "file-beta", + collectionId: collectionBeta.id, + parentId: null, + name: "beta.txt", + path: "/", + vespaDocId: "clf-beta", + originalName: "beta.txt", + mimeType: "text/plain", +} as any + +function createRepo() { + const collections = [collectionAlpha, collectionBeta] + const items = [projectsFolder, apiFolder, specFile, readmeFile, betaFile] + + return { + async getUserByEmail() { + return { id: 1 } + }, + async listUserScopedCollections() { + return collections + }, + async getCollectionById(collectionId: string) { + return ( + collections.find((collection) => collection.id === collectionId) ?? null + ) + }, + async getCollectionItemById(itemId: string) { + return items.find((item) => item.id === itemId) ?? null + }, + async listCollectionItems(collectionId: string) { + return items.filter((item) => item.collectionId === collectionId) + }, + async listCollectionItemsByIds(collectionId: string, itemIds: string[]) { + return items.filter( + (item) => + item.collectionId === collectionId && itemIds.includes(item.id), + ) + }, + async getCollectionLsProjection() { + return null + }, + async upsertCollectionLsProjection(params: any) { + return { + collectionId: params.collectionId, + projection: params.projection, + lsCollectionProjectionUpdatedAt: params.lsCollectionProjectionUpdatedAt, + createdAt: new Date("2025-01-01T00:00:00.000Z"), + updatedAt: new Date("2025-01-01T00:00:00.000Z"), + lastError: null, + } + }, + async recordCollectionLsProjectionError() { + return undefined + }, + } +} + +function createContext() { + return { + user: { + email: "tester@example.com", + numericId: 7, + workspaceNumericId: 9, + }, + agentPrompt: undefined, + } as any +} + +function buildKnowledgeBaseYql(processedSelections: { + collectionIds?: string[] + collectionFolderIds?: string[] + collectionFileIds?: string[] +}) { + return (sharedVespaService as any).HybridDefaultProfileForAgent( + 5, + Apps.KnowledgeBase, + null, + SearchModes.NativeRank, + null, + undefined, + [], + [Apps.KnowledgeBase], + [], + null, + [], + processedSelections, + [], + {}, + "tester@example.com", + {}, + "api contract", + ).yql as string +} + +async function runSearchAndCapture(params: any) { + const repo = createRepo() + let capturedSelections: any = null + let capturedProcessedSelections: any = null + let capturedYql = "" + + const searchExecutor = mock(async (options: any): Promise => { + capturedSelections = options.collectionSelections + capturedProcessedSelections = await extractCollectionVespaIds({ + collectionSelections: options.collectionSelections, + } as any) + capturedYql = buildKnowledgeBaseYql(capturedProcessedSelections) + + return [ + { + id: "fragment-1", + content: "hit", + source: { + docId: "clf-spec", + title: "spec.md", + url: "", + app: Apps.KnowledgeBase, + entity: KnowledgeBaseEntity.File, + }, + confidence: 0.9, + }, + ] + }) + + const result = await executeSearchKnowledgeBase(params, createContext(), { + repo: repo as any, + searchExecutor, + }) + + return { + result, + capturedSelections, + capturedProcessedSelections, + capturedYql, + searchExecutor, + } +} + +async function runLs(params: any) { + return executeLsKnowledgeBase(params, createContext(), createRepo() as any) +} + +describe("knowledge base target translation to Vespa", () => { + beforeEach(() => { + mockGetAllFolderIds.mockClear() + mockGetCollectionFilesVespaIds.mockClear() + }) + + test("collection targets remain collection IDs and become clId Vespa filters", async () => { + const outcome = await runSearchAndCapture({ + query: "alpha docs", + filters: { + targets: [ + { + type: "collection", + collectionId: collectionAlpha.id, + }, + ], + }, + }) + + expect(outcome.result.status).toBe("success") + expect(outcome.capturedSelections).toEqual([ + { + collectionIds: [collectionAlpha.id], + }, + ]) + expect(outcome.capturedProcessedSelections).toEqual({ + collectionIds: [collectionAlpha.id], + }) + expect(outcome.capturedYql).toContain(`clId contains '${collectionAlpha.id}'`) + }) + + test("folder and folder-path targets become clFd filters and expand descendant folders", async () => { + const outcome = await runSearchAndCapture({ + query: "api docs", + filters: { + targets: [ + { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects", + }, + ], + }, + }) + + expect(outcome.result.status).toBe("success") + expect(outcome.capturedSelections).toEqual([ + { + collectionFolderIds: [projectsFolder.id], + }, + ]) + expect(outcome.capturedProcessedSelections).toEqual({ + collectionFolderIds: [projectsFolder.id, apiFolder.id], + }) + expect(outcome.capturedYql).toContain(`clFd contains '${projectsFolder.id}'`) + expect(outcome.capturedYql).toContain(`clFd contains '${apiFolder.id}'`) + expect(mockGetAllFolderIds).toHaveBeenCalledWith([projectsFolder.id], {}) + }) + + test("file and file-path targets become docId filters using Vespa file doc IDs", async () => { + const outcome = await runSearchAndCapture({ + query: "spec", + filters: { + targets: [ + { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects/API/spec.md", + }, + ], + }, + }) + + expect(outcome.result.status).toBe("success") + expect(outcome.capturedSelections).toEqual([ + { + collectionFileIds: [specFile.id], + }, + ]) + expect(outcome.capturedProcessedSelections).toEqual({ + collectionFileIds: [specFile.vespaDocId], + }) + expect(outcome.capturedYql).toContain( + `docId contains '${specFile.vespaDocId}'`, + ) + expect(mockGetCollectionFilesVespaIds).toHaveBeenCalledWith( + [specFile.id], + {}, + ) + }) + + test("mixed targets are unioned before Vespa query construction", async () => { + const outcome = await runSearchAndCapture({ + query: "mixed", + filters: { + targets: [ + { + type: "collection", + collectionId: collectionBeta.id, + }, + { + type: "folder", + folderId: projectsFolder.id, + }, + { + type: "file", + fileId: specFile.id, + }, + ], + }, + }) + + expect(outcome.result.status).toBe("success") + expect(outcome.capturedSelections).toEqual([ + { + collectionIds: [collectionBeta.id], + collectionFolderIds: [projectsFolder.id], + collectionFileIds: [specFile.id], + }, + ]) + expect(outcome.capturedProcessedSelections).toEqual({ + collectionIds: [collectionBeta.id], + collectionFolderIds: [projectsFolder.id, apiFolder.id], + collectionFileIds: [specFile.vespaDocId], + }) + expect(outcome.capturedYql).toContain(`clId contains '${collectionBeta.id}'`) + expect(outcome.capturedYql).toContain(`clFd contains '${projectsFolder.id}'`) + expect(outcome.capturedYql).toContain( + `docId contains '${specFile.vespaDocId}'`, + ) + expect(outcome.capturedYql).toContain(" or ") + }) + + test("collection IDs returned by ls can be reused directly as collection targets", async () => { + const lsResult = await runLs({}) + expect(lsResult.status).toBe("success") + + const collectionEntry = (lsResult as any).data.entries.find( + (entry: any) => entry.type === "collection" && entry.id === collectionAlpha.id, + ) + expect(collectionEntry).toBeDefined() + + const outcome = await runSearchAndCapture({ + query: "alpha docs", + filters: { + targets: [ + { + type: "collection", + collectionId: collectionEntry.id, + }, + ], + }, + }) + + expect(outcome.capturedSelections).toEqual([ + { + collectionIds: [collectionAlpha.id], + }, + ]) + expect(outcome.capturedYql).toContain(`clId contains '${collectionAlpha.id}'`) + }) + + test("folder ls rows can be reused as folderId targets and path targets", async () => { + const lsResult = await runLs({ + target: { + type: "collection", + collectionId: collectionAlpha.id, + }, + depth: 2, + metadata: true, + }) + expect(lsResult.status).toBe("success") + + const folderEntry = (lsResult as any).data.entries.find( + (entry: any) => entry.type === "folder" && entry.id === projectsFolder.id, + ) + expect(folderEntry).toBeDefined() + expect(folderEntry.collection_id).toBe(collectionAlpha.id) + expect(folderEntry.path).toBe("/Projects") + + const folderIdOutcome = await runSearchAndCapture({ + query: "api docs", + filters: { + targets: [ + { + type: "folder", + folderId: folderEntry.id, + }, + ], + }, + }) + + expect(folderIdOutcome.capturedSelections).toEqual([ + { + collectionFolderIds: [projectsFolder.id], + }, + ]) + expect(folderIdOutcome.capturedYql).toContain( + `clFd contains '${projectsFolder.id}'`, + ) + + const pathOutcome = await runSearchAndCapture({ + query: "api docs", + filters: { + targets: [ + { + type: "path", + collectionId: folderEntry.collection_id, + path: folderEntry.path, + }, + ], + }, + }) + + expect(pathOutcome.capturedSelections).toEqual([ + { + collectionFolderIds: [projectsFolder.id], + }, + ]) + expect(pathOutcome.capturedYql).toContain(`clFd contains '${projectsFolder.id}'`) + }) + + test("file ls rows can be reused as fileId targets and path targets", async () => { + const lsResult = await runLs({ + target: { + type: "path", + collectionId: collectionAlpha.id, + path: "/Projects/API", + }, + depth: 1, + metadata: true, + }) + expect(lsResult.status).toBe("success") + + const fileEntry = (lsResult as any).data.entries.find( + (entry: any) => entry.type === "file" && entry.id === specFile.id, + ) + expect(fileEntry).toBeDefined() + expect(fileEntry.collection_id).toBe(collectionAlpha.id) + expect(fileEntry.path).toBe("/Projects/API/spec.md") + + const fileIdOutcome = await runSearchAndCapture({ + query: "spec", + filters: { + targets: [ + { + type: "file", + fileId: fileEntry.id, + }, + ], + }, + }) + + expect(fileIdOutcome.capturedSelections).toEqual([ + { + collectionFileIds: [specFile.id], + }, + ]) + expect(fileIdOutcome.capturedYql).toContain( + `docId contains '${specFile.vespaDocId}'`, + ) + + const pathOutcome = await runSearchAndCapture({ + query: "spec", + filters: { + targets: [ + { + type: "path", + collectionId: fileEntry.collection_id, + path: fileEntry.path, + }, + ], + }, + }) + + expect(pathOutcome.capturedSelections).toEqual([ + { + collectionFileIds: [specFile.id], + }, + ]) + expect(pathOutcome.capturedYql).toContain( + `docId contains '${specFile.vespaDocId}'`, + ) + }) +}) diff --git a/server/tests/messageAgentsFragments.test.ts b/server/tests/messageAgentsFragments.test.ts index 809e8aa69..7cfa38415 100644 --- a/server/tests/messageAgentsFragments.test.ts +++ b/server/tests/messageAgentsFragments.test.ts @@ -4,6 +4,7 @@ import { __messageAgentsHistoryInternals, __messageAgentsMetadataInternals, afterToolExecutionHook, + beforeToolExecutionHook, buildDelegatedAgentFragments, buildFinalSynthesisPayload, buildReviewPromptFromContext, @@ -150,6 +151,64 @@ describe("message-agents context tracking", () => { ) }) + test("excludedIds injection uses seen source docIds rather than fragment ids", async () => { + const context = createMockContext() + const chunkedFragment: MinimalAgentFragment = { + ...baseFragment, + id: "doc-1:0", + source: { + ...baseFragment.source, + docId: "doc-1", + }, + } + + await afterToolExecutionHook( + "searchGlobal", + { + status: "success", + metadata: { + contexts: [chunkedFragment], + }, + data: { + result: "Found ARR updates.", + }, + }, + { + toolCall: { id: "call-docid-1" } as any, + args: { query: "ARR" }, + state: { + context, + messages: [], + runId: createRunId("run-docid-1"), + traceId: createTraceId("trace-docid-1"), + currentAgentName: "xyne-agent", + turnCount: 1, + }, + agentName: "xyne-agent", + executionTime: 10, + status: "success", + }, + context.message.text, + [], + new Set(), + undefined, + context.turnCount + ) + + expect(Array.from(context.seenDocuments)).toEqual(["doc-1"]) + + const preparedArgs = await beforeToolExecutionHook( + "searchGlobal", + { + query: "ARR", + excludedIds: [], + }, + context + ) + + expect(preparedArgs.excludedIds).toEqual(["doc-1"]) + }, 15000) + test("afterToolExecutionHook enforces strict metadata constraints when no compliant docs exist", async () => { const context = createMockContext() const nonCompliantFragment: MinimalAgentFragment = { @@ -195,7 +254,7 @@ describe("message-agents context tracking", () => { expect(context.allFragments).toHaveLength(0) expect(context.currentTurnArtifacts.fragments).toHaveLength(0) - }) + }, 15000) test("buildReviewPromptFromContext includes plan, expectations, and image metadata", () => { const context = createMockContext()