feat: Assistant UI Toolbox

[heyglassy]

This PR implements a new Toolbox pattern that allows a developer to centralize the types and rendering definitions of their tools. An example implementation of this pattern is contained in this PR.

Preview DX.

export const toolbox = createToolbox<BackendTools>()({
  hi: frontendTool({
    parameters: z.object({
      name: z.string(),
    }),
    execute: async (args) => {
      return `Other: ${args.name}`;
    },

    render: function Render({ result }) {
      const [, setState] = useState(result);
      useEffect(() => {
        setState(result);
      }, [result]);
      return <div>Hi: {JSON.stringify(result)} confirmed</div>;
    },
  }),
// backend tool 
  weather: {
    render: ({ result }) => <div>Weather: {result?.weather}</div>,
  },
});

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Incorrect type assignment in AssistantRuntimeProvider	High	Define a tool with a custom UI and attempt to use the setUI function from useTool to update its UI.	The type of the callback in the setUI function from useTool is incorrectly defined. The prop type says that the setUI function takes a () => React.ReactNode, while the implementation in packages/react/src/model-context/tool.ts says that the type of the callback should be ComponentType<ToolCallContentPartProps<...>>. This will lead to type errors when trying to update the tool UI.
Missing prop type for render on FrontendTool.	Medium	Attempt to extract the render function from FrontendTool.	The line const { render, ...rest } = tool; in the useEffect callback extracts the render function. This may be a problem since render is not guaranteed to exist on the tool, and this could lead to a type error if render is undefined. Add a type check to ensure that the tool variable has the render prop before attempting to extract it.
useTool is not exposed correctly.	High	Define a tool and attempt to use the useTool hook to enable/disable the tool.	The useTool hook that is created in packages/react/src/model-context/tool.ts is not being passed down correctly through the toolbox prop. The enable/disable functions cannot be called.

_{Comments? Email us.}

[changeset-bot[bot]]

⚠️ No Changeset found

Latest commit: 0bc024729e1a1fce49e66e1a8747decd7d840629

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[vercel[bot]]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
assistant-ui	❌ Failed (Inspect)			Jun 19, 2025 9:00pm

[jazzberry-ai[bot]]

An error occured.

_{This error may be due to rate limits. If this error persists, please email us.}

[promptless[bot]]

📝 Documentation updates detected!

New suggestion: Document new Toolbox API for structured tool management

[coderabbitai[bot]]

Walkthrough

This set of changes enhances the typing and management of tools in the assistant-stream and react packages. The core tool types are refactored to support multiple schema formats with strong type inference for parameters and results. New utility functions facilitate creation and conversion of backend, frontend, and human tools. React integration is improved by adding support for rendering tool UIs and providing a typed toolbox factory with a useTool hook for UI management. Several modules update their filtering, mapping, and export logic to align with the revised tool types and structures. The changes also include safer invocation patterns and expanded type exports. Additionally, a new runtime dependency on zod is introduced to support schema validation.

---

📜 Recent review details

Configuration used: CodeRabbit UI Review profile: CHILL Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 9fdfe1a11d08130bb02cbbd593130e9a0eb24680 and d6afb704617e4c60353b11b81079f46e659b6ed5.

⛔ Files ignored due to path filters (1)

• pnpm-lock.yaml is excluded by !**/pnpm-lock.yaml

📒 Files selected for processing (4)

• packages/assistant-stream/package.json (1 hunks)
• packages/assistant-stream/src/core/tool/index.ts (1 hunks)
• packages/assistant-stream/src/core/tool/toolResultStream.ts (2 hunks)
• packages/react/src/model-context/makeAssistantVisible.tsx (2 hunks)

🚧 Files skipped from review as they are similar to previous changes (4)

• packages/assistant-stream/package.json
• packages/react/src/model-context/makeAssistantVisible.tsx
• packages/assistant-stream/src/core/tool/index.ts
• packages/assistant-stream/src/core/tool/toolResultStream.ts

