1import * as Context from "effect/Context";
2import * as Effect from "effect/Effect";
3
4export const TOOL_SEARCH = "tool_search";
5
6/** What a deferring client lists before it searches. Pi's extension declares the same three. */
7const DIRECT_TOOLS: ReadonlySet<string> = new Set([
8 "orchestrator_capabilities",
9 "delegate_task",
10 "task_status",
11 TOOL_SEARCH,
12]);
13
14const MAX_RESULTS = 10;
15// Sessions that ended are never reported here, so keep only the newest.
16const MAX_SESSIONS = 256;
17
18export interface ListedTool {
19 readonly name: string;
20 readonly description: string;
21}
22
23export class McpToolExposure extends Context.Service<
24 McpToolExposure,
25 {
26 /** Finds tools for the query and lists them to the session from now on. */
27 readonly search: (
28 providerSessionId: string,
29 query: string,
30 ) => Effect.Effect<ReadonlyArray<ListedTool>>;
31 readonly isListed: (providerSessionId: string, tool: string) => boolean;
32 }
33>()("t3/mcp/McpToolExposure") {}
34
35const score = (tool: ListedTool, terms: ReadonlyArray<string>) => {
36 const name = tool.name.toLowerCase();
37 const description = tool.description.toLowerCase();
38 let total = 0;
39 for (const term of terms) {
40 if (term === name) total += 100;
41 else if (name.includes(term)) total += 3;
42 else if (description.includes(term)) total += 1;
43 }
44 return total;
45};
46
47export const make = (
48 tools: () => ReadonlyArray<ListedTool>,
49 announceListChanged: () => Effect.Effect<void>,
50) => {
51 const loaded = new Map<string, Set<string>>();
52 const load = (providerSessionId: string, names: ReadonlyArray<string>) => {
53 const listed = loaded.get(providerSessionId) ?? new Set<string>();
54 loaded.delete(providerSessionId);
55 loaded.set(providerSessionId, listed);
56 for (const name of names) listed.add(name);
57 for (const oldest of loaded.keys()) {
58 if (loaded.size <= MAX_SESSIONS) break;
59 loaded.delete(oldest);
60 }
61 };
62 return McpToolExposure.of({
63 search: Effect.fn("McpToolExposure.search")(function* (providerSessionId, query) {
64 const terms = query
65 .toLowerCase()
66 .split(/[^a-z0-9_]+/u)
67 .filter(Boolean);
68 const found = tools()
69 .filter((tool) => tool.name !== TOOL_SEARCH)
70 .map((tool) => ({ tool, score: score(tool, terms) }))
71 .filter((match) => match.score > 0)
72 .sort((a, b) => b.score - a.score)
73 .slice(0, MAX_RESULTS)
74 .map((match) => match.tool);
75 if (found.length > 0) {
76 load(
77 providerSessionId,
78 found.map((tool) => tool.name),
79 );
80 yield* announceListChanged();
81 }
82 return found;
83 }),
84 isListed: (providerSessionId, tool) =>
85 DIRECT_TOOLS.has(tool) || loaded.get(providerSessionId)?.has(tool) === true,
86 });
87};