Skip to content

Commit bc13f76

Browse files
charlesliencharleslien
committed
pass in logger to agent-run.ts
1 parent eb7e170 commit bc13f76

File tree

8 files changed

+96
-4
lines changed

8 files changed

+96
-4
lines changed

‎backend/src/__tests__/agent-run.test.ts‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@ import {
1717
} from 'bun:test'
1818

1919
import { startAgentRun, finishAgentRun, addAgentStep } from '../agent-run'
20-
import { logger } from '../util/logger'
20+
21+
import type { Logger } from '@codebuff/agent-runtime'
2122

2223
describe('Agent Run Database Functions', () => {
24+
let logger: Logger
25+
2326
beforeEach(() => {
2427
// Setup spies for database operations
2528
spyOn(db, 'insert').mockReturnValue({
@@ -30,8 +33,12 @@ describe('Agent Run Database Functions', () => {
3033
where: mock(() => Promise.resolve()),
3134
})),
3235
} as any)
33-
// Mock logger
34-
spyOn(logger, 'error').mockImplementation(() => {})
36+
logger = {
37+
debug: mock(() => {}),
38+
info: mock(() => {}),
39+
warn: mock(() => {}),
40+
error: mock(() => {}),
41+
}
3542
})
3643

3744
afterEach(() => {
@@ -61,6 +68,7 @@ describe('Agent Run Database Functions', () => {
6168
userId: 'user-123',
6269
agentId: 'test-agent',
6370
ancestorRunIds: ['parent-run-1', 'parent-run-2'],
71+
logger,
6472
})
6573

6674
expect(result).toBe('generated-uuid')
@@ -84,6 +92,7 @@ describe('Agent Run Database Functions', () => {
8492
userId: 'user-123',
8593
agentId: 'test-agent',
8694
ancestorRunIds: [],
95+
logger,
8796
})
8897

8998
expect(result).toBe('custom-run-id')
@@ -105,6 +114,7 @@ describe('Agent Run Database Functions', () => {
105114
await startAgentRun({
106115
agentId: 'test-agent',
107116
ancestorRunIds: [],
117+
logger,
108118
})
109119

110120
expect(mockValues).toHaveBeenCalledWith({
@@ -125,6 +135,7 @@ describe('Agent Run Database Functions', () => {
125135
await startAgentRun({
126136
agentId: 'test-agent',
127137
ancestorRunIds: [],
138+
logger,
128139
})
129140

130141
expect(mockValues).toHaveBeenCalledWith(
@@ -142,6 +153,7 @@ describe('Agent Run Database Functions', () => {
142153
await startAgentRun({
143154
agentId: 'test-agent',
144155
ancestorRunIds: ['root-run', 'parent-run'],
156+
logger,
145157
})
146158

147159
expect(mockValues).toHaveBeenCalledWith(
@@ -161,6 +173,7 @@ describe('Agent Run Database Functions', () => {
161173
startAgentRun({
162174
agentId: 'test-agent',
163175
ancestorRunIds: [],
176+
logger,
164177
}),
165178
).rejects.toThrow('Database connection failed')
166179

@@ -190,6 +203,7 @@ describe('Agent Run Database Functions', () => {
190203
totalSteps: 5,
191204
directCredits: 150.5,
192205
totalCredits: 300.75,
206+
logger,
193207
})
194208

195209
expect(db.update).toHaveBeenCalledWith(schema.agentRun)
@@ -217,6 +231,7 @@ describe('Agent Run Database Functions', () => {
217231
directCredits: 75.25,
218232
totalCredits: 125.5,
219233
errorMessage: 'Agent execution failed',
234+
logger,
220235
})
221236

222237
expect(mockSet).toHaveBeenCalledWith({
@@ -242,6 +257,7 @@ describe('Agent Run Database Functions', () => {
242257
totalSteps: 2,
243258
directCredits: 50,
244259
totalCredits: 100,
260+
logger,
245261
})
246262

247263
expect(mockSet).toHaveBeenCalledWith(
@@ -266,6 +282,7 @@ describe('Agent Run Database Functions', () => {
266282
totalSteps: 5,
267283
directCredits: 150,
268284
totalCredits: 300,
285+
logger,
269286
}),
270287
).rejects.toThrow('Update failed')
271288

@@ -297,6 +314,7 @@ describe('Agent Run Database Functions', () => {
297314
messageId: 'msg-456',
298315
status: 'completed',
299316
startTime,
317+
logger,
300318
})
301319

302320
expect(result).toBe('step-uuid')
@@ -328,6 +346,7 @@ describe('Agent Run Database Functions', () => {
328346
stepNumber: 2,
329347
startTime,
330348
messageId: null,
349+
logger,
331350
})
332351

333352
expect(mockValues).toHaveBeenCalledWith({
@@ -359,6 +378,7 @@ describe('Agent Run Database Functions', () => {
359378
errorMessage: 'Step failed validation',
360379
startTime,
361380
messageId: null,
381+
logger,
362382
})
363383

364384
expect(mockValues).toHaveBeenCalledWith(
@@ -383,6 +403,7 @@ describe('Agent Run Database Functions', () => {
383403
status: 'running',
384404
startTime,
385405
messageId: null,
406+
logger,
386407
})
387408

388409
expect(mockValues).toHaveBeenCalledWith(
@@ -406,6 +427,7 @@ describe('Agent Run Database Functions', () => {
406427
credits: 0, // Zero credits
407428
startTime,
408429
messageId: null,
430+
logger,
409431
})
410432

411433
expect(mockValues).toHaveBeenCalledWith(
@@ -430,6 +452,7 @@ describe('Agent Run Database Functions', () => {
430452
stepNumber: 6,
431453
startTime,
432454
messageId: null,
455+
logger,
433456
}),
434457
).rejects.toThrow('Insert failed')
435458

@@ -457,6 +480,7 @@ describe('Agent Run Database Functions', () => {
457480
credits: 123.456789, // High precision number
458481
startTime: new Date(),
459482
messageId: null,
483+
logger,
460484
})
461485

462486
expect(mockValues).toHaveBeenCalledWith(
@@ -479,6 +503,7 @@ describe('Agent Run Database Functions', () => {
479503
stepNumber: 1,
480504
startTime: specificStartTime,
481505
messageId: null,
506+
logger,
482507
})
483508

484509
expect(mockValues).toHaveBeenCalledWith(

‎backend/src/agent-run.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@ import * as schema from '@codebuff/common/db/schema'
33
import { TEST_USER_ID } from '@codebuff/common/old-constants'
44
import { eq } from 'drizzle-orm'
55

6-
import { logger } from './util/logger'
6+
import type {
7+
AddAgentStepFn,
8+
FinishAgentRunFn,
9+
StartAgentRunFn,
10+
Logger,
11+
} from '@codebuff/agent-runtime'
712

813
/**
914
* Starts a new agent run and creates an entry in the agent_run table
@@ -13,11 +18,13 @@ export async function startAgentRun({
1318
userId,
1419
agentId,
1520
ancestorRunIds,
21+
logger,
1622
}: {
1723
runId?: string
1824
userId?: string
1925
agentId: string
2026
ancestorRunIds: string[]
27+
logger: Logger
2128
}): Promise<string> {
2229
if (userId === TEST_USER_ID) {
2330
return 'test-run-id'
@@ -44,6 +51,7 @@ export async function startAgentRun({
4451
throw error
4552
}
4653
}
54+
startAgentRun satisfies StartAgentRunFn
4755

4856
/**
4957
* Completes an agent run by updating its status and metrics
@@ -56,6 +64,7 @@ export async function finishAgentRun({
5664
directCredits,
5765
totalCredits,
5866
errorMessage,
67+
logger,
5968
}: {
6069
userId: string | undefined
6170
runId: string
@@ -64,6 +73,7 @@ export async function finishAgentRun({
6473
directCredits: number
6574
totalCredits: number
6675
errorMessage?: string
76+
logger: Logger
6777
}): Promise<void> {
6878
if (userId === TEST_USER_ID) {
6979
return
@@ -86,6 +96,7 @@ export async function finishAgentRun({
8696
throw error
8797
}
8898
}
99+
finishAgentRun satisfies FinishAgentRunFn
89100

90101
/**
91102
* Adds a completed step to the agent_step table
@@ -100,6 +111,7 @@ export async function addAgentStep({
100111
status = 'completed',
101112
errorMessage,
102113
startTime,
114+
logger,
103115
}: {
104116
userId: string | undefined
105117
agentRunId: string
@@ -110,6 +122,7 @@ export async function addAgentStep({
110122
status?: 'running' | 'completed' | 'skipped'
111123
errorMessage?: string
112124
startTime: Date
125+
logger: Logger
113126
}): Promise<string> {
114127
if (userId === TEST_USER_ID) {
115128
return 'test-step-id'
@@ -136,3 +149,4 @@ export async function addAgentStep({
136149
throw error
137150
}
138151
}
152+
addAgentStep satisfies AddAgentStepFn

‎backend/src/run-agent-step.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,7 @@ export const loopAgentSteps = async (
439439
userId,
440440
agentId: agentTemplate.id,
441441
ancestorRunIds: agentState.ancestorRunIds,
442+
logger,
442443
})
443444

444445
// Initialize message history with user prompt and instructions on first iteration
@@ -660,6 +661,7 @@ export const loopAgentSteps = async (
660661
messageId,
661662
status: 'completed',
662663
startTime,
664+
logger,
663665
})
664666
} else {
665667
logger.error('No runId found for agent state after finishing agent run')
@@ -689,6 +691,7 @@ export const loopAgentSteps = async (
689691
totalSteps,
690692
directCredits: currentAgentState.directCreditsUsed,
691693
totalCredits: currentAgentState.creditsUsed,
694+
logger,
692695
})
693696

694697
return {
@@ -721,6 +724,7 @@ export const loopAgentSteps = async (
721724
directCredits: currentAgentState.directCreditsUsed,
722725
totalCredits: currentAgentState.creditsUsed,
723726
errorMessage,
727+
logger,
724728
})
725729

726730
const errorObject = getErrorObject(error)

‎backend/src/run-programmatic-step.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ export async function runProgrammaticStep(
297297
status: 'completed',
298298
startTime,
299299
messageId: null,
300+
logger,
300301
})
301302
} else {
302303
logger.error('No runId found for agent state after finishing agent run')
@@ -345,6 +346,7 @@ export async function runProgrammaticStep(
345346
status: 'skipped',
346347
startTime,
347348
messageId: null,
349+
logger,
348350
})
349351
} else {
350352
logger.error('No runId found for agent state after failed agent run')

‎packages/agent-runtime/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export type * from './types'

‎packages/agent-runtime/src/types/database.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import type { Logger } from './logger'
2+
3+
export type StartAgentRunFn = (params: {
4+
runId?: string
5+
userId?: string
6+
agentId: string
7+
ancestorRunIds: string[]
8+
logger: Logger
9+
}) => Promise<string>
10+
export type FinishAgentRunFn = (params: {
11+
userId: string | undefined
12+
runId: string
13+
status: 'completed' | 'failed' | 'cancelled'
14+
totalSteps: number
15+
directCredits: number
16+
totalCredits: number
17+
errorMessage?: string
18+
logger: Logger
19+
}) => Promise<void>
20+
21+
export type AddAgentStepFn = (params: {
22+
userId: string | undefined
23+
agentRunId: string
24+
stepNumber: number
25+
credits?: number
26+
childRunIds?: string[]
27+
messageId: string | null
28+
status?: 'running' | 'completed' | 'skipped'
29+
errorMessage?: string
30+
startTime: Date
31+
logger: Logger
32+
}) => Promise<string>

‎packages/agent-runtime/src/types/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export type * from './database'
2+
export type * from './logger'

‎packages/agent-runtime/src/types/logger.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
export type LoggerFunction = (
2+
data: unknown,
3+
msg?: string,
4+
...args: unknown[]
5+
) => unknown
6+
7+
export type Logger = {
8+
debug: LoggerFunction
9+
info: LoggerFunction
10+
warn: LoggerFunction
11+
error: LoggerFunction
12+
}

0 commit comments

Comments
 (0)