add maxContextLength parameter

Author: charleslien

.agents/context-pruner.ts

@@ -1,4 +1,5 @@
 import { publisher } from './constants'
+
 import type { AgentDefinition, Message } from './types/agent-definition'
 
 const definition: AgentDefinition = {
@@ -11,19 +12,29 @@ const definition: AgentDefinition = {
 
   spawnerPrompt: `Spawn this agent between steps to prune context, starting with old tool results and then old messages.`,
 
-  inputSchema: {},
+  inputSchema: {
+    params: {
+      type: 'object',
+      properties: {
+        maxContextLength: {
+          type: 'number',
+        },
+      },
+      required: [],
+    },
+  },
 
   includeMessageHistory: true,
 
-  handleSteps: function* ({ agentState }) {
+  handleSteps: function* ({ agentState, params }) {
     const messages = agentState.messageHistory
 
     const countTokensJson = (obj: any): number => {
       // Very rough approximation
       return Math.ceil(JSON.stringify(obj).length / 3)
     }
 
-    const maxMessageTokens = 200_000
+    const maxMessageTokens: number = params?.maxContextLength ?? 200_000
     const numTerminalCommandsToKeep = 5
 
     // Remove the last assistant message if it contains the spawn call that invoked this context-pruner

.agents/factory/base.ts

@@ -19,6 +19,15 @@ export const base = (model: ModelName): Omit<SecretAgentDefinition, 'id'> => ({
       type: 'string',
       description: 'A coding task to complete',
     },
+    params: {
+      type: 'object',
+      properties: {
+        maxContextLength: {
+          type: 'number',
+        },
+      },
+      required: [],
+    },
   },
   outputMode: 'last_message',
   includeMessageHistory: false,
@@ -50,13 +59,14 @@ export const base = (model: ModelName): Omit<SecretAgentDefinition, 'id'> => ({
   instructionsPrompt: baseAgentUserInputPrompt(model),
   stepPrompt: baseAgentAgentStepPrompt(model),
 
-  handleSteps: function* ({ agentState }) {
+  handleSteps: function* ({ params }) {
     while (true) {
       // Run context-pruner before each step
       yield {
         toolName: 'spawn_agent_inline',
         input: {
           agent_type: 'context-pruner',
+          params,
         },
       } as any