Skip to content

Commit 8bb7715

Browse files
charlesliencharleslien
committed
inject sendAction
1 parent 3bfff7a commit 8bb7715

Some content is hidden

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

56 files changed

+555
-696
lines changed

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

Lines changed: 17 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,14 @@ import * as agentRegistry from '../templates/agent-registry'
2020
import * as websocketAction from '../websockets/websocket-action'
2121

2222
import type { AgentTemplate } from '../templates/types'
23+
import type { ServerAction } from '@codebuff/common/actions'
2324
import type {
2425
AgentRuntimeDeps,
2526
AgentRuntimeScopedDeps,
2627
} from '@codebuff/common/types/contracts/agent-runtime'
28+
import type { SendActionFn } from '@codebuff/common/types/contracts/client'
2729
import type { ProjectFileContext } from '@codebuff/common/util/file'
28-
import type { WebSocket } from 'ws'
30+
import type { Mock } from 'bun:test'
2931

3032
const mockFileContext: ProjectFileContext = {
3133
projectRoot: '/test',
@@ -84,34 +86,17 @@ const mockFileContext: ProjectFileContext = {
8486
},
8587
}
8688

87-
class MockWebSocket {
88-
sentActions: any[] = []
89-
90-
send(msg: string) {
91-
// Capture sent messages for verification
92-
try {
93-
const parsed = JSON.parse(msg)
94-
if (parsed.type === 'action') {
95-
this.sentActions.push(parsed.data)
96-
}
97-
} catch {}
98-
}
99-
100-
close() {}
101-
on(event: string, listener: (...args: any[]) => void) {}
102-
removeListener(event: string, listener: (...args: any[]) => void) {}
103-
}
104-
10589
describe('Cost Aggregation Integration Tests', () => {
10690
let mockLocalAgentTemplates: Record<string, any>
107-
let mockWebSocket: MockWebSocket
10891
let agentRuntimeImpl: AgentRuntimeDeps
10992
let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
11093

11194
beforeEach(async () => {
11295
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
113-
agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
114-
mockWebSocket = new MockWebSocket()
96+
agentRuntimeScopedImpl = {
97+
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
98+
sendAction: mock(() => {}),
99+
}
115100

116101
// Setup mock agent templates
117102
mockLocalAgentTemplates = {
@@ -258,7 +243,6 @@ describe('Cost Aggregation Integration Tests', () => {
258243
const result = await mainPrompt({
259244
...agentRuntimeImpl,
260245
...agentRuntimeScopedImpl,
261-
ws: mockWebSocket as unknown as WebSocket,
262246
action,
263247
userId: TEST_USER_ID,
264248
clientSessionId: 'test-session',
@@ -294,17 +278,18 @@ describe('Cost Aggregation Integration Tests', () => {
294278
await websocketAction.callMainPrompt({
295279
...agentRuntimeImpl,
296280
...agentRuntimeScopedImpl,
297-
ws: mockWebSocket as unknown as WebSocket,
298281
action,
299282
userId: TEST_USER_ID,
300283
promptId: 'test-prompt',
301284
clientSessionId: 'test-session',
302285
})
303286

304287
// Verify final cost is included in prompt response
305-
const promptResponse = mockWebSocket.sentActions.find(
306-
(action) => action.type === 'prompt-response',
307-
)
288+
const promptResponse = (
289+
agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
290+
).mock.calls
291+
.map((call) => call[0].action)
292+
.find((action: ServerAction) => action.type === 'prompt-response') as any
308293

309294
expect(promptResponse).toBeDefined()
310295
expect(promptResponse.promptId).toBe('test-prompt')
@@ -363,7 +348,6 @@ describe('Cost Aggregation Integration Tests', () => {
363348
const result = await mainPrompt({
364349
...agentRuntimeImpl,
365350
...agentRuntimeScopedImpl,
366-
ws: mockWebSocket as unknown as WebSocket,
367351
action,
368352
userId: TEST_USER_ID,
369353
clientSessionId: 'test-session',
@@ -420,7 +404,6 @@ describe('Cost Aggregation Integration Tests', () => {
420404
result = await mainPrompt({
421405
...agentRuntimeImpl,
422406
...agentRuntimeScopedImpl,
423-
ws: mockWebSocket as unknown as WebSocket,
424407
action,
425408
userId: TEST_USER_ID,
426409
clientSessionId: 'test-session',
@@ -470,7 +453,6 @@ describe('Cost Aggregation Integration Tests', () => {
470453
await mainPrompt({
471454
...agentRuntimeImpl,
472455
...agentRuntimeScopedImpl,
473-
ws: mockWebSocket as unknown as WebSocket,
474456
action,
475457
userId: TEST_USER_ID,
476458
clientSessionId: 'test-session',
@@ -511,17 +493,18 @@ describe('Cost Aggregation Integration Tests', () => {
511493
await websocketAction.callMainPrompt({
512494
...agentRuntimeImpl,
513495
...agentRuntimeScopedImpl,
514-
ws: mockWebSocket as unknown as WebSocket,
515496
action,
516497
userId: TEST_USER_ID,
517498
promptId: 'test-prompt',
518499
clientSessionId: 'test-session',
519500
})
520501

521502
// Server should have reset the malicious value and calculated correct cost
522-
const promptResponse = mockWebSocket.sentActions.find(
523-
(action) => action.type === 'prompt-response',
524-
)
503+
const promptResponse = (
504+
agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
505+
).mock.calls
506+
.map((call) => call[0].action)
507+
.find((action) => action.type === 'prompt-response') as any
525508

526509
expect(promptResponse).toBeDefined()
527510
expect(promptResponse.sessionState.mainAgentState.creditsUsed).toBeLessThan(

‎backend/src/__tests__/fast-rewrite.test.ts‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import path from 'path'
22

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,
@@ -11,11 +14,15 @@ import { createPatch } from 'diff'
1114

1215
import { rewriteWithOpenAI } from '../fast-rewrite'
1316

14-
import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
15-
16-
let agentRuntimeImpl: AgentRuntimeDeps
17+
import type {
18+
AgentRuntimeDeps,
19+
AgentRuntimeScopedDeps,
20+
} from '@codebuff/common/types/contracts/agent-runtime'
1721

1822
describe.skip('rewriteWithOpenAI', () => {
23+
let agentRuntimeImpl: AgentRuntimeDeps
24+
let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
25+
1926
beforeAll(() => {
2027
// Mock database interactions
2128
mockModule('pg-pool', () => ({
@@ -40,6 +47,7 @@ describe.skip('rewriteWithOpenAI', () => {
4047

4148
beforeEach(() => {
4249
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
50+
agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
4351
})
4452

4553
afterAll(() => {
@@ -54,6 +62,7 @@ describe.skip('rewriteWithOpenAI', () => {
5462

5563
const result = await rewriteWithOpenAI({
5664
...agentRuntimeImpl,
65+
...agentRuntimeScopedImpl,
5766
oldContent: originalContent,
5867
editSnippet,
5968
clientSessionId: 'clientSessionId',

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

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import { z } from 'zod/v4'
2727
import { withAppContext } from '../context/app-context'
2828
import { loopAgentSteps } from '../run-agent-step'
2929
import { clearAgentGeneratorCache } from '../run-programmatic-step'
30-
import { mockFileContext, MockWebSocket } from './test-utils'
30+
import { mockFileContext } from './test-utils'
3131

3232
import type { getAgentTemplate } from '../templates/agent-registry'
3333
import type { AgentTemplate } from '../templates/types'
@@ -38,7 +38,6 @@ import type {
3838
} from '@codebuff/common/types/contracts/agent-runtime'
3939
import type { ParamsOf } from '@codebuff/common/types/function-params'
4040
import type { AgentState } from '@codebuff/common/types/session-state'
41-
import type { WebSocket } from 'ws'
4241

4342
describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => {
4443
let mockTemplate: AgentTemplate
@@ -103,7 +102,10 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
103102

104103
beforeEach(() => {
105104
agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
106-
agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
105+
agentRuntimeScopedImpl = {
106+
...TEST_AGENT_RUNTIME_SCOPED_IMPL,
107+
sendAction: () => {},
108+
}
107109

108110
llmCallCount = 0
109111

@@ -210,7 +212,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
210212
const result = await runLoopAgentStepsWithContext({
211213
...agentRuntimeImpl,
212214
...agentRuntimeScopedImpl,
213-
ws: new MockWebSocket() as unknown as WebSocket,
214215
userInputId: 'test-user-input',
215216
agentType: 'test-agent',
216217
agentState: mockAgentState,
@@ -257,7 +258,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
257258
const result = await runLoopAgentStepsWithContext({
258259
...agentRuntimeImpl,
259260
...agentRuntimeScopedImpl,
260-
ws: new MockWebSocket() as unknown as WebSocket,
261261
userInputId: 'test-user-input',
262262
agentType: 'test-agent',
263263
agentState: mockAgentState,
@@ -306,7 +306,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
306306
const result = await runLoopAgentStepsWithContext({
307307
...agentRuntimeImpl,
308308
...agentRuntimeScopedImpl,
309-
ws: new MockWebSocket() as unknown as WebSocket,
310309
userInputId: 'test-user-input',
311310
agentType: 'test-agent',
312311
agentState: mockAgentState,
@@ -354,7 +353,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
354353
const result = await runLoopAgentStepsWithContext({
355354
...agentRuntimeImpl,
356355
...agentRuntimeScopedImpl,
357-
ws: new MockWebSocket() as unknown as WebSocket,
358356
userInputId: 'test-user-input',
359357
agentType: 'test-agent',
360358
agentState: mockAgentState,
@@ -395,7 +393,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
395393
const result = await runLoopAgentStepsWithContext({
396394
...agentRuntimeImpl,
397395
...agentRuntimeScopedImpl,
398-
ws: new MockWebSocket() as unknown as WebSocket,
399396
userInputId: 'test-user-input',
400397
agentType: 'test-agent',
401398
agentState: mockAgentState,
@@ -428,7 +425,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
428425
const result = await runLoopAgentStepsWithContext({
429426
...agentRuntimeImpl,
430427
...agentRuntimeScopedImpl,
431-
ws: new MockWebSocket() as unknown as WebSocket,
432428
userInputId: 'test-user-input',
433429
agentType: 'test-agent',
434430
agentState: mockAgentState,
@@ -463,7 +459,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
463459
const result = await runLoopAgentStepsWithContext({
464460
...agentRuntimeImpl,
465461
...agentRuntimeScopedImpl,
466-
ws: new MockWebSocket() as unknown as WebSocket,
467462
userInputId: 'test-user-input',
468463
agentType: 'test-agent',
469464
agentState: mockAgentState,
@@ -515,7 +510,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
515510
const result = await runLoopAgentStepsWithContext({
516511
...agentRuntimeImpl,
517512
...agentRuntimeScopedImpl,
518-
ws: new MockWebSocket() as unknown as WebSocket,
519513
userInputId: 'test-user-input',
520514
agentType: 'test-agent',
521515
agentState: mockAgentState,
@@ -565,7 +559,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
565559
await runLoopAgentStepsWithContext({
566560
...agentRuntimeImpl,
567561
...agentRuntimeScopedImpl,
568-
ws: new MockWebSocket() as unknown as WebSocket,
569562
userInputId: 'test-user-input',
570563
agentType: 'test-agent',
571564
agentState: mockAgentState,
@@ -653,7 +646,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
653646
const result = await runLoopAgentStepsWithContext({
654647
...agentRuntimeImpl,
655648
...agentRuntimeScopedImpl,
656-
ws: new MockWebSocket() as unknown as WebSocket,
657649
userInputId: 'test-user-input',
658650
agentType: 'test-agent',
659651
agentState: mockAgentState,
@@ -727,7 +719,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
727719
const result = await runLoopAgentStepsWithContext({
728720
...agentRuntimeImpl,
729721
...agentRuntimeScopedImpl,
730-
ws: new MockWebSocket() as unknown as WebSocket,
731722
userInputId: 'test-user-input',
732723
agentType: 'test-agent',
733724
agentState: mockAgentState,
@@ -774,7 +765,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
774765
const result = await runLoopAgentStepsWithContext({
775766
...agentRuntimeImpl,
776767
...agentRuntimeScopedImpl,
777-
ws: new MockWebSocket() as unknown as WebSocket,
778768
userInputId: 'test-user-input',
779769
agentType: 'test-agent',
780770
agentState: mockAgentState,
@@ -843,7 +833,6 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
843833
const result = await runLoopAgentStepsWithContext({
844834
...agentRuntimeImpl,
845835
...agentRuntimeScopedImpl,
846-
ws: new MockWebSocket() as unknown as WebSocket,
847836
userInputId: 'test-user-input',
848837
agentType: 'test-agent',
849838
agentState: mockAgentState,

‎backend/src/__tests__/main-prompt.integration.test.ts‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,9 @@ import type { RequestToolCallFn } from '@codebuff/common/types/contracts/client'
2828
import type { ParamsOf } from '@codebuff/common/types/function-params'
2929
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
3030
import type { ProjectFileContext } from '@codebuff/common/util/file'
31-
import type { WebSocket } from 'ws'
3231

3332
// --- Shared Mocks & Helpers ---
3433

35-
class MockWebSocket {
36-
send(msg: string) {}
37-
close() {}
38-
on(event: string, listener: (...args: any[]) => void) {}
39-
removeListener(event: string, listener: (...args: any[]) => void) {}
40-
}
41-
4234
const mockFileContext: ProjectFileContext = {
4335
projectRoot: '/test',
4436
cwd: '/test',
@@ -388,7 +380,6 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
388380
const { output, sessionState: finalSessionState } = await mainPrompt({
389381
...agentRuntimeImpl,
390382
...agentRuntimeScopedImpl,
391-
ws: new MockWebSocket() as unknown as WebSocket,
392383
action,
393384
userId: TEST_USER_ID,
394385
clientSessionId: 'test-session-delete-function-integration',
@@ -471,7 +462,6 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
471462
await mainPrompt({
472463
...agentRuntimeImpl,
473464
...agentRuntimeScopedImpl,
474-
ws: new MockWebSocket() as unknown as WebSocket,
475465
action,
476466
userId: TEST_USER_ID,
477467
clientSessionId: 'test-session-delete-function-integration',

0 commit comments

Comments
 (0)