fix(knowledge-tools): scope kb_* tools to composer-selected bases too

src/main/ai/runtime/aiSdk/params/__tests__/buildAgentParams.test.ts

@@ -1,10 +1,24 @@
 import type { ProviderOptions } from '@ai-sdk/provider-utils'
-import type { StopCondition, ToolSet } from 'ai'
-import { describe, expect, it } from 'vitest'
+import type { StopCondition, Tool, ToolSet } from 'ai'
+import { afterEach, describe, expect, it, vi } from 'vitest'
 
-import { makeModel } from '../../../../__tests__/fixtures'
+import { makeAssistant, makeModel } from '../../../../__tests__/fixtures'
+import { registry } from '../../../../tools/adapters/aiSdk/registry'
+import type { ToolEntry } from '../../../../tools/adapters/aiSdk/types'
 import type { CallOverrides } from '../../../../types/requests'
-import { applyCallOverrides, composeStopWhen } from '../buildAgentParams'
+
+vi.mock('@application', () => ({
+  application: {
+    get: (name: string) => {
+      if (name === 'KnowledgeService') return { hasAnyBase: () => true }
+      throw new Error(`unexpected service: ${name}`)
+    }
+  }
+}))
+
+const { applyCallOverrides, composeStopWhen, resolveKnowledgeBaseIds, resolveTools } = await import(
+  '../buildAgentParams'
+)
 
 /**
  * Covers the first-class per-request override merge that replaced the old
@@ -99,3 +113,62 @@ describe('composeStopWhen', () => {
     expect(await conditions[0]({ steps: new Array(19) } as never)).toBe(false)
   })
 })
+
+describe('resolveKnowledgeBaseIds', () => {
+  it('falls back to the assistant-bound bases when the request selects none', () => {
+    expect(resolveKnowledgeBaseIds(makeAssistant({ knowledgeBaseIds: ['kb-1'] }), undefined)).toEqual(['kb-1'])
+  })
+
+  it('trusts the request-selected bases when the assistant has no static binding', () => {
+    expect(resolveKnowledgeBaseIds(makeAssistant({ knowledgeBaseIds: [] }), ['kb-2'])).toEqual(['kb-2'])
+    expect(resolveKnowledgeBaseIds(undefined, ['kb-2'])).toEqual(['kb-2'])
+  })
+
+  it('drops request-selected bases outside the assistant scope instead of expanding it', () => {
+    // An assistant statically bound to `kb-public` must not become searchable for `kb-private`
+    // just because the renderer/IPC request asked for it — the assistant's own binding is the
+    // trust boundary, not whatever the composer UI happened to let the user pick.
+    expect(resolveKnowledgeBaseIds(makeAssistant({ knowledgeBaseIds: ['kb-public'] }), ['kb-private'])).toEqual([
+      'kb-public'
+    ])
+    expect(resolveKnowledgeBaseIds(makeAssistant({ knowledgeBaseIds: ['kb-1'] }), ['kb-1', 'kb-2'])).toEqual(['kb-1'])
+  })
+
+  it('returns an empty array when neither source selects a base', () => {
+    expect(resolveKnowledgeBaseIds(undefined, undefined)).toEqual([])
+    expect(resolveKnowledgeBaseIds(makeAssistant({ knowledgeBaseIds: [] }), undefined)).toEqual([])
+  })
+})
+
+describe('resolveTools knowledge-base wiring', () => {
+  const KB_GATED_TOOL_NAME = 'test-kb-gated-tool'
+
+  const kbGatedEntry: ToolEntry = {
+    name: KB_GATED_TOOL_NAME,
+    namespace: 'test',
+    description: 'test-only tool gated on knowledgeBaseIds',
+    defer: 'never',
+    tool: {} as Tool,
+    applies: (scope) => (scope.knowledgeBaseIds?.length ?? 0) > 0
+  }
+
+  afterEach(() => {
+    registry.deregister(KB_GATED_TOOL_NAME)
+  })
+
+  it('exposes a kb-gated tool when the effective knowledgeBaseIds is non-empty', async () => {
+    registry.register(kbGatedEntry)
+
+    const { tools } = await resolveTools({}, undefined, makeModel(), false, ['kb-1'])
+
+    expect(tools?.[KB_GATED_TOOL_NAME]).toBeDefined()
+  })
+
+  it('hides a kb-gated tool when the effective knowledgeBaseIds is empty', async () => {
+    registry.register(kbGatedEntry)
+
+    const { tools } = await resolveTools({}, undefined, makeModel(), false, [])
+
+    expect(tools?.[KB_GATED_TOOL_NAME]).toBeUndefined()
+  })
+})

src/main/ai/runtime/aiSdk/params/buildAgentParams.ts

@@ -75,8 +75,9 @@ export async function buildAgentParams(input: BuildAgentParamsInput): Promise<Bu
   applyHttpTrace(sdkConfig, request.chatId, model)
   const fileAttachments = collectFileAttachments(request.messages)
   const hasFileAttachments = fileAttachments.length > 0
+  const knowledgeBaseIds = resolveKnowledgeBaseIds(assistant, request.knowledgeBaseIds)
   const { tools, deferredEntries, mcpToolIds } = canModelConsumeTools(model)
-    ? await resolveTools(request, assistant, model, hasFileAttachments)
+    ? await resolveTools(request, assistant, model, hasFileAttachments, knowledgeBaseIds)
     : { tools: undefined, deferredEntries: [] as ToolEntry[], mcpToolIds: new Set<string>() }
   const capabilities = assistant ? resolveCapabilities(model, provider, assistant) : undefined
 
@@ -89,7 +90,8 @@ export async function buildAgentParams(input: BuildAgentParamsInput): Promise<Bu
     topicId: request.chatId,
     assistant,
     abortSignal: signal,
-    fileAttachments
+    fileAttachments,
+    knowledgeBaseIds
   }
 
   const scope: RequestScope = {
@@ -105,7 +107,8 @@ export async function buildAgentParams(input: BuildAgentParamsInput): Promise<Bu
     aiSdkProviderId,
     requestContext,
     mcpToolIds,
-    hasFileAttachments
+    hasFileAttachments,
+    knowledgeBaseIds
   }
 
   const features = extraFeatures?.length ? [...INTERNAL_FEATURES, ...extraFeatures] : INTERNAL_FEATURES
@@ -161,11 +164,12 @@ function canModelConsumeTools(model: Model): boolean {
  * sync the MCP entries into the registry, then materialise the active
  * `ToolSet` via `applies` predicates and defer exposition.
  */
