fix: add tool rule options (#44)

* fix: add tool rule options

* fix: remove unneeded file

* fix: types

* fix:  types

Author: gal-monday

__tests__/unit/tool-rules-filtering.test.ts

@@ -0,0 +1,220 @@
+/**
+ * Test to verify toolRules.allowOnlyApiGroups filtering works correctly.
+ */
+
+import { createServer, type AgentToolProtocolServer } from '../../packages/server/src/index';
+import { AgentToolProtocolClient } from '../../packages/client/src/index';
+import { runInRequestScope } from '../../packages/server/src/core/request-scope';
+
+describe('toolRules.allowOnlyApiGroups filtering', () => {
+	let server: AgentToolProtocolServer;
+	let client: AgentToolProtocolClient;
+
+	beforeAll(async () => {
+		server = createServer();
+
+		// Register "math" API group
+		server.use({
+			name: 'math',
+			type: 'custom',
+			functions: [
+				{
+					name: 'add',
+					description: 'Add two numbers',
+					inputSchema: {
+						type: 'object',
+						properties: {
+							a: { type: 'number' },
+							b: { type: 'number' },
+						},
+						required: ['a', 'b'],
+					},
+					handler: async (input: unknown) => {
+						const { a, b } = input as { a: number; b: number };
+						return { result: a + b };
+					},
+				},
+			],
+		});
+
+		// Register "text" API group
+		server.use({
+			name: 'text',
+			type: 'custom',
+			functions: [
+				{
+					name: 'uppercase',
+					description: 'Convert text to uppercase',
+					inputSchema: {
+						type: 'object',
+						properties: {
+							text: { type: 'string' },
+						},
+						required: ['text'],
+					},
+					handler: async (input: unknown) => {
+						const { text } = input as { text: string };
+						return { result: text.toUpperCase() };
+					},
+				},
+			],
+		});
+
+		await server.start();
+
+		client = new AgentToolProtocolClient({
+			server: server as any,
+		});
+
+		await client.init({ name: 'test-client', version: '1.0.0' });
+		await client.connect();
+	});
+
+	afterAll(async () => {
+		// Cleanup if needed
+	});
+
+	it('should allow math API when toolRules.allowOnlyApiGroups includes math', async () => {
+		const result = await client.execute('return await api.math.add({ a: 2, b: 3 });', {
+			toolRules: {
+				allowOnlyApiGroups: ['math'],
+			},
+		});
+
+		expect(result.status).toBe('completed');
+		expect((result.result as any)?.result).toBe(5);
+	});
+
+	it('should block text API when toolRules.allowOnlyApiGroups only includes math', async () => {
+		const result = await client.execute('return await api.text.uppercase({ text: "hello" });', {
+			toolRules: {
+				allowOnlyApiGroups: ['math'],
+			},
+		});
+
+		// When text API is blocked, the code should fail because api.text is undefined
+		expect(result.status).toBe('failed');
+	});
+
+	it('should allow both APIs when no toolRules are specified', async () => {
+		const mathResult = await client.execute('return await api.math.add({ a: 1, b: 2 });');
+		expect(mathResult.status).toBe('completed');
+		expect((mathResult.result as any)?.result).toBe(3);
+
+		const textResult = await client.execute('return await api.text.uppercase({ text: "hello" });');
+		expect(textResult.status).toBe('completed');
+		expect((textResult.result as any)?.result).toBe('HELLO');
+	});
+
+	it('should allow text API when toolRules.allowOnlyApiGroups includes text', async () => {
+		const result = await client.execute('return await api.text.uppercase({ text: "world" });', {
+			toolRules: {
+				allowOnlyApiGroups: ['text'],
+			},
+		});
+
+		expect(result.status).toBe('completed');
+		expect((result.result as any)?.result).toBe('WORLD');
+	});
+
+	it('should block math API when toolRules.allowOnlyApiGroups only includes text', async () => {
+		const result = await client.execute('return await api.math.add({ a: 5, b: 5 });', {
+			toolRules: {
+				allowOnlyApiGroups: ['text'],
+			},
+		});
+
+		expect(result.status).toBe('failed');
+	});
+
+	describe('explore_api filtering', () => {
+		it('should show all API groups at root when no toolRules', async () => {
+			const result = await client.exploreAPI('/');
+
+			// Should see 'custom' directory at root
+			expect(result).toHaveProperty('items');
+			expect(Array.isArray((result as any).items)).toBe(true);
+		});
+
+		it('should show both math and text when exploring /custom without toolRules', async () => {
+			const result = await client.exploreAPI('/custom');
+
+			const items = (result as any).items || [];
+			const itemNames = items.map((i: any) => i.name);
+
+			// Both APIs should be visible without filtering
+			expect(itemNames).toContain('math');
+			expect(itemNames).toContain('text');
+		});
+
+		it('should filter explore results when runInRequestScope is used with toolRules', async () => {
+			// Test that the filtering mechanism works when request scope is set
+			const result = await runInRequestScope(
+				{ toolRules: { allowOnlyApiGroups: ['math'] } },
+				async () => {
+					return await client.exploreAPI('/custom');
+				}
+			);
+
+			const items = (result as any).items || [];
+			const itemNames = items.map((i: any) => i.name);
+
+			// Only math should be visible when toolRules filters to math only
+			expect(itemNames).toContain('math');
+			expect(itemNames).not.toContain('text');
+		});
+
+		it('should filter explore results when toolRules passed directly to exploreAPI', async () => {
+			const result = await client.exploreAPI('/custom', {
+				toolRules: { allowOnlyApiGroups: ['math'] },
+			});
+
+			const items = (result as any).items || [];
+			const itemNames = items.map((i: any) => i.name);
+
+			// Only math should be visible when toolRules filters to math only
+			expect(itemNames).toContain('math');
+			expect(itemNames).not.toContain('text');
+		});
+	});
+
+	describe('search_api filtering', () => {
+		it('should return results from all APIs when no options', async () => {
+			const results = await client.searchAPI('add');
+
+			expect(Array.isArray(results)).toBe(true);
+			expect(results.length).toBeGreaterThan(0);
+		});
+
+		it('should find text API uppercase when no options', async () => {
+			const results = await client.searchAPI('uppercase');
+
+			expect(results.length).toBeGreaterThan(0);
+			expect(results.some((r) => r.functionName === 'uppercase')).toBe(true);
+		});
+
+		it('should only return math results when apiGroups filter includes only math', async () => {
+			// searchAPI accepts apiGroups option for filtering
+			const results = await client.searchAPI('uppercase', {
+				query: 'uppercase',
+				apiGroups: ['math'],
+			});
+
+			// When filtering by apiGroups, searching for "uppercase" with math-only filter
+			// should return no results because uppercase belongs to text API
+			expect(results.length).toBe(0);
+		});
+
+		it('should filter search results when toolRules passed directly to searchAPI', async () => {
+			// searchAPI with toolRules.allowOnlyApiGroups
+			const results = await client.searchAPI('uppercase', {
+				query: 'uppercase',
+				toolRules: { allowOnlyApiGroups: ['math'] },
+			});
+
+			// When filtering by toolRules, searching for "uppercase" with math-only filter
+			// should return no results because uppercase belongs to text API
+			expect(results.length).toBe(0);
+		});
+	});
+});

packages/client/src/client.ts

@@ -7,6 +7,7 @@ import type {
 	ClientToolDefinition,
 	ExploreResult,
 	ATPEvent,
+	ApiGroupRules,
 } from '@mondaydotcomorg/atp-protocol';
 import type { RuntimeAPIName } from '@mondaydotcomorg/atp-runtime';
 import { CallbackType } from '@mondaydotcomorg/atp-protocol';
@@ -246,8 +247,8 @@ export class AgentToolProtocolClient {
 	/**
 	 * Explores the API filesystem at the given path.
 	 */
-	async exploreAPI(path: string): Promise<ExploreResult> {
-		return await this.apiOps.exploreAPI(path);
+	async exploreAPI(path: string, options?: { toolRules?: ApiGroupRules }): Promise<ExploreResult> {
+		return await this.apiOps.exploreAPI(path, options);
 	}
 
 	/**

packages/client/src/core/api-operations.ts

@@ -1,4 +1,4 @@
-import type { SearchOptions, SearchResult, ExploreResult } from '@mondaydotcomorg/atp-protocol';
+import type { SearchOptions, SearchResult, ExploreResult, ApiGroupRules } from '@mondaydotcomorg/atp-protocol';
 import type { RuntimeAPIName } from '@mondaydotcomorg/atp-runtime';
 import type { ISession } from './session.js';
 import type { InProcessSession } from './in-process-session.js';
@@ -106,15 +106,15 @@ export class APIOperations {
 	/**
 	 * Explores the API filesystem at the given path.
 	 */
-	async exploreAPI(path: string): Promise<ExploreResult> {
+	async exploreAPI(path: string, options?: { toolRules?: ApiGroupRules }): Promise<ExploreResult> {
 		await this.session.ensureInitialized();
 
 		if (this.inProcessSession) {
-			return (await this.inProcessSession.explore(path)) as ExploreResult;
+			return (await this.inProcessSession.explore(path, options)) as ExploreResult;
 		}
 
 		const url = `${this.session.getBaseUrl()}/api/explore`;
-		const body = JSON.stringify({ path });
+		const body = JSON.stringify({ path, ...options });
 		const headers = await this.session.prepareHeaders('POST', url, body);
 
 		const response = await fetch(url, {

packages/client/src/core/in-process-session.ts

@@ -202,13 +202,13 @@ export class InProcessSession implements ISession {
 		return (await this.server.handleSearch(ctx)) as { results: unknown[] };
 	}
 
-	async explore(path: string): Promise<unknown> {
+	async explore(path: string, options?: Record<string, unknown>): Promise<unknown> {
 		await this.ensureInitialized();
 
 		const ctx = this.createContext({
 			method: 'POST',
 			path: '/api/explore',
-			body: { path },
+			body: { path, ...options },
 		});
 
 		return await this.server.handleExplore(ctx);

packages/protocol/src/types.ts

@@ -198,6 +198,8 @@ export interface ClientToolRules {
 	allowOnlyApiGroups?: string[];
 }
 
+export type ApiGroupRules = Pick<ClientToolRules, 'allowOnlyApiGroups' | 'blockApiGroups'>;
+
 /**
  * Tool/API metadata for security and risk management
  *
@@ -422,6 +424,7 @@ export interface SearchOptions {
 	maxResults?: number;
 	useEmbeddings?: boolean;
 	embeddingModel?: string;
+	toolRules?: ClientToolRules;
 }
 
 export interface SearchResult {
@@ -435,6 +438,7 @@ export interface SearchResult {
 
 export interface ExploreRequest {
 	path: string;
+	toolRules?: ClientToolRules;
 }
 
 export interface ExploreDirectoryResult {

packages/server/src/controllers/stream.controller.ts

@@ -93,6 +93,7 @@ export async function handleExecuteStream(
 			provenanceMode: request.config?.provenanceMode,
 			securityPolicies: request.config?.securityPolicies,
 			provenanceHints: request.config?.provenanceHints,
+			toolRules: request.config?.toolRules,
 		};
 
 		logger.info('Validating code for streaming execution', {

packages/server/src/handlers/execute.handler.ts

@@ -128,6 +128,7 @@ export async function handleExecute(
 		},
 		onToolCall,
 		eventCallback: requestConfig.eventCallback,
+		toolRules: requestConfig.toolRules,
 	};
 
 	// Verify provenance hints if provided

packages/server/src/handlers/explorer.handler.ts

@@ -1,18 +1,29 @@
 import type { RequestContext } from '../core/config.js';
 import type { ExplorerService } from '../explorer/index.js';
+import type { ApiGroupRules } from '@mondaydotcomorg/atp-protocol';
+import { runInRequestScope, getRequestScope } from '../core/request-scope.js';
 
 export async function handleExplore(
 	ctx: RequestContext,
 	explorerService: ExplorerService
 ): Promise<unknown> {
-	const body = ctx.body as { path?: string };
+	const body = ctx.body as { path?: string; toolRules?: ApiGroupRules };
 	const path = body.path || '/';
+	const { toolRules } = body;
 
-	const result = explorerService.explore(path);
+	const executeExplore = () => {
+		const result = explorerService.explore(path);
 
-	if (!result) {
-		ctx.throw(404, `Path not found: ${path}`);
+		if (!result) {
+			ctx.throw(404, `Path not found: ${path}`);
+		}
+
+		return result;
+	};
+
+	if (toolRules && !getRequestScope()?.toolRules) {
+		return runInRequestScope({ toolRules }, executeExplore);
 	}
 
-	return result;
+	return executeExplore();
 }

packages/server/src/handlers/search.handler.ts

@@ -1,19 +1,30 @@
 import type { RequestContext, ResolvedServerConfig } from '../core/config.js';
 import type { SearchEngine } from '../search/index.js';
+import { runInRequestScope, getRequestScope } from '../core/request-scope.js';
 
 export async function handleSearch(
 	ctx: RequestContext,
 	searchEngine: SearchEngine,
 	config: ResolvedServerConfig
 ): Promise<unknown> {
 	const searchOptions = ctx.body as any;
-	const results = await searchEngine.search(
-		searchOptions,
-		ctx.userId,
-		ctx.auth,
-		config.discovery.scopeFiltering
-	);
-	return { results };
+	const { toolRules, ...cleanOptions } = searchOptions;
+
+	const executeSearch = async () => {
+		const results = await searchEngine.search(
+			cleanOptions,
+			ctx.userId,
+			ctx.auth,
+			config.discovery.scopeFiltering
+		);
+		return { results };
+	};
+
+	if (toolRules && !getRequestScope()?.toolRules) {
+		return runInRequestScope({ toolRules }, executeSearch);
+	}
+
+	return executeSearch();
 }
 
 export async function handleSearchQuery(