Skip to content

Commit 7650ff0

Browse files
jahoomajahooma
committed
Deduplicate spawn agents code, add some tests
1 parent 33747cb commit 7650ff0

File tree

5 files changed

+751
-457
lines changed

5 files changed

+751
-457
lines changed

‎backend/src/__tests__/spawn-agents-permissions.test.ts‎

Lines changed: 231 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import { describe, expect, it, beforeEach, afterEach, mock, spyOn } from 'bun:test'
2-
import { getMatchingSpawn, handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
2+
import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
3+
import { handleSpawnAgentInline } from '../tools/handlers/tool/spawn-agent-inline'
4+
import { getMatchingSpawn } from '../tools/handlers/tool/spawn-agent-utils'
35
import { TEST_USER_ID } from '@codebuff/common/constants'
46
import { getInitialSessionState } from '@codebuff/common/types/session-state'
57
import { mockFileContext, MockWebSocket } from './test-utils'
@@ -14,6 +16,25 @@ describe('Spawn Agents Permissions', () => {
1416
let mockSendSubagentChunk: any
1517
let mockLoopAgentSteps: any
1618

19+
const createMockAgent = (id: string, spawnableAgents: string[] = []): AgentTemplate => ({
20+
id,
21+
displayName: `Mock ${id}`,
22+
outputMode: 'last_message' as const,
23+
inputSchema: {
24+
prompt: {
25+
safeParse: () => ({ success: true }),
26+
} as any,
27+
},
28+
spawnerPrompt: '',
29+
model: '',
30+
includeMessageHistory: true,
31+
toolNames: [],
32+
spawnableAgents,
33+
systemPrompt: '',
34+
instructionsPrompt: '',
35+
stepPrompt: '',
36+
})
37+
1738
beforeEach(() => {
1839
// Mock logger to reduce noise in tests
1940
spyOn(loggerModule.logger, 'debug').mockImplementation(() => {})
@@ -174,25 +195,6 @@ describe('Spawn Agents Permissions', () => {
174195
})
175196

176197
describe('handleSpawnAgents permission validation', () => {
177-
const createMockAgent = (id: string, spawnableAgents: string[] = []): AgentTemplate => ({
178-
id,
179-
displayName: `Mock ${id}`,
180-
outputMode: 'last_message' as const,
181-
inputSchema: {
182-
prompt: {
183-
safeParse: () => ({ success: true }),
184-
} as any,
185-
},
186-
spawnerPrompt: '',
187-
model: '',
188-
includeMessageHistory: true,
189-
toolNames: [],
190-
spawnableAgents,
191-
systemPrompt: '',
192-
instructionsPrompt: '',
193-
stepPrompt: '',
194-
})
195-
196198
const createSpawnToolCall = (agentType: string, prompt = 'test prompt'): CodebuffToolCall<'spawn_agents'> => ({
197199
toolName: 'spawn_agents' as const,
198200
toolCallId: 'test-tool-call-id',
@@ -436,4 +438,213 @@ describe('Spawn Agents Permissions', () => {
436438
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1) // Only thinker was spawned
437439
})
438440
})
441+
442+
describe('handleSpawnAgentInline permission validation', () => {
443+
const createInlineSpawnToolCall = (agentType: string, prompt = 'test prompt'): CodebuffToolCall<'spawn_agent_inline'> => ({
444+
toolName: 'spawn_agent_inline' as const,
445+
toolCallId: 'test-tool-call-id',
446+
input: {
447+
agent_type: agentType,
448+
prompt,
449+
},
450+
})
451+
452+
it('should allow spawning inline agent when agent is in spawnableAgents list', async () => {
453+
const parentAgent = createMockAgent('parent', ['thinker', 'reviewer'])
454+
const childAgent = createMockAgent('thinker')
455+
const ws = new MockWebSocket() as unknown as WebSocket
456+
const sessionState = getInitialSessionState(mockFileContext)
457+
const toolCall = createInlineSpawnToolCall('thinker')
458+
459+
const { result } = handleSpawnAgentInline({
460+
previousToolCallFinished: Promise.resolve(),
461+
toolCall,
462+
fileContext: mockFileContext,
463+
clientSessionId: 'test-session',
464+
userInputId: 'test-input',
465+
getLatestState: () => ({ messages: [] }),
466+
state: {
467+
ws,
468+
fingerprintId: 'test-fingerprint',
469+
userId: TEST_USER_ID,
470+
agentTemplate: parentAgent,
471+
localAgentTemplates: { thinker: childAgent },
472+
messages: [],
473+
agentState: sessionState.mainAgentState,
474+
},
475+
})
476+
477+
await result // Should not throw
478+
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
479+
})
480+
481+
it('should reject spawning inline agent when agent is not in spawnableAgents list', async () => {
482+
const parentAgent = createMockAgent('parent', ['thinker']) // Only allows thinker
483+
const childAgent = createMockAgent('reviewer')
484+
const ws = new MockWebSocket() as unknown as WebSocket
485+
const sessionState = getInitialSessionState(mockFileContext)
486+
const toolCall = createInlineSpawnToolCall('reviewer') // Try to spawn reviewer
487+
488+
const { result } = handleSpawnAgentInline({
489+
previousToolCallFinished: Promise.resolve(),
490+
toolCall,
491+
fileContext: mockFileContext,
492+
clientSessionId: 'test-session',
493+
userInputId: 'test-input',
494+
getLatestState: () => ({ messages: [] }),
495+
state: {
496+
ws,
497+
fingerprintId: 'test-fingerprint',
498+
userId: TEST_USER_ID,
499+
agentTemplate: parentAgent,
500+
localAgentTemplates: { reviewer: childAgent },
501+
messages: [],
502+
agentState: sessionState.mainAgentState,
503+
},
504+
})
505+
506+
await expect(result).rejects.toThrow('is not allowed to spawn child agent type reviewer')
507+
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
508+
})
509+
510+
it('should reject spawning inline agent when agent template is not found', async () => {
511+
const parentAgent = createMockAgent('parent', ['nonexistent'])
512+
const ws = new MockWebSocket() as unknown as WebSocket
513+
const sessionState = getInitialSessionState(mockFileContext)
514+
const toolCall = createInlineSpawnToolCall('nonexistent')
515+
516+
const { result } = handleSpawnAgentInline({
517+
previousToolCallFinished: Promise.resolve(),
518+
toolCall,
519+
fileContext: mockFileContext,
520+
clientSessionId: 'test-session',
521+
userInputId: 'test-input',
522+
getLatestState: () => ({ messages: [] }),
523+
state: {
524+
ws,
525+
fingerprintId: 'test-fingerprint',
526+
userId: TEST_USER_ID,
527+
agentTemplate: parentAgent,
528+
localAgentTemplates: {}, // Empty - agent not found
529+
messages: [],
530+
agentState: sessionState.mainAgentState,
531+
},
532+
})
533+
534+
await expect(result).rejects.toThrow('Agent type nonexistent not found')
535+
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
536+
})
537+
538+
it('should handle versioned inline agent permissions correctly', async () => {
539+
const parentAgent = createMockAgent('parent', ['codebuff/thinker@1.0.0'])
540+
const childAgent = createMockAgent('codebuff/thinker@1.0.0')
541+
const ws = new MockWebSocket() as unknown as WebSocket
542+
const sessionState = getInitialSessionState(mockFileContext)
543+
const toolCall = createInlineSpawnToolCall('codebuff/thinker@1.0.0')
544+
545+
const { result } = handleSpawnAgentInline({
546+
previousToolCallFinished: Promise.resolve(),
547+
toolCall,
548+
fileContext: mockFileContext,
549+
clientSessionId: 'test-session',
550+
userInputId: 'test-input',
551+
getLatestState: () => ({ messages: [] }),
552+
state: {
553+
ws,
554+
fingerprintId: 'test-fingerprint',
555+
userId: TEST_USER_ID,
556+
agentTemplate: parentAgent,
557+
localAgentTemplates: { 'codebuff/thinker@1.0.0': childAgent },
558+
messages: [],
559+
agentState: sessionState.mainAgentState,
560+
},
561+
})
562+
563+
await result // Should not throw
564+
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
565+
})
566+
567+
it('should allow spawning simple agent name inline when parent allows versioned agent', async () => {
568+
const parentAgent = createMockAgent('parent', ['codebuff/thinker@1.0.0'])
569+
const childAgent = createMockAgent('codebuff/thinker@1.0.0')
570+
const ws = new MockWebSocket() as unknown as WebSocket
571+
const sessionState = getInitialSessionState(mockFileContext)
572+
const toolCall = createInlineSpawnToolCall('thinker') // Simple name
573+
574+
const { result } = handleSpawnAgentInline({
575+
previousToolCallFinished: Promise.resolve(),
576+
toolCall,
577+
fileContext: mockFileContext,
578+
clientSessionId: 'test-session',
579+
userInputId: 'test-input',
580+
getLatestState: () => ({ messages: [] }),
581+
state: {
582+
ws,
583+
fingerprintId: 'test-fingerprint',
584+
userId: TEST_USER_ID,
585+
agentTemplate: parentAgent,
586+
localAgentTemplates: {
587+
'thinker': childAgent,
588+
'codebuff/thinker@1.0.0': childAgent, // Register with both keys
589+
},
590+
messages: [],
591+
agentState: sessionState.mainAgentState,
592+
},
593+
})
594+
595+
await result // Should not throw
596+
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
597+
})
598+
599+
it('should reject inline spawn when version mismatch exists', async () => {
600+
const parentAgent = createMockAgent('parent', ['codebuff/thinker@1.0.0'])
601+
const childAgent = createMockAgent('codebuff/thinker@2.0.0')
602+
const ws = new MockWebSocket() as unknown as WebSocket
603+
const sessionState = getInitialSessionState(mockFileContext)
604+
const toolCall = createInlineSpawnToolCall('codebuff/thinker@2.0.0')
605+
606+
const { result } = handleSpawnAgentInline({
607+
previousToolCallFinished: Promise.resolve(),
608+
toolCall,
609+
fileContext: mockFileContext,
610+
clientSessionId: 'test-session',
611+
userInputId: 'test-input',
612+
getLatestState: () => ({ messages: [] }),
613+
state: {
614+
ws,
615+
fingerprintId: 'test-fingerprint',
616+
userId: TEST_USER_ID,
617+
agentTemplate: parentAgent,
618+
localAgentTemplates: { 'codebuff/thinker@2.0.0': childAgent },
619+
messages: [],
620+
agentState: sessionState.mainAgentState,
621+
},
622+
})
623+
624+
await expect(result).rejects.toThrow('is not allowed to spawn child agent type')
625+
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
626+
})
627+
628+
it('should validate required state parameters for inline spawn', async () => {
629+
const parentAgent = createMockAgent('parent', ['thinker'])
630+
const toolCall = createInlineSpawnToolCall('thinker')
631+
632+
expect(() => {
633+
handleSpawnAgentInline({
634+
previousToolCallFinished: Promise.resolve(),
635+
toolCall,
636+
fileContext: mockFileContext,
637+
clientSessionId: 'test-session',
638+
userInputId: 'test-input',
639+
getLatestState: () => ({ messages: [] }),
640+
state: {
641+
// Missing required fields like ws, fingerprintId, etc.
642+
agentTemplate: parentAgent,
643+
localAgentTemplates: {},
644+
},
645+
})
646+
}).toThrow('Missing WebSocket in state')
647+
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
648+
})
649+
})
439650
})

0 commit comments

Comments
 (0)