@cap-js/agents 0.9.6 → 0.9.8

package/lib/protocol/persistence/cleanup.js

@@ -49,10 +49,19 @@ export async function triggerCleanup(serviceName) {
     LOG.debug(`triggerCleanup: srv.schedule not available (CDS < 9). Skipping cleanup scheduling.`)
     return
   }
-  serviceMap.set(serviceName, Date.now())
-  const MAX_TIMEOUT = 2_147_483_647
-  const delay = Math.min(ttlMs + MS_OF_A_DAY, MAX_TIMEOUT)
-  await srv.schedule("cleanupTasks", {}).after(delay)
+  const now = Date.now()
+  serviceMap.set(serviceName, now)
+  const delay = ttlMs + MS_OF_A_DAY
+  const scheduled = srv.schedule("cleanupTasks", {})
+  const taskName = `cleanupTasks-${srv.name}-${new Date().toDateString().substring(0, 10)}`
+  // .as in cds10
+  if (typeof scheduled.as === "function") await scheduled.as(taskName).after(delay)
+  // .asTask in cds9.9
+  else if (typeof scheduled.asTask === "function") await scheduled.asTask(taskName).after(delay)
+  // cds9 < 9.9 has no task naming API
+  else {
+    LOG.error("Cannot schedule cleanup of old tasks. Please update your @sap/cds version.")
+  }
 }
 
 /**

package/lib/telemetry/chat-tracing.js

@@ -308,6 +308,8 @@ export function setLLMSpanStartAttrs(
   if (params?.temperature != null)
     span.setAttribute("gen_ai.request.temperature", params.temperature)
   if (params?.max_tokens != null) span.setAttribute("gen_ai.request.max_tokens", params.max_tokens)
+  const modelParams = pickModelParams(params)
+  if (modelParams) span.setAttribute("gen_ai.request.model_params", JSON.stringify(modelParams))
   if (cds.context?.["agent.context.id"])
     span.setAttribute("gen_ai.conversation.id", cds.context["agent.context.id"])
   if (streaming) span.setAttribute("gen_ai.request.stream", true)
@@ -316,7 +318,10 @@ export function setLLMSpanStartAttrs(
 
   // MLflow inputs: chat messages for Chat tab
   if ((cds.env.agents?.mlflow || LOG._debug) && messages) {
-    const mlMessages = messages.map((m) => ({ role: toRole(m), content: extractText(m.content) }))
+    const mlMessages = messages.map((m) => ({
+      role: toRole(m),
+      content: extractText(m.content),
+    }))
     setSpanAttrs(span, mlflowAttrs("LLM", { model, provider, inputs: { messages: mlMessages } }))
     if (LOG._debug) {
       span.setAttribute("gen_ai.input.messages", JSON.stringify(messages.map((m) => m.content)))
@@ -367,7 +372,9 @@ export function setLLMSpanEndAttrs(span, { model, provider, tokenUsage, outputCo
       mlflowAttrs("LLM", {
         model,
         provider,
-        outputs: { choices: [{ message: { role: "assistant", content: outputContent } }] },
+        outputs: {
+          choices: [{ message: { role: "assistant", content: outputContent } }],
+        },
       }),
     )
   }

package/lib/telemetry/mlflow/evaluation.js

@@ -6,17 +6,36 @@ export async function createEvalRun({ name } = {}) {
   if (!exporter) return null
   const creds = cds.env.requires?.mlflow?.credentials || {}
   const experimentId = creds.MLFLOW_EXPERIMENT_ID || process.env.MLFLOW_EXPERIMENT_ID || "0"
-  return exporter.createRun(experimentId, name)
+  const tags = [
+    ...(await getSourceTags()),
+    { key: "mlflow.user", value: "@cap-js/agents evaluation" },
+  ]
+  return exporter.createRun(experimentId, name, tags)
 }
 
-export async function closeEvalRun(runId) {
+export async function closeEvalRun({ mlflowRunId: runId, metricKeys = [], prompts = [] }) {
   if (!runId) return
-  getMlflowExporter()?.closeRun(runId)
+  const exporter = getMlflowExporter()
+  if (!exporter) return
+  await logFinalMlflowMetrics(runId, metricKeys, exporter)
+  const uniquePrompts = []
+  for (const prompt of prompts) {
+    if (!uniquePrompts.some((p) => p.version === prompt.version && p.name === prompt.name)) {
+      uniquePrompts.push(prompt)
+    }
+  }
+  if (uniquePrompts.length)
+    await exporter.linkPromptVersionsToRun(runId, uniquePrompts).catch(() => {})
+  await exporter.closeRun(runId)
 }
 
 // Log a flat metrics object; null/undefined values are skipped.
-export async function logMlflowMetrics(runId, metrics) {
+export async function logMlflowMetrics({ mlflowRunId: runId, metricKeys }, metrics) {
   if (!runId) return
+  metricKeys ??= new Set()
+  for (const [key, value] of Object.entries(metrics)) {
+    if (value != null) metricKeys.add(key)
+  }
   const exporter = getMlflowExporter()
   if (!exporter) return
   await Promise.allSettled(
@@ -26,6 +45,42 @@ export async function logMlflowMetrics(runId, metrics) {
   )
 }
 
+export async function logMlflowRunMetadata(runId, metadata, exporter = getMlflowExporter()) {
+  if (!runId || !exporter || !metadata) return
+  const params = {
+    ...(metadata.model && { "llm.model": metadata.model }),
+    ...(metadata.provider && { "llm.provider": metadata.provider }),
+  }
+  for (const [key, value] of Object.entries(metadata.params ?? {})) {
+    if (value != null) params[`llm.param.${key}`] = _stringifyMlflowValue(value)
+  }
+
+  await Promise.allSettled(
+    Object.entries(params).map(([key, value]) => exporter.logParam(runId, key, value)),
+  )
+}
+
+const AVG_METRICS = { success_rate: 1, output_correctness: 1, latency_ms: 1 }
+
+export async function logFinalMlflowMetrics(runId, metricKeys, exporter = getMlflowExporter()) {
+  if (!runId || !exporter) return
+  const keys = Array.from(metricKeys ?? []).concat(["success_rate", "output_correctness"])
+  await Promise.allSettled(
+    keys.map(async (key) => {
+      const history = await exporter.getMetricHistory(runId, key)
+      const values = history
+        .filter((metric) => metric?.step !== 1)
+        .map((metric) => Number(metric?.value))
+        .filter(Number.isFinite)
+      if (!values.length) return
+
+      const total = values.reduce((sum, value) => sum + value, 0)
+      const aggregate = AVG_METRICS[key] ? total / values.length : total
+      await exporter.logMetric(runId, key, aggregate, { step: 1 })
+    }),
+  )
+}
+
 export async function postMlflowAssessment(
   traceId,
   score,
@@ -36,3 +91,27 @@ export async function postMlflowAssessment(
 ) {
   getMlflowExporter()?.postAssessment(traceId, score, rationale, assessmentName, sourceId, opts)
 }
+
+function _stringifyMlflowValue(value) {
+  return typeof value === "string" ? value : JSON.stringify(value)
+}
+
+function isCI() {
+  return process.env.GITHUB_ACTIONS === "true"
+}
+
+async function getSourceTags() {
+  const serverUrl = process.env.GITHUB_SERVER_URL?.replace(/\/$/, "")
+  const sourceName =
+    serverUrl && process.env.GITHUB_REPOSITORY
+      ? `${serverUrl}/${process.env.GITHUB_REPOSITORY}`
+      : undefined
+  const branch = process.env.GITHUB_HEAD_REF
+  const commit = process.env.GITHUB_SHA
+  return [
+    sourceName && { key: "mlflow.source.name", value: sourceName },
+    { key: "mlflow.source.type", value: isCI() ? "JOB" : "LOCAL" },
+    branch && { key: "mlflow.source.git.branch", value: branch },
+    commit && { key: "mlflow.source.git.commit", value: commit },
+  ].filter(Boolean)
+}

package/lib/telemetry/mlflow/exporter/DatabricksExporter.js

@@ -1,6 +1,28 @@
 import { MlflowExporter } from "./MlflowExporter.js"
 
 export class DatabricksExporter extends MlflowExporter {
+  async getMetricHistory(runId, key) {
+    const metrics = []
+    let pageToken
+    do {
+      const query = new URLSearchParams({
+        run_id: runId,
+        metric_key: key,
+        max_results: 1000,
+        ...(pageToken && { page_token: pageToken }),
+      }).toString()
+      // eslint-disable-next-line no-await-in-loop
+      const data = await this._fetch(
+        `/api/2.0/mlflow/metrics/get-history?${query}`,
+        undefined,
+        "GET",
+      )
+      metrics.push(...(data?.metrics ?? []))
+      pageToken = data?.next_page_token
+    } while (pageToken)
+    return metrics
+  }
+
   async postAssessment(
     traceId,
     score,
@@ -31,15 +53,27 @@ export class DatabricksExporter extends MlflowExporter {
     })
   }
 
+  async linkPromptVersionsToRun(runId, prompts = []) {
+    if (!prompts.length) return
+    await super.linkPromptVersionsToRun(runId, prompts)
+    await this._fetch("/api/2.0/mlflow/unity-catalog/prompt-versions/links-to-runs", {
+      prompt_versions: prompts,
+      run_ids: [runId],
+    })
+  }
+
   // Returns { tags: [{key,value}], latestVersion: {version, tags} | null }.
   async ensurePrompt(name, description) {
     let res = await this._fetch(
-      `/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}`,
+      `/api/2.0/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}`,
       undefined,
       "GET",
     )
     if (!res) {
-      res = await this._fetch("/mlflow/unity-catalog/prompts", { name, prompt: { description } })
+      res = await this._fetch("/api/2.0/mlflow/unity-catalog/prompts", {
+        name,
+        prompt: { description },
+      })
     }
     const latestVersion = await this._getLatestUcVersion(name)
     return { tags: res?.tags ?? [], latestVersion }
@@ -47,7 +81,7 @@ export class DatabricksExporter extends MlflowExporter {
 
   async createPromptVersion(name, description, tags = [], template = "") {
     const res = await this._fetch(
-      `/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/versions`,
+      `/api/2.0/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/versions`,
       {
         prompt_version: { template, description, tags },
       },
@@ -56,7 +90,7 @@ export class DatabricksExporter extends MlflowExporter {
   }
 
   async setRegisteredModelTag(name, key, value) {
-    await this._fetch(`/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/tags`, {
+    await this._fetch(`/api/2.0/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/tags`, {
       key,
       value,
     })
@@ -65,7 +99,7 @@ export class DatabricksExporter extends MlflowExporter {
   // Returns { version, tags } of the latest UC prompt version, or null.
   async _getLatestUcVersion(name) {
     const res = await this._fetch(
-      `/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/versions/search`,
+      `/api/2.0/mlflow/unity-catalog/prompts/${encodeURIComponent(name)}/versions/search`,
       { max_results: 1 },
     )
     const pv = res?.prompt_versions?.[0]

package/lib/telemetry/mlflow/exporter/MlflowExporter.js

@@ -29,12 +29,12 @@ export class MlflowExporter {
     }
   }
 
-  async createRun(experimentId, name) {
+  async createRun(experimentId, name, tags) {
     const data = await this._fetch("/api/2.0/mlflow/runs/create", {
       experiment_id: experimentId,
       run_name: name || `eval-${new Date().toISOString()}`,
       start_time: Date.now(),
-      tags: [{ key: "mlflow.source.type", value: "LOCAL" }],
+      tags,
     })
     return data?.run?.info?.run_id ?? null
   }
@@ -47,13 +47,52 @@ export class MlflowExporter {
     })
   }
 
-  async logMetric(runId, key, value) {
+  async linkPromptVersionsToRun(runId, prompts = []) {
+    if (!prompts.length) return
+    await this._fetch("/api/2.0/mlflow/runs/set-tag", {
+      run_id: runId,
+      key: "mlflow.linkedPrompts",
+      value: JSON.stringify(prompts),
+    })
+  }
+
+  async getMetricHistory(runId, key) {
+    const metrics = []
+    let pageToken
+    do {
+      const query = new URLSearchParams({
+        run_id: runId,
+        metric_key: key,
+        max_results: 1000,
+        ...(pageToken && { page_token: pageToken }),
+      }).toString()
+      // eslint-disable-next-line no-await-in-loop
+      const data = await this._fetch(
+        `/api/2.0/mlflow/metrics/get-history?${query}`,
+        undefined,
+        "GET",
+      )
+      metrics.push(...(data?.metrics ?? []))
+      pageToken = data?.next_page_token
+    } while (pageToken)
+    return metrics
+  }
+
+  async logMetric(runId, key, value, { step = 0 } = {}) {
     await this._fetch("/api/2.0/mlflow/runs/log-metric", {
       run_id: runId,
       key,
       value,
       timestamp: Date.now(),
-      step: 0,
+      step,
+    })
+  }
+
+  async logParam(runId, key, value) {
+    await this._fetch("/api/2.0/mlflow/runs/log-parameter", {
+      run_id: runId,
+      key,
+      value: String(value),
     })
   }
 

package/lib/telemetry/mlflow/index.js

@@ -7,10 +7,4 @@ export {
   RoutingSpanProcessor,
 } from "./tracing.js"
 export { postMlflowAssessment, createEvalRun, closeEvalRun } from "./evaluation.js"
-export {
-  syncPromptVersion,
-  syncSystemPrompt,
-  resolvePromptName,
-  linkedPromptsAttr,
-  hashPrompt,
-} from "./prompts.js"
+export { syncPromptVersion, resolvePromptName, linkedPromptsAttr, hashPrompt } from "./prompts.js"

package/lib/telemetry/mlflow/prompts.js

@@ -48,23 +48,27 @@ export function linkedPromptsAttr(promptName) {
 // Extracts the SystemMessage from prepared LLM messages and syncs it to MLflow.
 export function syncSystemPrompt(messages) {
   if (!cds.env.agents?.mlflow) return
-  const srvName = cds.context?.["agent.service"]
-  if (!srvName) return
-  const sysMsg = messages?.find((m) => m.type === "system")
-  if (!sysMsg) return
+  const msg = messages[0]
+  if (!msg) return
   const text =
-    typeof sysMsg.content === "string"
-      ? sysMsg.content
-      : Array.isArray(sysMsg.content)
-        ? sysMsg.content
+    typeof msg.content === "string"
+      ? msg.content
+      : Array.isArray(msg.content)
+        ? msg.content
             .filter((b) => b?.type === "text")
             .map((b) => b.text ?? "")
             .join("")
         : null
   if (!text) return
-  const srv = cds.services[srvName]
-  if (!srv) return
-  syncPromptVersion(resolvePromptName(srv), text).catch(() => {})
+  let name = msg.name
+  if (!name) {
+    const srvName = cds.context?.["agent.service"]
+    if (!srvName) return
+    const srv = cds.services[srvName]
+    if (!srv) return
+    name = resolvePromptName(srv)
+  }
+  syncPromptVersion(name, text).catch(() => {})
 }
 
 // AGENTS.md path relative to cds.root, or srv.name when no AGENTS.md exists.

package/lib/telemetry/mlflow/tracing.js

@@ -58,12 +58,12 @@ export function mlflowTraceAttrs() {
   const attrs = {
     // OTel semconv keys MLflow reads natively for user/session display
     "session.id": session,
-    "user.id": user,
     // mlflow.traceTag.* prefix for custom trace tags
-    "mlflow.traceTag.session": session,
-    "mlflow.traceTag.user": user,
     "mlflow.traceTag.tenant": tenant,
   }
+  if (cds.env.agents?.masking?.resolveInTraces) {
+    attrs["user.id"] = user
+  }
   return attrs
 }
 

package/lib/telemetry/span-masking.js

@@ -0,0 +1,165 @@
+import cds from "@sap/cds"
+import { decode, encode } from "@toon-format/toon"
+import { pseudonymizeToolResult } from "../masking/structured/index.js"
+import { resolveRemoteMcpTool } from "../masking/index.js"
+
+const LOG = cds.log("agents")
+const REGISTERED = Symbol.for("@cap-js/agents:trace-scrubber-registered")
+
+/**
+ * Returns true when PII should be replaced with pseudonym tokens in spans.
+ * Returns false when tokens should be resolved back to originals (resolveInTraces mode).
+ */
+function mustMaskValues() {
+  if (cds.env.agents?.masking?.resolveInTraces) {
+    LOG._debug &&
+      LOG.debug(
+        `Resolving pseudonymized values in OTEL spans because "cds.env.agents.masking.resolveInTraces" = true`,
+      )
+    return false
+  }
+  return true
+}
+
+function scrubValue(value, session) {
+  if (typeof value === "string") return session.scrubText(value)
+  if (Array.isArray(value)) return value.map((item) => scrubValue(item, session))
+  if (value !== null && typeof value === "object") {
+    for (const [k, v] of Object.entries(value)) value[k] = scrubValue(v, session)
+    return value
+  }
+  return value
+}
+
+function resolveValue(value, session) {
+  if (typeof value === "string") return session.resolveText(value)
+  if (Array.isArray(value)) return value.map((item) => resolveValue(item, session))
+  if (value !== null && typeof value === "object") {
+    for (const [k, v] of Object.entries(value)) value[k] = resolveValue(v, session)
+    return value
+  }
+  return value
+}
+
+function transformAttributes(attrs, session, mask) {
+  if (!attrs) return
+  const transform = mask ? scrubValue : resolveValue
+  for (const [key, value] of Object.entries(attrs)) attrs[key] = transform(value, session)
+}
+
+function parseJson(value) {
+  if (typeof value !== "string") return undefined
+  try {
+    return JSON.parse(value)
+  } catch {
+    return undefined
+  }
+}
+
+function scrubToolOutputs(span, session) {
+  const attrs = span.attributes
+  if (attrs?.["mlflow.spanType"] !== "TOOL") return
+  const rawOutputs = parseJson(attrs["mlflow.spanOutputs"])
+  // rawOutputs is the parsed value of the JSON-stringified tool outputs.
+  // Only TOON-encoded strings (the success path) can be scrubbed; skip objects (error case).
+  if (typeof rawOutputs !== "string") return
+  const inputs = parseJson(attrs["mlflow.spanInputs"]) ?? {}
+
+  const remote = resolveRemoteMcpTool(attrs["gen_ai.tool.call.id"])
+  const effectiveToolName = remote?.bareName ?? attrs["gen_ai.tool.call.id"]
+  const effectiveSrv = remote?.remoteSrv
+
+  const srvName = cds.context?.["agent.service"]
+  const defaultSrv = cds.services?.[srvName]
+  const args = {
+    toolName: effectiveToolName,
+    cql: inputs.cql,
+    strictMasking: true,
+    session,
+    srv: effectiveSrv ?? defaultSrv,
+    model: cds.context?.model ?? effectiveSrv?.model ?? defaultSrv?.model,
+  }
+
+  // Decode TOON or JSON, pseudonymize decoded object, re-encode
+  let scrubbed = rawOutputs
+  try {
+    const decoded = decode(rawOutputs)
+    pseudonymizeToolResult({ decoded, ...args })
+    scrubbed = encode(decoded)
+  } catch {
+    try {
+      const decoded = JSON.parse(rawOutputs)
+      pseudonymizeToolResult({ decoded, ...args })
+      scrubbed = JSON.stringify(decoded)
+    } catch {
+      // unparseable — leave unchanged
+    }
+  }
+  attrs["mlflow.spanOutputs"] = JSON.stringify(scrubbed)
+  if (attrs["gen_ai.tool.call.result"] === rawOutputs) attrs["gen_ai.tool.call.result"] = scrubbed
+}
+
+export class MaskingSpanProcessor {
+  onStart() {}
+
+  onEnd(span) {
+    if (!cds.env.agents?.masking) return
+    const session = cds.context?.["agent.pseudonyms"]
+    if (!session) return
+
+    const mask = mustMaskValues()
+    if (mask) scrubToolOutputs(span, session)
+
+    transformAttributes(span.attributes, session, mask)
+    for (const event of span.events ?? []) transformAttributes(event.attributes, session, mask)
+    if (span.status?.message) {
+      span.status.message = mask
+        ? session.scrubText(span.status.message)
+        : session.resolveText(span.status.message)
+    }
+  }
+
+  forceFlush() {
+    return Promise.resolve()
+  }
+
+  shutdown() {
+    return Promise.resolve()
+  }
+}
+
+function registerFirst(delegate, processor) {
+  // @opentelemetry/sdk-trace-base >= 2.0: internal MultiSpanProcessor holds processors in _spanProcessors
+  if (Array.isArray(delegate._activeSpanProcessor?._spanProcessors)) {
+    delegate._activeSpanProcessor._spanProcessors.unshift(processor)
+    return true
+  }
+  // @opentelemetry/sdk-trace-base ^1.x: BasicTracerProvider exposes _registeredSpanProcessors directly
+  if (Array.isArray(delegate._registeredSpanProcessors)) {
+    delegate._registeredSpanProcessors.unshift(processor)
+    return true
+  }
+  // @opentelemetry/sdk-trace-base ^1.x fallback: public addSpanProcessor API (appends, not prepends)
+  if (typeof delegate.addSpanProcessor === "function") {
+    delegate.addSpanProcessor(processor)
+    return true
+  }
+  return false
+}
+
+export async function setupTraceScrubbing() {
+  if (!cds.env.agents?.masking) return
+  try {
+    const { trace } = await import("@opentelemetry/api")
+    const provider = trace.getTracerProvider()
+    const delegate = provider.getDelegate?.() || provider
+    if (delegate[REGISTERED]) return
+    if (!registerFirst(delegate, new MaskingSpanProcessor())) {
+      LOG.warn("Trace scrubbing: no TracerProvider with span processor support")
+      return
+    }
+    delegate[REGISTERED] = true
+  } catch (err) {
+    LOG.warn("Trace scrubbing setup failed", { error: err.message })
+  }
+}

