diff --git a/apps/studio/lib/ai/tool-filter.test.ts b/apps/studio/lib/ai/tool-filter.test.ts index 3dae04865f2..c419c119e47 100644 --- a/apps/studio/lib/ai/tool-filter.test.ts +++ b/apps/studio/lib/ai/tool-filter.test.ts @@ -35,6 +35,7 @@ describe('tool allowance by opt-in level', () => { list_policies: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, // Log tools get_advisors: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, + get_logs: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, } as unknown as ToolSet const filtered = filterToolsByOptInLevel(mockTools, optInLevel as any) @@ -61,6 +62,7 @@ describe('tool allowance by opt-in level', () => { expect(tools).not.toContain('list_extensions') expect(tools).not.toContain('list_edge_functions') expect(tools).not.toContain('list_branches') + expect(tools).not.toContain('get_logs') expect(tools).not.toContain('execute_sql') }) @@ -76,6 +78,7 @@ describe('tool allowance by opt-in level', () => { expect(tools).toContain('list_policies') expect(tools).toContain('search_docs') expect(tools).not.toContain('get_advisors') + expect(tools).not.toContain('get_logs') expect(tools).not.toContain('execute_sql') }) @@ -91,6 +94,7 @@ describe('tool allowance by opt-in level', () => { expect(tools).toContain('list_policies') expect(tools).toContain('search_docs') expect(tools).toContain('get_advisors') + expect(tools).toContain('get_logs') expect(tools).not.toContain('execute_sql') }) @@ -106,6 +110,7 @@ describe('tool allowance by opt-in level', () => { expect(tools).toContain('list_policies') expect(tools).toContain('search_docs') expect(tools).toContain('get_advisors') + expect(tools).toContain('get_logs') expect(tools).not.toContain('execute_sql') }) }) @@ -125,6 +130,7 @@ describe('filterToolsByOptInLevel', () => { search_docs: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, // Log tools get_advisors: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, + get_logs: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, // Unknown tool - should be filtered out entirely some_other_tool: { execute: vitest.fn().mockResolvedValue({ status: 'success' }) }, } as unknown as ToolSet @@ -180,6 +186,7 @@ describe('filterToolsByOptInLevel', () => { 'list_branches', 'list_policies', 'get_advisors', + 'get_logs', ]) }) @@ -193,13 +200,14 @@ describe('filterToolsByOptInLevel', () => { 'list_branches', 'list_policies', 'get_advisors', + 'get_logs', ]) }) it('should stub log tools for schema opt-in level', async () => { const tools = filterToolsByOptInLevel(mockTools, 'schema') - await expectStubsFor(tools, ['get_advisors']) + await expectStubsFor(tools, ['get_advisors', 'get_logs']) }) // No execute_sql tool, so nothing additional to stub for schema_and_log opt-in level diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 2a32b8638b3..33d5f07c7fa 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -25,6 +25,7 @@ export const toolSetValidationSchema = z.record( 'list_branches', 'search_docs', 'get_advisors', + 'get_logs', // Local tools 'display_query', @@ -105,6 +106,7 @@ export const TOOL_CATEGORY_MAP: Record = { // Log tools - MCP and local get_advisors: TOOL_CATEGORIES.LOG, + get_logs: TOOL_CATEGORIES.LOG, } /**