Skip to content

Commit 4691ec5

Browse files
charlesliencharleslien
committed
sdk: implement startAgentRun
1 parent 5f21a90 commit 4691ec5

File tree

5 files changed

+82
-43
lines changed

5 files changed

+82
-43
lines changed

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

Lines changed: 32 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import db from '@codebuff/common/db'
22
import * as schema from '@codebuff/common/db/schema'
3+
import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
34
import {
45
mockModule,
56
clearMockedModules,
@@ -18,12 +19,20 @@ import {
1819

1920
import { startAgentRun, finishAgentRun, addAgentStep } from '../agent-run'
2021

21-
import type { Logger } from '@codebuff/common/types/contracts/logger'
22+
import type {
23+
AgentRuntimeDeps,
24+
AgentRuntimeScopedDeps,
25+
} from '@codebuff/common/types/contracts/agent-runtime'
2226

2327
describe('Agent Run Database Functions', () => {
24-
let logger: Logger
28+
let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps
2529

2630
beforeEach(() => {
31+
agentRuntimeImpl = {
32+
...TEST_AGENT_RUNTIME_IMPL,
33+
}
34+
agentRuntimeImpl.logger.error = mock(() => {})
35+
2736
// Setup spies for database operations
2837
spyOn(db, 'insert').mockReturnValue({
2938
values: mock(() => Promise.resolve({ id: 'test-run-id' })),
@@ -33,12 +42,6 @@ describe('Agent Run Database Functions', () => {
3342
where: mock(() => Promise.resolve()),
3443
})),
3544
} as any)
36-
logger = {
37-
debug: mock(() => {}),
38-
info: mock(() => {}),
39-
warn: mock(() => {}),
40-
error: mock(() => {}),
41-
}
4245
})
4346

4447
afterEach(() => {
@@ -65,10 +68,10 @@ describe('Agent Run Database Functions', () => {
6568
spyOn(crypto, 'randomUUID').mockReturnValue('generated-uuid')
6669

6770
const result = await startAgentRun({
71+
...agentRuntimeImpl,
6872
userId: 'user-123',
6973
agentId: 'test-agent',
7074
ancestorRunIds: ['parent-run-1', 'parent-run-2'],
71-
logger,
7275
})
7376

7477
expect(result).toBe('generated-uuid')
@@ -88,11 +91,11 @@ describe('Agent Run Database Functions', () => {
8891
spyOn(db, 'insert').mockReturnValue({ values: mockValues } as any)
8992

9093
const result = await startAgentRun({
94+
...agentRuntimeImpl,
9195
runId: 'custom-run-id',
9296
userId: 'user-123',
9397
agentId: 'test-agent',
9498
ancestorRunIds: [],
95-
logger,
9699
})
97100

98101
expect(result).toBe('custom-run-id')
@@ -112,9 +115,9 @@ describe('Agent Run Database Functions', () => {
112115
spyOn(crypto, 'randomUUID').mockReturnValue('generated-uuid')
113116

114117
await startAgentRun({
118+
...agentRuntimeImpl,
115119
agentId: 'test-agent',
116120
ancestorRunIds: [],
117-
logger,
118121
})
119122

120123
expect(mockValues).toHaveBeenCalledWith({
@@ -133,9 +136,9 @@ describe('Agent Run Database Functions', () => {
133136
spyOn(crypto, 'randomUUID').mockReturnValue('generated-uuid')
134137

135138
await startAgentRun({
139+
...agentRuntimeImpl,
136140
agentId: 'test-agent',
137141
ancestorRunIds: [],
138-
logger,
139142
})
140143

141144
expect(mockValues).toHaveBeenCalledWith(
@@ -151,9 +154,9 @@ describe('Agent Run Database Functions', () => {
151154
spyOn(crypto, 'randomUUID').mockReturnValue('generated-uuid')
152155

153156
await startAgentRun({
157+
...agentRuntimeImpl,
154158
agentId: 'test-agent',
155159
ancestorRunIds: ['root-run', 'parent-run'],
156-
logger,
157160
})
158161

159162
expect(mockValues).toHaveBeenCalledWith(
@@ -171,13 +174,13 @@ describe('Agent Run Database Functions', () => {
171174

172175
expect(
173176
startAgentRun({
177+
...agentRuntimeImpl,
174178
agentId: 'test-agent',
175179
ancestorRunIds: [],
176-
logger,
177180
}),
178181
).rejects.toThrow('Database connection failed')
179182

180-
expect(logger.error).toHaveBeenCalledWith(
183+
expect(agentRuntimeImpl.logger.error).toHaveBeenCalledWith(
181184
{
182185
error: mockError,
183186
runId: undefined,
@@ -197,13 +200,13 @@ describe('Agent Run Database Functions', () => {
197200
spyOn(db, 'update').mockReturnValue({ set: mockSet } as any)
198201

199202
await finishAgentRun({
203+
...agentRuntimeImpl,
200204
userId: undefined,
201205
runId: 'test-run-id',
202206
status: 'completed',
203207
totalSteps: 5,
204208
directCredits: 150.5,
205209
totalCredits: 300.75,
206-
logger,
207210
})
208211

209212
expect(db.update).toHaveBeenCalledWith(schema.agentRun)
@@ -224,14 +227,14 @@ describe('Agent Run Database Functions', () => {
224227
spyOn(db, 'update').mockReturnValue({ set: mockSet } as any)
225228

226229
await finishAgentRun({
230+
...agentRuntimeImpl,
227231
userId: undefined,
228232
runId: 'test-run-id',
229233
status: 'failed',
230234
totalSteps: 3,
231235
directCredits: 75.25,
232236
totalCredits: 125.5,
233237
errorMessage: 'Agent execution failed',
234-
logger,
235238
})
236239

237240
expect(mockSet).toHaveBeenCalledWith({
@@ -251,13 +254,13 @@ describe('Agent Run Database Functions', () => {
251254
mockSet.mockReturnValue({ where: mockWhere })
252255

253256
await finishAgentRun({
257+
...agentRuntimeImpl,
254258
userId: undefined,
255259
runId: 'test-run-id',
256260
status: 'cancelled',
257261
totalSteps: 2,
258262
directCredits: 50,
259263
totalCredits: 100,
260-
logger,
261264
})
262265

263266
expect(mockSet).toHaveBeenCalledWith(
@@ -276,17 +279,17 @@ describe('Agent Run Database Functions', () => {
276279

277280
expect(
278281
finishAgentRun({
282+
...agentRuntimeImpl,
279283
userId: undefined,
280284
runId: 'test-run-id',
281285
status: 'completed',
282286
totalSteps: 5,
283287
directCredits: 150,
284288
totalCredits: 300,
285-
logger,
286289
}),
287290
).rejects.toThrow('Update failed')
288291

289-
expect(logger.error).toHaveBeenCalledWith(
292+
expect(agentRuntimeImpl.logger.error).toHaveBeenCalledWith(
290293
{
291294
error: mockError,
292295
runId: 'test-run-id',
@@ -306,6 +309,7 @@ describe('Agent Run Database Functions', () => {
306309
const startTime = new Date('2023-01-01T10:00:00Z')
307310

308311
const result = await addAgentStep({
312+
...agentRuntimeImpl,
309313
userId: undefined,
310314
agentRunId: 'run-123',
311315
stepNumber: 1,
@@ -314,7 +318,6 @@ describe('Agent Run Database Functions', () => {
314318
messageId: 'msg-456',
315319
status: 'completed',
316320
startTime,
317-
logger,
318321
})
319322

320323
expect(result).toBe('step-uuid')
@@ -341,12 +344,12 @@ describe('Agent Run Database Functions', () => {
341344
const startTime = new Date('2023-01-01T10:00:00Z')
342345

343346
await addAgentStep({
347+
...agentRuntimeImpl,
344348
userId: undefined,
345349
agentRunId: 'run-123',
346350
stepNumber: 2,
347351
startTime,
348352
messageId: null,
349-
logger,
350353
})
351354

352355
expect(mockValues).toHaveBeenCalledWith({
@@ -371,14 +374,14 @@ describe('Agent Run Database Functions', () => {
371374
const startTime = new Date('2023-01-01T10:00:00Z')
372375

373376
await addAgentStep({
377+
...agentRuntimeImpl,
374378
userId: undefined,
375379
agentRunId: 'run-123',
376380
stepNumber: 3,
377381
status: 'skipped',
378382
errorMessage: 'Step failed validation',
379383
startTime,
380384
messageId: null,
381-
logger,
382385
})
383386

384387
expect(mockValues).toHaveBeenCalledWith(
@@ -397,13 +400,13 @@ describe('Agent Run Database Functions', () => {
397400
const startTime = new Date('2023-01-01T10:00:00Z')
398401

399402
await addAgentStep({
403+
...agentRuntimeImpl,
400404
userId: undefined,
401405
agentRunId: 'run-123',
402406
stepNumber: 4,
403407
status: 'running',
404408
startTime,
405409
messageId: null,
406-
logger,
407410
})
408411

409412
expect(mockValues).toHaveBeenCalledWith(
@@ -421,13 +424,13 @@ describe('Agent Run Database Functions', () => {
421424
const startTime = new Date('2023-01-01T10:00:00Z')
422425

423426
await addAgentStep({
427+
...agentRuntimeImpl,
424428
userId: undefined,
425429
agentRunId: 'run-123',
426430
stepNumber: 5,
427431
credits: 0, // Zero credits
428432
startTime,
429433
messageId: null,
430-
logger,
431434
})
432435

433436
expect(mockValues).toHaveBeenCalledWith(
@@ -447,16 +450,16 @@ describe('Agent Run Database Functions', () => {
447450

448451
expect(
449452
addAgentStep({
453+
...agentRuntimeImpl,
450454
userId: undefined,
451455
agentRunId: 'run-123',
452456
stepNumber: 6,
453457
startTime,
454458
messageId: null,
455-
logger,
456459
}),
457460
).rejects.toThrow('Insert failed')
458461

459-
expect(logger.error).toHaveBeenCalledWith(
462+
expect(agentRuntimeImpl.logger.error).toHaveBeenCalledWith(
460463
{
461464
error: mockError,
462465
agentRunId: 'run-123',
@@ -474,13 +477,13 @@ describe('Agent Run Database Functions', () => {
474477
spyOn(crypto, 'randomUUID').mockReturnValue('step-uuid')
475478

476479
await addAgentStep({
480+
...agentRuntimeImpl,
477481
userId: undefined,
478482
agentRunId: 'run-123',
479483
stepNumber: 1,
480484
credits: 123.456789, // High precision number
481485
startTime: new Date(),
482486
messageId: null,
483-
logger,
484487
})
485488

486489
expect(mockValues).toHaveBeenCalledWith(
@@ -498,12 +501,12 @@ describe('Agent Run Database Functions', () => {
498501
const specificStartTime = new Date('2023-01-01T10:30:45.123Z')
499502

500503
await addAgentStep({
504+
...agentRuntimeImpl,
501505
userId: undefined,
502506
agentRunId: 'run-123',
503507
stepNumber: 1,
504508
startTime: specificStartTime,
505509
messageId: null,
506-
logger,
507510
})
508511

509512
expect(mockValues).toHaveBeenCalledWith(

‎common/src/types/contracts/database.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,13 @@ export type FetchAgentFromDatabaseFn = (params: {
5656
}) => Promise<AgentTemplate | null>
5757

5858
export type StartAgentRunFn = (params: {
59+
apiKey: string
5960
runId?: string
6061
userId?: string
6162
agentId: string
6263
ancestorRunIds: string[]
6364
logger: Logger
64-
}) => Promise<string>
65+
}) => Promise<string | null>
6566

6667
export type FinishAgentRunFn = (params: {
6768
userId: string | undefined

‎packages/agent-runtime/src/run-agent-step.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,8 @@ export async function loopAgentSteps(
445445
typeof getMCPToolData,
446446
'toolNames' | 'mcpServers' | 'writeTo'
447447
> &
448-
ParamsOf<CheckLiveUserInputFn>,
448+
ParamsOf<CheckLiveUserInputFn> &
449+
ParamsExcluding<StartAgentRunFn, 'runId' | 'agentId' | 'ancestorRunIds'>,
449450
): Promise<{
450451
agentState: AgentState
451452
output: AgentOutput
@@ -482,11 +483,10 @@ export async function loopAgentSteps(
482483
const runId = crypto.randomUUID()
483484
agentState.runId = runId
484485
await startAgentRun({
486+
...params,
485487
runId,
486-
userId,
487488
agentId: agentTemplate.id,
488489
ancestorRunIds: agentState.ancestorRunIds,
489-
logger,
490490
})
491491

492492
// Initialize message history with user prompt and instructions on first iteration

‎sdk/src/impl/agent-runtime.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import { trackEvent } from '@codebuff/common/analytics'
22

3-
import { fetchAgentFromDatabase, getUserInfoFromApiKey } from './database'
3+
import {
4+
fetchAgentFromDatabase,
5+
getUserInfoFromApiKey,
6+
startAgentRun,
7+
} from './database'
48

59
import type {
610
AgentRuntimeDeps,
@@ -9,7 +13,6 @@ import type {
913

1014
export const CLI_AGENT_RUNTIME_IMPL: Omit<
1115
AgentRuntimeDeps & AgentRuntimeScopedDeps,
12-
| 'startAgentRun'
1316
| 'finishAgentRun'
1417
| 'addAgentStep'
1518
| 'consumeCreditsWithFallback'
@@ -27,7 +30,7 @@ export const CLI_AGENT_RUNTIME_IMPL: Omit<
2730
// Database
2831
getUserInfoFromApiKey,
2932
fetchAgentFromDatabase,
30-
// startAgentRun: StartAgentRunFn
33+
startAgentRun,
3134
// finishAgentRun: FinishAgentRunFn
3235
// addAgentStep: AddAgentStepFn
3336

0 commit comments

Comments
 (0)