-async function resolveTools(
+export async function resolveTools(
   request: BuildAgentParamsInput['request'],
   assistant: Assistant | undefined,
   model: Model,
-  hasFileAttachments: boolean
+  hasFileAttachments: boolean,
+  knowledgeBaseIds: readonly string[]
 ): Promise<{
   tools: ToolSet | undefined
   deferredEntries: ToolEntry[]
@@ -184,7 +188,13 @@ async function resolveTools(
   }
 
   const hasAnyKnowledgeBase = resolveHasAnyKnowledgeBase()
-  const activeEntries = registry.selectActive({ assistant, mcpToolIds, hasFileAttachments, hasAnyKnowledgeBase })
+  const activeEntries = registry.selectActive({
+    assistant,
+    mcpToolIds,
+    hasFileAttachments,
+    hasAnyKnowledgeBase,
+    knowledgeBaseIds
+  })
   let tools: ToolSet | undefined
   if (activeEntries.length > 0) {
     tools = {}
@@ -218,6 +228,20 @@ function resolveHasAnyKnowledgeBase(): boolean {
   }
 }
 
+/**
+ * Effective knowledge base scope for this request. When the assistant has its own static binding,
+ * that binding IS the scope — the composer's per-turn selection can never expand it, since main
+ * cannot trust a renderer/IPC-supplied id list to stay within the assistant's configured bases (the
+ * composer UI happens to restrict picks to that set today, but that's a UI nicety, not a security
+ * boundary). Only an assistant with no static binding lets the per-turn selection define the scope —
+ * which is the actual gap this resolves: composer-only, ad-hoc knowledge base use.
+ */
+export function resolveKnowledgeBaseIds(assistant: Assistant | undefined, requestIds: string[] | undefined): string[] {
+  const assistantIds = assistant?.knowledgeBaseIds ?? []
+  if (assistantIds.length > 0) return assistantIds
+  return Array.from(new Set(requestIds ?? []))
+}
+
 /**
  * Assemble `AgentOptions`: capability-driven providerOptions overlaid with
  * the user's customParameters (split into AI-SDK standard params vs

src/main/ai/streamManager/context/PersistentChatContextProvider.ts

@@ -275,7 +275,14 @@ export class PersistentChatContextProvider implements ChatContextProvider {
       const history = this.buildHistory(userMessage.id)
       const models_ = assistantPlaceholders.map(({ model, placeholder, rootSpan }) => ({
         modelId: model.id,
-        request: this.buildStreamRequest(req.topicId, assistantId, model.id, history, placeholder.id),
+        request: this.buildStreamRequest(
+          req.topicId,
+          assistantId,
+          model.id,
+          history,
+          placeholder.id,
+          req.knowledgeBaseIds
+        ),
         rootSpan
       }))
       // Author the turn span's input attributes here, where the built request payload is available.
@@ -368,7 +375,7 @@ export class PersistentChatContextProvider implements ChatContextProvider {
         models: [
           {
             modelId: model.id,
-            request: this.buildStreamRequest(req.topicId, assistantId, model.id, history, anchor.id),
+            request: this.buildStreamRequest(req.topicId, assistantId, model.id, history, anchor.id, undefined),
             rootSpan
           }
         ],
@@ -438,7 +445,7 @@ export class PersistentChatContextProvider implements ChatContextProvider {
         models: [
           {
             modelId: model.id,
-            request: this.buildStreamRequest(req.topicId, assistantId, model.id, history, placeholder.id),
+            request: this.buildStreamRequest(req.topicId, assistantId, model.id, history, placeholder.id, undefined),
             rootSpan
           }
         ],
@@ -471,15 +478,17 @@ export class PersistentChatContextProvider implements ChatContextProvider {
     assistantId: string | undefined,
     uniqueModelId: UniqueModelId,
     history: CherryUIMessage[],
-    messageId: string
+    messageId: string,
+    knowledgeBaseIds: string[] | undefined
   ): AiStreamRequest {
     return {
       chatId: topicId,
       trigger: 'submit-message',
       assistantId,
       uniqueModelId,
       messages: history,
-      messageId
+      messageId,
+      knowledgeBaseIds
     }
   }
 }

src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeListTool.ts

@@ -10,8 +10,9 @@
  * Both modes live in the shared `knowledgeLookup` core so the Claude Code MCP bridge runs identical
  * logic; this file is just the AI-SDK `tool()` wrapper.
  *
- * Scope: when `assistant.knowledgeBaseIds` is non-empty, only those bases are reachable; when empty,
- * all user bases are.
+ * Scope: when the effective scope (the assistant's static binding when non-empty, else the composer's
+ * per-turn selection — see `resolveKnowledgeBaseIds`) is non-empty, only those bases are reachable;
+ * when empty, all user bases are.
  */
 
 import {
@@ -45,7 +46,7 @@ const kbListTool = tool({
   strict: true,
   execute: async (input, options) => {
     const { request } = getToolCallContext(options)
-    return listOrOutlineKnowledge(input, request.assistant?.knowledgeBaseIds ?? [])
+    return listOrOutlineKnowledge(input, request.knowledgeBaseIds ?? [])
   },
   toModelOutput: ({ input, output }) => knowledgeListModelOutput(output, input)
 })
@@ -58,9 +59,9 @@ export function createKbListToolEntry(): ToolEntry {
     defer: 'never',
     tool: kbListTool,
     // Discovery entry point, always inlined (defer: 'never') — but gated identically to kb_search /
-    // kb_read / kb_manage: a base must exist AND be bound to this assistant. Listing every base while
-    // none are bound (no kb_read / kb_search to act on them) is a discovery dead-end and widens the
-    // per-assistant scope, so kb_list shares the siblings' gate.
-    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.assistant?.knowledgeBaseIds?.length ?? 0) > 0
+    // kb_read / kb_manage: a base must exist AND be in scope (bound to this assistant, or selected
+    // this turn). Listing every base while none are in scope (no kb_read / kb_search to act on them)
+    // is a discovery dead-end and widens the scope, so kb_list shares the siblings' gate.
+    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.knowledgeBaseIds?.length ?? 0) > 0
   }
 }

