feat(ai): stub non-core tool schemas + tool_search to reduce context window usage by 2witstudios · Pull Request #1301 · 2witstudios/PageSpace

2witstudios · 2026-05-11T20:22:06Z

Summary

Sending all 37 full Zod schemas to Claude every request costs ~7–18k tokens per turn in tool overhead
Non-core tools (29 of 37) now use `z.object({}).passthrough()` stub schemas; the AI calls `tool_search` to fetch full parameter schemas on demand before using them — same pattern as Claude Code's own deferred tools / ToolSearch
Core 8 tools (`list_drives`, `list_pages`, `read_page`, `get_page_details`, `create_page`, `replace_lines`, `regex_search`, `multi_drive_search`) keep full schemas since they're used in almost every session
`tool_search` is scoped to the enabled tool set: in read-only mode write tools are absent; when web search is off `web_search` is absent — so the AI can never discover unavailable tools
Estimated saving: ~7,000–16,000 tokens per request

Changes

New files:

`stub-tools.ts` — wraps non-core tools with passthrough stubs; real `safeParse` validation inside `execute` returns a helpful `tool_search` error on invalid args
`tool-search-tool.ts` — introspects the live registry via `z.toJSONSchema`; supports `select:name1,name2` and keyword queries; accepts a `ToolSet` so it can be scoped at call time
`ai-evals/tool-discovery/` — `riteway ai` prompt evals for behavioral testing (run with `riteway ai 'apps/web/ai-evals/tool-discovery/*.sudo'`)

Modified files:

`chat/route.ts` — uses `pageSpaceToolsStubbed` + builds `enabledFullSchemaTools` (intersection of `filteredTools` keys and full-schema `pageSpaceTools`) + injects `tool_search` from that scoped set
`system-prompt.ts` — adds `TOOL_DISCOVERY_PROMPT` section teaching the AI which tools are core and how to use `tool_search`
`ai-tools.ts` — exports `pageSpaceToolsStubbed`

Test plan

23 unit tests passing: `stub-tools.test.ts` (7), `tool-search-tool.test.ts` (8), `ai-tools.test.ts` (4 new + 4 existing)
Fixed mocks in `stream-socket-events.test.ts` and `mcp-scope.test.ts` (both tests mock `@/lib/ai/core` — added `pageSpaceToolsStubbed: {}` and mocked `tool-search-tool` module)
Run prompt evals: `riteway ai 'apps/web/ai-evals/tool-discovery/*.sudo'`
Smoke test AI chat in the UI — verify calendar/task/agent tools still work (AI should call `tool_search` first, then the tool with correct params)

🤖 Generated with Claude Code