package/lib/telemetry/tool-tracing.js

@@ -50,9 +50,6 @@ export function _patchToolsProto(proto) {
         span.setAttribute("gen_ai.operation.name", "execute_tool")
         span.setAttribute("gen_ai.provider.name", "langchain")
         span.setAttribute("gen_ai.tool.call.id", toolName)
-        const inputs = args?.args ?? args
-        setSpanAttrs(span, mlflowAttrs("TOOL", { inputs, functionName: toolName }))
-        if (LOG._debug) span.setAttribute("gen_ai.tool.call.arguments", JSON.stringify(inputs))
       }
       const taskId = config?.configurable?._taskId || cds.context?.["agent.task.id"]
       let outcome
@@ -64,9 +61,12 @@ export function _patchToolsProto(proto) {
 
         const semanticError = result?.artifact?.isError === true
         outcome = semanticError ? "error" : "success"
-        const outputs = result?.content ?? result
+        let outputs = result?.content ?? result
 
         if (span) {
+          const inputs = args?.args ?? args
+          setSpanAttrs(span, mlflowAttrs("TOOL", { inputs, functionName: toolName }))
+          if (LOG._debug) span.setAttribute("gen_ai.tool.call.arguments", JSON.stringify(inputs))
           span.setAttribute("gen_ai.tool.call.outcome", outcome)
           if (semanticError) {
             const msg = typeof outputs === "string" ? outputs : "tool returned isError"

package/lib/utils/markdown.js

@@ -1,5 +1,5 @@
 import cds from "@sap/cds"
-import { slugified } from "./utils.js"
+import { effectiveDefinition, slugified } from "./utils.js"
 const { path, fs } = cds.utils
 const LOG = cds.log("agents")
 
@@ -7,19 +7,12 @@ const LOG = cds.log("agents")
  * Absolute filesystem directory of the `.cds` source file that defines `srv`.
  */
 export function serviceSourceDir(srv) {
-  const file = srv?.definition?.$location?.file
+  const file = srv?.definition?.["@source"] || srv?.definition?.$location?.file
   if (!file) return undefined
-  // $location.file is relative to cds.root
+  // Both @source and $location.file are relative to cds.root
   return path.dirname(path.join(cds.root, file))
 }
 
-/**
- * Get the effective service definition (feature-toggled if available, else base).
- */
-function effectiveDefinition(srv) {
-  return cds.context?.model?.definitions?.[srv.name] || srv?.definition
-}
-
 /**
  * Resolve a path from a `@agent.directory` / `@agent.card` annotation against the
  * `.cds` source file's directory. Absolute paths are returned as-is.

package/lib/utils/resilience.js

@@ -40,7 +40,7 @@ export function timeout(ms = DEFAULT_TIMEOUT_MS) {
 }
 
 // Retries on 5xx / network errors; bails immediately on 4xx.
-export function retry(retries = DEFAULT_RETRIES) {
+export function retry(retries = DEFAULT_RETRIES, { shouldRetry } = {}) {
   if (retries < 0) throw new Error("Number of retries must be greater or equal to 0.")
 
   return ({ fn }) =>
@@ -50,6 +50,7 @@ export function retry(retries = DEFAULT_RETRIES) {
           return await fn(arg) // eslint-disable-line no-await-in-loop
         } catch (error) {
           const status = error?.response?.status
+          if (shouldRetry && !shouldRetry(error)) throw error
           if (status == null)
             LOG.debug("HTTP request failed without a response status. Rethrowing.")
           else if (`${status}`.startsWith("4"))