src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeManageTool.ts

@@ -2,9 +2,11 @@
  * Knowledge base write tool — the destructive companion to the read tools.
  *
  * One tool with an `action`: add a new source (file / url / note), or delete /
- * re-index existing documents addressed by their Concept ID. Per-request
- * `assistant.knowledgeBaseIds` flows in via RequestContext and scopes which bases
- * are reachable. Every action mutates the base, so the tool is approval-gated
+ * re-index existing documents addressed by their Concept ID. The effective knowledge
+ * base scope (the assistant's static binding when non-empty, else the composer's
+ * per-turn selection — see `resolveKnowledgeBaseIds`) flows in via
+ * `RequestContext.knowledgeBaseIds` and scopes which bases are reachable. Every
+ * action mutates the base, so the tool is approval-gated
  * (`needsApproval: true`) — Cherry surfaces the approval card before it runs. The
  * mutation itself lives in the shared `knowledgeLookup` core so the Claude Code
  * MCP bridge runs identical logic (gated there by Claude Code's own permission
@@ -43,7 +45,7 @@ const kbManageTool = tool({
   needsApproval: true,
   execute: async (input, options) => {
     const { request } = getToolCallContext(options)
-    return manageKnowledge(input, request.assistant?.knowledgeBaseIds ?? [])
+    return manageKnowledge(input, request.knowledgeBaseIds ?? [])
   },
   toModelOutput: ({ output }) => knowledgeManageModelOutput(output)
 })
@@ -55,7 +57,7 @@ export function createKbManageToolEntry(): ToolEntry {
     description: 'Add, delete, or re-index documents in a knowledge base (requires approval)',
     defer: 'never',
     tool: kbManageTool,
-    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.assistant?.knowledgeBaseIds?.length ?? 0) > 0
+    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.knowledgeBaseIds?.length ?? 0) > 0
   }
 }
 

