diff --git a/src/main/ai/runtime/aiSdk/params/__tests__/buildAgentParams.test.ts b/src/main/ai/runtime/aiSdk/params/__tests__/buildAgentParams.test.ts index 2dbe04dd4bd..ee9790d4217 100644 --- a/src/main/ai/runtime/aiSdk/params/__tests__/buildAgentParams.test.ts +++ b/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() + }) +}) diff --git a/src/main/ai/runtime/aiSdk/params/buildAgentParams.ts b/src/main/ai/runtime/aiSdk/params/buildAgentParams.ts index fabe1f1c310..215fbd6ca47 100644 --- a/src/main/ai/runtime/aiSdk/params/buildAgentParams.ts +++ b/src/main/ai/runtime/aiSdk/params/buildAgentParams.ts @@ -75,8 +75,9 @@ export async function buildAgentParams(input: BuildAgentParamsInput): Promise 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() } const capabilities = assistant ? resolveCapabilities(model, provider, assistant) : undefined @@ -89,7 +90,8 @@ export async function buildAgentParams(input: BuildAgentParamsInput): Promise 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 diff --git a/src/main/ai/streamManager/context/PersistentChatContextProvider.ts b/src/main/ai/streamManager/context/PersistentChatContextProvider.ts index d3438126d7d..73b01b8513c 100644 --- a/src/main/ai/streamManager/context/PersistentChatContextProvider.ts +++ b/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,7 +478,8 @@ export class PersistentChatContextProvider implements ChatContextProvider { assistantId: string | undefined, uniqueModelId: UniqueModelId, history: CherryUIMessage[], - messageId: string + messageId: string, + knowledgeBaseIds: string[] | undefined ): AiStreamRequest { return { chatId: topicId, @@ -479,7 +487,8 @@ export class PersistentChatContextProvider implements ChatContextProvider { assistantId, uniqueModelId, messages: history, - messageId + messageId, + knowledgeBaseIds } } } diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeListTool.ts b/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeListTool.ts index e247880c42b..c9cde97be8f 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeListTool.ts +++ b/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 } } diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeManageTool.ts b/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeManageTool.ts index 6927f3c8c99..7b08a974631 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeManageTool.ts +++ b/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 } } diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeReadTool.ts b/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeReadTool.ts index d08e598093e..c76ef677f22 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeReadTool.ts +++ b/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 } } diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeSearchTool.ts b/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeSearchTool.ts index ea07d393fba..923098c121b 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/KnowledgeSearchTool.ts +++ b/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 } } diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeListTool.test.ts b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeListTool.test.ts index 4d3dd82deb6..4dfc00f7c83 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeListTool.test.ts +++ b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeListTool.test.ts @@ -157,14 +157,14 @@ function makeProcessingFileItem(id: string): KnowledgeItem { type ListArgs = { query?: string | null; groupId?: string | null; baseId?: string | null; maxDepth?: number | null } -function callExecute(args: ListArgs, ctx: { assistant?: Assistant } = {}): Promise { +function callExecute(args: ListArgs, ctx: { knowledgeBaseIds?: string[] } = {}): Promise { const execute = entry.tool.execute as (args: ListArgs, options: ToolExecutionOptions) => Promise return execute(args, { toolCallId: 'tc-1', messages: [], experimental_context: { requestId: 'req-1', - assistant: ctx.assistant, + knowledgeBaseIds: ctx.knowledgeBaseIds ?? [], abortSignal: new AbortController().signal } } as ToolExecutionOptions) @@ -192,7 +192,7 @@ describe('kb_list', () => { ]) knowledgeServiceListRootItems.mockReturnValue([]) - const result = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const result = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ id: string }> expect(result.map((b) => b.id)).toEqual(['kb-1']) @@ -204,7 +204,7 @@ describe('kb_list', () => { knowledgeServiceListBases.mockReturnValue([makeBase({ id: 'kb-1' }), makeBase({ id: 'kb-2' })]) knowledgeServiceListRootItems.mockReturnValue([]) - const result = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: [] }) })) as Array<{ + const result = (await callExecute({}, { knowledgeBaseIds: [] })) as Array<{ id: string }> expect(result.map((b) => b.id).sort()).toEqual(['kb-1', 'kb-2']) @@ -217,10 +217,9 @@ describe('kb_list', () => { ]) knowledgeServiceListRootItems.mockReturnValue([]) - const result = (await callExecute( - { groupId: 'g1' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1', 'kb-2'] }) } - )) as Array<{ id: string }> + const result = (await callExecute({ groupId: 'g1' }, { knowledgeBaseIds: ['kb-1', 'kb-2'] })) as Array<{ + id: string + }> expect(result.map((b) => b.id)).toEqual(['kb-1']) }) @@ -233,7 +232,7 @@ describe('kb_list', () => { const result = (await callExecute( { query: null, groupId: null }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1', 'kb-2'] }) } + { knowledgeBaseIds: ['kb-1', 'kb-2'] } )) as Array<{ id: string }> // null groupId must NOT collapse to `base.groupId === null`; both bases come back. expect(result.map((b) => b.id).sort()).toEqual(['kb-1', 'kb-2']) @@ -250,10 +249,9 @@ describe('kb_list', () => { return [] }) - const result = (await callExecute( - { query: 'RUST' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1', 'kb-2', 'kb-3'] }) } - )) as Array<{ id: string }> + const result = (await callExecute({ query: 'RUST' }, { knowledgeBaseIds: ['kb-1', 'kb-2', 'kb-3'] })) as Array<{ + id: string + }> expect(result.map((b) => b.id).sort()).toEqual(['kb-1', 'kb-3']) }) @@ -267,7 +265,7 @@ describe('kb_list', () => { makeProcessingFileItem('i5') ]) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ sampleSources: string[] itemCount: number }> @@ -284,7 +282,7 @@ describe('kb_list', () => { knowledgeServiceListBases.mockReturnValue([makeBase({ id: 'kb-1' })]) knowledgeServiceListRootItems.mockReturnValue([makeNoteItem('n1', 'a'.repeat(200))]) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ sampleSources: string[] }> expect(base.sampleSources).toHaveLength(1) @@ -298,7 +296,7 @@ describe('kb_list', () => { const items = Array.from({ length: 12 }, (_, idx) => makeFileItem(`i${idx}`, `file-${idx}.md`)) knowledgeServiceListRootItems.mockReturnValue(items) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ sampleSources: string[] }> expect(base.sampleSources).toHaveLength(8) @@ -309,7 +307,7 @@ describe('kb_list', () => { makeBase({ id: 'kb-1', status: 'failed', error: 'missing_embedding_model' }) ]) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ id: string status: string sampleSources: string[] @@ -328,7 +326,7 @@ describe('kb_list', () => { throw new Error('boom') }) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ id: string sampleSources: string[] itemCount?: number @@ -345,7 +343,7 @@ describe('kb_list', () => { knowledgeServiceListBases.mockReturnValue([makeBase({ id: 'kb-1' })]) knowledgeServiceListRootItems.mockReturnValue([]) - const [base] = (await callExecute({}, { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) })) as Array<{ + const [base] = (await callExecute({}, { knowledgeBaseIds: ['kb-1'] })) as Array<{ itemCount?: number itemsUnavailable?: boolean }> @@ -370,10 +368,7 @@ describe('kb_list', () => { it('outlines an in-scope base, forwarding maxDepth and mapping itemType → type', async () => { knowledgeServiceGetOrganizationTree.mockReturnValue(orgTree()) - const result = await callExecute( - { baseId: 'kb-1', maxDepth: 2 }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - ) + const result = await callExecute({ baseId: 'kb-1', maxDepth: 2 }, { knowledgeBaseIds: ['kb-1'] }) expect(knowledgeServiceGetOrganizationTree).toHaveBeenCalledWith('kb-1', { maxDepth: 2 }) expect(result).toEqual({ @@ -390,10 +385,7 @@ describe('kb_list', () => { }) it('returns an error and does not traverse when the base is outside the assistant scope', async () => { - const result = (await callExecute( - { baseId: 'kb-other' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - )) as { error: string } + const result = (await callExecute({ baseId: 'kb-other' }, { knowledgeBaseIds: ['kb-1'] })) as { error: string } expect(result.error).toContain('kb-other') expect(knowledgeServiceGetOrganizationTree).not.toHaveBeenCalled() @@ -404,10 +396,7 @@ describe('kb_list', () => { throw DataApiErrorFactory.notFound('Knowledge base', 'kb-gone') }) - const result = (await callExecute( - { baseId: 'kb-gone' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-gone'] }) } - )) as { error: string } + const result = (await callExecute({ baseId: 'kb-gone' }, { knowledgeBaseIds: ['kb-gone'] })) as { error: string } expect(result.error).toContain('kb-gone') expect(result.error).toContain('kb_list') @@ -476,32 +465,46 @@ describe('kb_list', () => { }) describe('applies', () => { - it('applies only when a base exists AND one is bound to the assistant (matches kb_search/kb_read)', () => { + it('applies only when a base exists AND one is in the effective scope (matches kb_search/kb_read)', () => { const applies = entry.applies! // No base in the system → never applies, even with bound ids. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: false + hasAnyKnowledgeBase: false, + knowledgeBaseIds: ['kb-1'] }) ).toBe(false) - // A base exists but none bound (or no assistant) → does NOT apply: listing every base would be a - // discovery dead-end (no kb_read / kb_search to act on them) and widen the per-assistant scope. - expect(applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true })).toBe(false) + // A base exists but the effective scope is empty → does NOT apply: listing every base would be a + // discovery dead-end (no kb_read / kb_search to act on them) and widen the scope. + expect( + applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true, knowledgeBaseIds: [] }) + ).toBe(false) expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: [] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: [] }) ).toBe(false) - // A base exists AND is bound → applies. + // A base exists AND is bound to the assistant → applies. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-1'] + }) + ).toBe(true) + // Assistant has no static binding, but the composer selected one for this turn → applies. + expect( + applies({ + assistant: makeAssistant({ knowledgeBaseIds: [] }), + mcpToolIds: new Set(), + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-selected-this-turn'] }) ).toBe(true) }) diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeManageTool.test.ts b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeManageTool.test.ts index e3cf202da7d..e24e079e80e 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeManageTool.test.ts +++ b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeManageTool.test.ts @@ -42,14 +42,14 @@ type ManageArgs = { conceptIds?: string[] } -function callExecute(args: ManageArgs, ctx: { assistant?: Assistant } = {}): Promise { +function callExecute(args: ManageArgs, ctx: { knowledgeBaseIds?: string[] } = {}): Promise { const execute = entry.tool.execute as (args: ManageArgs, options: ToolExecutionOptions) => Promise return execute(args, { toolCallId: 'tc-1', messages: [], experimental_context: { requestId: 'req-1', - assistant: ctx.assistant, + knowledgeBaseIds: ctx.knowledgeBaseIds ?? [], abortSignal: new AbortController().signal } } as ToolExecutionOptions) @@ -80,7 +80,7 @@ describe('kb_manage', () => { it('returns an error and does not mutate when the base is outside the assistant scope', async () => { const result = (await callExecute( { baseId: 'kb-other', action: 'delete', conceptIds: ['docs/a.md'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('kb-other') @@ -90,7 +90,7 @@ describe('kb_manage', () => { it('adds a file by absolute path, deriving the source name from the basename', async () => { const result = await callExecute( { baseId: 'kb-1', action: 'add', type: 'file', path: '/Users/me/docs/report.pdf' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(addItems).toHaveBeenCalledWith('kb-1', [ @@ -102,7 +102,7 @@ describe('kb_manage', () => { it('rejects a non-absolute file path via schema validation and does not add', async () => { const result = (await callExecute( { baseId: 'kb-1', action: 'add', type: 'file', path: 'relative/report.pdf' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('Invalid knowledge item to add') @@ -115,7 +115,7 @@ describe('kb_manage', () => { it('adds a url, using the url as its source', async () => { const result = await callExecute( { baseId: 'kb-1', action: 'add', type: 'url', url: 'https://example.com/post' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(addItems).toHaveBeenCalledWith('kb-1', [ @@ -127,7 +127,7 @@ describe('kb_manage', () => { it('adds a note, deriving the source from the first line when no title is given', async () => { const result = await callExecute( { baseId: 'kb-1', action: 'add', type: 'note', content: 'First line\nsecond line' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(addItems).toHaveBeenCalledWith('kb-1', [ @@ -139,7 +139,7 @@ describe('kb_manage', () => { it('returns a steer and does not add when a required add field is missing', async () => { const result = (await callExecute( { baseId: 'kb-1', action: 'add', type: 'file' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('path') @@ -151,7 +151,7 @@ describe('kb_manage', () => { const result = await callExecute( { baseId: 'kb-1', action: 'delete', conceptIds: ['docs/a.md', 'docs/gone.md'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(deleteConcepts).toHaveBeenCalledWith('kb-1', ['docs/a.md', 'docs/gone.md']) @@ -163,7 +163,7 @@ describe('kb_manage', () => { const result = await callExecute( { baseId: 'kb-1', action: 'refresh', conceptIds: ['docs/a.md'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(refreshConcepts).toHaveBeenCalledWith('kb-1', ['docs/a.md']) @@ -171,10 +171,9 @@ describe('kb_manage', () => { }) it('returns a steer and does not delete when conceptIds are missing', async () => { - const result = (await callExecute( - { baseId: 'kb-1', action: 'delete' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - )) as { error: string } + const result = (await callExecute({ baseId: 'kb-1', action: 'delete' }, { knowledgeBaseIds: ['kb-1'] })) as { + error: string + } expect(result.error).toContain('conceptIds') expect(deleteConcepts).not.toHaveBeenCalled() @@ -185,7 +184,7 @@ describe('kb_manage', () => { const result = (await callExecute( { baseId: 'kb-gone', action: 'delete', conceptIds: ['docs/a.md'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-gone'] }) } + { knowledgeBaseIds: ['kb-gone'] } )) as { error: string } expect(result.error).toContain('kb-gone') @@ -216,31 +215,45 @@ describe('kb_manage', () => { }) describe('applies', () => { - it('returns true only when a base exists AND at least one is bound to the assistant', () => { + it('returns true only when a base exists AND at least one is in the effective scope', () => { const applies = entry.applies! // No base in the system → never applies, even with bound ids. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: false + hasAnyKnowledgeBase: false, + knowledgeBaseIds: ['kb-1'] }) ).toBe(false) - // A base exists but none bound to this assistant → does not apply. - expect(applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true })).toBe(false) + // A base exists but the effective scope is empty → does not apply. + expect( + applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true, knowledgeBaseIds: [] }) + ).toBe(false) expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: [] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: [] }) ).toBe(false) - // A base exists AND is bound → applies. + // A base exists AND is bound to the assistant → applies. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-1'] + }) + ).toBe(true) + // Assistant has no static binding, but the composer selected one for this turn → applies. + expect( + applies({ + assistant: makeAssistant({ knowledgeBaseIds: [] }), + mcpToolIds: new Set(), + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-selected-this-turn'] }) ).toBe(true) }) diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeReadTool.test.ts b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeReadTool.test.ts index d455db57e0a..02df734b319 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeReadTool.test.ts +++ b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeReadTool.test.ts @@ -41,14 +41,14 @@ type ReadArgs = { maxMatches?: number } -function callExecute(args: ReadArgs, ctx: { assistant?: Assistant } = {}): Promise { +function callExecute(args: ReadArgs, ctx: { knowledgeBaseIds?: string[] } = {}): Promise { const execute = entry.tool.execute as (args: ReadArgs, options: ToolExecutionOptions) => Promise return execute(args, { toolCallId: 'tc-1', messages: [], experimental_context: { requestId: 'req-1', - assistant: ctx.assistant, + knowledgeBaseIds: ctx.knowledgeBaseIds ?? [], abortSignal: new AbortController().signal } } as ToolExecutionOptions) @@ -86,7 +86,7 @@ describe('kb_read', () => { it('returns an error and does not read when the base is outside the assistant scope', async () => { const result = (await callExecute( { baseId: 'kb-other', conceptId: 'docs/intro.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('kb-other') @@ -99,7 +99,7 @@ describe('kb_read', () => { const result = await callExecute( { baseId: 'kb-1', conceptId: 'docs/intro.md', charStart: 0, charEnd: 11 }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(readConcept).toHaveBeenCalledWith('kb-1', 'docs/intro.md', { charStart: 0, charEnd: 11 }) @@ -118,10 +118,7 @@ describe('kb_read', () => { it('reads unscoped when the assistant has no knowledge scope', async () => { readConcept.mockResolvedValue(conceptContent()) - await callExecute( - { baseId: 'kb-1', conceptId: 'docs/intro.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: [] }) } - ) + await callExecute({ baseId: 'kb-1', conceptId: 'docs/intro.md' }, { knowledgeBaseIds: [] }) expect(readConcept).toHaveBeenCalledWith('kb-1', 'docs/intro.md', { charStart: undefined, charEnd: undefined }) }) @@ -131,7 +128,7 @@ describe('kb_read', () => { const result = (await callExecute( { baseId: 'kb-1', conceptId: 'docs/gone.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('docs/gone.md') @@ -145,7 +142,7 @@ describe('kb_read', () => { const result = (await callExecute( { baseId: 'kb-1', conceptId: 'docs/intro.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('docs/intro.md') @@ -158,10 +155,9 @@ describe('kb_read', () => { // NOT_FOUND — it must not be reported as a bad conceptId (it would send the model re-checking ids). readConcept.mockRejectedValue(DataApiErrorFactory.notFound('KnowledgeBase', 'kb-gone')) - const result = (await callExecute( - { baseId: 'kb-gone', conceptId: 'docs/intro.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: [] }) } - )) as { error: string } + const result = (await callExecute({ baseId: 'kb-gone', conceptId: 'docs/intro.md' }, { knowledgeBaseIds: [] })) as { + error: string + } expect(result.error).toContain('kb-gone') expect(result.error).toContain('kb_list') @@ -173,7 +169,7 @@ describe('kb_read', () => { const result = (await callExecute( { baseId: 'kb-1', conceptId: 'docs/intro.md' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toBe('vector store down') @@ -191,7 +187,7 @@ describe('kb_read', () => { const result = await callExecute( { baseId: 'kb-1', conceptId: 'docs/intro.md', pattern: 'match', ignoreCase: false, maxMatches: 10 }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } ) expect(grepConcept).toHaveBeenCalledWith('kb-1', 'docs/intro.md', { @@ -220,7 +216,7 @@ describe('kb_read', () => { const result = (await callExecute( { baseId: 'kb-1', conceptId: 'docs/intro.md', pattern: '(' }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } + { knowledgeBaseIds: ['kb-1'] } )) as { error: string } expect(result.error).toContain('Invalid kb_read regular expression') @@ -277,31 +273,45 @@ describe('kb_read', () => { }) describe('applies', () => { - it('returns true only when a base exists AND at least one is bound to the assistant', () => { + it('returns true only when a base exists AND at least one is in the effective scope', () => { const applies = entry.applies! // No base in the system → never applies, even with bound ids. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: false + hasAnyKnowledgeBase: false, + knowledgeBaseIds: ['kb-1'] }) ).toBe(false) - // A base exists but none bound to this assistant → does not apply. - expect(applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true })).toBe(false) + // A base exists but the effective scope is empty → does not apply. + expect( + applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true, knowledgeBaseIds: [] }) + ).toBe(false) expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: [] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: [] }) ).toBe(false) - // A base exists AND is bound → applies. + // A base exists AND is bound to the assistant → applies. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-1'] + }) + ).toBe(true) + // Assistant has no static binding, but the composer selected one for this turn → applies. + expect( + applies({ + assistant: makeAssistant({ knowledgeBaseIds: [] }), + mcpToolIds: new Set(), + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-selected-this-turn'] }) ).toBe(true) }) diff --git a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeSearchTool.test.ts b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeSearchTool.test.ts index cb9951f3b62..068529ba042 100644 --- a/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeSearchTool.test.ts +++ b/src/main/ai/tools/adapters/aiSdk/builtin/__tests__/KnowledgeSearchTool.test.ts @@ -36,7 +36,7 @@ function makeAssistant(overrides: Partial = {}): Assistant { function callExecute( args: { query: string; baseIds: string[] }, - ctx: { assistant?: Assistant; abortSignal?: AbortSignal } = {} + ctx: { knowledgeBaseIds?: string[]; abortSignal?: AbortSignal } = {} ): Promise { const execute = entry.tool.execute as ( args: { query: string; baseIds: string[] }, @@ -47,7 +47,7 @@ function callExecute( messages: [], experimental_context: { requestId: 'req-1', - assistant: ctx.assistant, + knowledgeBaseIds: ctx.knowledgeBaseIds ?? [], abortSignal: ctx.abortSignal ?? new AbortController().signal } } as ToolExecutionOptions) @@ -68,19 +68,13 @@ describe('kb_search', () => { }) it('returns [] and does not search when every requested baseId is outside the assistant scope', async () => { - const result = await callExecute( - { query: 'foo', baseIds: ['kb-other'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - ) + const result = await callExecute({ query: 'foo', baseIds: ['kb-other'] }, { knowledgeBaseIds: ['kb-1'] }) expect(result).toEqual([]) expect(knowledgeServiceSearch).not.toHaveBeenCalled() }) it('warns about dropped baseIds even when the whole set is out of scope (warn before the early return)', async () => { - const result = await callExecute( - { query: 'foo', baseIds: ['kb-other', 'kb-gone'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - ) + const result = await callExecute({ query: 'foo', baseIds: ['kb-other', 'kb-gone'] }, { knowledgeBaseIds: ['kb-1'] }) expect(result).toEqual([]) expect(knowledgeServiceSearch).not.toHaveBeenCalled() // The all-dropped case must still surface the rejection — the warn fires before the empty-target @@ -93,17 +87,14 @@ describe('kb_search', () => { it('drops out-of-scope baseIds but still searches the in-scope ones', async () => { knowledgeServiceSearch.mockResolvedValue([]) - await callExecute( - { query: 'q', baseIds: ['kb-1', 'kb-other'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }) } - ) + await callExecute({ query: 'q', baseIds: ['kb-1', 'kb-other'] }, { knowledgeBaseIds: ['kb-1'] }) expect(knowledgeServiceSearch).toHaveBeenCalledTimes(1) expect(knowledgeServiceSearch).toHaveBeenCalledWith('kb-1', 'q') }) it('trusts the requested baseIds when assistant scope is empty (future toggle path)', async () => { knowledgeServiceSearch.mockResolvedValue([]) - await callExecute({ query: 'q', baseIds: ['kb-1', 'kb-2'] }, { assistant: makeAssistant({ knowledgeBaseIds: [] }) }) + await callExecute({ query: 'q', baseIds: ['kb-1', 'kb-2'] }, { knowledgeBaseIds: [] }) expect(knowledgeServiceSearch).toHaveBeenCalledTimes(2) expect(knowledgeServiceSearch).toHaveBeenCalledWith('kb-1', 'q') expect(knowledgeServiceSearch).toHaveBeenCalledWith('kb-2', 'q') @@ -111,10 +102,7 @@ describe('kb_search', () => { it('queries every requested base when all are in-scope', async () => { knowledgeServiceSearch.mockResolvedValue([]) - await callExecute( - { query: 'how does X work', baseIds: ['kb-1', 'kb-2'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1', 'kb-2'] }) } - ) + await callExecute({ query: 'how does X work', baseIds: ['kb-1', 'kb-2'] }, { knowledgeBaseIds: ['kb-1', 'kb-2'] }) expect(knowledgeServiceSearch).toHaveBeenCalledTimes(2) expect(knowledgeServiceSearch).toHaveBeenCalledWith('kb-1', 'how does X work') expect(knowledgeServiceSearch).toHaveBeenCalledWith('kb-2', 'how does X work') @@ -137,7 +125,7 @@ describe('kb_search', () => { const result = (await callExecute( { query: 'q', baseIds: ['kb-1', 'kb-2'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['kb-1', 'kb-2'] }) } + { knowledgeBaseIds: ['kb-1', 'kb-2'] } )) as Array<{ id: number; content: string; score: number }> expect(result).toEqual([ @@ -154,7 +142,7 @@ describe('kb_search', () => { }) const result = (await callExecute( { query: 'q', baseIds: ['broken', 'good'] }, - { assistant: makeAssistant({ knowledgeBaseIds: ['broken', 'good'] }) } + { knowledgeBaseIds: ['broken', 'good'] } )) as Array<{ id: number; content: string }> expect(result).toEqual([{ id: 1, content: 'ok', score: 0.7 }]) }) @@ -192,31 +180,45 @@ describe('kb_search', () => { }) describe('applies', () => { - it('returns true only when a base exists AND at least one is bound to the assistant', () => { + it('returns true only when a base exists AND at least one is in the effective scope', () => { const applies = entry.applies! // No base in the system → never applies, even with bound ids. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: false + hasAnyKnowledgeBase: false, + knowledgeBaseIds: ['kb-1'] }) ).toBe(false) - // A base exists but none bound to this assistant → does not apply. - expect(applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true })).toBe(false) + // A base exists but the effective scope is empty → does not apply. + expect( + applies({ assistant: undefined, mcpToolIds: new Set(), hasAnyKnowledgeBase: true, knowledgeBaseIds: [] }) + ).toBe(false) expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: [] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: [] }) ).toBe(false) - // A base exists AND is bound → applies. + // A base exists AND is bound to the assistant → applies. expect( applies({ assistant: makeAssistant({ knowledgeBaseIds: ['kb-1'] }), mcpToolIds: new Set(), - hasAnyKnowledgeBase: true + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-1'] + }) + ).toBe(true) + // Assistant has no static binding, but the composer selected one for this turn → applies. + expect( + applies({ + assistant: makeAssistant({ knowledgeBaseIds: [] }), + mcpToolIds: new Set(), + hasAnyKnowledgeBase: true, + knowledgeBaseIds: ['kb-selected-this-turn'] }) ).toBe(true) }) diff --git a/src/main/ai/tools/adapters/aiSdk/context.ts b/src/main/ai/tools/adapters/aiSdk/context.ts index 3c4ec045ef6..b08c72774e8 100644 --- a/src/main/ai/tools/adapters/aiSdk/context.ts +++ b/src/main/ai/tools/adapters/aiSdk/context.ts @@ -22,6 +22,14 @@ export interface RequestContext { /** Attachments the `read_file` tool may read this request (filename → entry allow-list). */ readonly fileAttachments?: ReadonlyArray + + /** + * Effective knowledge base scope for this request, resolved by `resolveKnowledgeBaseIds`: the + * assistant's static `knowledgeBaseIds` binding when non-empty, otherwise the composer's per-turn + * `/` picker selection. The `kb_*` tools read this instead of `assistant.knowledgeBaseIds` + * directly. Defaults to empty. + */ + readonly knowledgeBaseIds?: readonly string[] } /** Per-call context: {@link RequestContext} + AI SDK's per-`execute` fields. */ diff --git a/src/main/ai/tools/adapters/aiSdk/types.ts b/src/main/ai/tools/adapters/aiSdk/types.ts index caf2049633a..57a9561eeba 100644 --- a/src/main/ai/tools/adapters/aiSdk/types.ts +++ b/src/main/ai/tools/adapters/aiSdk/types.ts @@ -14,6 +14,10 @@ export interface ToolApplyScope { readonly hasFileAttachments?: boolean /** True when the user has at least one knowledge base — gates the `kb_*` tools. Defaults to false. */ readonly hasAnyKnowledgeBase?: boolean + /** + * Effective knowledge base scope for this request; see `resolveKnowledgeBaseIds`. Defaults to empty. + */ + readonly knowledgeBaseIds?: readonly string[] } /** diff --git a/src/main/ai/tools/knowledgeLookup.ts b/src/main/ai/tools/knowledgeLookup.ts index cb956587698..a735d700a2a 100644 --- a/src/main/ai/tools/knowledgeLookup.ts +++ b/src/main/ai/tools/knowledgeLookup.ts @@ -3,10 +3,12 @@ * * Single source of truth shared by the AI-SDK builtin tools (`kb_search` / * `kb_list`) and the Claude Code in-process MCP bridge. `allowedIds` scopes - * which bases are reachable: in the AI-SDK path it is the assistant's - * `knowledgeBaseIds`; an empty array means "no scope" (all user bases), - * which is what the Claude Code agent path passes since agents have no - * per-assistant knowledge scope. + * which bases are reachable: in the AI-SDK path it is the scope resolved by + * `resolveKnowledgeBaseIds` (the assistant's own bound bases take precedence + * when non-empty; only when the assistant has none does the composer's + * per-turn selection define the scope); an empty array means "no scope" + * (all user bases), which is what the Claude Code agent path passes since + * agents have no per-assistant knowledge scope. * * `searchKnowledge` never throws: an infrastructure failure (every targeted * base errored) returns `{ error }` so it is distinguishable from "ran fine, @@ -143,7 +145,7 @@ function isConceptLookupError(output: object): output is KnowledgeLookupError { export async function searchKnowledge( query: string, baseIds: string[], - allowedIds: string[] + allowedIds: readonly string[] ): Promise { const targetIds = allowedIds.length > 0 ? baseIds.filter((id) => allowedIds.includes(id)) : baseIds @@ -228,7 +230,7 @@ async function readConcept( baseId: string, conceptId: string, range: { charStart?: number; charEnd?: number }, - allowedIds: string[] + allowedIds: readonly string[] ): Promise { if (allowedIds.length > 0 && !allowedIds.includes(baseId)) { logger.warn('kb_read targeted a base outside the assistant scope', { baseId, allowedIds }) @@ -279,7 +281,7 @@ async function grepConcept( baseId: string, conceptId: string, options: { pattern: string; ignoreCase?: boolean; maxMatches?: number }, - allowedIds: string[] + allowedIds: readonly string[] ): Promise { if (allowedIds.length > 0 && !allowedIds.includes(baseId)) { logger.warn('kb_read (grep mode) targeted a base outside the assistant scope', { baseId, allowedIds }) @@ -304,7 +306,10 @@ async function grepConcept( * One tool with two modes (see KNOWLEDGE_READ_DESCRIPTION); both cores above share the `{ error }` * contract, so this only routes by the presence of `pattern`. (`!= null` also covers undefined.) */ -export async function readOrGrepConcept(input: KbReadInput, allowedIds: string[]): Promise { +export async function readOrGrepConcept( + input: KbReadInput, + allowedIds: readonly string[] +): Promise { if (input.pattern != null) { return grepConcept( input.baseId, @@ -366,7 +371,11 @@ function conceptLookupError( * throws: an out-of-scope base or a service error returns `{ error }`; a missing * base maps to a clear "not found" message. `allowedIds` scopes reachable bases. */ -function readTree(baseId: string, options: { maxDepth?: number }, allowedIds: string[]): KnowledgeTreeResultOrError { +function readTree( + baseId: string, + options: { maxDepth?: number }, + allowedIds: readonly string[] +): KnowledgeTreeResultOrError { if (allowedIds.length > 0 && !allowedIds.includes(baseId)) { logger.warn('kb_list (outline mode) targeted a base outside the assistant scope', { baseId, allowedIds }) return { error: `Knowledge base "${baseId}" is not available to this assistant.` } @@ -403,7 +412,7 @@ function readTree(baseId: string, options: { maxDepth?: number }, allowedIds: st */ export async function listOrOutlineKnowledge( input: { query?: string | null; groupId?: string | null; baseId?: string | null; maxDepth?: number | null }, - allowedIds: string[] + allowedIds: readonly string[] ): Promise { if (input.baseId != null) { return readTree(input.baseId, { maxDepth: input.maxDepth ?? undefined }, allowedIds) @@ -437,7 +446,7 @@ type ManageKnowledgeInput = { */ export async function manageKnowledge( input: ManageKnowledgeInput, - allowedIds: string[] + allowedIds: readonly string[] ): Promise { if (allowedIds.length > 0 && !allowedIds.includes(input.baseId)) { logger.warn('kb_manage targeted a base outside the assistant scope', { baseId: input.baseId, allowedIds }) @@ -564,7 +573,7 @@ function deriveNoteSource(content: string, title?: string | null): string { async function listKnowledgeBases( query: string | null | undefined, groupId: string | null | undefined, - allowedIds: string[] + allowedIds: readonly string[] ): Promise { try { const knowledgeService = application.get('KnowledgeService') diff --git a/src/main/ai/types/requests.ts b/src/main/ai/types/requests.ts index 67114d559b5..42038bab254 100644 --- a/src/main/ai/types/requests.ts +++ b/src/main/ai/types/requests.ts @@ -42,6 +42,12 @@ export interface AiBaseRequest { /** "providerId::modelId" */ uniqueModelId?: UniqueModelId mcpToolIds?: string[] + /** + * Knowledge bases selected for this turn. Scope is resolved by `resolveKnowledgeBaseIds`: the + * assistant's own bound bases take precedence when non-empty (these ids are then ignored); only + * when the assistant has none does this selection define the scope. + */ + knowledgeBaseIds?: string[] requestOptions?: AiTransportOptions /** Per-request overrides (in-process only; assistant-less callers like the API gateway). */ callOverrides?: CallOverrides diff --git a/src/renderer/pages/home/HomePage.tsx b/src/renderer/pages/home/HomePage.tsx index 034d0e3b75f..9bf31c69531 100644 --- a/src/renderer/pages/home/HomePage.tsx +++ b/src/renderer/pages/home/HomePage.tsx @@ -376,7 +376,8 @@ const HomePage: FC = () => { trigger: 'submit-message', topicId: topic.id, userMessageParts: options?.userMessageParts ?? [{ type: 'text', text }], - mentionedModelIds: options?.mentionedModels + mentionedModelIds: options?.mentionedModels, + knowledgeBaseIds: options?.knowledgeBaseIds }) const rendererTopic = mapApiTopicToRendererTopic(topic) setDraftAssistantSelectionState(undefined) diff --git a/src/renderer/pages/home/useChatRuntimeState.ts b/src/renderer/pages/home/useChatRuntimeState.ts index d2d45623d45..d5595e60bbd 100644 --- a/src/renderer/pages/home/useChatRuntimeState.ts +++ b/src/renderer/pages/home/useChatRuntimeState.ts @@ -230,7 +230,8 @@ export function useChatRuntimeState({ topicId: conversation.topicId, parentAnchorId: conversation.parentAnchorId ?? undefined, userMessageParts: options?.userMessageParts ?? [{ type: 'text', text }], - mentionedModelIds: options?.mentionedModels + mentionedModelIds: options?.mentionedModels, + knowledgeBaseIds: options?.knowledgeBaseIds }), refreshMetadata: ({ topicId }) => invalidateCache(['/topics', `/topics/${topicId}`]) }) diff --git a/src/shared/ai/transport/stream.ts b/src/shared/ai/transport/stream.ts index 8a9f39afee6..5df81597c35 100644 --- a/src/shared/ai/transport/stream.ts +++ b/src/shared/ai/transport/stream.ts @@ -134,6 +134,13 @@ export type AiStreamOpenRequest = { topicId: string /** UniqueModelIds selected by the composer model selector — Main dispatches one execution per model. */ mentionedModelIds?: UniqueModelId[] + /** + * Knowledge bases selected via the composer `/` picker for this turn. Scope is resolved by + * `resolveKnowledgeBaseIds`: the assistant's own bound bases take precedence when non-empty + * (these ids are then ignored); only when the assistant has none does this selection define + * the scope. + */ + knowledgeBaseIds?: string[] } & ( | { /** Brand-new user turn: create the user msg + N assistant placeholders. */ diff --git a/src/shared/ipc/schemas/ai.ts b/src/shared/ipc/schemas/ai.ts index 8945b7a9a20..3681132e6ab 100644 --- a/src/shared/ipc/schemas/ai.ts +++ b/src/shared/ipc/schemas/ai.ts @@ -114,7 +114,8 @@ export const aiRequestSchemas = { input: z.intersection( z.object({ topicId: z.string().min(1), - mentionedModelIds: z.array(z.custom()).optional() + mentionedModelIds: z.array(z.custom()).optional(), + knowledgeBaseIds: z.array(z.string()).optional() }), z.discriminatedUnion('trigger', [ z.object({