refactor: clean up message context management #2952
Conversation
|
This pull request has been ignored for the connected project Preview Branches by Supabase. |
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughIntroduces a BaseContext abstraction, adds six concrete context classes (File, Highlight, Error, Branch, Image, AgentRule), centralizes prompt/label generation in a context index, refactors provider to delegate formatting to context classes, updates tests and test data, and removes shared prompt constants. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor UI
participant Provider as provider.ts
participant CtxIndex as contexts/index
participant Class as ContextClass (e.g., FileContext)
participant LLM as Downstream
UI->>Provider: request hydrate/composed prompt (opts)
Provider->>CtxIndex: getContextPrompt(context)
CtxIndex->>CtxIndex: resolve type -> class (CONTEXT_CLASSES)
CtxIndex->>Class: Class.getPrompt(context)
Class-->>CtxIndex: prompt string
CtxIndex-->>Provider: prompt string
Provider-->>LLM: send composed prompt
LLM-->>UI: response
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)
✨ Finishing touches
🧪 Generate unit tests
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
packages/ai/test/prompt/prompt.test.ts (1)
161-180: Align the test helper call with the branch-aware signatureOnce
getHighlightsContentaccepts abranchId, keep the test exercising the real path by threading that value too:- const prompt = HighlightContext.getHighlightsContent('test.txt', [ - { - path: 'test.txt', - start: 1, - end: 2, - content: 'test', - type: MessageContextType.HIGHLIGHT, - displayName: 'test.txt', - branchId: 'test', - }, - { - path: 'test.txt', - start: 3, - end: 4, - content: 'test2', - type: MessageContextType.HIGHLIGHT, - displayName: 'test.txt', - branchId: 'test', - }, - ]); + const prompt = HighlightContext.getHighlightsContent( + 'test.txt', + [ + { + path: 'test.txt', + start: 1, + end: 2, + content: 'test', + type: MessageContextType.HIGHLIGHT, + displayName: 'test.txt', + branchId: 'test', + }, + { + path: 'test.txt', + start: 3, + end: 4, + content: 'test2', + type: MessageContextType.HIGHLIGHT, + displayName: 'test.txt', + branchId: 'test', + }, + ], + 'test', + );packages/ai/src/prompt/provider.ts (1)
136-142: Consolidate duplicate helpers into FileContext
Remove the exportedgetLanguageFromFilePathandgetBranchContentfrom packages/ai/src/prompt/provider.ts and instead reuse the corresponding static methods on FileContext (make them public if needed).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
context-architecture.md(1 hunks)packages/ai/src/contexts/classes/agent-rule-context.ts(1 hunks)packages/ai/src/contexts/classes/branch-context.ts(1 hunks)packages/ai/src/contexts/classes/error-context.ts(1 hunks)packages/ai/src/contexts/classes/file-context.ts(1 hunks)packages/ai/src/contexts/classes/highlight-context.ts(1 hunks)packages/ai/src/contexts/classes/image-context.ts(1 hunks)packages/ai/src/contexts/index.ts(1 hunks)packages/ai/src/contexts/models/base.ts(1 hunks)packages/ai/src/prompt/provider.ts(2 hunks)packages/ai/test/prompt/data/file.txt(1 hunks)packages/ai/test/prompt/data/highlights.txt(1 hunks)packages/ai/test/prompt/data/user.txt(2 hunks)packages/ai/test/prompt/prompt.test.ts(3 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Do not use the any type unless necessary
Files:
packages/ai/test/prompt/prompt.test.tspackages/ai/src/contexts/classes/image-context.tspackages/ai/src/contexts/index.tspackages/ai/src/contexts/classes/error-context.tspackages/ai/src/contexts/models/base.tspackages/ai/src/contexts/classes/branch-context.tspackages/ai/src/contexts/classes/agent-rule-context.tspackages/ai/src/contexts/classes/highlight-context.tspackages/ai/src/contexts/classes/file-context.tspackages/ai/src/prompt/provider.ts
{apps,packages}/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
Avoid using the any type unless absolutely necessary
Files:
packages/ai/test/prompt/prompt.test.tspackages/ai/src/contexts/classes/image-context.tspackages/ai/src/contexts/index.tspackages/ai/src/contexts/classes/error-context.tspackages/ai/src/contexts/models/base.tspackages/ai/src/contexts/classes/branch-context.tspackages/ai/src/contexts/classes/agent-rule-context.tspackages/ai/src/contexts/classes/highlight-context.tspackages/ai/src/contexts/classes/file-context.tspackages/ai/src/prompt/provider.ts
🧬 Code graph analysis (9)
packages/ai/src/contexts/classes/image-context.ts (2)
packages/ui/src/components/icons/index.tsx (1)
Icons(138-3667)packages/models/src/chat/message/context.ts (1)
ImageMessageContext(38-41)
packages/ai/src/contexts/index.ts (7)
packages/ai/src/contexts/classes/file-context.ts (1)
FileContext(7-87)packages/ai/src/contexts/classes/highlight-context.ts (1)
HighlightContext(7-51)packages/ai/src/contexts/classes/error-context.ts (1)
ErrorContext(7-41)packages/ai/src/contexts/classes/branch-context.ts (1)
BranchContext(6-28)packages/ai/src/contexts/classes/image-context.ts (1)
ImageContext(5-29)packages/ai/src/contexts/classes/agent-rule-context.ts (1)
AgentRuleContext(7-30)packages/models/src/chat/message/context.ts (1)
MessageContext(53-59)
packages/ai/src/contexts/classes/error-context.ts (4)
packages/ai/src/contexts/index.ts (2)
ErrorContext(35-35)BaseContext(33-33)packages/models/src/chat/message/context.ts (1)
ErrorMessageContext(43-46)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)packages/ai/src/prompt/constants/context.ts (1)
CONTEXT_PROMPTS(26-32)
packages/ai/src/contexts/models/base.ts (2)
packages/ai/src/contexts/index.ts (1)
BaseContext(33-33)packages/models/src/chat/message/context.ts (1)
MessageContext(53-59)
packages/ai/src/contexts/classes/branch-context.ts (3)
packages/ui/src/components/icons/index.tsx (1)
Icons(138-3667)packages/models/src/chat/message/context.ts (1)
BranchMessageContext(18-21)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)
packages/ai/src/contexts/classes/agent-rule-context.ts (3)
packages/models/src/chat/message/context.ts (1)
AgentRuleMessageContext(48-51)packages/ai/src/prompt/constants/context.ts (1)
CONTEXT_PROMPTS(26-32)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)
packages/ai/src/contexts/classes/highlight-context.ts (5)
packages/ui/src/components/icons/index.tsx (1)
Icons(138-3667)packages/models/src/chat/message/context.ts (1)
HighlightMessageContext(29-36)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)packages/ai/src/prompt/constants/format.ts (1)
CODE_FENCE(1-4)packages/ai/src/prompt/constants/context.ts (1)
CONTEXT_PROMPTS(26-32)
packages/ai/src/contexts/classes/file-context.ts (5)
packages/models/src/chat/message/context.ts (2)
FileMessageContext(23-27)HighlightMessageContext(29-36)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)packages/ai/src/prompt/constants/format.ts (1)
CODE_FENCE(1-4)packages/ai/src/prompt/constants/context.ts (1)
CONTEXT_PROMPTS(26-32)packages/ai/src/contexts/classes/highlight-context.ts (1)
HighlightContext(7-51)
packages/ai/src/prompt/provider.ts (2)
packages/ai/src/contexts/classes/file-context.ts (1)
FileContext(7-87)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)
🪛 markdownlint-cli2 (0.18.1)
context-architecture.md
3-3: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
376-376: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (15)
packages/ai/test/prompt/data/highlights.txt (1)
2-7: Highlight path formatting looks goodEmbedding the line span inside the
<path>text matches the new HighlightContext expectations. Update reads well.packages/ai/test/prompt/data/file.txt (1)
8-13: File highlight markup matches new conventionThe consolidated
<path>value with#L1:L2aligns with the refreshed highlight serialization. Looks consistent with other fixtures.packages/ai/test/prompt/data/user.txt (1)
8-29: User context fixture aligns with new structureHighlight path formatting and the separate
<instruction>block both match the refactored context model. Fixture updates look correct.packages/ai/src/contexts/models/base.ts (1)
8-26: BaseContext scaffold looks solidThe abstract base cleanly defines the static contract and guards with throwers; subclasses can override as needed. No concerns.
packages/ai/src/contexts/classes/agent-rule-context.ts (3)
7-10: LGTM! Clean metadata definition.The static metadata follows the BaseContext interface correctly with appropriate icon and display name choices.
12-14: LGTM! Straightforward prompt formatting.The prompt correctly combines path and content with a newline separator.
16-18: LGTM! Proper label fallback.The label method appropriately falls back to
pathwhendisplayNameis not available.packages/ai/src/contexts/classes/error-context.ts (5)
7-10: LGTM! Proper metadata definition.The static metadata correctly implements the BaseContext interface with an appropriate icon choice.
12-16: LGTM! Clean prompt formatting.The method correctly combines branch display and error content with proper XML wrapping.
18-20: LGTM! Appropriate label fallback.The label method properly falls back to 'Error' when displayName is not available.
25-36: LGTM! Proper empty array handling.The early return for empty arrays prevents unnecessary XML wrapping and prefix generation. The loop correctly aggregates error prompts (getPrompt already includes trailing newline from line 15).
38-40: LGTM! Clean helper method.The private helper correctly encapsulates branch content formatting.
packages/ai/src/prompt/provider.ts (3)
14-18: LGTM! Context class imports are well-structured.The new imports support the refactored architecture that centralizes prompt generation logic into dedicated context classes.
20-25: LGTM! Interface extension supports message truncation logic.The new index properties enable the function to determine when to truncate file context for older messages in the conversation.
65-134: LGTM! Clean delegation to centralized context classes.The refactoring successfully moves prompt generation logic to dedicated context classes while preserving the original behavior. The type filtering and method calls are well-structured.
context-architecture.md
Outdated
| @@ -0,0 +1,388 @@ | |||
| # Context Base Class Architecture | |||
|
|
|||
| *NOTE: Treat all the code here as suggestions and validate the correct code usage. Typecheck should be ran under packages/ai* | |||
There was a problem hiding this comment.
Fix markdown: Use proper heading instead of emphasis.
Line 3 uses emphasis (italics) instead of a proper heading, which violates markdown best practices.
As per static analysis.
Apply this diff:
# Context Base Class Architecture
-*NOTE: Treat all the code here as suggestions and validate the correct code usage. Typecheck should be ran under packages/ai*
+> **NOTE**: Treat all the code here as suggestions and validate the correct code usage. Typecheck should be run under packages/ai
## OverviewCommittable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 markdownlint-cli2 (0.18.1)
3-3: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
🤖 Prompt for AI Agents
In context-architecture.md around line 3, the note is formatted with emphasis
(italics) instead of a proper markdown heading; replace the italicized line with
an appropriate markdown heading (e.g., prefix with one or more # characters such
as "###") so it becomes a heading rather than emphasized text, preserving the
same text content and punctuation.
context-architecture.md
Outdated
| ``` | ||
| packages/ai/src/contexts/ | ||
| ├── models/ | ||
| │ └── base.ts | ||
| ├── classes/ | ||
| │ ├── file-context.ts | ||
| │ ├── highlight-context.ts | ||
| │ ├── error-context.ts | ||
| │ ├── branch-context.ts | ||
| │ ├── image-context.ts | ||
| │ └── agent-rule-context.ts | ||
| └── index.ts | ||
| ``` |
There was a problem hiding this comment.
Specify language for fenced code block.
The directory structure code block at line 376 should specify a language identifier for proper rendering.
As per static analysis.
Apply this diff:
## Directory Structure
-```
+```text
packages/ai/src/contexts/
├── models/
│ └── base.ts🧰 Tools
🪛 markdownlint-cli2 (0.18.1)
376-376: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 Prompt for AI Agents
In context-architecture.md around lines 376 to 388, the fenced code block
showing the directory structure lacks a language identifier; update the opening
fence from ``` to ```text so the block is rendered with the text language (i.e.,
replace the existing triple-backtick start with ```text and keep the closing
triple-backticks unchanged).
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
packages/ai/src/contexts/classes/file-context.ts (1)
33-54: Propagate branch id to highlight filtering.Line 43 calls
getHighlightsForFilewithout the file's branch id, which may cause highlights from different branches to be incorrectly mixed into the file content.Apply the suggested diff from the previous review to pass
file.branchIdthrough to the highlight context filtering.packages/ai/src/contexts/classes/agent-rule-context.ts (1)
24-30: Add separator between agent rules.Line 27 concatenates agent rule prompts without any separator, causing multiple rules to run together.
Apply the suggested diff from the previous review to add a newline after each rule prompt.
🧹 Nitpick comments (1)
packages/ai/src/prompt/provider.ts (1)
59-128: LGTM! Correct delegation to context classes.The function properly delegates to the new context class methods for all content generation. The refactoring successfully moves prompt logic out of the provider.
Lines 66-71 use
.map((c) => c)which is redundant since the filter already returns the correct type. Consider simplifying:- const files = context.filter((c) => c.type === MessageContextType.FILE).map((c) => c); - const highlights = context.filter((c) => c.type === MessageContextType.HIGHLIGHT).map((c) => c); - const errors = context.filter((c) => c.type === MessageContextType.ERROR).map((c) => c); - const agentRules = context.filter((c) => c.type === MessageContextType.AGENT_RULE).map((c) => c); - const images = context.filter((c) => c.type === MessageContextType.IMAGE).map((c) => c); - const branches = context.filter((c) => c.type === MessageContextType.BRANCH).map((c) => c); + const files = context.filter((c) => c.type === MessageContextType.FILE); + const highlights = context.filter((c) => c.type === MessageContextType.HIGHLIGHT); + const errors = context.filter((c) => c.type === MessageContextType.ERROR); + const agentRules = context.filter((c) => c.type === MessageContextType.AGENT_RULE); + const images = context.filter((c) => c.type === MessageContextType.IMAGE); + const branches = context.filter((c) => c.type === MessageContextType.BRANCH);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
packages/ai/src/contexts/classes/agent-rule-context.ts(1 hunks)packages/ai/src/contexts/classes/error-context.ts(1 hunks)packages/ai/src/contexts/classes/file-context.ts(1 hunks)packages/ai/src/contexts/classes/highlight-context.ts(1 hunks)packages/ai/src/prompt/constants/context.ts(0 hunks)packages/ai/src/prompt/constants/index.ts(0 hunks)packages/ai/src/prompt/provider.ts(2 hunks)
💤 Files with no reviewable changes (2)
- packages/ai/src/prompt/constants/index.ts
- packages/ai/src/prompt/constants/context.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/ai/src/contexts/classes/highlight-context.ts
🧰 Additional context used
📓 Path-based instructions (2)
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Do not use the any type unless necessary
Files:
packages/ai/src/contexts/classes/error-context.tspackages/ai/src/contexts/classes/file-context.tspackages/ai/src/contexts/classes/agent-rule-context.tspackages/ai/src/prompt/provider.ts
{apps,packages}/**/*.{ts,tsx}
📄 CodeRabbit inference engine (CLAUDE.md)
Avoid using the any type unless absolutely necessary
Files:
packages/ai/src/contexts/classes/error-context.tspackages/ai/src/contexts/classes/file-context.tspackages/ai/src/contexts/classes/agent-rule-context.tspackages/ai/src/prompt/provider.ts
🧬 Code graph analysis (4)
packages/ai/src/contexts/classes/error-context.ts (2)
packages/models/src/chat/message/context.ts (1)
ErrorMessageContext(43-46)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)
packages/ai/src/contexts/classes/file-context.ts (5)
packages/ui/src/components/icons/index.tsx (1)
Icons(138-3667)packages/models/src/chat/message/context.ts (2)
FileMessageContext(23-27)HighlightMessageContext(29-36)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)packages/ai/src/prompt/constants/format.ts (1)
CODE_FENCE(1-4)packages/ai/src/contexts/classes/highlight-context.ts (1)
HighlightContext(7-53)
packages/ai/src/contexts/classes/agent-rule-context.ts (3)
packages/ui/src/components/icons/index.tsx (1)
Icons(138-3667)packages/models/src/chat/message/context.ts (1)
AgentRuleMessageContext(48-51)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)
packages/ai/src/prompt/provider.ts (4)
packages/ai/src/contexts/classes/file-context.ts (1)
FileContext(7-91)packages/ai/src/prompt/helpers.ts (1)
wrapXml(1-3)packages/ai/src/contexts/classes/error-context.ts (1)
ErrorContext(6-55)packages/ai/src/contexts/classes/agent-rule-context.ts (1)
AgentRuleContext(6-31)
🔇 Additional comments (13)
packages/ai/src/contexts/classes/file-context.ts (5)
1-10: LGTM! Clean class declaration and metadata.The imports are well-organized and the static metadata (contextType, displayName, icon) follows the established pattern across context classes.
12-14: LGTM! Clear instructional prefixes.The prefix messages appropriately guide the AI on how to interpret the context.
16-28: LGTM! Clean prompt and label generation.Both methods correctly format their output with appropriate fallbacks.
59-76: LGTM! Correct truncated content generation.The method appropriately strips content while preserving path and branch metadata for truncated contexts.
78-90: LGTM! Well-designed helper methods.The dynamic require pattern correctly avoids circular dependencies, and the language extraction safely handles edge cases.
packages/ai/src/contexts/classes/agent-rule-context.ts (2)
1-11: LGTM! Standard context class structure.The imports, metadata, and prefix follow the established pattern across all context classes.
13-19: LGTM! Appropriate prompt and label generation.The simple format is suitable for agent rules, with XML wrapping deferred to the aggregate method.
packages/ai/src/contexts/classes/error-context.ts (5)
1-9: LGTM! Standard context class setup.The imports and metadata are consistent with the other context classes.
11-24: LGTM! Comprehensive error debugging guidance.The prefix provides clear, actionable guidance for common Next.js + Bun issues and emphasizes the correct package manager commands.
26-34: LGTM! Clean prompt and label methods.Both methods correctly format error context with appropriate XML structure and fallbacks.
39-50: LGTM! Correct error aggregation.The method properly aggregates multiple errors, and the trailing newline in
getPromptensures proper separation between errors.
52-54: LGTM! Standard branch helper.The branch content wrapper is consistent with other context classes.
packages/ai/src/prompt/provider.ts (1)
1-19: LGTM! Clean refactored imports.The imports correctly use the new context classes instead of direct type imports, aligning with the new architecture.
![ellipsis-dev[bot]](https://avatars.githubusercontent.com/in/64358?s=60&v=4){kind=small}
apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-pills/helpers.tsx
Outdated
Show resolved
Hide resolved
Kitenite
changed the title
Description
Related Issues
Type of Change
Testing
Screenshots (if applicable)
Additional Notes
Important
Refactor message context handling by introducing dedicated context classes for modularity and updating tests accordingly.
File,Highlight,Error,Branch,Image,Agent Rule.agent-rule-context.ts,branch-context.ts,error-context.ts).agent-rule-context.test.ts,branch-context.test.ts,error-context.test.ts.index.tsandpackage.json.This description was created by
for 2a52504. You can customize this summary. It will automatically update as commits are pushed.
Summary by CodeRabbit
New Features
Refactor
Documentation
Tests
Chores