Skip to content

Commit c2978f5

Browse files
charlesliencharleslien
committed
inject requestToolCall
1 parent 5240d28 commit c2978f5

26 files changed

+626
-394
lines changed

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

Lines changed: 31 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
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 {
3+
TEST_AGENT_RUNTIME_IMPL,
4+
TEST_AGENT_RUNTIME_SCOPED_IMPL,
5+
} from '@codebuff/common/testing/impl/agent-runtime'
36
import { getInitialSessionState } from '@codebuff/common/types/session-state'
47
import {
58
spyOn,
@@ -17,7 +20,10 @@ import * as agentRegistry from '../templates/agent-registry'
1720
import * as websocketAction from '../websockets/websocket-action'
1821

1922
import type { AgentTemplate } from '../templates/types'
20-
import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
23+
import type {
24+
AgentRuntimeDeps,
25+
AgentRuntimeScopedDeps,
26+
} from '@codebuff/common/types/contracts/agent-runtime'
2127
import type { ProjectFileContext } from '@codebuff/common/util/file'
2228
import type { WebSocket } from 'ws'
2329

@@ -100,9 +106,11 @@ describe('Cost Aggregation Integration Tests', () => {
100106
let mockLocalAgentTemplates: Record<string, any>
101107
let mockWebSocket: MockWebSocket
102108
let agentRuntimeImpl: AgentRuntimeDeps
109+
let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
103110

104111
beforeEach(async () => {
105112
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
113+
agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
106114
mockWebSocket = new MockWebSocket()
107115

108116
// Setup mock agent templates
@@ -179,32 +187,30 @@ describe('Cost Aggregation Integration Tests', () => {
179187
}
180188

181189
// Mock tool call execution
182-
spyOn(websocketAction, 'requestToolCall').mockImplementation(
183-
async (ws, userInputId, toolName, input) => {
184-
if (toolName === 'write_file') {
185-
return {
186-
output: [
187-
{
188-
type: 'json',
189-
value: {
190-
message: `File ${input.path} created successfully`,
191-
},
192-
},
193-
],
194-
}
195-
}
190+
agentRuntimeScopedImpl.requestToolCall = async ({ toolName, input }) => {
191+
if (toolName === 'write_file') {
196192
return {
197193
output: [
198194
{
199195
type: 'json',
200196
value: {
201-
message: 'Tool executed successfully',
197+
message: `File ${input.path} created successfully`,
202198
},
203199
},
204200
],
205201
}
206-
},
207-
)
202+
}
203+
return {
204+
output: [
205+
{
206+
type: 'json',
207+
value: {
208+
message: 'Tool executed successfully',
209+
},
210+
},
211+
],
212+
}
213+
}
208214

209215
// Mock file reading
210216
spyOn(websocketAction, 'requestFiles').mockImplementation(
@@ -251,6 +257,7 @@ describe('Cost Aggregation Integration Tests', () => {
251257

252258
const result = await mainPrompt({
253259
...agentRuntimeImpl,
260+
...agentRuntimeScopedImpl,
254261
ws: mockWebSocket as unknown as WebSocket,
255262
action,
256263
userId: TEST_USER_ID,
@@ -286,6 +293,7 @@ describe('Cost Aggregation Integration Tests', () => {
286293
// Call through websocket action handler to test full integration
287294
await websocketAction.callMainPrompt({
288295
...agentRuntimeImpl,
296+
...agentRuntimeScopedImpl,
289297
ws: mockWebSocket as unknown as WebSocket,
290298
action,
291299
userId: TEST_USER_ID,
@@ -354,6 +362,7 @@ describe('Cost Aggregation Integration Tests', () => {
354362

355363
const result = await mainPrompt({
356364
...agentRuntimeImpl,
365+
...agentRuntimeScopedImpl,
357366
ws: mockWebSocket as unknown as WebSocket,
358367
action,
359368
userId: TEST_USER_ID,
@@ -410,6 +419,7 @@ describe('Cost Aggregation Integration Tests', () => {
410419
try {
411420
result = await mainPrompt({
412421
...agentRuntimeImpl,
422+
...agentRuntimeScopedImpl,
413423
ws: mockWebSocket as unknown as WebSocket,
414424
action,
415425
userId: TEST_USER_ID,
@@ -459,6 +469,7 @@ describe('Cost Aggregation Integration Tests', () => {
459469

460470
await mainPrompt({
461471
...agentRuntimeImpl,
472+
...agentRuntimeScopedImpl,
462473
ws: mockWebSocket as unknown as WebSocket,
463474
action,
464475
userId: TEST_USER_ID,
@@ -499,6 +510,7 @@ describe('Cost Aggregation Integration Tests', () => {
499510
// Call through websocket action to test server-side reset
500511
await websocketAction.callMainPrompt({
501512
...agentRuntimeImpl,
513+
...agentRuntimeScopedImpl,
502514
ws: mockWebSocket as unknown as WebSocket,
503515
action,
504516
userId: TEST_USER_ID,

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
1+
import {
2+
TEST_AGENT_RUNTIME_IMPL,
3+
TEST_AGENT_RUNTIME_SCOPED_IMPL,
4+
} from '@codebuff/common/testing/impl/agent-runtime'
25
import {
36
getInitialAgentState,
47
getInitialSessionState,
@@ -181,6 +184,7 @@ describe('Cost Aggregation System', () => {
181184

182185
const result = handleSpawnAgents({
183186
...TEST_AGENT_RUNTIME_IMPL,
187+
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
184188
previousToolCallFinished: Promise.resolve(),
185189
toolCall: mockToolCall,
186190
fileContext: mockFileContext,
@@ -261,6 +265,7 @@ describe('Cost Aggregation System', () => {
261265

262266
const result = handleSpawnAgents({
263267
...TEST_AGENT_RUNTIME_IMPL,
268+
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
264269
previousToolCallFinished: Promise.resolve(),
265270
toolCall: mockToolCall,
266271
fileContext: mockFileContext,
@@ -418,6 +423,7 @@ describe('Cost Aggregation System', () => {
418423

419424
const result = handleSpawnAgents({
420425
...TEST_AGENT_RUNTIME_IMPL,
426+
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
421427
previousToolCallFinished: Promise.resolve(),
422428
toolCall: mockToolCall,
423429
fileContext: mockFileContext,

‎backend/src/__tests__/loop-agent-steps.test.ts‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import * as analytics from '@codebuff/common/analytics'
22
import db from '@codebuff/common/db'
33
import { TEST_USER_ID } from '@codebuff/common/old-constants'
4-
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
4+
import {
5+
TEST_AGENT_RUNTIME_IMPL,
6+
TEST_AGENT_RUNTIME_SCOPED_IMPL,
7+
} from '@codebuff/common/testing/impl/agent-runtime'
58
import {
69
clearMockedModules,
710
mockModule,
@@ -29,17 +32,20 @@ import { mockFileContext, MockWebSocket } from './test-utils'
2932
import type { getAgentTemplate } from '../templates/agent-registry'
3033
import type { AgentTemplate } from '../templates/types'
3134
import type { StepGenerator } from '@codebuff/common/types/agent-template'
32-
import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
35+
import type {
36+
AgentRuntimeDeps,
37+
AgentRuntimeScopedDeps,
38+
} from '@codebuff/common/types/contracts/agent-runtime'
3339
import type { ParamsOf } from '@codebuff/common/types/function-params'
3440
import type { AgentState } from '@codebuff/common/types/session-state'
3541
import type { WebSocket } from 'ws'
3642

37-
let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL }
38-
3943
describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => {
4044
let mockTemplate: AgentTemplate
4145
let mockAgentState: AgentState
4246
let llmCallCount: number
47+
let agentRuntimeImpl: AgentRuntimeDeps
48+
let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
4349

4450
const runLoopAgentStepsWithContext = async (
4551
options: ParamsOf<typeof loopAgentSteps>,
@@ -96,6 +102,9 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
96102
})
97103

98104
beforeEach(() => {
105+
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
106+
agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
107+
99108
llmCallCount = 0
100109

101110
// Setup spies for database operations
@@ -200,6 +209,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
200209

201210
const result = await runLoopAgentStepsWithContext({
202211
...agentRuntimeImpl,
212+
...agentRuntimeScopedImpl,
203213
ws: new MockWebSocket() as unknown as WebSocket,
204214
userInputId: 'test-user-input',
205215
agentType: 'test-agent',
@@ -246,6 +256,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
246256

247257
const result = await runLoopAgentStepsWithContext({
248258
...agentRuntimeImpl,
259+
...agentRuntimeScopedImpl,
249260
ws: new MockWebSocket() as unknown as WebSocket,
250261
userInputId: 'test-user-input',
251262
agentType: 'test-agent',
@@ -294,6 +305,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
294305

295306
const result = await runLoopAgentStepsWithContext({
296307
...agentRuntimeImpl,
308+
...agentRuntimeScopedImpl,
297309
ws: new MockWebSocket() as unknown as WebSocket,
298310
userInputId: 'test-user-input',
299311
agentType: 'test-agent',
@@ -341,6 +353,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
341353

342354
const result = await runLoopAgentStepsWithContext({
343355
...agentRuntimeImpl,
356+
...agentRuntimeScopedImpl,
344357
ws: new MockWebSocket() as unknown as WebSocket,
345358
userInputId: 'test-user-input',
346359
agentType: 'test-agent',
@@ -381,6 +394,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
381394

382395
const result = await runLoopAgentStepsWithContext({
383396
...agentRuntimeImpl,
397+
...agentRuntimeScopedImpl,
384398
ws: new MockWebSocket() as unknown as WebSocket,
385399
userInputId: 'test-user-input',
386400
agentType: 'test-agent',
@@ -413,6 +427,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
413427

414428
const result = await runLoopAgentStepsWithContext({
415429
...agentRuntimeImpl,
430+
...agentRuntimeScopedImpl,
416431
ws: new MockWebSocket() as unknown as WebSocket,
417432
userInputId: 'test-user-input',
418433
agentType: 'test-agent',
@@ -447,6 +462,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
447462

448463
const result = await runLoopAgentStepsWithContext({
449464
...agentRuntimeImpl,
465+
...agentRuntimeScopedImpl,
450466
ws: new MockWebSocket() as unknown as WebSocket,
451467
userInputId: 'test-user-input',
452468
agentType: 'test-agent',
@@ -498,6 +514,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
498514

499515
const result = await runLoopAgentStepsWithContext({
500516
...agentRuntimeImpl,
517+
...agentRuntimeScopedImpl,
501518
ws: new MockWebSocket() as unknown as WebSocket,
502519
userInputId: 'test-user-input',
503520
agentType: 'test-agent',
@@ -547,6 +564,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
547564

548565
await runLoopAgentStepsWithContext({
549566
...agentRuntimeImpl,
567+
...agentRuntimeScopedImpl,
550568
ws: new MockWebSocket() as unknown as WebSocket,
551569
userInputId: 'test-user-input',
552570
agentType: 'test-agent',
@@ -634,6 +652,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
634652

635653
const result = await runLoopAgentStepsWithContext({
636654
...agentRuntimeImpl,
655+
...agentRuntimeScopedImpl,
637656
ws: new MockWebSocket() as unknown as WebSocket,
638657
userInputId: 'test-user-input',
639658
agentType: 'test-agent',
@@ -707,6 +726,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
707726

708727
const result = await runLoopAgentStepsWithContext({
709728
...agentRuntimeImpl,
729+
...agentRuntimeScopedImpl,
710730
ws: new MockWebSocket() as unknown as WebSocket,
711731
userInputId: 'test-user-input',
712732
agentType: 'test-agent',
@@ -753,6 +773,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
753773

754774
const result = await runLoopAgentStepsWithContext({
755775
...agentRuntimeImpl,
776+
...agentRuntimeScopedImpl,
756777
ws: new MockWebSocket() as unknown as WebSocket,
757778
userInputId: 'test-user-input',
758779
agentType: 'test-agent',
@@ -821,6 +842,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
821842

822843
const result = await runLoopAgentStepsWithContext({
823844
...agentRuntimeImpl,
845+
...agentRuntimeScopedImpl,
824846
ws: new MockWebSocket() as unknown as WebSocket,
825847
userInputId: 'test-user-input',
826848
agentType: 'test-agent',

0 commit comments

Comments
 (0)