Skip to content

Commit 1eb069d

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 1eb069d

File tree

11 files changed

+146
-88
lines changed

11 files changed

+146
-88
lines changed

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-programmatic-step.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ export function clearAgentGeneratorCache() {
3838
}
3939
runIdToStepAll.clear()
4040
// Clean up QuickJS sandboxes
41-
sandboxManager.dispose()
41+
sandboxManager.dispose({ logger })
4242
}
4343

4444
// Function to handle programmatic agents
@@ -363,7 +363,7 @@ export async function runProgrammaticStep(
363363
if (endTurn) {
364364
if (sandbox) {
365365
// Clean up QuickJS sandbox if execution is complete
366-
sandboxManager.removeSandbox({ runId: agentState.runId })
366+
sandboxManager.removeSandbox({ runId: agentState.runId, logger })
367367
}
368368
delete runIdToGenerator[agentState.runId]
369369
runIdToStepAll.delete(agentState.runId)

backend/src/util/__tests__/messages.test.ts

Lines changed: 36 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import {
88
spyOn,
99
} from 'bun:test'
1010

11-
import { logger } from '../logger'
1211
import {
1312
trimMessagesToFitTokenLimit,
1413
messagesWithSystem,
@@ -40,6 +39,14 @@ describe('messagesWithSystem', () => {
4039
})
4140
})
4241

42+
// Mock logger for tests
43+
const logger = {
44+
debug: () => {},
45+
info: () => {},
46+
warn: () => {},
47+
error: () => {},
48+
}
49+
4350
describe('trimMessagesToFitTokenLimit', () => {
4451
beforeEach(() => {
4552
// Mock countTokensJson to just count characters
@@ -211,11 +218,12 @@ describe('trimMessagesToFitTokenLimit', () => {
211218
it('handles all features working together correctly', () => {
212219
const maxTotalTokens = 3000
213220
const systemTokens = 0
214-
const result = trimMessagesToFitTokenLimit(
215-
testMessages,
221+
const result = trimMessagesToFitTokenLimit({
222+
messages: testMessages,
216223
systemTokens,
217224
maxTotalTokens,
218-
)
225+
logger,
226+
})
219227

220228
// Should have replacement message for omitted content
221229
expect(result.length).toBeGreaterThan(0)
@@ -236,11 +244,12 @@ describe('trimMessagesToFitTokenLimit', () => {
236244
it('subtracts system tokens from total tokens', () => {
237245
const maxTotalTokens = 10_000
238246
const systemTokens = 7_000
239-
const result = trimMessagesToFitTokenLimit(
240-
testMessages,
247+
const result = trimMessagesToFitTokenLimit({
248+
messages: testMessages,
241249
systemTokens,
242250
maxTotalTokens,
243-
)
251+
logger,
252+
})
244253

245254
// Should have replacement message for omitted content
246255
expect(result.length).toBeGreaterThan(0)
@@ -261,11 +270,12 @@ describe('trimMessagesToFitTokenLimit', () => {
261270
it('does not simplify if under token limit', () => {
262271
const maxTotalTokens = 10_000
263272
const systemTokens = 100
264-
const result = trimMessagesToFitTokenLimit(
265-
testMessages,
273+
const result = trimMessagesToFitTokenLimit({
274+
messages: testMessages,
266275
systemTokens,
267276
maxTotalTokens,
268-
)
277+
logger,
278+
})
269279

270280
// All messages should be unchanged
271281
expect(result).toHaveLength(testMessages.length)
@@ -282,7 +292,7 @@ describe('trimMessagesToFitTokenLimit', () => {
282292
it('handles empty messages array', () => {
283293
const maxTotalTokens = 200
284294
const systemTokens = 100
285-
const result = trimMessagesToFitTokenLimit([], systemTokens, maxTotalTokens)
295+
const result = trimMessagesToFitTokenLimit({ messages: [], systemTokens, maxTotalTokens, logger })
286296

287297
expect(result).toEqual([])
288298
})
@@ -305,7 +315,7 @@ describe('trimMessagesToFitTokenLimit', () => {
305315
},
306316
]
307317

308-
const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
318+
const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
309319

310320
// Should contain the kept messages
311321
const keptMessages = result.filter(
@@ -335,7 +345,7 @@ describe('trimMessagesToFitTokenLimit', () => {
335345
},
336346
] as Message[]
337347

338-
const result = trimMessagesToFitTokenLimit(messages, 0, 10000)
348+
const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 10000, logger })
339349

340350
// Should be unchanged when under token limit
341351
expect(result).toHaveLength(2)
@@ -351,7 +361,7 @@ describe('trimMessagesToFitTokenLimit', () => {
351361
{ role: 'user', content: 'Keep this', keepDuringTruncation: true },
352362
]
353363

354-
const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
364+
const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
355365

356366
// Should only have one replacement message for consecutive removals
357367
const replacementMessages = result.filter(
@@ -381,7 +391,7 @@ describe('trimMessagesToFitTokenLimit', () => {
381391
{ role: 'user', content: 'C'.repeat(100) }, // Might be kept
382392
]
383393

384-
const result = trimMessagesToFitTokenLimit(messages, 0, 2000)
394+
const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 2000, logger })
385395

386396
// Should preserve the keepDuringTruncation message
387397
const keptMessage = result.find(
@@ -405,7 +415,7 @@ describe('trimMessagesToFitTokenLimit', () => {
405415
{ role: 'user', content: 'C'.repeat(800) }, // Large message to force truncation
406416
]
407417

408-
const result = trimMessagesToFitTokenLimit(messages, 0, 500)
418+
const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 500, logger })
409419

410420
// Should keep both marked messages
411421
const keptMessages = result.filter(
@@ -428,7 +438,7 @@ describe('trimMessagesToFitTokenLimit', () => {
428438

429439
describe('getPreviouslyReadFiles', () => {
430440
it('returns empty array when no messages provided', () => {
431-
const result = getPreviouslyReadFiles([])
441+
const result = getPreviouslyReadFiles({ messages: [], logger })
432442
expect(result).toEqual([])
433443
})
434444

@@ -449,7 +459,7 @@ describe('getPreviouslyReadFiles', () => {
449459
} satisfies CodebuffToolMessage<'write_file'>,
450460
]
451461

452-
const result = getPreviouslyReadFiles(messages)
462+
const result = getPreviouslyReadFiles({ messages, logger })
453463
expect(result).toEqual([])
454464
})
455465

@@ -481,7 +491,7 @@ describe('getPreviouslyReadFiles', () => {
481491
} satisfies CodebuffToolMessage<'read_files'>,
482492
]
483493

484-
const result = getPreviouslyReadFiles(messages)
494+
const result = getPreviouslyReadFiles({ messages, logger })
485495
expect(result).toEqual([
486496
{
487497
path: 'src/test.ts',
@@ -518,7 +528,7 @@ describe('getPreviouslyReadFiles', () => {
518528
} satisfies CodebuffToolMessage<'find_files'>,
519529
]
520530

521-
const result = getPreviouslyReadFiles(messages)
531+
const result = getPreviouslyReadFiles({ messages, logger })
522532
expect(result).toEqual([
523533
{
524534
path: 'components/Button.tsx',
@@ -573,7 +583,7 @@ describe('getPreviouslyReadFiles', () => {
573583
},
574584
]
575585

576-
const result = getPreviouslyReadFiles(messages)
586+
const result = getPreviouslyReadFiles({ messages, logger })
577587
expect(result).toEqual([
578588
{ path: 'file1.ts', content: 'content 1' },
579589
{ path: 'file2.ts', content: 'content 2' },
@@ -611,7 +621,7 @@ describe('getPreviouslyReadFiles', () => {
611621
} satisfies CodebuffToolMessage<'read_files'>,
612622
]
613623

614-
const result = getPreviouslyReadFiles(messages)
624+
const result = getPreviouslyReadFiles({ messages, logger })
615625
expect(result).toEqual([
616626
{ path: 'small-file.ts', content: 'small content' },
617627
{ path: 'another-small-file.ts', content: 'another small content' },
@@ -633,7 +643,7 @@ describe('getPreviouslyReadFiles', () => {
633643
},
634644
]
635645

636-
const result = getPreviouslyReadFiles(messages)
646+
const result = getPreviouslyReadFiles({ messages, logger })
637647
expect(result).toEqual([])
638648
expect(mockLoggerError).toHaveBeenCalled()
639649

@@ -660,7 +670,7 @@ describe('getPreviouslyReadFiles', () => {
660670
} satisfies CodebuffToolMessage<'find_files'>,
661671
]
662672

663-
const result = getPreviouslyReadFiles(messages)
673+
const result = getPreviouslyReadFiles({ messages, logger })
664674
expect(result).toEqual([])
665675
})
666676

@@ -690,7 +700,7 @@ describe('getPreviouslyReadFiles', () => {
690700
} satisfies CodebuffToolMessage<'read_files'>,
691701
]
692702

693-
const result = getPreviouslyReadFiles(messages)
703+
const result = getPreviouslyReadFiles({ messages, logger })
694704
expect(result).toEqual([{ path: 'test.ts', content: 'test content' }])
695705
})
696706

@@ -712,7 +722,7 @@ describe('getPreviouslyReadFiles', () => {
712722
} satisfies CodebuffToolMessage<'read_files'>,
713723
]
714724

715-
const result = getPreviouslyReadFiles(messages)
725+
const result = getPreviouslyReadFiles({ messages, logger })
716726
expect(result).toEqual([])
717727
})
718728
})

0 commit comments

Comments
 (0)