Skip to content

Commit 0affa94

Browse files
charlesliencharleslien
committed
pass in logger to ai sdk
1 parent 194e307 commit 0affa94

16 files changed

+240
-130
lines changed

‎backend/src/__tests__/live-user-inputs.test.ts‎

Lines changed: 124 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -133,60 +133,80 @@ describe('live-user-inputs', () => {
133133
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
134134
setSessionConnected('session-1', true)
135135

136-
const isLive = checkLiveUserInput('user-1', 'input-123', 'session-1')
136+
const isLive = checkLiveUserInput({
137+
userId: 'user-1',
138+
userInputId: 'input-123',
139+
clientSessionId: 'session-1',
140+
})
137141
expect(isLive).toBe(true)
138142
})
139143

140144
it('should return true for user input with matching prefix', () => {
141145
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
142146
setSessionConnected('session-1', true)
143147

144-
const isLive = checkLiveUserInput(
145-
'user-1',
146-
'input-123-async-agent',
147-
'session-1',
148-
)
148+
const isLive = checkLiveUserInput({
149+
userId: 'user-1',
150+
userInputId: 'input-123-async-agent',
151+
clientSessionId: 'session-1',
152+
})
149153
expect(isLive).toBe(true)
150154
})
151155

152156
it('should return false for non-existent user', () => {
153157
setSessionConnected('session-1', true)
154158

155-
const isLive = checkLiveUserInput(
156-
'user-nonexistent',
157-
'input-123',
158-
'session-1',
159-
)
159+
const isLive = checkLiveUserInput({
160+
userId: 'user-nonexistent',
161+
userInputId: 'input-123',
162+
clientSessionId: 'session-1',
163+
})
160164
expect(isLive).toBe(false)
161165
})
162166

163167
it('should return false for undefined user', () => {
164168
setSessionConnected('session-1', true)
165169

166-
const isLive = checkLiveUserInput(undefined, 'input-123', 'session-1')
170+
const isLive = checkLiveUserInput({
171+
userId: undefined,
172+
userInputId: 'input-123',
173+
clientSessionId: 'session-1',
174+
})
167175
expect(isLive).toBe(false)
168176
})
169177

170178
it('should return false for disconnected session', () => {
171179
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
172180
setSessionConnected('session-1', false)
173181

174-
const isLive = checkLiveUserInput('user-1', 'input-123', 'session-1')
182+
const isLive = checkLiveUserInput({
183+
userId: 'user-1',
184+
userInputId: 'input-123',
185+
clientSessionId: 'session-1',
186+
})
175187
expect(isLive).toBe(false)
176188
})
177189

178190
it('should return false for non-matching user input', () => {
179191
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
180192
setSessionConnected('session-1', true)
181193

182-
const isLive = checkLiveUserInput('user-1', 'input-456', 'session-1')
194+
const isLive = checkLiveUserInput({
195+
userId: 'user-1',
196+
userInputId: 'input-456',
197+
clientSessionId: 'session-1',
198+
})
183199
expect(isLive).toBe(false)
184200
})
185201

186202
it('should return true when live user input check is disabled', () => {
187203
disableLiveUserInputCheck()
188204

189-
const isLive = checkLiveUserInput('user-1', 'input-123', 'session-1')
205+
const isLive = checkLiveUserInput({
206+
userId: 'user-1',
207+
userInputId: 'input-123',
208+
clientSessionId: 'session-1',
209+
})
190210
expect(isLive).toBe(true)
191211
})
192212
})
@@ -196,7 +216,11 @@ describe('live-user-inputs', () => {
196216
setSessionConnected('session-1', true)
197217
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
198218

199-
const isLive = checkLiveUserInput('user-1', 'input-123', 'session-1')
219+
const isLive = checkLiveUserInput({
220+
userId: 'user-1',
221+
userInputId: 'input-123',
222+
clientSessionId: 'session-1',
223+
})
200224
expect(isLive).toBe(true)
201225
})
202226

@@ -205,11 +229,23 @@ describe('live-user-inputs', () => {
205229
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
206230

207231
// First verify it's connected
208-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(true)
232+
expect(
233+
checkLiveUserInput({
234+
userId: 'user-1',
235+
userInputId: 'input-123',
236+
clientSessionId: 'session-1',
237+
}),
238+
).toBe(true)
209239

210240
// Then disconnect
211241
setSessionConnected('session-1', false)
212-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(false)
242+
expect(
243+
checkLiveUserInput({
244+
userId: 'user-1',
245+
userInputId: 'input-123',
246+
clientSessionId: 'session-1',
247+
}),
248+
).toBe(false)
213249
})
214250

