Add classes for BaseAgent and RootAgent + add Agent Streamer #2922
Conversation
feat: Add streamer to manage agent streaming
|
@samschooler is attempting to deploy a commit to the Onlook Team on Vercel. A member of the Team first needs to authorize it. |
|
Caution Review failedThe pull request is closed. WalkthroughRefactors chat streaming to use a RootAgent + AgentStreamer abstraction: adds BaseAgent, RootAgent, AgentStreamer, re-exports agents from package entrypoints, replaces direct streamText usage in the API route, and moves streaming configuration, telemetry, and error handling into the new agent-based flow. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant API as /api/chat (route.ts)
participant AgentClass as RootAgent
participant Streamer as AgentStreamer
participant Provider as LLM Provider
Client->>API: POST chat messages
API->>AgentClass: RootAgent.create(chatType)
API->>Streamer: new AgentStreamer(agent, conversationId)
API->>Streamer: streamText(messages, { streamTextConfig, toUI... })
Streamer->>AgentClass: agent.streamText(streamTextConfig + messages)
AgentClass->>Provider: stream with model, tools, systemPrompt
Provider-->>AgentClass: streaming parts
AgentClass-->>Streamer: stream result
Streamer-->>API: UI message stream response (ids, metadata)
API-->>Client: SSE/streamed response
Note over API,Streamer: telemetry & error handlers triggered on error/finish
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (4)
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. Comment |
samschooler
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
apps/web/client/src/app/api/chat/route.ts (1)
41-41: Await streamResponse so errors are caught by the surrounding try/catch.Without
await, thrown errors insidestreamResponsewon’t be caught by thePOSThandler’s catch block.Apply this diff:
- return streamResponse(req, user.id); + return await streamResponse(req, user.id);
🧹 Nitpick comments (6)
packages/ai/src/agents/streamer.ts (2)
24-33: Stabilize createdAt per streamed messagecreatedAt changes per part; set it once to keep a consistent timestamp for the message.
- const result = await this.agent.streamText(convertToStreamMessages(messages), streamTextConfig); + const result = await this.agent.streamText(convertToStreamMessages(messages), streamTextConfig); + const createdAt = new Date(); @@ - messageMetadata: ({ part }) => { + messageMetadata: ({ part }) => { return { - createdAt: new Date(), + createdAt, conversationId, context: [], checkpoints: [], finishReason: part.type === 'finish-step' ? part.finishReason : undefined, usage: part.type === 'finish-step' ? part.usage : undefined, } satisfies ChatMetadata; },
5-5: Nit: remove unused importUIMessageStreamOnFinishCallback isn’t used.
-import type { streamText, UIMessageStreamOnFinishCallback, UIMessageStreamOptions } from "ai"; +import type { streamText, UIMessageStreamOptions } from "ai";packages/ai/src/agents/models/base.ts (1)
2-2: Nit: drop unused type importTool isn’t referenced in this file.
-import { type ModelMessage, type Tool, type ToolSet } from 'ai'; +import { type ModelMessage, type ToolSet } from 'ai';apps/web/client/src/app/api/chat/route.ts (1)
3-5: Remove unused imports.
convertToStreamMessagesandChatMetadataare not used in this file.Apply this diff:
-import { AgentStreamer, convertToStreamMessages, RootAgent } from '@onlook/ai'; +import { AgentStreamer, RootAgent } from '@onlook/ai'; ... -import { ChatType, type ChatMessage, type ChatMetadata } from '@onlook/models'; +import { ChatType, type ChatMessage } from '@onlook/models';packages/ai/src/agents/classes/root.ts (2)
1-1: Remove unused type imports from 'ai'.
JSONValue,LanguageModel,ModelMessage, andToolSetare not used in this file.Apply this diff:
-import { type JSONValue, type LanguageModel, type ModelMessage, type ToolSet } from 'ai';
8-17: Optional: Mark id as override if BaseAgent declares it.If
BaseAgentdefinesid, useoverride readonly id = 'root-agent'for stricter typing.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
apps/web/client/src/app/api/chat/route.ts(3 hunks)packages/ai/src/agents/classes/index.ts(1 hunks)packages/ai/src/agents/classes/root.ts(1 hunks)packages/ai/src/agents/index.ts(1 hunks)packages/ai/src/agents/models/base.ts(1 hunks)packages/ai/src/agents/models/index.ts(1 hunks)packages/ai/src/agents/streamer.ts(1 hunks)packages/ai/src/index.ts(1 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Do not use the any type unless necessary
Files:
packages/ai/src/agents/index.tspackages/ai/src/index.tspackages/ai/src/agents/classes/index.tsapps/web/client/src/app/api/chat/route.tspackages/ai/src/agents/models/base.tspackages/ai/src/agents/models/index.tspackages/ai/src/agents/streamer.tspackages/ai/src/agents/classes/root.ts
{apps,packages}/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
Avoid using the any type unless absolutely necessary
Files:
packages/ai/src/agents/index.tspackages/ai/src/index.tspackages/ai/src/agents/classes/index.tsapps/web/client/src/app/api/chat/route.tspackages/ai/src/agents/models/base.tspackages/ai/src/agents/models/index.tspackages/ai/src/agents/streamer.tspackages/ai/src/agents/classes/root.ts
apps/web/client/src/**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
apps/web/client/src/**/*.{ts,tsx}: Use path aliases @/* and ~/* for imports that map to apps/web/client/src/*
Avoid hardcoded user-facing text; use next-intl messages/hooks insteadUse path aliases @/* and ~/* for imports mapping to src/*
Files:
apps/web/client/src/app/api/chat/route.ts
apps/web/client/src/app/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
Default to Server Components; add 'use client' only when using events, state/effects, browser APIs, or client-only libs
Files:
apps/web/client/src/app/api/chat/route.ts
🧠 Learnings (2)
📚 Learning: 2025-09-14T01:44:21.209Z
Learnt from: CR
PR: onlook-dev/onlook#0
File: AGENTS.md:0-0
Timestamp: 2025-09-14T01:44:21.209Z
Learning: Applies to apps/web/client/src/server/api/root.ts : Export all tRPC routers from apps/web/client/src/server/api/root.ts
Applied to files:
packages/ai/src/agents/classes/index.ts
📚 Learning: 2025-09-16T19:22:52.461Z
Learnt from: CR
PR: onlook-dev/onlook#0
File: CLAUDE.md:0-0
Timestamp: 2025-09-16T19:22:52.461Z
Learning: Applies to apps/web/client/src/server/api/root.ts : Export all tRPC routers from src/server/api/root.ts
Applied to files:
packages/ai/src/agents/classes/index.ts
🧬 Code graph analysis (4)
apps/web/client/src/app/api/chat/route.ts (4)
packages/ai/src/agents/classes/root.ts (1)
RootAgent(7-57)packages/ai/src/agents/streamer.ts (1)
AgentStreamer(7-37)apps/web/client/src/app/api/chat/helpers/stream.ts (1)
repairToolCall(47-87)apps/web/client/src/app/api/chat/helpers/usage.ts (1)
decrementUsage(67-89)
packages/ai/src/agents/models/base.ts (2)
packages/models/src/llm/index.ts (1)
ModelConfig(34-39)packages/ai/src/agents/streamer.ts (1)
streamText(16-36)
packages/ai/src/agents/streamer.ts (3)
packages/models/src/chat/message/message.ts (2)
ChatMessage(18-18)ChatMetadata(6-13)packages/ai/src/agents/models/base.ts (1)
streamText(31-52)packages/ai/src/stream/index.ts (1)
convertToStreamMessages(5-23)
packages/ai/src/agents/classes/root.ts (5)
packages/models/src/llm/index.ts (1)
ModelConfig(34-39)packages/ai/src/tools/toolset.ts (1)
getToolSetFromType(66-68)packages/ai/src/prompt/provider.ts (2)
getCreatePageSystemPrompt(29-33)getAskModeSystemPrompt(41-45)packages/ai/src/agents/models/base.ts (1)
getSystemPrompt(15-17)packages/ai/src/chat/providers.ts (1)
initModel(14-51)
🔇 Additional comments (9)
packages/ai/src/agents/models/index.ts (1)
1-1: LGTM: clean public re‑exportExposing BaseAgent via the models index makes sense and keeps imports tidy.
packages/ai/src/agents/classes/index.ts (1)
1-1: LGTM: targeted re‑exportRe-exporting RootAgent from the classes index keeps consumer imports clean.
packages/ai/src/index.ts (1)
1-1: Verify no export name collisions from new export-starAdding export * from './agents' widens the public API and might collide with existing exports (types or values) from other modules already re-exported here. Please verify no duplicate export names and consider a minor version bump if API surface expands.
Run this read-only check to surface potential duplicate exported identifiers across packages/ai/src (heuristic; may include false positives due to re-exports):
packages/ai/src/agents/index.ts (1)
1-3: Public agents API re-export looks good.Consolidated exports are clear and minimal. No issues spotted.
packages/ai/src/agents/classes/root.ts (3)
23-33: Confirm desired prompt for ChatType.FIX.
FIXis mapped to a CREATE/FIX model below but falls into the default system prompt path here. Verify whether FIX should use a specialized prompt (e.g., EDIT) for consistency.
35-38: Factory method is appropriate.Asynchronous
createneatly encapsulates model selection and agent construction.
40-55: OPEN_AI_GPT_5 model and token configuration verified
OPENROUTER_MODELS.OPEN_AI_GPT_5 is defined in the enum and MODEL_MAX_TOKENS has a 400 000 entry for it.apps/web/client/src/app/api/chat/route.ts (2)
78-126: Import path stability confirmed: The@onlook/aipackage’s top-level export re-exports bothRootAgentandAgentStreamer, so the import path remains stable.
71-73: Runtime supports Array.prototype.findLast
package.json requires Node.js ≥18.0.0 and Array.prototype.findLast is available since Node 18.
2 participants
Description
Encapsulate the root agent in a class in order to have separation of concerns in the chat streaming route.
Type of Change
Testing
Verified with QA + existing tests
Important
Refactor to encapsulate root agent functionality into
BaseAgent,RootAgent, andAgentStreamerclasses for improved separation of concerns in chat streaming.BaseAgentclass inmodels/base.tswith methods for system prompt and toolset management.RootAgentclass inclasses/root.tsextendingBaseAgent, with methods for system prompt based onChatTypeand model configuration.AgentStreamerclass instreamer.tsto handle message streaming usingBaseAgent.streamResponseinroute.tsto useRootAgentandAgentStreamerfor handling chat streaming.index.tsfiles to include new classes and modules.This description was created by
for 8aeab47. You can customize this summary. It will automatically update as commits are pushed.
Summary by CodeRabbit
Refactor
New Features