Merge pull request #18706 from Budibase/feat/handle-gemini-503

feat: handle gemini 503 better

Author: Dakuan

packages/server/src/ai/tools/budibase/knowledgeFiles.ts

@@ -30,6 +30,33 @@ interface RankedMatch {
   file: KnowledgeBaseFile
 }
 
+const GEMINI_RETRIEVAL_UNAVAILABLE_MESSAGE =
+  "Gemini knowledge retrieval is temporarily unavailable (upstream 503). Budibase is operating normally; please retry shortly."
+const GEMINI_UPSTREAM_EVENT = "ai.gemini.upstream_unavailable"
+
+const isGeminiRetrievalUnavailable = (error: unknown): boolean => {
+  if (!error || typeof error !== "object") {
+    return false
+  }
+
+  const maybeStatus = (error as { status?: unknown }).status
+  if (maybeStatus === 503) {
+    return true
+  }
+
+  const maybeMessage = String((error as { message?: unknown }).message || "")
+    .toLowerCase()
+    .trim()
+  if (!maybeMessage) {
+    return false
+  }
+
+  return (
+    maybeMessage.includes("upstream unavailable") ||
+    maybeMessage.includes("service unavailable")
+  )
+}
+
 const toEpochMillis = (value?: string | number) => {
   if (value == null) {
     return 0
@@ -274,11 +301,31 @@ export const createKnowledgeSearchTool = (
       }
 
       const agent = await sdk.ai.agents.getOrThrow(agentId)
-      const result = await sdk.ai.rag.retrieveContextForAgent(agent, question)
-      return {
-        context: result.text,
-        sources: result.sources,
-        chunks: result.chunks,
+      try {
+        const result = await sdk.ai.rag.retrieveContextForAgent(agent, question)
+        return {
+          context: result.text,
+          sources: result.sources,
+          chunks: result.chunks,
+        }
+      } catch (error: any) {
+        if (isGeminiRetrievalUnavailable(error)) {
+          console.error("[AI_UPSTREAM] Gemini unavailable", {
+            event: GEMINI_UPSTREAM_EVENT,
+            provider: "gemini",
+            path: "knowledge_retrieval",
+            upstreamStatus: error?.status,
+            agentId,
+            errorMessage: error?.message,
+          })
+          throw new Error(GEMINI_RETRIEVAL_UNAVAILABLE_MESSAGE)
+        }
+        console.error("Failed to retrieve agent knowledge context", {
+          agentId,
+          status: error?.status,
+          message: error?.message,
+        })
+        throw error
       }
     },
   }),

packages/server/src/ai/tools/budibase/tests/knowledgeFiles.spec.ts

@@ -2,20 +2,35 @@ import {
   KnowledgeBaseFileStatus,
   type KnowledgeBaseFile,
 } from "@budibase/types"
+
+import { features } from "@budibase/backend-core"
+import sdk from "../../../../sdk"
+import {
+  createKnowledgeFilesTool,
+  createKnowledgeSearchTool,
+} from "../knowledgeFiles"
+
+jest.mock("@budibase/backend-core", () => ({
+  features: {
+    isEnabled: jest.fn(),
+  },
+}))
+
 jest.mock("../../../../sdk", () => ({
   __esModule: true,
   default: {
     ai: {
+      agents: {
+        getOrThrow: jest.fn(),
+      },
       rag: {
         listFilesForAgent: jest.fn(),
+        retrieveContextForAgent: jest.fn(),
       },
     },
   },
 }))
 
-import sdk from "../../../../sdk"
-import { createKnowledgeFilesTool } from "../knowledgeFiles"
-
 const executeTool = async (
   agentId: string,
   input: {
@@ -35,6 +50,21 @@ const executeTool = async (
   })
 }
 
+const executeSearchTool = async (
+  agentId: string,
+  input: { question: string }
+) => {
+  const toolDef = createKnowledgeSearchTool(agentId)
+  if (!toolDef.tool.execute) {
+    throw new Error("tool.execute is not a function")
+  }
+
+  return await toolDef.tool.execute(input, {
+    toolCallId: "test-tool-call",
+    messages: [],
+  })
+}
+
 describe("AI Tools - Knowledge files", () => {
   beforeEach(() => {
     jest.restoreAllMocks()
@@ -319,3 +349,38 @@ describe("AI Tools - Knowledge files", () => {
     expect(result.files[0].matchedBy).toBe("filename-contains")
   })
 })
+
+describe("AI Tools - Knowledge search", () => {
+  beforeEach(() => {
+    jest.restoreAllMocks()
+    jest.mocked(features.isEnabled).mockResolvedValue(true)
+  })
+
+  it("hard-fails with a clear message when Gemini retrieval is unavailable", async () => {
+    jest
+      .spyOn(sdk.ai.agents, "getOrThrow")
+      .mockResolvedValue({ _id: "agent_1" } as any)
+    jest
+      .spyOn(sdk.ai.rag, "retrieveContextForAgent")
+      .mockRejectedValue({ status: 503, message: "upstream unavailable" })
+
+    await expect(
+      executeSearchTool("agent_1", { question: "What is policy?" })
+    ).rejects.toThrow(
+      "Gemini knowledge retrieval is temporarily unavailable (upstream 503). Budibase is operating normally; please retry shortly."
+    )
+  })
+
+  it("preserves non-provider retrieval errors", async () => {
+    jest
+      .spyOn(sdk.ai.agents, "getOrThrow")
+      .mockResolvedValue({ _id: "agent_1" } as any)
+    jest
+      .spyOn(sdk.ai.rag, "retrieveContextForAgent")
+      .mockRejectedValue(new Error("Failed to map source metadata"))
+
+    await expect(
+      executeSearchTool("agent_1", { question: "What is policy?" })
+    ).rejects.toThrow("Failed to map source metadata")
+  })
+})

packages/server/src/api/controllers/ai/files.ts

@@ -22,9 +22,12 @@ import {
   RequiredKeys,
 } from "@budibase/types"
 import sdk from "../../../sdk"
+
 import { fetchSharePointSitesByConnection } from "../../../sdk/workspace/ai/knowledgeSources/sharepointConnection"
 import { getSharePointSiteIds, getSharePointSources } from "./sharepoint"
 
+const GEMINI_UPSTREAM_EVENT = "ai.gemini.upstream_unavailable"
+
 const normalizeUpload = (fileInput: any) => {
   if (!fileInput) {
     return undefined
@@ -171,6 +174,23 @@ export async function uploadAgentFile(
     ctx.body = { file: updated }
     ctx.status = 201
   } catch (error: any) {
+    const normalizedMessage = String(error?.message || "").toLowerCase()
+    const isGeminiUpstreamUnavailable =
+      error?.status === 503 ||
+      error?.statusCode === 503 ||
+      normalizedMessage.includes("upstream unavailable") ||
+      normalizedMessage.includes("service unavailable")
+
+    if (isGeminiUpstreamUnavailable) {
+      console.error("[AI_UPSTREAM] Gemini unavailable", {
+        event: GEMINI_UPSTREAM_EVENT,
+        provider: "gemini",
+        path: "knowledge_ingest",
+        upstreamStatus: error?.status,
+        agentId,
+        errorMessage: error?.message,
+      })
+    }
     console.error("Failed to upload agent file", error)
     throw new HTTPError(
       error?.message || "Failed to process uploaded file",