215251
it('should handle multiple sessions independently', () => {
@@ -218,8 +254,20 @@ describe('live-user-inputs', () => {
218254

219255
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
220256

221-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(true)
222-
expect(checkLiveUserInput('user-1', 'input-123', 'session-2')).toBe(false)
257+
expect(
258+
checkLiveUserInput({
259+
userId: 'user-1',
260+
userInputId: 'input-123',
261+
clientSessionId: 'session-1',
262+
}),
263+
).toBe(true)
264+
expect(
265+
checkLiveUserInput({
266+
userId: 'user-1',
267+
userInputId: 'input-123',
268+
clientSessionId: 'session-2',
269+
}),
270+
).toBe(false)
223271
})
224272
})
225273

@@ -250,14 +298,26 @@ describe('live-user-inputs', () => {
250298
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
251299

252300
// Verify input is live
253-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(true)
301+
expect(
302+
checkLiveUserInput({
303+
userId: 'user-1',
304+
userInputId: 'input-123',
305+
clientSessionId: 'session-1',
306+
}),
307+
).toBe(true)
254308
expect(getLiveUserInputIds('user-1')).toEqual(['input-123'])
255309

256310
// End user input
257311
cancelUserInput({ userId: 'user-1', userInputId: 'input-123', logger })
258312

259313
// Verify input is no longer live
260-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(false)
314+
expect(
315+
checkLiveUserInput({
316+
userId: 'user-1',
317+
userInputId: 'input-123',
318+
clientSessionId: 'session-1',
319+
}),
320+
).toBe(false)
261321
expect(getLiveUserInputIds('user-1')).toBeUndefined()
262322
})
263323

@@ -267,13 +327,25 @@ describe('live-user-inputs', () => {
267327
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
268328

269329
// Verify input is live
270-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(true)
330+
expect(
331+
checkLiveUserInput({
332+
userId: 'user-1',
333+
userInputId: 'input-123',
334+
clientSessionId: 'session-1',
335+
}),
336+
).toBe(true)
271337

272338
// Disconnect session
273339
setSessionConnected('session-1', false)
274340

275341
// Input should no longer be considered live
276-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(false)
342+
expect(
343+
checkLiveUserInput({
344+
userId: 'user-1',
345+
userInputId: 'input-123',
346+
clientSessionId: 'session-1',
347+
}),
348+
).toBe(false)
277349

278350
// But input ID should still exist (for potential reconnection)
279351
expect(getLiveUserInputIds('user-1')).toEqual(['input-123'])
@@ -285,15 +357,39 @@ describe('live-user-inputs', () => {
285357
startUserInput({ userId: 'user-1', userInputId: 'input-123' })
286358
startUserInput({ userId: 'user-1', userInputId: 'input-456' })
287359

288-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(true)
289-
expect(checkLiveUserInput('user-1', 'input-456', 'session-1')).toBe(true)
360+
expect(
361+
checkLiveUserInput({
362+
userId: 'user-1',
363+
userInputId: 'input-123',
364+
clientSessionId: 'session-1',
365+
}),
366+
).toBe(true)
367+
expect(
368+
checkLiveUserInput({
369+
userId: 'user-1',
370+
userInputId: 'input-456',
371+
clientSessionId: 'session-1',
372+
}),
373+
).toBe(true)
290374
expect(getLiveUserInputIds('user-1')).toEqual(['input-123', 'input-456'])
291375

292376
// Cancel one input
293377
cancelUserInput({ userId: 'user-1', userInputId: 'input-123', logger })
294378

295-
expect(checkLiveUserInput('user-1', 'input-123', 'session-1')).toBe(false)
296-
expect(checkLiveUserInput('user-1', 'input-456', 'session-1')).toBe(true)
379+
expect(
380+
checkLiveUserInput({
381+
userId: 'user-1',
382+
userInputId: 'input-123',
383+
clientSessionId: 'session-1',
384+
}),
385+
).toBe(false)
386+
expect(
387+
checkLiveUserInput({
388+
userId: 'user-1',
389+
userInputId: 'input-456',
390+
clientSessionId: 'session-1',
391+
}),
392+
).toBe(true)
297393
expect(getLiveUserInputIds('user-1')).toEqual(['input-456'])
298394
})
299395
})

‎backend/src/admin/grade-runs.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { closeXml } from '@codebuff/common/util/xml'
44
import { promptAiSdk } from '../llm-apis/vercel-ai-sdk/ai-sdk'
55

66
import type { Relabel, GetRelevantFilesTrace } from '@codebuff/bigquery'
7+
import type { Logger } from '@codebuff/types/logger'
78

