Skip to content

Commit 518e3db

Browse files
Refactor utility functions and middleware to pass logger as parameter
Refactored the following files to accept logger as a parameter instead of importing it directly: - util/simplify-tool-results.ts: simplifyTerminalCommandResults now takes logger param - util/messages.ts: All message utility functions now accept logger param - util/check-auth.ts: checkAuth and checkAdmin refactored with logger param - util/quickjs-sandbox.ts: dispose methods now require logger param - Updated all callers in production code and test files - Updated plan documentation with clearer instructions for test logger naming 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent 8387a3e commit 518e3db

14 files changed

+270
-204
lines changed

backend/src/__tests__/run-programmatic-step.test.ts

Lines changed: 57 additions & 53 deletions
Large diffs are not rendered by default.

backend/src/__tests__/sandbox-generator.test.ts

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,3 @@
1-
import {
2-
clearMockedModules,
3-
mockModule,
4-
} from '@codebuff/common/testing/mock-modules'
51
import {
62
getInitialAgentState,
73
type AgentState,
@@ -18,8 +14,16 @@ import * as requestContext from '../websockets/request-context'
1814
import * as websocketAction from '../websockets/websocket-action'
1915

2016
import type { AgentTemplate } from '../templates/types'
17+
import type { Logger } from '@codebuff/types/logger'
2118
import type { WebSocket } from 'ws'
2219

20+
const logger: Logger = {
21+
debug: () => {},
22+
error: () => {},
23+
info: () => {},
24+
warn: () => {},
25+
}
26+
2327
describe('QuickJS Sandbox Generator', () => {
2428
let mockAgentState: AgentState
2529
let mockParams: any
@@ -41,17 +45,6 @@ describe('QuickJS Sandbox Generator', () => {
4145
'mock-uuid-0000-0000-0000-000000000000' as `${string}-${string}-${string}-${string}-${string}`,
4246
)
4347

44-
// Mock logger
45-
mockModule('@codebuff/backend/util/logger', () => ({
46-
logger: {
47-
debug: () => {},
48-
error: () => {},
49-
info: () => {},
50-
warn: () => {},
51-
},
52-
withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
53-
}))
54-
5548
// Reuse common test data structure
5649
mockAgentState = {
5750
...getInitialAgentState(),
@@ -85,6 +78,7 @@ describe('QuickJS Sandbox Generator', () => {
8578

8679
// Common params structure
8780
mockParams = {
81+
agentState: mockAgentState,
8882
template: mockTemplate,
8983
prompt: 'Test prompt',
9084
params: { testParam: 'value' },
@@ -101,12 +95,12 @@ describe('QuickJS Sandbox Generator', () => {
10195
localAgentTemplates: {},
10296
stepsComplete: false,
10397
stepNumber: 1,
98+
logger,
10499
}
105100
})
106101

107102
afterEach(() => {
108103
clearAgentGeneratorCache()
109-
clearMockedModules()
110104
})
111105

112106
test('should execute string-based generator in QuickJS sandbox', async () => {
@@ -126,7 +120,7 @@ describe('QuickJS Sandbox Generator', () => {
126120
mockParams.template = mockTemplate
127121
mockParams.localAgentTemplates = { 'test-vm-agent': mockTemplate }
128122

129-
const result = await runProgrammaticStep(mockAgentState, mockParams)
123+
const result = await runProgrammaticStep(mockParams)
130124

131125
expect(result.agentState.output).toEqual({
132126
message: 'Hello from QuickJS sandbox!',
@@ -155,7 +149,7 @@ describe('QuickJS Sandbox Generator', () => {
155149
mockParams.params = {}
156150
mockParams.localAgentTemplates = { 'test-vm-agent-error': mockTemplate }
157151

158-
const result = await runProgrammaticStep(mockAgentState, mockParams)
152+
const result = await runProgrammaticStep(mockParams)
159153

160154
expect(result.endTurn).toBe(true)
161155
expect(result.agentState.output?.error).toContain(

backend/src/api/usage.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ async function usageHandler(
4848
fingerprintId,
4949
authToken,
5050
clientSessionId,
51+
logger,
5152
})
5253
if (authResult) {
5354
const errorMessage =

backend/src/find-files/request-files-prompt.ts

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -344,16 +344,17 @@ async function getRelevantFiles(params: {
344344
logger,
345345
} = params
346346
const bufferTokens = 100_000
347-
const messagesWithPrompt = getMessagesSubset(
348-
[
347+
const messagesWithPrompt = getMessagesSubset({
348+
messages: [
349349
...messages,
350350
{
351351
role: 'user' as const,
352352
content: userPrompt,
353353
},
354354
],
355-
bufferTokens,
356-
)
355+
otherTokens: bufferTokens,
356+
logger,
357+
})
357358
const start = performance.now()
358359
let codebuffMessages = messagesWithSystem({ messages: messagesWithPrompt, system })
359360

@@ -435,16 +436,17 @@ async function getRelevantFilesForTraining(params: {
435436
logger,
436437
} = params
437438
const bufferTokens = 100_000
438-
const messagesWithPrompt = getMessagesSubset(
439-
[
439+
const messagesWithPrompt = getMessagesSubset({
440+
messages: [
440441
...messages,
441442
{
442443
role: 'user' as const,
443444
content: userPrompt,
444445
},
445446
],
446-
bufferTokens,
447-
)
447+
otherTokens: bufferTokens,
448+
logger,
449+
})
448450
const start = performance.now()
449451
let response = await promptAiSdk({
450452
messages: messagesWithSystem({ messages: messagesWithPrompt, system }),

backend/src/run-agent-step.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -574,7 +574,8 @@ export const loopAgentSteps = async (
574574
agentState: programmaticAgentState,
575575
endTurn,
576576
stepNumber,
577-
} = await runProgrammaticStep(currentAgentState, {
577+
} = await runProgrammaticStep({
578+
agentState: currentAgentState,
578579
userId,
579580
userInputId,
580581
clientSessionId,
@@ -585,10 +586,11 @@ export const loopAgentSteps = async (
585586
template: agentTemplate,
586587
localAgentTemplates,
587588
prompt: currentPrompt,
588-
params: currentParams,
589+
toolCallParams: currentParams,
589590
system,
590591
stepsComplete: shouldEndTurn,
591592
stepNumber: totalSteps,
593+
logger,
592594
})
593595
currentAgentState = programmaticAgentState
594596
totalSteps = stepNumber

backend/src/run-programmatic-step.ts

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import type {
2222
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
2323
import type { AgentState } from '@codebuff/common/types/session-state'
2424
import type { ProjectFileContext } from '@codebuff/common/util/file'
25+
import type { Logger } from '@codebuff/types/logger'
2526
import type { WebSocket } from 'ws'
2627

2728
// Global sandbox manager for QuickJS contexts
@@ -38,16 +39,33 @@ export function clearAgentGeneratorCache() {
3839
}
3940
runIdToStepAll.clear()
4041
// Clean up QuickJS sandboxes
41-
sandboxManager.dispose()
42+
sandboxManager.dispose({ logger })
4243
}
4344

4445
// Function to handle programmatic agents
45-
export async function runProgrammaticStep(
46-
agentState: AgentState,
47-
{
46+
export async function runProgrammaticStep(params: {
47+
agentState: AgentState
48+
template: AgentTemplate
49+
prompt: string | undefined
50+
toolCallParams: Record<string, any> | undefined
51+
system: string | undefined
52+
userId: string | undefined
53+
userInputId: string
54+
clientSessionId: string
55+
fingerprintId: string
56+
onResponseChunk: (chunk: string | PrintModeEvent) => void
57+
fileContext: ProjectFileContext
58+
ws: WebSocket
59+
localAgentTemplates: Record<string, AgentTemplate>
60+
stepsComplete: boolean
61+
stepNumber: number
62+
logger: Logger
63+
}): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
64+
const {
65+
agentState,
4866
template,
4967
prompt,
50-
params,
68+
toolCallParams,
5169
system,
5270
userId,
5371
userInputId,
@@ -58,24 +76,10 @@ export async function runProgrammaticStep(
5876
ws,
5977
localAgentTemplates,
6078
stepsComplete,
61-
stepNumber,
62-
}: {
63-
template: AgentTemplate
64-
prompt: string | undefined
65-
params: Record<string, any> | undefined
66-
system: string | undefined
67-
userId: string | undefined
68-
userInputId: string
69-
clientSessionId: string
70-
fingerprintId: string
71-
onResponseChunk: (chunk: string | PrintModeEvent) => void
72-
fileContext: ProjectFileContext
73-
ws: WebSocket
74-
localAgentTemplates: Record<string, AgentTemplate>
75-
stepsComplete: boolean
76-
stepNumber: number
77-
},
78-
): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
79+
logger,
80+
} = params
81+
let { stepNumber } = params
82+
7983
if (!template.handleSteps) {
8084
throw new Error('No step handler found for agent template ' + template.id)
8185
}
@@ -119,11 +123,12 @@ export async function runProgrammaticStep(
119123
initialInput: {
120124
agentState,
121125
prompt,
122-
params,
126+
params: toolCallParams,
123127
logger: streamingLogger,
124128
},
125129
config: undefined, // config
126-
logger: streamingLogger, // pass the streaming logger instance for internal use
130+
sandboxLogger: streamingLogger, // pass the streaming logger instance for internal use
131+
logger,
127132
})
128133
} else {
129134
// Initialize native generator
@@ -363,7 +368,7 @@ export async function runProgrammaticStep(
363368
if (endTurn) {
364369
if (sandbox) {
365370
// Clean up QuickJS sandbox if execution is complete
366-
sandboxManager.removeSandbox({ runId: agentState.runId })
371+
sandboxManager.removeSandbox({ runId: agentState.runId, logger })
367372
}
368373
delete runIdToGenerator[agentState.runId]
369374
runIdToStepAll.delete(agentState.runId)

0 commit comments

Comments
 (0)