Skip to content

Commit d690688

Browse files
charlesliencharleslien
committed
temporarily revert ai sdk changes
1 parent 94251da commit d690688

26 files changed

+357
-316
lines changed

‎backend/src/__tests__/cost-aggregation.integration.test.ts‎

Lines changed: 79 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { TEST_USER_ID } from '@codebuff/common/old-constants'
2-
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
2+
import { testAgentRuntimeImpl } from '@codebuff/common/testing/impl/agent-runtime'
33
import { getInitialSessionState } from '@codebuff/common/types/session-state'
44
import {
55
spyOn,
@@ -12,12 +12,12 @@ import {
1212
} from 'bun:test'
1313

1414
import * as messageCostTracker from '../llm-apis/message-cost-tracker'
15+
import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
1516
import { mainPrompt } from '../main-prompt'
1617
import * as agentRegistry from '../templates/agent-registry'
1718
import * as websocketAction from '../websockets/websocket-action'
1819

1920
import type { AgentTemplate } from '../templates/types'
20-
import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
2121
import type { ProjectFileContext } from '@codebuff/common/util/file'
2222
import type { WebSocket } from 'ws'
2323

@@ -99,7 +99,6 @@ class MockWebSocket {
9999
describe('Cost Aggregation Integration Tests', () => {
100100
let mockLocalAgentTemplates: Record<string, any>
101101
let mockWebSocket: MockWebSocket
102-
let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
103102

104103
beforeEach(async () => {
105104
mockWebSocket = new MockWebSocket()
@@ -151,31 +150,33 @@ describe('Cost Aggregation Integration Tests', () => {
151150
// Mock LLM streaming
152151
let callCount = 0
153152
const creditHistory: number[] = []
154-
agentRuntimeImpl.promptAiSdkStream = async function* (options) {
155-
callCount++
156-
const credits = callCount === 1 ? 10 : 7 // Main agent vs subagent costs
157-
creditHistory.push(credits)
158-
159-
if (options.onCostCalculated) {
160-
await options.onCostCalculated(credits)
161-
}
162-
163-
// Simulate different responses based on call
164-
if (callCount === 1) {
165-
// Main agent spawns a subagent
166-
yield {
167-
type: 'text' as const,
168-
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write a simple hello world file"}]}\n</codebuff_tool_call>',
153+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(
154+
async function* (options) {
155+
callCount++
156+
const credits = callCount === 1 ? 10 : 7 // Main agent vs subagent costs
157+
creditHistory.push(credits)
158+
159+
if (options.onCostCalculated) {
160+
await options.onCostCalculated(credits)
169161
}
170-
} else {
171-
// Subagent writes a file
172-
yield {
173-
type: 'text' as const,
174-
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "hello.txt", "instructions": "Create hello world file", "content": "Hello, World!"}\n</codebuff_tool_call>',
162+
163+
// Simulate different responses based on call
164+
if (callCount === 1) {
165+
// Main agent spawns a subagent
166+
yield {
167+
type: 'text' as const,
168+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write a simple hello world file"}]}\n</codebuff_tool_call>',
169+
}
170+
} else {
171+
// Subagent writes a file
172+
yield {
173+
type: 'text' as const,
174+
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "hello.txt", "instructions": "Create hello world file", "content": "Hello, World!"}\n</codebuff_tool_call>',
175+
}
175176
}
176-
}
177-
return 'mock-message-id'
178-
}
177+
return 'mock-message-id'
178+
},
179+
)
179180

180181
// Mock tool call execution
181182
spyOn(websocketAction, 'requestToolCall').mockImplementation(
@@ -230,7 +231,6 @@ describe('Cost Aggregation Integration Tests', () => {
230231

231232
afterEach(() => {
232233
mock.restore()
233-
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
234234
})
235235

236236
it('should correctly aggregate costs across the entire main prompt flow', async () => {
@@ -250,7 +250,7 @@ describe('Cost Aggregation Integration Tests', () => {
250250
}
251251

252252
const result = await mainPrompt({
253-
...agentRuntimeImpl,
253+
...testAgentRuntimeImpl,
254254
ws: mockWebSocket as unknown as WebSocket,
255255
action,
256256
userId: TEST_USER_ID,
@@ -285,7 +285,7 @@ describe('Cost Aggregation Integration Tests', () => {
285285

286286
// Call through websocket action handler to test full integration
287287
await websocketAction.callMainPrompt({
288-
...agentRuntimeImpl,
288+
...testAgentRuntimeImpl,
289289
ws: mockWebSocket as unknown as WebSocket,
290290
action,
291291
userId: TEST_USER_ID,
@@ -308,35 +308,37 @@ describe('Cost Aggregation Integration Tests', () => {
308308
it('should handle multi-level subagent hierarchies correctly', async () => {
309309
// Mock a more complex scenario with nested subagents
310310
let callCount = 0
311-
agentRuntimeImpl.promptAiSdkStream = async function* (options) {
312-
callCount++
311+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(
312+
async function* (options) {
313+
callCount++
313314

314-
if (options.onCostCalculated) {
315-
await options.onCostCalculated(5) // Each call costs 5 credits
316-
}
317-
318-
if (callCount === 1) {
319-
// Main agent spawns first-level subagent
320-
yield {
321-
type: 'text' as const,
322-
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Create files"}]}\n</codebuff_tool_call>',
323-
}
324-
} else if (callCount === 2) {
325-
// First-level subagent spawns second-level subagent
326-
yield {
327-
type: 'text' as const,
328-
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write specific file"}]}\n</codebuff_tool_call>',
315+
if (options.onCostCalculated) {
316+
await options.onCostCalculated(5) // Each call costs 5 credits
329317
}
330-
} else {
331-
// Second-level subagent does actual work
332-
yield {
333-
type: 'text' as const,
334-
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "nested.txt", "instructions": "Create nested file", "content": "Nested content"}\n</codebuff_tool_call>',
318+
319+
if (callCount === 1) {
320+
// Main agent spawns first-level subagent
321+
yield {
322+
type: 'text' as const,
323+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Create files"}]}\n</codebuff_tool_call>',
324+
}
325+
} else if (callCount === 2) {
326+
// First-level subagent spawns second-level subagent
327+
yield {
328+
type: 'text' as const,
329+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write specific file"}]}\n</codebuff_tool_call>',
330+
}
331+
} else {
332+
// Second-level subagent does actual work
333+
yield {
334+
type: 'text' as const,
335+
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "nested.txt", "instructions": "Create nested file", "content": "Nested content"}\n</codebuff_tool_call>',
336+
}
335337
}
336-
}
337338

338-
return 'mock-message-id'
339-
}
339+
return 'mock-message-id'
340+
},
341+
)
340342

341343
const sessionState = getInitialSessionState(mockFileContext)
342344
sessionState.mainAgentState.stepsRemaining = 10
@@ -353,7 +355,7 @@ describe('Cost Aggregation Integration Tests', () => {
353355
}
354356

355357
const result = await mainPrompt({
356-
...agentRuntimeImpl,
358+
...testAgentRuntimeImpl,
357359
ws: mockWebSocket as unknown as WebSocket,
358360
action,
359361
userId: TEST_USER_ID,
@@ -371,27 +373,29 @@ describe('Cost Aggregation Integration Tests', () => {
371373
it('should maintain cost integrity when subagents fail', async () => {
372374
// Mock scenario where subagent fails after incurring partial costs
373375
let callCount = 0
374-
agentRuntimeImpl.promptAiSdkStream = async function* (options) {
375-
callCount++
376+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(
377+
async function* (options) {
378+
callCount++
376379

377-
if (options.onCostCalculated) {
378-
await options.onCostCalculated(6) // Each call costs 6 credits
379-
}
380+
if (options.onCostCalculated) {
381+
await options.onCostCalculated(6) // Each call costs 6 credits
382+
}
380383

381-
if (callCount === 1) {
382-
// Main agent spawns subagent
383-
yield {
384-
type: 'text' as const,
385-
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "This will fail"}]}\n</codebuff_tool_call>',
384+
if (callCount === 1) {
385+
// Main agent spawns subagent
386+
yield {
387+
type: 'text' as const,
388+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "This will fail"}]}\n</codebuff_tool_call>',
389+
}
390+
} else {
391+
// Subagent fails after incurring cost
392+
yield { type: 'text' as const, text: 'Some response' }
393+
throw new Error('Subagent execution failed')
386394
}
387-
} else {
388-
// Subagent fails after incurring cost
389-
yield { type: 'text' as const, text: 'Some response' }
390-
throw new Error('Subagent execution failed')
391-
}
392395

393-
return 'mock-message-id'
394-
}
396+
return 'mock-message-id'
397+
},
398+
)
395399

396400
const sessionState = getInitialSessionState(mockFileContext)
397401
sessionState.mainAgentState.agentType = 'base'
@@ -409,7 +413,7 @@ describe('Cost Aggregation Integration Tests', () => {
409413
let result
410414
try {
411415
result = await mainPrompt({
412-
...agentRuntimeImpl,
416+
...testAgentRuntimeImpl,
413417
ws: mockWebSocket as unknown as WebSocket,
414418
action,
415419
userId: TEST_USER_ID,
@@ -458,7 +462,7 @@ describe('Cost Aggregation Integration Tests', () => {
458462
}
459463

460464
await mainPrompt({
461-
...agentRuntimeImpl,
465+
...testAgentRuntimeImpl,
462466
ws: mockWebSocket as unknown as WebSocket,
463467
action,
464468
userId: TEST_USER_ID,
@@ -498,7 +502,7 @@ describe('Cost Aggregation Integration Tests', () => {
498502

499503
// Call through websocket action to test server-side reset
500504
await websocketAction.callMainPrompt({
501-
...agentRuntimeImpl,
505+
...testAgentRuntimeImpl,
502506
ws: mockWebSocket as unknown as WebSocket,
503507
action,
504508
userId: TEST_USER_ID,

‎backend/src/__tests__/cost-aggregation.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
1+
import { testAgentRuntimeImpl } from '@codebuff/common/testing/impl/agent-runtime'
22
import {
33
getInitialAgentState,
44
getInitialSessionState,
@@ -180,7 +180,7 @@ describe('Cost Aggregation System', () => {
180180
}
181181

182182
const result = handleSpawnAgents({
183-
...TEST_AGENT_RUNTIME_IMPL,
183+
...testAgentRuntimeImpl,
184184
previousToolCallFinished: Promise.resolve(),
185185
toolCall: mockToolCall,
186186
fileContext: mockFileContext,
@@ -260,7 +260,7 @@ describe('Cost Aggregation System', () => {
260260
}
261261

262262
const result = handleSpawnAgents({
263-
...TEST_AGENT_RUNTIME_IMPL,
263+
...testAgentRuntimeImpl,
264264
previousToolCallFinished: Promise.resolve(),
265265
toolCall: mockToolCall,
266266
fileContext: mockFileContext,
@@ -417,7 +417,7 @@ describe('Cost Aggregation System', () => {
417417
}
418418

419419
const result = handleSpawnAgents({
420-
...TEST_AGENT_RUNTIME_IMPL,
420+
...testAgentRuntimeImpl,
421421
previousToolCallFinished: Promise.resolve(),
422422
toolCall: mockToolCall,
423423
fileContext: mockFileContext,

‎backend/src/__tests__/generate-diffs-prompt.test.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,16 @@
1-
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
21
import { expect, describe, it } from 'bun:test'
32

43
import { parseAndGetDiffBlocksSingleFile } from '../generate-diffs-prompt'
54

5+
import type { Logger } from '@codebuff/common/types/contracts/logger'
6+
7+
const logger: Logger = {
8+
debug: () => {},
9+
info: () => {},
10+
warn: () => {},
11+
error: () => {},
12+
}
13+
614
describe('parseAndGetDiffBlocksSingleFile', () => {
715
it('should parse diff blocks with newline before closing marker', () => {
816
const oldContent = 'function test() {\n return true;\n}\n'
@@ -18,9 +26,9 @@ function test() {
1826
>>>>>>> REPLACE`
1927

2028
const result = parseAndGetDiffBlocksSingleFile({
21-
...TEST_AGENT_RUNTIME_IMPL,
2229
newContent,
2330
oldFileContent: oldContent,
31+
logger,
2432
})
2533
console.log(JSON.stringify({ result }))
2634

@@ -47,9 +55,9 @@ function test() {
4755
}>>>>>>> REPLACE`
4856

4957
const result = parseAndGetDiffBlocksSingleFile({
50-
...TEST_AGENT_RUNTIME_IMPL,
5158
newContent,
5259
oldFileContent: oldContent,
60+
logger,
5361
})
5462

5563
expect(result.diffBlocks.length).toBe(1)
@@ -100,9 +108,9 @@ function subtract(a, b) {
100108
>>>>>>> REPLACE`
101109

102110
const result = parseAndGetDiffBlocksSingleFile({
103-
...TEST_AGENT_RUNTIME_IMPL,
104111
newContent,
105112
oldFileContent: oldContent,
113+
logger,
106114
})
107115

108116
expect(result.diffBlocks.length).toBe(2)
@@ -128,9 +136,9 @@ function subtract(a, b) {
128136
>>>>>>> REPLACE`
129137

130138
const result = parseAndGetDiffBlocksSingleFile({
131-
...TEST_AGENT_RUNTIME_IMPL,
132139
newContent,
133140
oldFileContent: oldContent,
141+
logger,
134142
})
135143

136144
expect(result.diffBlocks.length).toBe(1)

0 commit comments

Comments
 (0)