⏰ Context from checks skipped due to timeout of 90000ms (2)

• GitHub Check: Trag Review
• GitHub Check: Analyze (javascript-typescript)

✨ Finishing Touches

• [ ] 📝 Generate Docstrings

🧪 Generate Unit Tests

• [ ] Create PR with Unit Tests
• [ ] Post Copyable Unit Tests in Comment
• [ ] Commit Unit Tests in branch create-toolbox-types-and-proxy-object

---

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

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai auto-generate unit tests to generate unit tests for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[jazzberry-ai[bot]]

An error occured.

_{This error may be due to rate limits. If this error persists, please email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
HumanTools not registered	Medium	Create a toolbox with a HumanTool and verify that it is registered.	HumanTools are not registered with the runtime because the useEffect hook in AssistantRuntimeProvider.tsx only registers tools that have an execute function. This prevents the UI for HumanTools from being displayed. The fix is to register any valid Tool, not only tools with an execute function.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Tool UIs not updated when render functions change.	Medium	1. Define a toolbox with a tool that has a render function. 2. Render the Assistant UI with the toolbox. 3. Dynamically update the render function for the tool. 4. Observe that the tool UI does not update with the new render function.	The useEffect hook in AssistantRuntimeProvider that registers render functions with the useToolUIs store only depends on the toolbox object reference. If the toolbox object reference remains the same, but the render functions within the toolbox change, the effect will not re-run, and the tool UIs will not be updated. This can lead to stale UIs. The dependecy array should include useToolUIs and toolbox.tools.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
HumanTool/frontendTool without execute is not registered	Medium	Create a toolbox with a HumanTool that does not have an execute function. Observe that the tool is not registered with the runtime.	The filter in AssistantRuntimeProvider.tsx incorrectly filters out HumanTool instances that do not have an execute function.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Found a bug in packages/react/src/model-context/tool.ts. The Toolbox type definition had default types for BackendTools and FrontendAndHumanTools which caused Typescript errors when defining toolboxes with only FrontendTools. The default type has been removed to fix this. Severity: Medium Example test case: Define a toolbox with only frontend tools without explicitly passing the BackendTools generic type. Description: The default type for BackendTools ({[key: string]: BackendTool;}) caused Typescript to require a render property on all tools, even frontend tools which do not have a render property. This was fixed by removing the default type for the BackendTools and FrontendAndHumanTools generic types in the Toolbox type definition.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Missing name property in tool definitions	High	Define a tool without a name property.	The name property is now required for all tool types (BackendTool, FrontendTool, HumanTool). Existing code that defines tools without a name property will need to be updated, or it will cause a type error.
Missing disable and enable functions in useTool hook	Medium	Attempt to disable or enable a tool using the useTool hook.	The disable and enable functions in the useTool hook are commented out, which means that the functionality to disable or enable tools is currently not available. This reduces the dynamic control over tools at the component level.
Type Conversion Issue in fromLanguageModelTools	High	Use fromLanguageModelTools to convert tools and verify that the tool type can be accessed by the runtime.	The fromLanguageModelTools function in react-edge no longer infers or sets the type property. This means code that consumes those tools may not be able to properly determine the type of tool.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
BackendTool execute argument type bypass	High	Define a BackendTool with a Zod schema for parameters. Implement the execute function with a specific type for the arguments. Pass incorrect arguments to the tool at runtime.	The execute function in BackendTool has `args: InferArgsFromParameters<TParameters>

_{Comments? Email us.}

[jazzberry-ai[bot]]

An error occured.

_{This error may be due to rate limits. If this error persists, please email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
useTool Proxy Missing hasOwnProperty Check	Low-Medium	Access a non-existent property on the object returned by useTool.	Accessing non-existent properties on the object returned by useTool will silently return undefined, which can lead to debugging issues.
AssistantRuntimeProvider toolbox Prop Update Issue	Medium	Update the toolbox prop after the initial render.	Updates to the toolbox prop after initial render are not reflected in the model context, which can cause UI inconsistencies.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name: Incorrect Type for useTool.setUI

Severity: Medium

Example test case:

1. Define a backend tool with a render function.
2. Attempt to use the useTool hook to change the UI. Observe that the typechecker prevents usage.

Description: The useTool hook's setUI function is overly restrictive. Specifically, the logic assumes that backend tools can only set the UI, whereas the frontend/human tools can enable, disable and set the UI. This isn't necessarily true - a backend tool could have it's UI changed.

Name: Missing name prop type enforce

Severity: Low

Example test case:

1. Create a backend tool without the name prop.
2. Observe there is no type error.

Description: The tool-types.ts file requires the name field to be present, but this is not enforced at type level when creating a new object.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Potential Type Inference Issue in createToolbox	Medium	Define a backend tool with an optional execute function.	The type inference for the render property in the createToolbox function relies on NonNullable<BackendTools[K]["execute"]>. If execute is optional, this could lead to type errors or unexpected behavior.
useTool Hook and Backend Tools	Medium	Try to use the enable or disable functions from the useTool hook with a backend tool.	The useTool hook appears to be designed primarily for frontend and human tools. It is unclear how this hook is intended to be used with backend tools. This could lead to confusion or errors if developers try to use enable or disable with backend tools.
fromLanguageModelTools and Tool Types	Low	Use fromLanguageModelTools to convert a language model tool, then check if the resulting tool object has the type property.	The fromLanguageModelTools function removes the type property from the tool object. This might be problematic if other parts of the system rely on the type property to determine how to handle the tool.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Missing name property for tools	High	Create a tool without a name property.	The BackendTool and FrontendTool types now require a name property. This could break existing code if tools are not updated to include this property.
Incomplete tool enabling/disabling	Medium	Attempt to enable or disable a tool using useTool.	The useTool hook suggests that tools can be enabled or disabled, but the functionality is commented out and not implemented.
Type assignment issue in withPromiseOrValue	Low	Examine the src/core/utils/withPromiseOrValue.ts file.	The type assignment in withPromiseOrValue.ts may cause runtime errors due to incorrect type assumptions.
Missing module 'ai'	Medium	Attempt to build the assistant-stream package	The TypeScript compiler is unable to locate the 'ai' module, used in tool-types.ts.

_{Comments? Email us.}

[jazzberry-ai[bot]]

Bug Report

Name	Severity	Example test case	Description
Conflicting Tool Parameter Types	Medium	Define a tool with a StandardSchemaV1 parameter and attempt to use it with a runtime expecting JSONSchema7.	The FrontendTool, HumanTool and BackendTool types have different possible values for TParameters. This could lead to type errors if a tool is used in a context where the expected parameter type differs.
Missing experimental_onSchemaValidationError	Low	Create a FrontendTool and attempt to define experimental_onSchemaValidationError.	While the BackendTool has an experimental_onSchemaValidationError field, the FrontendTool and HumanTool interfaces do not. This could lead to inconsistent handling of schema validation errors in different tool types.
Unsafe Type Cast in toAISDKTool	Medium	Use a BackendTool with properties that don't perfectly align with AITool and rely on type safety.	The toAISDKTool function uses any in casting the result to AITool, which circumvents type checking and could introduce runtime errors if the properties don't match the expected AITool type.
Incorrect Type Inference for useTool	Low	Use the useTool hook with a complex tool parameter or return type and observe unexpected type behavior.	The type definition for the return value of useTool is complex and may fail in certain edge cases, especially with complicated parameter types or return types for the execute function. There is also use of any to defeat the use of types in the first place
Stale toolbox values	Medium	Change the toolbox prop after initial render.	useEffect hook uses an empty dependency array [], which means it will only run once on initial render. This could cause issues if the toolbox prop changes after the initial render, as the component will not re-register the tools or update the UI.

_{Comments? Email us.}