src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeReadTool.ts

@@ -8,8 +8,10 @@
  *     match's line, offsets, and snippet — for a precise lookup when semantic search is too fuzzy.
  *
  * The model passes a `conceptId` + `baseId` from a `kb_search` hit (or a `kb_list` outline).
- * Per-request `assistant.knowledgeBaseIds` flows in via RequestContext and scopes which bases are
- * reachable. Both modes live in the shared `knowledgeLookup` core so the Claude Code MCP bridge runs
+ * The effective knowledge base scope (the assistant's static binding when non-empty, else the
+ * composer's per-turn selection — see `resolveKnowledgeBaseIds`) flows in via
+ * `RequestContext.knowledgeBaseIds` and scopes which bases are reachable. Both modes live in the
+ * shared `knowledgeLookup` core so the Claude Code MCP bridge runs
  * identical logic; this file is just the AI-SDK `tool()` wrapper.
  */
 
@@ -39,7 +41,7 @@ const kbReadTool = tool({
   strict: true,
   execute: async (input, options) => {
     const { request } = getToolCallContext(options)
-    return readOrGrepConcept(input, request.assistant?.knowledgeBaseIds ?? [])
+    return readOrGrepConcept(input, request.knowledgeBaseIds ?? [])
   },
   toModelOutput: ({ output }) => knowledgeReadModelOutput(output)
 })
@@ -51,7 +53,7 @@ export function createKbReadToolEntry(): ToolEntry {
     description: 'Read a knowledge base document by its Concept ID, or grep within it',
     defer: 'always',
     tool: kbReadTool,
-    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.assistant?.knowledgeBaseIds?.length ?? 0) > 0
+    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.knowledgeBaseIds?.length ?? 0) > 0
   }
 }
 

src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeSearchTool.ts

@@ -2,10 +2,11 @@
  * Knowledge base search tool — agentic.
  *
  * The model picks the query and target `baseIds` (typically after `kb_list`).
- * Per-request `assistant.knowledgeBaseIds` flows in via RequestContext and
- * scopes which base IDs are accepted. The search itself lives in the shared
- * `knowledgeLookup` core so the Claude Code MCP bridge runs identical logic;
- * this file is just the AI-SDK `tool()` wrapper.
+ * The effective knowledge base scope (the assistant's static binding when non-empty, else the
+ * composer's per-turn selection — see `resolveKnowledgeBaseIds`) flows in via
+ * `RequestContext.knowledgeBaseIds` and scopes which base IDs are accepted. The search itself lives
+ * in the shared `knowledgeLookup` core so the Claude Code MCP
+ * bridge runs identical logic; this file is just the AI-SDK `tool()` wrapper.
  */
 
 import { KB_SEARCH_TOOL_NAME, kbSearchInputSchema, kbSearchOutputSchema } from '@shared/ai/builtinTools'
@@ -33,7 +34,7 @@ const kbSearchTool = tool({
   strict: true,
   execute: async ({ query, baseIds }, options) => {
     const { request } = getToolCallContext(options)
-    return searchKnowledge(query, baseIds, request.assistant?.knowledgeBaseIds ?? [])
+    return searchKnowledge(query, baseIds, request.knowledgeBaseIds ?? [])
   },
   toModelOutput: ({ output }) => knowledgeSearchModelOutput(output)
 })
@@ -45,7 +46,7 @@ export function createKbSearchToolEntry(): ToolEntry {
     description: "Search the user's private knowledge base",
     defer: 'always',
     tool: kbSearchTool,
-    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.assistant?.knowledgeBaseIds?.length ?? 0) > 0
+    applies: (scope) => scope.hasAnyKnowledgeBase === true && (scope.knowledgeBaseIds?.length ?? 0) > 0
   }
 }