Skip to content

Commit a7ba3b5

Browse files
daniel-lxsdaniel-lxs
authored
feat: wire RooMessage storage into Task.ts and all providers (#11386)
1 parent 4e659b4 commit a7ba3b5

File tree

96 files changed

+2089
-2206
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

96 files changed

+2089
-2206
lines changed

‎src/api/index.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ import OpenAI from "openai"
33

44
import { isRetiredProvider, type ProviderSettings, type ModelInfo } from "@roo-code/types"
55

6+
import type { RooMessage } from "../core/task-persistence/rooMessage"
7+
68
import { ApiStream } from "./transform/stream"
79

810
import {
@@ -89,11 +91,7 @@ export interface ApiHandlerCreateMessageMetadata {
8991
}
9092

9193
export interface ApiHandler {
92-
createMessage(
93-
systemPrompt: string,
94-
messages: Anthropic.Messages.MessageParam[],
95-
metadata?: ApiHandlerCreateMessageMetadata,
96-
): ApiStream
94+
createMessage(systemPrompt: string, messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata): ApiStream
9795

9896
getModel(): { id: string; info: ModelInfo }
9997

‎src/api/providers/__tests__/anthropic-vertex.spec.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// npx vitest run src/api/providers/__tests__/anthropic-vertex.spec.ts
23

34
import { AnthropicVertexHandler } from "../anthropic-vertex"
@@ -184,7 +185,7 @@ describe("AnthropicVertexHandler", () => {
184185
})
185186

186187
describe("createMessage", () => {
187-
const mockMessages: Anthropic.Messages.MessageParam[] = [
188+
const mockMessages: RooMessage[] = [
188189
{
189190
role: "user",
190191
content: "Hello",
@@ -244,15 +245,20 @@ describe("AnthropicVertexHandler", () => {
244245
)
245246
})
246247

247-
it("should call convertToAiSdkMessages with the messages", async () => {
248+
it("should pass messages directly to streamText as ModelMessage[]", async () => {
248249
mockStreamText.mockReturnValue(createMockStreamResult([]))
249250

250251
const stream = handler.createMessage(systemPrompt, mockMessages)
251252
for await (const _chunk of stream) {
252253
// consume
253254
}
254255

255-
expect(convertToAiSdkMessages).toHaveBeenCalledWith(mockMessages)
256+
// Messages are now already in ModelMessage format, passed directly to streamText
257+
expect(mockStreamText).toHaveBeenCalledWith(
258+
expect.objectContaining({
259+
messages: mockMessages,
260+
}),
261+
)
256262
})
257263

258264
it("should pass tools through AI SDK conversion pipeline", async () => {

‎src/api/providers/__tests__/azure.spec.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
23
const { mockStreamText, mockGenerateText, mockCreateAzure } = vi.hoisted(() => ({
34
mockStreamText: vi.fn(),
@@ -132,7 +133,7 @@ describe("AzureHandler", () => {
132133

133134
describe("createMessage", () => {
134135
const systemPrompt = "You are a helpful assistant."
135-
const messages: Anthropic.Messages.MessageParam[] = [
136+
const messages: RooMessage[] = [
136137
{
137138
role: "user",
138139
content: [
@@ -376,7 +377,7 @@ describe("AzureHandler", () => {
376377

377378
describe("tools", () => {
378379
const systemPrompt = "You are a helpful assistant."
379-
const messages: Anthropic.Messages.MessageParam[] = [
380+
const messages: RooMessage[] = [
380381
{
381382
role: "user",
382383
content: [{ type: "text" as const, text: "Use a tool" }],

‎src/api/providers/__tests__/base-provider.spec.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
import { Anthropic } from "@anthropic-ai/sdk"
23

34
import type { ModelInfo } from "@roo-code/types"
@@ -7,7 +8,7 @@ import type { ApiStream } from "../../transform/stream"
78

89
// Create a concrete implementation for testing
910
class TestProvider extends BaseProvider {
10-
createMessage(_systemPrompt: string, _messages: Anthropic.Messages.MessageParam[]): ApiStream {
11+
createMessage(_systemPrompt: string, _messages: RooMessage[]): ApiStream {
1112
throw new Error("Not implemented")
1213
}
1314

‎src/api/providers/__tests__/baseten.spec.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// npx vitest run src/api/providers/__tests__/baseten.spec.ts
23

34
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
@@ -101,7 +102,7 @@ describe("BasetenHandler", () => {
101102

102103
describe("createMessage", () => {
103104
const systemPrompt = "You are a helpful assistant."
104-
const messages: Anthropic.Messages.MessageParam[] = [
105+
const messages: RooMessage[] = [
105106
{
106107
role: "user",
107108
content: [
@@ -281,7 +282,7 @@ describe("BasetenHandler", () => {
281282

282283
describe("tool handling", () => {
283284
const systemPrompt = "You are a helpful assistant."
284-
const messages: Anthropic.Messages.MessageParam[] = [
285+
const messages: RooMessage[] = [
285286
{
286287
role: "user",
287288
content: [{ type: "text" as const, text: "Hello!" }],
@@ -389,7 +390,7 @@ describe("BasetenHandler", () => {
389390

390391
describe("error handling", () => {
391392
const systemPrompt = "You are a helpful assistant."
392-
const messages: Anthropic.Messages.MessageParam[] = [
393+
const messages: RooMessage[] = [
393394
{
394395
role: "user",
395396
content: [{ type: "text" as const, text: "Hello!" }],

‎src/api/providers/__tests__/bedrock.spec.ts‎

Lines changed: 20 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// Mock TelemetryService before other imports
23
const mockCaptureException = vi.fn()
34

@@ -490,17 +491,14 @@ describe("AwsBedrockHandler", () => {
490491
it("should properly pass image content through to streamText via AI SDK messages", async () => {
491492
setupMockStreamText()
492493

493-
const messages: Anthropic.Messages.MessageParam[] = [
494+
const messages: any[] = [
494495
{
495496
role: "user",
496497
content: [
497498
{
498499
type: "image",
499-
source: {
500-
type: "base64",
501-
data: mockImageData,
502-
media_type: "image/jpeg",
503-
},
500+
image: `data:image/jpeg;base64,${mockImageData}`,
501+
mimeType: "image/jpeg",
504502
},
505503
{
506504
type: "text",
@@ -530,7 +528,7 @@ describe("AwsBedrockHandler", () => {
530528
expect(userMsg).toBeDefined()
531529
expect(Array.isArray(userMsg.content)).toBe(true)
532530

533-
// The AI SDK convertToAiSdkMessages converts images to { type: "image", image: "data:...", mimeType: "..." }
531+
// Messages are already in AI SDK ImagePart format
534532
const imagePart = userMsg.content.find((p: { type: string }) => p.type === "image")
535533
expect(imagePart).toBeDefined()
536534
expect(imagePart.image).toContain("data:image/jpeg;base64,")
@@ -544,29 +542,23 @@ describe("AwsBedrockHandler", () => {
544542
it("should handle multiple images in a single message", async () => {
545543
setupMockStreamText()
546544

547-
const messages: Anthropic.Messages.MessageParam[] = [
545+
const messages: any[] = [
548546
{
549547
role: "user",
550548
content: [
551549
{
552550
type: "image",
553-
source: {
554-
type: "base64",
555-
data: mockImageData,
556-
media_type: "image/jpeg",
557-
},
551+
image: `data:image/jpeg;base64,${mockImageData}`,
552+
mimeType: "image/jpeg",
558553
},
559554
{
560555
type: "text",
561556
text: "First image",
562557
},
563558
{
564559
type: "image",
565-
source: {
566-
type: "base64",
567-
data: mockImageData,
568-
media_type: "image/png",
569-
},
560+
image: `data:image/png;base64,${mockImageData}`,
561+
mimeType: "image/png",
570562
},
571563
{
572564
type: "text",
@@ -761,7 +753,7 @@ describe("AwsBedrockHandler", () => {
761753
awsBedrock1MContext: true,
762754
})
763755

764-
const messages: Anthropic.Messages.MessageParam[] = [
756+
const messages: RooMessage[] = [
765757
{
766758
role: "user",
767759
content: "Test message",
@@ -794,7 +786,7 @@ describe("AwsBedrockHandler", () => {
794786
awsBedrock1MContext: false,
795787
})
796788

797-
const messages: Anthropic.Messages.MessageParam[] = [
789+
const messages: RooMessage[] = [
798790
{
799791
role: "user",
800792
content: "Test message",
@@ -828,7 +820,7 @@ describe("AwsBedrockHandler", () => {
828820
awsBedrock1MContext: true,
829821
})
830822

831-
const messages: Anthropic.Messages.MessageParam[] = [
823+
const messages: RooMessage[] = [
832824
{
833825
role: "user",
834826
content: "Test message",
@@ -881,7 +873,7 @@ describe("AwsBedrockHandler", () => {
881873
awsBedrock1MContext: true,
882874
})
883875

884-
const messages: Anthropic.Messages.MessageParam[] = [
876+
const messages: RooMessage[] = [
885877
{
886878
role: "user",
887879
content: "Test message",
@@ -1013,7 +1005,7 @@ describe("AwsBedrockHandler", () => {
10131005
awsBedrockServiceTier: "PRIORITY",
10141006
})
10151007

1016-
const messages: Anthropic.Messages.MessageParam[] = [
1008+
const messages: RooMessage[] = [
10171009
{
10181010
role: "user",
10191011
content: "Test message",
@@ -1050,7 +1042,7 @@ describe("AwsBedrockHandler", () => {
10501042
awsBedrockServiceTier: "FLEX",
10511043
})
10521044

1053-
const messages: Anthropic.Messages.MessageParam[] = [
1045+
const messages: RooMessage[] = [
10541046
{
10551047
role: "user",
10561048
content: "Test message",
@@ -1087,7 +1079,7 @@ describe("AwsBedrockHandler", () => {
10871079
awsBedrockServiceTier: "PRIORITY", // Try to apply PRIORITY tier
10881080
})
10891081

1090-
const messages: Anthropic.Messages.MessageParam[] = [
1082+
const messages: RooMessage[] = [
10911083
{
10921084
role: "user",
10931085
content: "Test message",
@@ -1122,7 +1114,7 @@ describe("AwsBedrockHandler", () => {
11221114
// No awsBedrockServiceTier specified
11231115
})
11241116

1125-
const messages: Anthropic.Messages.MessageParam[] = [
1117+
const messages: RooMessage[] = [
11261118
{
11271119
role: "user",
11281120
content: "Test message",
@@ -1192,7 +1184,7 @@ describe("AwsBedrockHandler", () => {
11921184
awsRegion: "us-east-1",
11931185
})
11941186

1195-
const messages: Anthropic.Messages.MessageParam[] = [
1187+
const messages: RooMessage[] = [
11961188
{
11971189
role: "user",
11981190
content: "Hello",
@@ -1267,7 +1259,7 @@ describe("AwsBedrockHandler", () => {
12671259
awsRegion: "us-east-1",
12681260
})
12691261

1270-
const messages: Anthropic.Messages.MessageParam[] = [
1262+
const messages: RooMessage[] = [
12711263
{
12721264
role: "user",
12731265
content: "Hello",

‎src/api/providers/__tests__/deepseek.spec.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
23
const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({
34
mockStreamText: vi.fn(),
@@ -173,7 +174,7 @@ describe("DeepSeekHandler", () => {
173174

174175
describe("createMessage", () => {
175176
const systemPrompt = "You are a helpful assistant."
176-
const messages: Anthropic.Messages.MessageParam[] = [
177+
const messages: RooMessage[] = [
177178
{
178179
role: "user",
179180
content: [
@@ -400,7 +401,7 @@ describe("DeepSeekHandler", () => {
400401

401402
describe("reasoning content with deepseek-reasoner", () => {
402403
const systemPrompt = "You are a helpful assistant."
403-
const messages: Anthropic.Messages.MessageParam[] = [
404+
const messages: RooMessage[] = [
404405
{
405406
role: "user",
406407
content: [
@@ -570,7 +571,7 @@ describe("DeepSeekHandler", () => {
570571

571572
describe("tool handling", () => {
572573
const systemPrompt = "You are a helpful assistant."
573-
const messages: Anthropic.Messages.MessageParam[] = [
574+
const messages: RooMessage[] = [
574575
{
575576
role: "user",
576577
content: [{ type: "text" as const, text: "Hello!" }],

‎src/api/providers/__tests__/fireworks.spec.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// npx vitest run src/api/providers/__tests__/fireworks.spec.ts
23

34
// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
@@ -363,7 +364,7 @@ describe("FireworksHandler", () => {
363364

364365
describe("createMessage", () => {
365366
const systemPrompt = "You are a helpful assistant."
366-
const messages: Anthropic.Messages.MessageParam[] = [
367+
const messages: RooMessage[] = [
367368
{
368369
role: "user",
369370
content: [
@@ -730,7 +731,7 @@ describe("FireworksHandler", () => {
730731

731732
describe("tool handling", () => {
732733
const systemPrompt = "You are a helpful assistant."
733-
const messages: Anthropic.Messages.MessageParam[] = [
734+
const messages: RooMessage[] = [
734735
{
735736
role: "user",
736737
content: [{ type: "text" as const, text: "Hello!" }],

‎src/api/providers/__tests__/gemini.spec.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
// npx vitest run src/api/providers/__tests__/gemini.spec.ts
23

34
import { NoOutputGeneratedError } from "ai"
@@ -102,7 +103,7 @@ describe("GeminiHandler", () => {
102103
})
103104

104105
describe("createMessage", () => {
105-
const mockMessages: Anthropic.Messages.MessageParam[] = [
106+
const mockMessages: RooMessage[] = [
106107
{
107108
role: "user",
108109
content: "Hello",
@@ -377,7 +378,7 @@ describe("GeminiHandler", () => {
377378
})
378379

379380
describe("error telemetry", () => {
380-
const mockMessages: Anthropic.Messages.MessageParam[] = [
381+
const mockMessages: RooMessage[] = [
381382
{
382383
role: "user",
383384
content: "Hello",

‎src/api/providers/__tests__/lite-llm.spec.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type { RooMessage } from "../../../core/task-persistence/rooMessage"
12
const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({
23
mockStreamText: vi.fn(),
34
mockGenerateText: vi.fn(),
@@ -257,7 +258,7 @@ describe("LiteLLMHandler", () => {
257258
})
258259

259260
const systemPrompt = "You are a helpful assistant"
260-
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
261+
const messages: RooMessage[] = [{ role: "user", content: "Hello" }]
261262

262263
const generator = handler.createMessage(systemPrompt, messages)
263264
for await (const _chunk of generator) {

0 commit comments

Comments
 (0)