Merge branch 'main' into brandon/tui

Author: brandonkachen

.agents/base2/base2-gpt-5.ts

@@ -15,7 +15,6 @@ const definition: SecretAgentDefinition = {
     'researcher-docs',
     'commander',
     'reviewer-gpt-5',
-    'editor-gpt-5',
     'context-pruner',
   ),
 }

.agents/base2/base2.ts

@@ -42,7 +42,6 @@ export const createBase2: (
       'commander',
       'generate-plan',
       'reviewer',
-      'editor',
       'context-pruner',
     ),
 
@@ -56,18 +55,6 @@ In between layers, you are encouraged to use the read_files tool to read files t
 
 Continue to spawn layers of agents until have completed the user's request or require more information from the user.
 
-## Example layers
-
-The user asks you to implement a new feature. You respond in multiple steps:
-
-1. Spawn file-pickers with different prompts to find relevant files; spawn a find-all-referencer to find more relevant files and answer questions about the codebase; spawn 1 docs researcher to find relevant docs.
-1a. Read all the relevant files using the read_files tool.
-2. Spawn one more file picker and one more find-all-referencer with different prompts to find relevant files.
-2a. Read all the relevant files using the read_files tool.
-3. Spawn a generate-plan agent to generate a plan for the changes.
-4. Use the str_replace or write_file tool to make the changes.
-5. Spawn a reviewer to review the changes.
-
 ## Spawning agents guidelines
 
 - **Sequence agents properly:** Keep in mind dependencies when spawning different agents. Don't spawn agents in parallel that depend on each other. Be conservative sequencing agents so they can build on each other's insights:
@@ -77,7 +64,6 @@ The user asks you to implement a new feature. You respond in multiple steps:
   - Reviewers should be spawned after you have made your edits.
 - **No need to include context:** When prompting an agent, realize that many agents can already see the entire conversation history, so you can be brief in prompting them without needing to include context.
 - **Don't spawn reviewers for trivial changes or quick follow-ups:** You should spawn the reviewer for most changes, but not for little changes or simple follow-ups.
-- **Don't spawn editors unless asked to parallelize or use multiple agents:** The editor performs worse at editing and is not to be used most of the time.
 
 # Core Mandates
 
@@ -130,9 +116,23 @@ The following is the state of the git repository at the start of the conversatio
 ${PLACEHOLDER.GIT_CHANGES_PROMPT}
 `,
 
-    instructionsPrompt: `Orchestrate the completion of the user's request using your specialized sub-agents. Take your time and be comprehensive.`,
+    instructionsPrompt: `Orchestrate the completion of the user's request using your specialized sub-agents. Take your time and be comprehensive.
+    
+## Example response
+
+The user asks you to implement a new feature. You respond in multiple steps:
+
+1. Spawn file-pickers with different prompts to find relevant files; spawn a find-all-referencer to find more relevant files and answer questions about the codebase; spawn 1 docs researcher to find relevant docs.
+1a. Read all the relevant files using the read_files tool.
+2. Spawn one more file picker and one more find-all-referencer with different prompts to find relevant files.
+2a. Read all the relevant files using the read_files tool.
+3. Spawn a generate-plan agent to generate a plan for the changes.
+4. Use the str_replace or write_file tool to make the changes.
+5. Spawn a reviewer to review the changes.
+6. Fix any issues raised by the reviewer.
+7. Inform the user that you have completed the task in one sentence without a final summary.`,
 
