fix: add methods to fetch builtIn and builtInWrite tools (#581)

Author: jguoamz

runtimes/runtimes/agent.test.ts

@@ -1,4 +1,4 @@
-import { Agent } from '../server-interface'
+import { Agent, ToolClassification } from '../server-interface'
 import { newAgent } from './agent'
 import * as assert from 'assert'
 
@@ -178,4 +178,30 @@ describe('Agent Tools', () => {
         AGENT.removeTool(SOME_TOOL_SPEC.name)
         await assert.rejects(() => AGENT.runTool(SOME_TOOL_SPEC.name, { test: 'test' }), /not found/)
     })
+
+    it('should track built-in tools', () => {
+        AGENT.addTool(SOME_TOOL_SPEC, SOME_TOOL_HANDLER, ToolClassification.BuiltIn)
+        AGENT.addTool({ ...SOME_TOOL_SPEC, name: 'regular' }, SOME_TOOL_HANDLER)
+
+        const builtInTools = AGENT.getBuiltInToolNames()
+        assert.equal(builtInTools.length, 1)
+        assert.equal(builtInTools[0], SOME_TOOL_SPEC.name)
+    })
+
+    it('should track built-in write tools', () => {
+        AGENT.addTool(SOME_TOOL_SPEC, SOME_TOOL_HANDLER, ToolClassification.BuiltInCanWrite)
+        AGENT.addTool({ ...SOME_TOOL_SPEC, name: 'readOnly' }, SOME_TOOL_HANDLER, ToolClassification.BuiltIn)
+
+        const builtInTools = AGENT.getBuiltInToolNames()
+        const builtInWriteTools = AGENT.getBuiltInWriteToolNames()
+
+        assert.equal(builtInTools.length, 2)
+        assert.equal(builtInWriteTools.length, 1)
+        assert.equal(builtInWriteTools[0], SOME_TOOL_SPEC.name)
+    })
+
+    it('should return empty arrays for built-in tools when none are added', () => {
+        assert.equal(AGENT.getBuiltInToolNames().length, 0)
+        assert.equal(AGENT.getBuiltInWriteToolNames().length, 0)
+    })
 })

runtimes/runtimes/agent.ts

@@ -6,6 +6,7 @@ import {
     GetToolsOptions,
     InferSchema,
     ObjectSchema,
+    ToolClassification,
     Tools,
     ToolSpec,
 } from '../server-interface'
@@ -21,11 +22,14 @@ type Tool<T, R> = {
 export const newAgent = (): Agent => {
     const tools: Record<string, Tool<any, any>> = {}
     const ajv = new Ajv({ strictSchema: false })
+    const builtInToolNames: string[] = []
+    const builtInWriteToolNames: string[] = []
 
     return {
         addTool: <T extends InferSchema<S['inputSchema']>, S extends ToolSpec>(
             spec: S,
-            handler: (input: T, token?: CancellationToken) => Promise<any>
+            handler: (input: T, token?: CancellationToken) => Promise<any>,
+            toolClassification?: ToolClassification
         ) => {
             const validator = ajv.compile(spec.inputSchema)
             const tool = {
@@ -40,6 +44,15 @@ export const newAgent = (): Agent => {
             }
 
             tools[spec.name] = tool
+            if (
+                toolClassification === ToolClassification.BuiltIn ||
+                toolClassification === ToolClassification.BuiltInCanWrite
+            ) {
+                builtInToolNames.push(spec.name)
+                if (toolClassification === ToolClassification.BuiltInCanWrite) {
+                    builtInWriteToolNames.push(spec.name)
+                }
+            }
         },
 
         runTool: async (toolName: string, input: any, token?: CancellationToken, updates?: WritableStream) => {
@@ -91,5 +104,13 @@ export const newAgent = (): Agent => {
         removeTool: (name: string) => {
             delete tools[name]
         },
+
+        getBuiltInToolNames: () => {
+            return builtInToolNames
+        },
+
+        getBuiltInWriteToolNames: () => {
+            return builtInWriteToolNames
+        },
     }
 }

runtimes/server-interface/agent.ts

@@ -112,6 +112,12 @@ export type BedrockTools = {
     toolSpecification: Omit<ToolSpec, 'inputSchema'> & { inputSchema: { json: ToolSpec['inputSchema'] } }
 }[]
 
+export enum ToolClassification {
+    BuiltIn = 'builtIn',
+    BuiltInCanWrite = 'builtInCanWrite',
+    MCP = 'mcp',
+}
+
 export type Agent = {
     /**
      * Add a tool to the local tool repository. Tools with the same name will be overwritten.
@@ -123,7 +129,8 @@ export type Agent = {
      */
     addTool: <T extends InferSchema<S['inputSchema']>, S extends ToolSpec, R>(
         spec: S,
-        handler: (input: T, token?: CancellationToken, updates?: WritableStream) => Promise<R>
+        handler: (input: T, token?: CancellationToken, updates?: WritableStream) => Promise<R>,
+        toolClassification?: ToolClassification
     ) => void
 
     /**
@@ -150,4 +157,16 @@ export type Agent = {
      * @param toolName The name of the tool to remove
      */
     removeTool: (toolName: string) => void
+
+    /**
+     * Get the list of built-in tool names in the local tool repository.
+     * @returns The list of built-in tool names
+     */
+    getBuiltInToolNames: () => string[]
+
+    /**
+     * Get the list of built-in write tool names in the local tool repository.
+     * @returns The list of built-in write tool names
+     */
+    getBuiltInWriteToolNames: () => string[]
 }