For "last_message" output mode, return the entire assistant response, including tools and results

Author: jahooma

common/src/types/session-state.ts

@@ -48,7 +48,7 @@ export const AgentOutputSchema = z.discriminatedUnion('type', [
   }),
   z.object({
     type: z.literal('lastMessage'),
-    value: z.any(),
+    value: z.array(z.any()), // Array of assistant and tool messages from the last turn, including tool results
   }),
   z.object({
     type: z.literal('allMessages'),

packages/agent-runtime/src/__tests__/cost-aggregation.test.ts

@@ -159,7 +159,7 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 75, // First subagent uses 75 credits
           },
-          output: { type: 'lastMessage', value: 'Sub-agent 1 response' },
+          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] },
         })
         .mockResolvedValueOnce({
           agentState: {
@@ -169,7 +169,7 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 100, // Second subagent uses 100 credits
           },
-          output: { type: 'lastMessage', value: 'Sub-agent 2 response' },
+          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] },
         })
 
       const mockToolCall = {
@@ -223,7 +223,7 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: 50, // Successful agent
           },
-          output: { type: 'lastMessage', value: 'Successful response' },
+          output: { type: 'lastMessage', value: [assistantMessage('Successful response')] },
         })
         .mockRejectedValueOnce(
           (() => {
@@ -370,7 +370,7 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: subAgent1Cost,
           } as AgentState,
-          output: { type: 'lastMessage', value: 'Sub-agent 1 response' },
+          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] },
         })
         .mockResolvedValueOnce({
           agentState: {
@@ -381,7 +381,7 @@ describe('Cost Aggregation System', () => {
             stepsRemaining: 10,
             creditsUsed: subAgent2Cost,
           } as AgentState,
-          output: { type: 'lastMessage', value: 'Sub-agent 2 response' },
+          output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] },
         })
 
       const mockToolCall = {

packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts

@@ -52,7 +52,7 @@ describe('Spawn Agents Message History', () => {
             assistantMessage('Mock agent response'),
           ],
         },
-        output: { type: 'lastMessage', value: 'Mock agent response' },
+        output: { type: 'lastMessage', value: [assistantMessage('Mock agent response')] },
       }
     })
 

packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts

@@ -85,7 +85,7 @@ describe('Spawn Agents Permissions', () => {
           ...options.agentState,
           messageHistory: [assistantMessage('Mock agent response')],
         },
-        output: { type: 'lastMessage', value: 'Mock agent response' },
+        output: { type: 'lastMessage', value: [assistantMessage('Mock agent response')] },
       }
     })
   })

packages/agent-runtime/src/__tests__/subagent-streaming.test.ts

@@ -96,7 +96,7 @@ describe('Subagent Streaming', () => {
           ...options.agentState,
           messageHistory: [assistantMessage('Test response from subagent')],
         },
-        output: { type: 'lastMessage', value: 'Test response from subagent' },
+        output: { type: 'lastMessage', value: [assistantMessage('Test response from subagent')] },
       }
     })
 

packages/agent-runtime/src/util/agent-output.ts

@@ -1,10 +1,49 @@
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
-import type { AssistantMessage } from '@codebuff/common/types/messages/codebuff-message'
+import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type {
   AgentState,
   AgentOutput,
 } from '@codebuff/common/types/session-state'
 
+/**
+ * Get the last assistant turn messages, which includes the last assistant message
+ * and any subsequent tool messages that are responses to its tool calls.
+ */
+function getLastAssistantTurnMessages(messageHistory: Message[]): Message[] {
+  // Find the index of the last assistant message
+  let lastAssistantIndex = -1
+  for (let i = messageHistory.length - 1; i >= 0; i--) {
+    if (messageHistory[i].role === 'assistant') {
+      lastAssistantIndex = i
+      break
+    }
+  }
+
+  for (let i = lastAssistantIndex; i >= 0; i--) {
+    if (messageHistory[i].role === 'assistant') {
+      lastAssistantIndex = i
+    } else break
+  }
+
+  if (lastAssistantIndex === -1) {
+    return []
+  }
+
+  // Collect the assistant message and all subsequent tool messages
+  const result: Message[] = []
+  for (let i = lastAssistantIndex; i < messageHistory.length; i++) {
+    const message = messageHistory[i]
+    if (message.role === 'assistant' || message.role === 'tool') {
+      result.push(message)
+    } else {
+      // Stop if we hit a user or system message
+      break
+    }
+  }
+
+  return result
+}
+
 export function getAgentOutput(
   agentState: AgentState,
   agentTemplate: AgentTemplate,
@@ -16,19 +55,18 @@ export function getAgentOutput(
     }
   }
   if (agentTemplate.outputMode === 'last_message') {
-    const assistantMessages = agentState.messageHistory.filter(
-      (message): message is AssistantMessage => message.role === 'assistant',
+    const lastTurnMessages = getLastAssistantTurnMessages(
+      agentState.messageHistory,
     )
-    const lastAssistantMessage = assistantMessages[assistantMessages.length - 1]
-    if (!lastAssistantMessage) {
+    if (lastTurnMessages.length === 0) {
       return {
         type: 'error',
         message: 'No response from agent',
       }
     }
     return {
       type: 'lastMessage',
-      value: lastAssistantMessage.content,
+      value: lastTurnMessages,
     }
   }
   if (agentTemplate.outputMode === 'all_messages') {