-    stepPrompt: `Don't forget to spawn agents that could help, especially: the file-picker and find-all-referencer to get codebase context, the generate-plan agent to create a plan, and the reviewer to review changes. No need to provide any final summary.`,
+    stepPrompt: `Don't forget to spawn agents that could help, especially: the file-picker and find-all-referencer to get codebase context, the generate-plan agent to create a plan, and the reviewer to review changes.`,
 
     handleSteps: function* ({ prompt, params }) {
       let steps = 0

.agents/file-explorer/file-picker-max.ts

@@ -0,0 +1,43 @@
+import { ToolCall } from 'types/agent-definition'
+import { publisher } from '../constants'
+
+import {
+  PLACEHOLDER,
+  type SecretAgentDefinition,
+} from '../types/secret-agent-definition'
+
+const definition: SecretAgentDefinition = {
+  id: 'file-picker-max',
+  displayName: 'Fletcher the File Fetcher',
+  publisher,
+  model: 'anthropic/claude-haiku-4.5',
+  spawnerPrompt:
+    'Spawn to find relevant files in a codebase related to the prompt. Cannot do string searches on the codebase.',
+  inputSchema: {
+    prompt: {
+      type: 'string',
+      description: 'A coding task to complete',
+    },
+  },
+  outputMode: 'last_message',
+  includeMessageHistory: false,
+  toolNames: [],
+  spawnableAgents: [],
+
+  systemPrompt: `You are an expert at finding relevant files in a codebase. ${PLACEHOLDER.FILE_TREE_PROMPT_SMALL}`,
+  instructionsPrompt: `Instructions:
+- Don't use any tools.
+- Provide a short report of the locations in the codebase that could be helpful. Focus on the files that are most relevant to the user prompt.
+In your report, please give a very concise analysis that includes the full paths of files that are relevant and (briefly) how they could be useful.
+  `.trim(),
+
+  handleSteps: function* ({ agentState, prompt, params }) {
+    yield {
+      toolName: 'find_files',
+      input: { prompt: prompt ?? '' },
+    } satisfies ToolCall
+    yield 'STEP_ALL'
+  },
+}
+
+export default definition

backend/src/__tests__/cost-aggregation.integration.test.ts

@@ -20,12 +20,14 @@ import * as agentRegistry from '../templates/agent-registry'
 import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '../templates/types'
+import type { ServerAction } from '@codebuff/common/actions'
 import type {
   AgentRuntimeDeps,
   AgentRuntimeScopedDeps,
 } from '@codebuff/common/types/contracts/agent-runtime'
+import type { SendActionFn } from '@codebuff/common/types/contracts/client'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
+import type { Mock } from 'bun:test'
 
 const mockFileContext: ProjectFileContext = {
   projectRoot: '/test',
@@ -84,34 +86,17 @@ const mockFileContext: ProjectFileContext = {
   },
 }
 
-class MockWebSocket {
-  sentActions: any[] = []
-
-  send(msg: string) {
-    // Capture sent messages for verification
-    try {
-      const parsed = JSON.parse(msg)
-      if (parsed.type === 'action') {
-        this.sentActions.push(parsed.data)
-      }
-    } catch {}
-  }
-
-  close() {}
-  on(event: string, listener: (...args: any[]) => void) {}
-  removeListener(event: string, listener: (...args: any[]) => void) {}
-}
-
 describe('Cost Aggregation Integration Tests', () => {
   let mockLocalAgentTemplates: Record<string, any>
-  let mockWebSocket: MockWebSocket
   let agentRuntimeImpl: AgentRuntimeDeps
   let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
 
   beforeEach(async () => {
     agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
-    agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
-    mockWebSocket = new MockWebSocket()
+    agentRuntimeScopedImpl = {
+      ...TEST_AGENT_RUNTIME_SCOPED_IMPL,
+      sendAction: mock(() => {}),
+    }
 
     // Setup mock agent templates
     mockLocalAgentTemplates = {
@@ -258,7 +243,6 @@ describe('Cost Aggregation Integration Tests', () => {
     const result = await mainPrompt({
       ...agentRuntimeImpl,
       ...agentRuntimeScopedImpl,
-      ws: mockWebSocket as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
@@ -294,17 +278,18 @@ describe('Cost Aggregation Integration Tests', () => {
     await websocketAction.callMainPrompt({
       ...agentRuntimeImpl,
       ...agentRuntimeScopedImpl,
-      ws: mockWebSocket as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
       promptId: 'test-prompt',
       clientSessionId: 'test-session',
     })
 
     // Verify final cost is included in prompt response
-    const promptResponse = mockWebSocket.sentActions.find(
-      (action) => action.type === 'prompt-response',
-    )
+    const promptResponse = (
+      agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
+    ).mock.calls
+      .map((call) => call[0].action)
+      .find((action: ServerAction) => action.type === 'prompt-response') as any
 
     expect(promptResponse).toBeDefined()
     expect(promptResponse.promptId).toBe('test-prompt')
@@ -363,7 +348,6 @@ describe('Cost Aggregation Integration Tests', () => {
     const result = await mainPrompt({
       ...agentRuntimeImpl,
       ...agentRuntimeScopedImpl,
-      ws: mockWebSocket as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
@@ -420,7 +404,6 @@ describe('Cost Aggregation Integration Tests', () => {
       result = await mainPrompt({
         ...agentRuntimeImpl,
         ...agentRuntimeScopedImpl,
-        ws: mockWebSocket as unknown as WebSocket,
         action,
         userId: TEST_USER_ID,
         clientSessionId: 'test-session',
@@ -470,7 +453,6 @@ describe('Cost Aggregation Integration Tests', () => {
     await mainPrompt({
       ...agentRuntimeImpl,
       ...agentRuntimeScopedImpl,
-      ws: mockWebSocket as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
       clientSessionId: 'test-session',
@@ -511,17 +493,18 @@ describe('Cost Aggregation Integration Tests', () => {
     await websocketAction.callMainPrompt({
       ...agentRuntimeImpl,
       ...agentRuntimeScopedImpl,
-      ws: mockWebSocket as unknown as WebSocket,
       action,
       userId: TEST_USER_ID,
       promptId: 'test-prompt',
       clientSessionId: 'test-session',
     })
 
     // Server should have reset the malicious value and calculated correct cost
-    const promptResponse = mockWebSocket.sentActions.find(
-      (action) => action.type === 'prompt-response',
-    )
+    const promptResponse = (
+      agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
+    ).mock.calls
+      .map((call) => call[0].action)
+      .find((action) => action.type === 'prompt-response') as any
 
     expect(promptResponse).toBeDefined()
     expect(promptResponse.sessionState.mainAgentState.creditsUsed).toBeLessThan(

backend/src/__tests__/cost-aggregation.test.ts

@@ -22,7 +22,6 @@ import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
 
 import type { AgentState } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
-import type { WebSocket } from 'ws'
 
 const mockFileContext: ProjectFileContext = {
   projectRoot: '/test',
@@ -137,7 +136,6 @@ describe('Cost Aggregation System', () => {
       }
 
       const mockValidatedState = {
-        ws: new MockWebSocket() as unknown as WebSocket,
         fingerprintId: 'test-fingerprint',
         userId: 'test-user',
         agentTemplate: mockAgentTemplate,
@@ -212,7 +210,6 @@ describe('Cost Aggregation System', () => {
       }
 
       const mockValidatedState = {
-        ws: new MockWebSocket() as unknown as WebSocket,
         fingerprintId: 'test-fingerprint',
         userId: 'test-user',
         agentTemplate: mockAgentTemplate,
@@ -371,7 +368,6 @@ describe('Cost Aggregation System', () => {
 
       // Mock subagent spawning that adds their costs
       const mockValidatedState = {
-        ws: new MockWebSocket() as unknown as WebSocket,
         fingerprintId: 'test-fingerprint',
         userId: 'test-user',
         agentTemplate: mockAgentTemplate,

backend/src/__tests__/fast-rewrite.test.ts

@@ -1,7 +1,10 @@
 import path from 'path'
 
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
-import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
+import {
+  TEST_AGENT_RUNTIME_IMPL,
+  TEST_AGENT_RUNTIME_SCOPED_IMPL,
+} from '@codebuff/common/testing/impl/agent-runtime'
 import {
   clearMockedModules,
   mockModule,
@@ -11,11 +14,15 @@ import { createPatch } from 'diff'
 
 import { rewriteWithOpenAI } from '../fast-rewrite'
 
-import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
-
-let agentRuntimeImpl: AgentRuntimeDeps
+import type {
+  AgentRuntimeDeps,
+  AgentRuntimeScopedDeps,
+} from '@codebuff/common/types/contracts/agent-runtime'
 
 describe.skip('rewriteWithOpenAI', () => {
+  let agentRuntimeImpl: AgentRuntimeDeps
+  let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
+
   beforeAll(() => {
     // Mock database interactions
     mockModule('pg-pool', () => ({
@@ -40,6 +47,7 @@ describe.skip('rewriteWithOpenAI', () => {
 
   beforeEach(() => {
     agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
+    agentRuntimeScopedImpl = { ...TEST_AGENT_RUNTIME_SCOPED_IMPL }
   })
 
   afterAll(() => {
@@ -54,6 +62,7 @@ describe.skip('rewriteWithOpenAI', () => {
 
     const result = await rewriteWithOpenAI({
       ...agentRuntimeImpl,
+      ...agentRuntimeScopedImpl,
       oldContent: originalContent,
       editSnippet,
       clientSessionId: 'clientSessionId',