change tool handlers to have required state fields

Author: charleslien

packages/agent-runtime/src/tools/handlers/handler-function-type.ts

@@ -23,59 +23,61 @@ import type { ProjectFileContext } from '@codebuff/common/util/file'
 type PresentOrAbsent<K extends PropertyKey, V> =
   | { [P in K]: V }
   | { [P in K]: never }
+export type State = {
+  agentStepId: string
+  clientSessionId: string
+  userInputId: string
+  creditsUsed?: number | Promise<number>
+  fingerprintId: string
+  userId: string | undefined
+  repoId: string | undefined
+  agentTemplate: AgentTemplate
+  localAgentTemplates: Record<string, AgentTemplate>
+  sendSubagentChunk: SendSubagentChunkFn
+  agentState: AgentState
+  prompt: string | undefined
+  fullResponse: string | undefined
+  agentContext: Record<
+    string,
+    {
+      logs: string[]
+      objective?: string | undefined
+      status?:
+        | 'NOT_STARTED'
+        | 'IN_PROGRESS'
+        | 'COMPLETE'
+        | 'ABORTED'
+        | undefined
+      plan?: string | undefined
+    }
+  >
+  messages: Message[]
+  system: string
+  logger: Logger
+} & FileProcessingState
 
 export type CodebuffToolHandlerFunction<T extends ToolName = ToolName> = (
   params: {
     previousToolCallFinished: Promise<void>
     toolCall: CodebuffToolCall<T>
 
-    runId: string
     agentStepId: string
+    ancestorRunIds: string[]
+    apiKey: string
     clientSessionId: string
-    userInputId: string
-    repoUrl: string | undefined
-    repoId: string | undefined
     fileContext: ProjectFileContext
-    apiKey: string
-
+    fullResponse: string
+    repoId: string | undefined
+    repoUrl: string | undefined
+    runId: string
     signal: AbortSignal
+    state: State
+    userInputId: string
 
-    ancestorRunIds: string[]
-
-    fullResponse: string
     fetch: typeof globalThis.fetch
-
-    writeToClient: (chunk: string | PrintModeEvent) => void
+    getLatestState: () => State
     trackEvent: TrackEventFn
-
-    getLatestState: () => any
-    state: {
-      creditsUsed?: number | Promise<number>
-      fingerprintId: string
-      userId: string | undefined
-      repoId: string | undefined
-      agentTemplate: AgentTemplate
-      localAgentTemplates: Record<string, AgentTemplate>
-      sendSubagentChunk: SendSubagentChunkFn
-      agentState: AgentState
-      agentContext: Record<
-        string,
-        {
-          logs: string[]
-          objective?: string | undefined
-          status?:
-            | 'NOT_STARTED'
-            | 'IN_PROGRESS'
-            | 'COMPLETE'
-            | 'ABORTED'
-            | undefined
-          plan?: string | undefined
-        }
-      >
-      messages: Message[]
-      system: string
-      logger: Logger
-    } & FileProcessingState
+    writeToClient: (chunk: string | PrintModeEvent) => void
   } & PresentOrAbsent<
     'requestClientToolCall',
     (
@@ -89,7 +91,6 @@ export type CodebuffToolHandlerFunction<T extends ToolName = ToolName> = (
   state?: Partial<
     {
       creditsUsed?: number | Promise<number>
-      fingerprintId: string
       userId: string | undefined
       repoId: string | undefined
       agentTemplate: AgentTemplate

packages/agent-runtime/src/tools/handlers/tool/add-subgoal.ts

@@ -10,7 +10,7 @@ import type { Subgoal } from '@codebuff/common/types/session-state'
 export const handleAddSubgoal = ((params: {
   previousToolCallFinished: Promise<void>
   toolCall: CodebuffToolCall<'add_subgoal'>
-  state: { agentContext?: Record<string, Subgoal> }
+  state: { agentContext: Record<string, Subgoal> }
 }): {
   result: Promise<CodebuffToolOutput<'add_subgoal'>>
   state: { agentContext: Record<string, Subgoal> }

packages/agent-runtime/src/tools/handlers/tool/create-plan.ts

@@ -3,10 +3,7 @@ import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import { getFileProcessingValues, postStreamProcessing } from './write-file'
 
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
-import type {
-  FileProcessingState,
-  OptionalFileProcessingState,
-} from './write-file'
+import type { FileProcessingState } from './write-file'
 import type {
   ClientToolCall,
   CodebuffToolCall,
@@ -27,13 +24,13 @@ export const handleCreatePlan = ((params: {
 
   getLatestState: () => FileProcessingState
   state: {
-    agentStepId?: string
-    clientSessionId?: string
-    fingerprintId?: string
-    userId?: string
-    userInputId?: string
-    repoId?: string
-  } & OptionalFileProcessingState
+    agentStepId: string
+    clientSessionId: string
+    fingerprintId: string
+    userId: string | undefined
+    userInputId: string
+    repoId: string | undefined
+  } & FileProcessingState
 }): {
   result: Promise<CodebuffToolOutput<'create_plan'>>
   state: FileProcessingState

packages/agent-runtime/src/tools/handlers/tool/find-files.ts

@@ -37,10 +37,10 @@ export const handleFindFiles = ((
     userInputId: string
 
     state: {
-      fingerprintId?: string
-      userId?: string
-      repoId?: string
-      messages?: Message[]
+      fingerprintId: string
+      userId: string | undefined
+      repoId: string | undefined
+      messages: Message[]
     }
   } & ParamsExcluding<
     typeof requestRelevantFiles,

packages/agent-runtime/src/tools/handlers/tool/read-docs.ts

@@ -1,6 +1,6 @@
-import { fetchContext7LibraryDocumentation } from '../../../llm-api/context7-api'
 import { callDocsSearchAPI } from '../../../llm-api/codebuff-web-api'
 
+import type { fetchContext7LibraryDocumentation } from '../../../llm-api/context7-api'
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 import type {
   CodebuffToolCall,
@@ -20,9 +20,9 @@ export const handleReadDocs = ((
     userInputId: string
 
     state: {
-      userId?: string
-      fingerprintId?: string
-      repoId?: string
+      fingerprintId: string
+      userId: string | undefined
+      repoId: string | undefined
     }
   } & ParamsExcluding<
     typeof fetchContext7LibraryDocumentation,
@@ -43,7 +43,7 @@ export const handleReadDocs = ((
     fetch,
   } = params
   const { libraryTitle, topic, max_tokens } = toolCall.input
-  const { userId, fingerprintId, repoId } = state
+  const { userId, repoId, fingerprintId } = state
 
   const docsStartTime = Date.now()
   const docsContext = {

packages/agent-runtime/src/tools/handlers/tool/read-files.ts

@@ -20,10 +20,10 @@ export const handleReadFiles = ((
     fileContext: ProjectFileContext
 
     state: {
-      userId?: string
-      fingerprintId?: string
-      repoId?: string
-      messages?: Message[]
+      userId: string | undefined
+      fingerprintId: string
+      repoId: string | undefined
+      messages: Message[]
     }
   } & ParamsExcluding<typeof getFileReadingUpdates, 'requestedFiles'>,
 ): {

packages/agent-runtime/src/tools/handlers/tool/set-output.ts

@@ -1,11 +1,27 @@
 import { getAgentTemplate } from '../../../templates/agent-registry'
 
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
+import type { CodebuffToolCall } from '@codebuff/common/tools/list'
+import type {
+  AgentTemplate,
+  Logger,
+} from '@codebuff/common/types/agent-template'
+import type { FetchAgentFromDatabaseFn } from '@codebuff/common/types/contracts/database'
+import type { AgentState } from '@codebuff/common/types/session-state'
 
 type ToolName = 'set_output'
-export const handleSetOutput: CodebuffToolHandlerFunction<ToolName> = (
-  params,
-) => {
+export const handleSetOutput = ((params: {
+  previousToolCallFinished: Promise<void>
+  toolCall: CodebuffToolCall<ToolName>
+  state: {
+    agentState: AgentState
+    localAgentTemplates: Record<string, AgentTemplate>
+  }
+  fetchAgentFromDatabase: FetchAgentFromDatabaseFn
+  databaseAgentCache: Map<string, AgentTemplate | null>
+  logger: Logger
+  apiKey: string
+}) => {
   const { previousToolCallFinished, toolCall, state, logger } = params
   const output = toolCall.input
   const { agentState, localAgentTemplates } = state
@@ -70,4 +86,4 @@ export const handleSetOutput: CodebuffToolHandlerFunction<ToolName> = (
     })(),
     state: { agentState: agentState },
   }
-}
+}) satisfies CodebuffToolHandlerFunction<ToolName>

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts

@@ -32,13 +32,13 @@ export const handleSpawnAgentInline = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
-      fingerprintId?: string
-      userId?: string
-      agentTemplate?: AgentTemplate
-      localAgentTemplates?: Record<string, AgentTemplate>
-      messages?: Message[]
-      agentState?: AgentState
-      system?: string
+      fingerprintId: string
+      userId: string | undefined
+      agentTemplate: AgentTemplate
+      localAgentTemplates: Record<string, AgentTemplate>
+      messages: Message[]
+      agentState: AgentState
+      system: string
     }
     logger: Logger
   } & ParamsExcluding<

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-utils.ts

@@ -27,13 +27,13 @@ export interface SpawnAgentParams {
 }
 
 export interface BaseSpawnState {
-  fingerprintId?: string
-  userId?: string
-  agentTemplate?: AgentTemplate
-  localAgentTemplates?: Record<string, AgentTemplate>
-  messages?: Message[]
-  agentState?: AgentState
-  system?: string
+  fingerprintId: string
+  userId: string | undefined
+  agentTemplate: AgentTemplate
+  localAgentTemplates: Record<string, AgentTemplate>
+  messages: Message[]
+  agentState: AgentState
+  system: string
 }
 
 export interface SpawnContext {

packages/agent-runtime/src/tools/handlers/tool/spawn-agents.ts

@@ -39,14 +39,14 @@ export const handleSpawnAgents = ((
 
     getLatestState: () => { messages: Message[] }
     state: {
-      fingerprintId?: string
-      userId?: string
-      agentTemplate?: AgentTemplate
-      localAgentTemplates?: Record<string, AgentTemplate>
-      sendSubagentChunk?: SendSubagentChunk
-      messages?: Message[]
-      agentState?: AgentState
-      system?: string
+      fingerprintId: string
+      userId: string | undefined
+      agentTemplate: AgentTemplate
+      localAgentTemplates: Record<string, AgentTemplate>
+      sendSubagentChunk: SendSubagentChunk
+      messages: Message[]
+      agentState: AgentState
+      system: string
     }
     logger: Logger
   } & ParamsExcluding<