History

feat(execution): allow custom tool discovery providers

Saatvik Arya committed 5 months ago

Commit c598d641 parent eefb8b55 files · +156 −16Browse files

Unified · Split

5 files
  1. packages/core/execution/README.md +19 −0
  2. packages/core/execution/src/engine.ts +34 −15
  3. packages/core/execution/src/index.ts +5 −0
  4. packages/core/execution/src/tool-invoker.test.ts +79 −1
  5. packages/core/execution/src/tool-invoker.ts +19 −0
packages/core/execution/README.md+19−0
@@ -51,6 +51,25 @@ console.log(result); · 50 unmodified lines
51// { result: 12, logs: [...] }51// { result: 12, logs: [...] }
52```52```
53 53
54## Custom tool discovery
55
56`tools.search(...)` uses Executor's built-in lexical tool discovery by default. Hosts can provide their own implementation, such as an indexed or semantic search provider, without replacing the sandbox runtime:
57
58```ts
59import { createExecutionEngine, type ToolDiscoveryProvider } from "@executor-js/execution";
60
61const toolDiscoveryProvider: ToolDiscoveryProvider = {
62 searchTools: ({ query, namespace, limit, offset }) =>
63 mySearchIndex.searchTools({ query, namespace, limit, offset }),
64};
65
66const engine = createExecutionEngine({
67 executor,
68 codeExecutor: makeQuickJsExecutor(),
69 toolDiscoveryProvider,
70});
71```
72
54## Pause/resume for elicitation73## Pause/resume for elicitation
55 74
56When the host doesn't support inline elicitation, use `executeWithPause` to intercept the first request as a pause point:75When the host doesn't support inline elicitation, use `executeWithPause` to intercept the first request as a pause point:
packages/core/execution/src/engine.ts+34−15
@@ -12,10 +12,11 @@ import { CodeExecutionError } from "@executor-js/codemode-core"; · 11 unmodified lines
12import type { CodeExecutor, ExecuteResult, SandboxToolInvoker } from "@executor-js/codemode-core";12import type { CodeExecutor, ExecuteResult, SandboxToolInvoker } from "@executor-js/codemode-core";
13 13
14import {14import {
15 defaultToolDiscoveryProvider,
15 makeExecutorToolInvoker,16 makeExecutorToolInvoker,
16 searchTools,
17 listExecutorSources,17 listExecutorSources,
18 describeTool,18 describeTool,
19 type ToolDiscoveryProvider,
19} from "./tool-invoker";20} from "./tool-invoker";
20import { ExecutionToolError } from "./errors";21import { ExecutionToolError } from "./errors";
21import { buildExecuteDescription } from "./description";22import { buildExecuteDescription } from "./description";
@@ -27,6 +28,7 @@ import { buildExecuteDescription } from "./description"; · 5 unmodified lines
27export type ExecutionEngineConfig<E extends Cause.YieldableError = CodeExecutionError> = {28export type ExecutionEngineConfig<E extends Cause.YieldableError = CodeExecutionError> = {
28 readonly executor: Executor;29 readonly executor: Executor;
29 readonly codeExecutor: CodeExecutor<E>;30 readonly codeExecutor: CodeExecutor<E>;
31 readonly toolDiscoveryProvider?: ToolDiscoveryProvider;
30};32};
31 33
32export type ExecutionResult =34export type ExecutionResult =
@@ -186,7 +188,11 @@ const readOptionalOffset = (value: unknown, toolName: string): number | Executio · 153 unmodified lines
186 return Math.floor(value);188 return Math.floor(value);
187};189};
188 190
189const makeFullInvoker = (executor: Executor, invokeOptions: InvokeOptions): SandboxToolInvoker => {191const makeFullInvoker = (
192 executor: Executor,
193 invokeOptions: InvokeOptions,
194 toolDiscoveryProvider: ToolDiscoveryProvider,
195): SandboxToolInvoker => {
190 const base = makeExecutorToolInvoker(executor, { invokeOptions });196 const base = makeExecutorToolInvoker(executor, { invokeOptions });
191 return {197 return {
192 invoke: ({ path, args }) => {198 invoke: ({ path, args }) => {
@@ -226,14 +232,19 @@ const makeFullInvoker = (executor: Executor, invokeOptions: InvokeOptions): Sand · 33 unmodified lines
226 return Effect.fail(offset);232 return Effect.fail(offset);
227 }233 }
228 234
229 return searchTools(executor, args.query ?? "", limit, {235 return toolDiscoveryProvider
230 namespace: args.namespace,236 .searchTools({
231 offset,237 executor,
232 }).pipe(238 query: args.query ?? "",
233 Effect.withSpan("mcp.tool.dispatch", {239 limit,
234 attributes: { "mcp.tool.name": path, "executor.tool.builtin": true },240 namespace: args.namespace,
235 }),241 offset,
236 );242 })
243 .pipe(
244 Effect.withSpan("mcp.tool.dispatch", {
245 attributes: { "mcp.tool.name": path, "executor.tool.builtin": true },
246 }),
247 );
237 }248 }
238 if (path === "executor.sources.list") {249 if (path === "executor.sources.list") {
239 if (args !== undefined && !isRecord(args)) {250 if (args !== undefined && !isRecord(args)) {
@@ -364,7 +375,7 @@ export type ExecutionEngine<E extends Cause.YieldableError = CodeExecutionError> · 124 unmodified lines
364export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecutionError>(375export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecutionError>(
365 config: ExecutionEngineConfig<E>,376 config: ExecutionEngineConfig<E>,
366): ExecutionEngine<E> => {377): ExecutionEngine<E> => {
367 const { executor, codeExecutor } = config;378 const { executor, codeExecutor, toolDiscoveryProvider = defaultToolDiscoveryProvider } = config;
368 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();379 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();
369 let nextId = 0;380 let nextId = 0;
370 381
@@ -433,7 +444,11 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 62 unmodified lines
433 return yield* Deferred.await(responseDeferred);444 return yield* Deferred.await(responseDeferred);
434 });445 });
435 446
436 const invoker = makeFullInvoker(executor, { onElicitation: elicitationHandler });447 const invoker = makeFullInvoker(
448 executor,
449 { onElicitation: elicitationHandler },
450 toolDiscoveryProvider,
451 );
437 fiber = yield* Effect.forkDetach(452 fiber = yield* Effect.forkDetach(
438 codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec")),453 codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec")),
439 );454 );
@@ -477,9 +492,13 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 37 unmodified lines
477 "mcp.execute.mode": "inline",492 "mcp.execute.mode": "inline",
478 "mcp.execute.code_length": code.length,493 "mcp.execute.code_length": code.length,
479 });494 });
480 const invoker = makeFullInvoker(executor, {495 const invoker = makeFullInvoker(
481 onElicitation: options.onElicitation,496 executor,
482 });497 {
498 onElicitation: options.onElicitation,
499 },
500 toolDiscoveryProvider,
501 );
483 return yield* codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec"));502 return yield* codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec"));
484 });503 });
485 504
packages/core/execution/src/index.ts+5−0
@@ -12,8 +12,13 @@ export { · 11 unmodified lines
12export { buildExecuteDescription } from "./description";12export { buildExecuteDescription } from "./description";
13export { ExecutionToolError } from "./errors";13export { ExecutionToolError } from "./errors";
14export {14export {
15 defaultToolDiscoveryProvider,
15 makeExecutorToolInvoker,16 makeExecutorToolInvoker,
16 searchTools,17 searchTools,
17 listExecutorSources,18 listExecutorSources,
18 describeTool,19 describeTool,
20 type ToolDiscoveryInput,
21 type ToolDiscoveryProvider,
22 type PagedResult,
23 type ToolDiscoveryResult,
19} from "./tool-invoker";24} from "./tool-invoker";
packages/core/execution/src/tool-invoker.test.ts+79−1
@@ -12,7 +12,12 @@ import { · 11 unmodified lines
12import { makeTestConfig, typeCheckOutputTypeScript } from "@executor-js/sdk/testing";12import { makeTestConfig, typeCheckOutputTypeScript } from "@executor-js/sdk/testing";
13import { makeQuickJsExecutor } from "@executor-js/runtime-quickjs";13import { makeQuickJsExecutor } from "@executor-js/runtime-quickjs";
14import { createExecutionEngine } from "./engine";14import { createExecutionEngine } from "./engine";
15import { describeTool, makeExecutorToolInvoker, searchTools } from "./tool-invoker";15import {
16 describeTool,
17 makeExecutorToolInvoker,
18 searchTools,
19 type ToolDiscoveryProvider,
20} from "./tool-invoker";
16 21
17const codeExecutor = makeQuickJsExecutor();22const codeExecutor = makeQuickJsExecutor();
18 23
@@ -334,6 +339,79 @@ describe("tool discovery", () => { · 315 unmodified lines
334 }),339 }),
335 );340 );
336 341
342 it.effect("lets execution hosts provide custom tool discovery", () =>
343 Effect.gen(function* () {
344 const executor = yield* makeSearchExecutor();
345 const calls: Array<{
346 readonly query: string;
347 readonly namespace?: string;
348 readonly limit: number;
349 readonly offset: number;
350 }> = [];
351 const provider: ToolDiscoveryProvider = {
352 searchTools: ({ query, namespace, limit, offset }) =>
353 Effect.sync(() => {
354 calls.push({ query, namespace, limit, offset });
355 return {
356 items: [
357 {
358 path: "custom.searchResult",
359 name: "searchResult",
360 description: "Provided by the host",
361 sourceId: "custom",
362 score: 999,
363 },
364 ],
365 total: 1,
366 hasMore: false,
367 nextOffset: null,
368 };
369 }),
370 };
371 const engine = createExecutionEngine({
372 executor,
373 codeExecutor,
374 toolDiscoveryProvider: provider,
375 });
376
377 const result = yield* engine.execute(
378 [
379 "return await tools.search({",
380 ' query: "calendar events",',
381 ' namespace: "calendar",',
382 " limit: 7,",
383 " offset: 2,",
384 "});",
385 ].join("\n"),
386 { onElicitation: acceptAll },
387 );
388
389 expect(result.error).toBeUndefined();
390 expect(result.result).toEqual({
391 items: [
392 {
393 path: "custom.searchResult",
394 name: "searchResult",
395 description: "Provided by the host",
396 sourceId: "custom",
397 score: 999,
398 },
399 ],
400 total: 1,
401 hasMore: false,
402 nextOffset: null,
403 });
404 expect(calls).toEqual([
405 {
406 query: "calendar events",
407 namespace: "calendar",
408 limit: 7,
409 offset: 2,
410 },
411 ]);
412 }),
413 );
414
337 it.effect("supports executor-scoped source listing and tool search", () =>415 it.effect("supports executor-scoped source listing and tool search", () =>
338 Effect.gen(function* () {416 Effect.gen(function* () {
339 const executor = yield* makeSearchExecutor();417 const executor = yield* makeSearchExecutor();
packages/core/execution/src/tool-invoker.ts+19−0
@@ -249,6 +249,20 @@ export type ExecutorSourceListItem = { · 248 unmodified lines
249 readonly toolCount: number;249 readonly toolCount: number;
250};250};
251 251
252export type ToolDiscoveryInput = {
253 readonly executor: Executor;
254 readonly query: string;
255 readonly namespace?: string;
256 readonly limit: number;
257 readonly offset: number;
258};
259
260export interface ToolDiscoveryProvider {
261 readonly searchTools: (
262 input: ToolDiscoveryInput,
263 ) => Effect.Effect<PagedResult<ToolDiscoveryResult>, ExecutionToolError>;
264}
265
252/**266/**
253 * Page of results from a list-style discovery tool. Shared by267 * Page of results from a list-style discovery tool. Shared by
254 * `tools.search` and `tools.executor.sources.list` so the model sees one268 * `tools.search` and `tools.executor.sources.list` so the model sees one
@@ -515,6 +529,11 @@ export const searchTools = Effect.fn("executor.tools.search")(function* ( · 260 unmodified lines
515 return page;529 return page;
516});530});
517 531
532export const defaultToolDiscoveryProvider: ToolDiscoveryProvider = {
533 searchTools: ({ executor, query, namespace, limit, offset }) =>
534 searchTools(executor, query, limit, { namespace, offset }),
535};
536
518/** What `tools.executor.sources.list()` calls inside the sandbox. */537/** What `tools.executor.sources.list()` calls inside the sandbox. */
519export const listExecutorSources = Effect.fn("executor.sources.list")(function* (538export const listExecutorSources = Effect.fn("executor.sources.list")(function* (
520 executor: Executor,539 executor: Executor,

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t