89
const PROMPT = `
910
You are an evaluator system, measuring how well various models perform at selecting the most relevant files for a given user request.
@@ -96,11 +97,12 @@ function extractResponse(response: string): {
9697
}
9798
}
9899

99-
export async function gradeRun(tracesAndRelabels: {
100+
export async function gradeRun(params: {
100101
trace: GetRelevantFilesTrace
101102
relabels: Relabel[]
103+
logger: Logger
102104
}) {
103-
const { trace, relabels } = tracesAndRelabels
105+
const { trace, relabels, logger } = params
104106
const messages = trace.payload.messages
105107

106108
const originalOutput = trace.payload.output
@@ -152,6 +154,7 @@ export async function gradeRun(tracesAndRelabels: {
152154
// type: 'enabled',
153155
// budget_tokens: 10000,
154156
// },
157+
logger,
155158
})
156159

157160
const { scores } = extractResponse(response)

‎backend/src/admin/relabelRuns.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ export async function relabelForUserHandler(params: {
192192
fingerprintId: 'relabel-trace-api',
193193
userInputId: 'relabel-trace-api',
194194
userId: TEST_USER_ID,
195+
logger,
195196
})
196197

197198
// Create relabel record
@@ -436,6 +437,7 @@ export async function relabelWithClaudeWithFullFileContext(params: {
436437
userInputId: 'relabel-trace-api',
437438
userId: TEST_USER_ID,
438439
maxOutputTokens: 1000,
440+
logger,
439441
})
440442

441443
const relabel = {

‎backend/src/fast-rewrite.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ Please output just the complete updated file content with the edit applied and n
147147
fingerprintId,
148148
userInputId,
149149
userId,
150+
logger,
150151
})
151152

152153
return parseMarkdownCodeBlock(response) + '\n'

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,7 @@ async function getRelevantFilesForTraining(params: {
461461
model: models.openrouter_claude_sonnet_4,
462462
userId,
463463
chargeUser: false,
464+
logger,
464465
})
465466

466467
const end = performance.now()

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,7 @@ Provide a new set of SEARCH/REPLACE changes to make the intended edit from the o
174174
fingerprintId,
175175
userInputId,
176176
userId,
177+
logger,
177178
})
178179
const {
179180
diffBlocks: newDiffBlocks,

‎backend/src/get-documentation-for-query.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,7 @@ ${closeXml('user_query')}
202202
),
203203
}),
204204
timeout: 5_000,
205+
logger,
205206
})
206207
return {
207208
libraries: response.libraries,
@@ -266,6 +267,7 @@ ${closeXml('documentation_chunks')}
266267
relevant_chunks: z.array(z.number()),
267268
}),
268269
timeout: 20_000,
270+
logger,
269271
})
270272
const geminiDuration = Date.now() - geminiStartTime
271273

‎backend/src/live-user-inputs.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,13 @@ export function cancelUserInput(params: {
4141
}
4242
}
4343

44-
export function checkLiveUserInput(
45-
userId: string | undefined,
46-
userInputId: string,
47-
sessionId: string,
48-
): boolean {
44+
export function checkLiveUserInput(params: {
45+
userId: string | undefined
46+
userInputId: string
47+
clientSessionId: string
48+
}): boolean {
49+
const { userId, userInputId, clientSessionId } = params
50+
4951
if (!liveUserInputCheckEnabled) {
5052
return true
5153
}
@@ -54,7 +56,7 @@ export function checkLiveUserInput(
5456
}
5557

5658
// Check if WebSocket is still connected for this session
57-
if (!sessionConnections[sessionId]) {
59+
if (!sessionConnections[clientSessionId]) {
5860
return false
5961
}
6062

‎backend/src/llm-apis/gemini-with-fallbacks.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ export async function promptFlashWithFallbacks(params: {
6666
...geminiOptions,
6767
messages,
6868
model: useFinetunedModel,
69+
logger,
6970
})
7071
} catch (error) {
7172
logger.warn(
@@ -77,7 +78,7 @@ export async function promptFlashWithFallbacks(params: {
7778

7879
try {
7980
// First try Gemini
80-
return await promptAiSdk({ ...geminiOptions, messages })
81+
return await promptAiSdk({ ...geminiOptions, messages, logger })
8182
} catch (error) {
8283
logger.warn(
8384
{ error },
@@ -95,6 +96,7 @@ export async function promptFlashWithFallbacks(params: {
9596
experimental: openrouterModels.openrouter_claude_3_5_haiku,
9697
ask: openrouterModels.openrouter_claude_3_5_haiku,
9798
}[costMode ?? 'normal'],
99+
logger,
98100
})
99101
}
100102
}

0 commit comments

Comments
 (0)