…ontext window usage Sending all 37 full Zod schemas every request cost ~7–18k tokens per turn. Non-core tools now use passthrough stub schemas; the AI calls tool_search("keyword") to get full parameter schemas on demand before using extended tools (same pattern as Claude Code's own ToolSearch). - New stub-tools.ts: wraps non-core tools with z.object({}).passthrough() inputSchema; real safeParse validation runs inside execute so invalid args return a helpful error - New tool-search-tool.ts: introspects the live registry via z.toJSONSchema; supports select:name1,name2 and keyword queries - chat/route.ts: uses pageSpaceToolsStubbed + injects tool_search referencing full registry - system-prompt.ts: adds TOOL_DISCOVERY_PROMPT teaching the 8 core tools and tool_search usage - Tests: 23 unit tests across stub-tools, tool-search-tool, and ai-tools - Prompt evals: riteway ai .sudo files in ai-evals/tool-discovery/ for behavioral testing Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

coderabbitai · 2026-05-11T20:22:20Z

📝 Walkthrough

Walkthrough

Implements a tool discovery system where AI models use core tools directly and dynamically discover non-core tool schemas via tool_search. Core tools (list_drives, read_page, etc.) have unmodified schemas; non-core tools are stubbed to accept any input, validate requests, and redirect to tool_search on mismatch. The chat route injects tool_search backed by the full tool registry alongside the stubbed tools.

Changes

Tool Discovery & Schema Resolution

Layer / File(s)	Summary
Tool Discovery Spec & AI Evaluation `apps/web/ai-evals/tool-discovery/spec.mdc`, `apps/web/ai-evals/tool-discovery/core-tools-no-search.sudo`, `apps/web/ai-evals/tool-discovery/non-core-requires-search.sudo`	Specification defines core tools (list_drives, read_page, etc.) that use no tool_search, and extended tools (calendar, tasks, agents) that require tool_search first. Evaluation prompts demonstrate expected behavior for workspace listing and calendar event creation scenarios.
Tool Stubbing Mechanism `apps/web/src/lib/ai/core/stub-tools.ts`	`CORE_TOOL_NAMES` set identifies eight core tool names. `wrapWithStub` replaces non-core tool schemas with permissive objects and validates arguments against the original schema; returns error with tool_search reference on mismatch. `buildStubbedTools` applies stubs to all non-core tools in a registry.
Tool Search Implementation `apps/web/src/lib/ai/tools/tool-search-tool.ts`	`createToolSearchTool` factory creates a tool that queries a full tool registry. Supports `select:name1,name2` syntax and case-insensitive keyword matching against tool names and descriptions. Returns matched tools with name, description, and JSON-schema representation of their input parameters.
Stubbed Tools Export `apps/web/src/lib/ai/core/ai-tools.ts`	Imports `buildStubbedTools` and exports `pageSpaceToolsStubbed` derived from `pageSpaceTools` with non-core tools replaced by stubs.
System Prompt Tool Discovery Guidance `apps/web/src/lib/ai/core/system-prompt.ts`	Adds `TOOL_DISCOVERY_PROMPT` constant instructing the model to call `tool_search(...)` to determine parameter schemas for non-core tools before invoking them. Integrated into `buildSystemPrompt` output.
Chat Route Tool Injection `apps/web/src/app/api/ai/chat/route.ts`	Builds initial tool set from `pageSpaceToolsStubbed` with feature toggles applied, then injects `createToolSearchTool(pageSpaceTools)` to provide dynamic access to full tool schemas for non-core tools.
Stub Tools Tests `apps/web/src/lib/ai/core/__tests__/stub-tools.test.ts`	Validates `CORE_TOOL_NAMES` content, core tool passthrough by reference identity, non-core stub schema replacement with permissive object, stub `execute` delegation on valid arguments, error message mentioning tool_search on validation failure, and empty registry handling.
Tool Search Tests `apps/web/src/lib/ai/tools/__tests__/tool-search-tool.test.ts`	Verifies tool-search tool description, `select:` query returning exact matches, unknown selects returning empty results, keyword matching against tool names and descriptions, JSON schema `type: "object"` conversion, and empty keyword results.
AI Tools Integration Tests `apps/web/src/lib/ai/core/__tests__/ai-tools.test.ts`	Validates `pageSpaceToolsStubbed` export, keyset matching with `pageSpaceTools`, core tool identity preservation, and non-core tool stub object replacement.

Sequence Diagram

sequenceDiagram
  participant Model as AI Model
  participant StubTools as Stubbed Tool Set
  participant ToolSearch as tool_search
  participant FullTools as Full Tool Registry
  participant CoreTool as Core Tool (e.g., list_drives)
  participant NonCoreTool as Non-core Tool (e.g., create_calendar_event)
  
  Model->>StubTools: Call list_drives with valid args
  StubTools->>CoreTool: Execute directly (no stub)
  CoreTool-->>Model: Return workspace list
  
  Model->>StubTools: Call create_calendar_event with incomplete args
  StubTools->>StubTools: Validate args against original schema
  StubTools-->>Model: Error: use tool_search first
  
  Model->>ToolSearch: Query tool_search("calendar")
  ToolSearch->>FullTools: Look up matching tools
  FullTools-->>ToolSearch: Return calendar tool with full schema
  ToolSearch-->>Model: Return create_calendar_event schema
  
  Model->>StubTools: Call create_calendar_event with valid args from schema
  StubTools->>NonCoreTool: Execute with validated args
  NonCoreTool-->>Model: Return event confirmation

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

2witstudios/PageSpace#1016: Modifies apps/web/src/lib/ai/core/ai-tools.ts to update tool registry composition; overlaps with this PR's export of pageSpaceToolsStubbed from the same file.
2witstudios/PageSpace#829: Updates apps/web/src/app/api/ai/chat/route.ts to inject additional tools and modifies system-prompt behavior; shares the same API route tooling and prompt integration pattern.
2witstudios/PageSpace#776: Renames tool descriptor properties to inputSchema; this PR reads and wraps tools via that same inputSchema property.

Poem

🐰 Stumbling through schemas with search in the air,
Core tools run swift, while extended tools care,
Query the discovery, find what you need,
Validation and wisdom plant knowledge like seed!
Build it with stubs, watch the AI tools grow,
Tool search reveals secrets, let the model know. 🔍

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat(ai): stub non-core tool schemas + tool_search to reduce context window usage' clearly and specifically summarizes the main objective of the changeset: implementing tool schema stubbing and a tool discovery mechanism to optimize LLM context usage.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch pu/tool-disc

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

Generate code and open pull requests
Plan features and break down work
Investigate incidents and troubleshoot customer tickets together
Automate recurring tasks and respond to alerts with triggers
Summarize progress and report instantly

Built for teams:

Shared memory across your entire org—no repeating context
Per-thread sandboxes to safely plan and execute work
Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

apps/web/src/app/api/ai/chat/route.ts (1)

527-532: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update logging to reflect stubbed tools usage.

The debug log on line 528 references pageSpaceTools.length, but the code now uses pageSpaceToolsStubbed (line 515). While the key counts are identical, the log message should reflect the actual tool set in use for accurate debugging context.

📊 Proposed fix

-    loggers.ai.debug('AI Page Chat API: Tools built from baseline + runtime toggles', {
-      totalTools: Object.keys(pageSpaceTools).length,
+    loggers.ai.debug('AI Page Chat API: Tools built from stubbed baseline + runtime toggles', {
+      totalTools: Object.keys(pageSpaceToolsStubbed).length,
       filteredTools: Object.keys(filteredTools).length,
       isReadOnly: readOnlyMode,
       webSearchEnabled: webSearchMode

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@apps/web/src/app/api/ai/chat/route.ts` around lines 527 - 532, The debug
message in loggers.ai.debug currently reports counts using pageSpaceTools but
the runtime uses pageSpaceToolsStubbed; update the logging call
(loggers.ai.debug in this file) to use Object.keys(pageSpaceToolsStubbed).length
for totalTools so the log reflects the actual tool set in use alongside
filteredTools, readOnlyMode and webSearchMode.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@apps/web/src/app/api/ai/chat/route.ts`:
- Around line 527-532: The debug message in loggers.ai.debug currently reports
counts using pageSpaceTools but the runtime uses pageSpaceToolsStubbed; update
the logging call (loggers.ai.debug in this file) to use
Object.keys(pageSpaceToolsStubbed).length for totalTools so the log reflects the
actual tool set in use alongside filteredTools, readOnlyMode and webSearchMode.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: edd8f665-958d-401c-9db4-462a202e80fe

📥 Commits

Reviewing files that changed from the base of the PR and between 64388fa and 31a79e2.

📒 Files selected for processing (11)

apps/web/ai-evals/tool-discovery/core-tools-no-search.sudo
apps/web/ai-evals/tool-discovery/non-core-requires-search.sudo
apps/web/ai-evals/tool-discovery/spec.mdc
apps/web/src/app/api/ai/chat/route.ts
apps/web/src/lib/ai/core/__tests__/ai-tools.test.ts
apps/web/src/lib/ai/core/__tests__/stub-tools.test.ts
apps/web/src/lib/ai/core/ai-tools.ts
apps/web/src/lib/ai/core/stub-tools.ts
apps/web/src/lib/ai/core/system-prompt.ts
apps/web/src/lib/ai/tools/__tests__/tool-search-tool.test.ts
apps/web/src/lib/ai/tools/tool-search-tool.ts

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 31a79e2595

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-05-11T20:27:22Z

+    // return complete parameter schemas on demand.
+    filteredTools = {
+      ...filteredTools,
+      tool_search: createToolSearchTool(pageSpaceTools),


Restrict tool_search to currently enabled tools

tool_search is wired against pageSpaceTools (the full registry), not the runtime-filtered toolset, so in read-only mode or when web search is disabled it still advertises schemas for tools that were intentionally removed from filteredTools. In those sessions the model can discover (and then attempt) unavailable tools like web_search/write tools, which causes unavailable-tool execution errors instead of a normal response path. Build tool_search from the same filtered availability set (with full schemas for only those allowed tools) to avoid exposing tools that cannot be executed in this request.

Useful? React with 👍 / 👎.

Fixed in f57a557. tool_search now receives enabledFullSchemaTools — the intersection of filteredTools keys and pageSpaceTools full schemas. In read-only mode, write tools are absent from filteredTools so they don't appear in tool_search results. Same for web_search when web search is disabled.

…lSearchTool in chat route tests The two chat route tests mock @/lib/ai/core but didn't include pageSpaceToolsStubbed (newly exported) or mock the new @/lib/ai/tools/tool-search-tool import. This caused all lifecycle spy assertions to fail with 0 calls in CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

tool_search was wired to pageSpaceTools (full registry), meaning in read-only mode or with web search disabled it could advertise unavailable tools. Now builds enabledFullSchemaTools by intersecting the filtered key set with pageSpaceTools, so tool_search only returns schemas for tools that can actually execute in this request. Also fixes the debug log to count pageSpaceToolsStubbed instead of pageSpaceTools. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…t in route Aligns tool-search-tool.ts with the AI SDK's ToolSet type (Record<string, CoreTool>) instead of the looser Record<string, Tool>. Adds explicit as ToolSet cast on Object.fromEntries in route.ts since TypeScript cannot infer the tuple value type. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…window usage (#1301) * feat(ai): stub non-core tool schemas + tool_search meta-tool to cut context window usage Sending all 37 full Zod schemas every request cost ~7–18k tokens per turn. Non-core tools now use passthrough stub schemas; the AI calls tool_search("keyword") to get full parameter schemas on demand before using extended tools (same pattern as Claude Code's own ToolSearch). - New stub-tools.ts: wraps non-core tools with z.object({}).passthrough() inputSchema; real safeParse validation runs inside execute so invalid args return a helpful error - New tool-search-tool.ts: introspects the live registry via z.toJSONSchema; supports select:name1,name2 and keyword queries - chat/route.ts: uses pageSpaceToolsStubbed + injects tool_search referencing full registry - system-prompt.ts: adds TOOL_DISCOVERY_PROMPT teaching the 8 core tools and tool_search usage - Tests: 23 unit tests across stub-tools, tool-search-tool, and ai-tools - Prompt evals: riteway ai .sudo files in ai-evals/tool-discovery/ for behavioral testing Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(tests): add missing mocks for pageSpaceToolsStubbed and createToolSearchTool in chat route tests The two chat route tests mock @/lib/ai/core but didn't include pageSpaceToolsStubbed (newly exported) or mock the new @/lib/ai/tools/tool-search-tool import. This caused all lifecycle spy assertions to fail with 0 calls in CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ai): scope tool_search to enabled tools + fix stale log count tool_search was wired to pageSpaceTools (full registry), meaning in read-only mode or with web search disabled it could advertise unavailable tools. Now builds enabledFullSchemaTools by intersecting the filtered key set with pageSpaceTools, so tool_search only returns schemas for tools that can actually execute in this request. Also fixes the debug log to count pageSpaceToolsStubbed instead of pageSpaceTools. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(ai): use ToolSet type in createToolSearchTool + explicit cast in route Aligns tool-search-tool.ts with the AI SDK's ToolSet type (Record<string, CoreTool>) instead of the looser Record<string, Tool>. Adds explicit as ToolSet cast on Object.fromEntries in route.ts since TypeScript cannot infer the tuple value type. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

coderabbitai Bot reviewed May 11, 2026

View reviewed changes

chatgpt-codex-connector Bot reviewed May 11, 2026

View reviewed changes

2witstudios and others added 3 commits May 11, 2026 15:47

2witstudios merged commit 35b83eb into master May 11, 2026
10 checks passed

2witstudios mentioned this pull request May 12, 2026

fix(ai): move tool discovery to global assistant + enforce page agent enabledTools #1306

Merged

3 tasks done

coderabbitai Bot mentioned this pull request May 12, 2026

feat(ai): execute_tool dispatcher — drop 29 non-core tool schemas from model context #1310

Merged

5 tasks done

2witstudios deleted the pu/tool-disc branch May 14, 2026 19:21

This was referenced May 24, 2026

feat(ai): add tool_docs meta-tool with per-category usage guides #1411

Closed

feat(ai-agents): per-agent tool exposure mode (upfront vs search) #1462

Merged

coderabbitai Bot mentioned this pull request Jun 2, 2026

Agent Code Execution — bash/code in chat via Fly Sprites (epic) #1487

Merged

coderabbitai Bot mentioned this pull request Jul 28, 2026

feat(ai): unified agent skill system — discoverable, loadable, discardable capabilities #2244

Merged

Navigation Menu

Repository navigation

feat(ai): stub non-core tool schemas + tool_search to reduce context window usage - #1301

feat(ai): stub non-core tool schemas + tool_search to reduce context window usage#1301
2witstudios merged 4 commits into
masterfrom
pu/tool-disc

2witstudios commented May 11, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented May 11, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

Possibly related PRs

Poem

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot May 11, 2026

Uh oh!

2witstudios May 11, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

2witstudios commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Test plan

Uh oh!

coderabbitai Bot commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

Possibly related PRs

Poem

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot May 11, 2026

Choose a reason for hiding this comment

Uh oh!

2witstudios May 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

2witstudios commented May 11, 2026 •

edited

Loading

coderabbitai Bot commented May 11, 2026 •

edited

Loading