From 2b1f301848ed48b6808de26218eae20a73d5e039 Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Tue, 15 Jul 2025 16:14:33 -0700 Subject: [PATCH] fix: Better filtering of invalid messages (#417) * fix: Better filtering of invalid messages * cr * cr --- apps/open-swe/src/__tests__/tokens.test.ts | 430 +++++++++++++++++- .../planner/nodes/generate-message/index.ts | 14 +- .../planner/nodes/generate-plan/index.ts | 13 +- .../src/graphs/planner/nodes/proposed-plan.ts | 2 +- .../src/graphs/planner/nodes/take-action.ts | 4 + .../nodes/generate-message/index.ts | 12 +- .../programmer/nodes/progress-plan-step.ts | 12 +- .../programmer/nodes/summarize-history.ts | 53 +-- .../graphs/programmer/nodes/take-action.ts | 4 + .../graphs/reviewer/nodes/initialize-state.ts | 2 +- .../reviewer/nodes/take-review-action.ts | 3 + .../src/graphs/shared/initialize-sandbox.ts | 2 +- apps/open-swe/src/tools/shell.ts | 2 +- apps/open-swe/src/utils/message/content.ts | 17 + apps/open-swe/src/utils/tokens.ts | 123 ++++- .../src/components/gen-ui/task-summary.tsx | 8 +- .../web/src/components/thread/messages/ai.tsx | 10 +- packages/shared/src/open-swe/tools.ts | 4 +- 18 files changed, 605 insertions(+), 110 deletions(-) diff --git a/apps/open-swe/src/__tests__/tokens.test.ts b/apps/open-swe/src/__tests__/tokens.test.ts index 5532a393..6863dc73 100644 --- a/apps/open-swe/src/__tests__/tokens.test.ts +++ b/apps/open-swe/src/__tests__/tokens.test.ts @@ -6,13 +6,13 @@ import { MAX_INTERNAL_TOKENS, } from "../utils/tokens.js"; -describe("calculateConversationHistoryTokenCount", () => { - it("should return 0 for empty messages array", () => { +describe("calculateConversationHistoryTokenCount", async () => { + it("should return 0 for empty messages array", async () => { const result = calculateConversationHistoryTokenCount([]); expect(result).toBe(0); }); - it("should calculate token count for human messages", () => { + it("should calculate token count for human messages", async () => { const messages = [ new HumanMessage({ content: "This is a test message with exactly 10 words in it.", @@ -25,7 +25,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(result).toBe(13); }); - it("should calculate token count for AI messages with usage metadata", () => { + it("should calculate token count for AI messages with usage metadata", async () => { const messages = [ new AIMessage({ content: "AI response", @@ -41,7 +41,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(result).toBe(20); }); - it("should calculate token count for AI messages without usage metadata", () => { + it("should calculate token count for AI messages without usage metadata", async () => { const messages = [ new AIMessage({ content: "This is an AI response with no usage metadata.", @@ -54,7 +54,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(result).toBe(12); }); - it("should calculate token count for AI messages with tool calls", () => { + it("should calculate token count for AI messages with tool calls", async () => { const messages = [ new AIMessage({ content: "Using a tool", @@ -75,7 +75,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(result).toBeGreaterThan(0); }); - it("should calculate token count for tool messages", () => { + it("should calculate token count for tool messages", async () => { const messages = [ new ToolMessage({ content: "Result of tool execution with some data.", @@ -89,7 +89,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(result).toBe(10); }); - it("should exclude hidden messages when option is provided", () => { + it("should exclude hidden messages when option is provided", async () => { const messages = [ new HumanMessage({ content: "Visible message", @@ -110,7 +110,7 @@ describe("calculateConversationHistoryTokenCount", () => { expect(resultWithOption).toBe(4); // "Visible message" is ~4 tokens }); - it("should exclude messages from the end when option is provided", () => { + it("should exclude messages from the end when option is provided", async () => { const messages = [ new HumanMessage({ content: "First message" }), new HumanMessage({ content: "Second message" }), @@ -127,22 +127,107 @@ describe("calculateConversationHistoryTokenCount", () => { // First two messages should be ~7 tokens expect(resultWithOption).toBe(7); }); + + it("should not separate AI messages with tool calls from their tool messages when excluding from end", async () => { + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll help you with that", + tool_calls: [ + { + name: "test_tool", + args: { param: "value" }, + id: "call_123", + }, + ], + }); + + const toolMessage = new ToolMessage({ + content: "Tool result", + tool_call_id: "call_123", + }); + + const messages = [ + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage, + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 2 messages from the end, which would normally cut between AI and tool message + const result = calculateConversationHistoryTokenCount(messages, { + excludeCountFromEnd: 2, + }); + + // Should only count the first human message since we can't separate AI/tool pair + const expectedResult = calculateConversationHistoryTokenCount([ + new HumanMessage({ content: "First message" }), + ]); + + expect(result).toBe(expectedResult); + }); + + it("should preserve multiple tool messages following an AI message", async () => { + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll use multiple tools", + tool_calls: [ + { + name: "tool1", + args: { param: "value1" }, + id: "call_1", + }, + { + name: "tool2", + args: { param: "value2" }, + id: "call_2", + }, + ], + }); + + const toolMessage1 = new ToolMessage({ + content: "Tool 1 result", + tool_call_id: "call_1", + }); + + const toolMessage2 = new ToolMessage({ + content: "Tool 2 result", + tool_call_id: "call_2", + }); + + const messages = [ + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage1, + toolMessage2, + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 3 messages from the end, which would cut in the middle of tool messages + const result = calculateConversationHistoryTokenCount(messages, { + excludeCountFromEnd: 3, + }); + + // Should only count the first human message + const expectedResult = calculateConversationHistoryTokenCount([ + new HumanMessage({ content: "First message" }), + ]); + + expect(result).toBe(expectedResult); + }); }); -describe("getMessagesSinceLastSummary", () => { - it("should return all messages when there is no summary message", () => { +describe("getMessagesSinceLastSummary", async () => { + it("should return all messages when there is no summary message", async () => { const messages = [ new HumanMessage({ content: "Message 1" }), new AIMessage({ content: "Message 2" }), new HumanMessage({ content: "Message 3" }), ]; - const result = getMessagesSinceLastSummary(messages); + const result = await getMessagesSinceLastSummary(messages); expect(result).toHaveLength(3); expect(result).toEqual(messages); }); - it("should return messages after the last summary message", () => { + it("should return messages after the last summary message", async () => { const summaryMessage = new AIMessage({ content: "Summary of conversation", additional_kwargs: { summary_message: true }, @@ -155,13 +240,13 @@ describe("getMessagesSinceLastSummary", () => { new AIMessage({ content: "Message 4" }), ]; - const result = getMessagesSinceLastSummary(messages); + const result = await getMessagesSinceLastSummary(messages); expect(result).toHaveLength(2); expect(result[0].content).toBe("Message 3"); expect(result[1].content).toBe("Message 4"); }); - it("should exclude hidden messages when option is provided", () => { + it("should exclude hidden messages when option is provided", async () => { const summaryMessage = new AIMessage({ content: "Summary of conversation", additional_kwargs: { summary_message: true }, @@ -177,7 +262,7 @@ describe("getMessagesSinceLastSummary", () => { new AIMessage({ content: "Another visible message" }), ]; - const result = getMessagesSinceLastSummary(messages, { + const result = await getMessagesSinceLastSummary(messages, { excludeHiddenMessages: true, }); @@ -186,7 +271,7 @@ describe("getMessagesSinceLastSummary", () => { expect(result[1].content).toBe("Another visible message"); }); - it("should exclude messages from the end when option is provided", () => { + it("should exclude messages from the end when option is provided", async () => { const summaryMessage = new AIMessage({ content: "Summary of conversation", additional_kwargs: { summary_message: true }, @@ -199,7 +284,7 @@ describe("getMessagesSinceLastSummary", () => { new HumanMessage({ content: "Message 3" }), ]; - const result = getMessagesSinceLastSummary(messages, { + const result = await getMessagesSinceLastSummary(messages, { excludeCountFromEnd: 1, }); @@ -208,7 +293,7 @@ describe("getMessagesSinceLastSummary", () => { expect(result[1].content).toBe("Message 2"); }); - it("should handle both excludeHiddenMessages and excludeCountFromEnd options", () => { + it("should handle both excludeHiddenMessages and excludeCountFromEnd options", async () => { const summaryMessage = new AIMessage({ content: "Summary of conversation", additional_kwargs: { summary_message: true }, @@ -225,7 +310,7 @@ describe("getMessagesSinceLastSummary", () => { new HumanMessage({ content: "Message 4" }), ]; - const result = getMessagesSinceLastSummary(messages, { + const result = await getMessagesSinceLastSummary(messages, { excludeHiddenMessages: true, excludeCountFromEnd: 1, }); @@ -235,7 +320,304 @@ describe("getMessagesSinceLastSummary", () => { expect(result[1].content).toBe("Message 3"); }); - it("should return empty array if all messages are before the summary", () => { + it("should not separate AI messages with tool calls from their tool messages when excluding from end", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll help you with that", + tool_calls: [ + { + name: "test_tool", + args: { param: "value" }, + id: "call_123", + }, + ], + }); + + const toolMessage = new ToolMessage({ + content: "Tool result", + tool_call_id: "call_123", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage, + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 2 messages from the end, which would normally cut between AI and tool message + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 2, + }); + + // Should only include the first human message since we can't separate AI/tool pair + expect(result).toHaveLength(1); + expect(result[0].content).toBe("First message"); + }); + + it("should preserve multiple tool messages following an AI message in getMessagesSinceLastSummary", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll use multiple tools", + tool_calls: [ + { + name: "tool1", + args: { param: "value1" }, + id: "call_1", + }, + { + name: "tool2", + args: { param: "value2" }, + id: "call_2", + }, + ], + }); + + const toolMessage1 = new ToolMessage({ + content: "Tool 1 result", + tool_call_id: "call_1", + }); + + const toolMessage2 = new ToolMessage({ + content: "Tool 2 result", + tool_call_id: "call_2", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage1, + toolMessage2, + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 3 messages from the end, which would cut in the middle of tool messages + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 3, + }); + + // Should only include the first human message + expect(result).toHaveLength(1); + expect(result[0].content).toBe("First message"); + }); + + it("should exclude entire AI/tool group when cut point would separate them", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll use a tool", + tool_calls: [ + { + name: "test_tool", + args: { param: "value" }, + id: "call_123", + }, + ], + }); + + const toolMessage = new ToolMessage({ + content: "Tool result", + tool_call_id: "call_123", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage, + ]; + + // Try to exclude 1 message from the end (just the tool message) + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 1, + }); + + // Should exclude the entire AI/tool group to maintain integrity + expect(result).toHaveLength(1); + expect(result[0].content).toBe("First message"); + }); + + it("should preserve AI message with multiple tool calls and their corresponding tool messages", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithMultipleToolCalls = new AIMessage({ + content: "I'll use multiple tools to help you", + tool_calls: [ + { + name: "search_tool", + args: { query: "example query" }, + id: "call_search_123", + }, + { + name: "calculator_tool", + args: { expression: "2 + 2" }, + id: "call_calc_456", + }, + { + name: "file_tool", + args: { filename: "test.txt" }, + id: "call_file_789", + }, + ], + }); + + const searchToolMessage = new ToolMessage({ + content: "Search results found", + tool_call_id: "call_search_123", + }); + + const calculatorToolMessage = new ToolMessage({ + content: "Result: 4", + tool_call_id: "call_calc_456", + }); + + const fileToolMessage = new ToolMessage({ + content: "File contents: Hello world", + tool_call_id: "call_file_789", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithMultipleToolCalls, + searchToolMessage, + calculatorToolMessage, + fileToolMessage, + new HumanMessage({ content: "After all tools" }), + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 4 messages from the end, which would cut in the middle of the tool messages + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 4, + }); + + // Should include the first human message and the complete AI/tool group since we can't separate them + expect(result).toHaveLength(5); + expect(result[0].content).toBe("First message"); + expect(result[1].content).toBe("I'll use multiple tools to help you"); + expect(result[2].content).toBe("Search results found"); + expect(result[3].content).toBe("Result: 4"); + expect(result[4].content).toBe("File contents: Hello world"); + }); + + it("should include complete AI/tool group when exclusion doesn't break the group", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithMultipleToolCalls = new AIMessage({ + content: "I'll use two tools", + tool_calls: [ + { + name: "tool1", + args: { param: "value1" }, + id: "call_1", + }, + { + name: "tool2", + args: { param: "value2" }, + id: "call_2", + }, + ], + }); + + const tool1Message = new ToolMessage({ + content: "Tool 1 result", + tool_call_id: "call_1", + }); + + const tool2Message = new ToolMessage({ + content: "Tool 2 result", + tool_call_id: "call_2", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithMultipleToolCalls, + tool1Message, + tool2Message, + new HumanMessage({ content: "After tools" }), + new HumanMessage({ content: "Second to last" }), + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 2 messages from the end (just the last two human messages) + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 2, + }); + + // Should include the first human message, the AI message, both tool messages, and the "After tools" message + expect(result).toHaveLength(5); + expect(result[0].content).toBe("First message"); + expect(result[1].content).toBe("I'll use two tools"); + expect(result[2].content).toBe("Tool 1 result"); + expect(result[3].content).toBe("Tool 2 result"); + expect(result[4].content).toBe("After tools"); + }); + + it("should handle case where AI/tool group can be included entirely", async () => { + const summaryMessage = new AIMessage({ + content: "Summary of conversation", + additional_kwargs: { summary_message: true }, + }); + + const aiMessageWithToolCalls = new AIMessage({ + content: "I'll use a tool", + tool_calls: [ + { + name: "test_tool", + args: { param: "value" }, + id: "call_123", + }, + ], + }); + + const toolMessage = new ToolMessage({ + content: "Tool result", + tool_call_id: "call_123", + }); + + const messages = [ + summaryMessage, + new HumanMessage({ content: "First message" }), + aiMessageWithToolCalls, + toolMessage, + new HumanMessage({ content: "After tool message" }), + new HumanMessage({ content: "Last message" }), + ]; + + // Try to exclude 2 messages from the end (the last two human messages) + const result = await getMessagesSinceLastSummary(messages, { + excludeCountFromEnd: 2, + }); + + // Should include the first human message and the complete AI/tool group + expect(result).toHaveLength(3); + expect(result[0].content).toBe("First message"); + expect(result[1].content).toBe("I'll use a tool"); + expect(result[2].content).toBe("Tool result"); + }); + + it("should return empty array if all messages are before the summary", async () => { const messages = [ new HumanMessage({ content: "Message 1" }), new AIMessage({ content: "Message 2" }), @@ -245,13 +627,13 @@ describe("getMessagesSinceLastSummary", () => { }), ]; - const result = getMessagesSinceLastSummary(messages); + const result = await getMessagesSinceLastSummary(messages); expect(result).toHaveLength(0); }); }); -describe("MAX_INTERNAL_TOKENS constant", () => { - it("should be defined as 60,000", () => { +describe("MAX_INTERNAL_TOKENS constant", async () => { + it("should be defined as 60,000", async () => { expect(MAX_INTERNAL_TOKENS).toBe(60_000); }); }); diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts index f5ebeabd..7856ff26 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts @@ -17,12 +17,12 @@ import { import { SYSTEM_PROMPT } from "./prompt.js"; import { getRepoAbsolutePath } from "@open-swe/shared/git"; import { getMissingMessages } from "../../../../utils/github/issue-messages.js"; -import { filterHiddenMessages } from "../../../../utils/message/filter-hidden.js"; import { getPlansFromIssue } from "../../../../utils/github/issue-task.js"; import { createSearchTool } from "../../../../tools/search.js"; import { formatCustomRulesPrompt } from "../../../../utils/custom-rules.js"; import { createPlannerNotesTool } from "../../../../tools/planner-notes.js"; import { getMcpTools } from "../../../../utils/mcp-client.js"; +import { filterMessagesWithoutContent } from "../../../../utils/message/content.js"; const logger = createLogger(LogLevel.INFO, "GeneratePlanningMessageNode"); @@ -73,6 +73,15 @@ export async function generateAction( getMissingMessages(state, config), getPlansFromIssue(state, config), ]); + + const inputMessages = filterMessagesWithoutContent([ + ...state.messages, + ...missingMessages, + ]); + if (!inputMessages.length) { + throw new Error("No messages to process."); + } + const response = await modelWithTools .withConfig({ tags: ["nostream"] }) .invoke([ @@ -83,8 +92,7 @@ export async function generateAction( taskPlan: latestTaskPlan ?? state.taskPlan, }), }, - ...filterHiddenMessages(state.messages), - ...missingMessages, + ...inputMessages, ]); logger.info("Generated planning message", { diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-plan/index.ts b/apps/open-swe/src/graphs/planner/nodes/generate-plan/index.ts index 4fbd309c..a79ef25b 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-plan/index.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-plan/index.ts @@ -13,12 +13,12 @@ import { isFollowupRequest, } from "../../utils/followup.js"; import { stopSandbox } from "../../../../utils/sandbox.js"; -import { filterHiddenMessages } from "../../../../utils/message/filter-hidden.js"; import { z } from "zod"; import { formatCustomRulesPrompt } from "../../../../utils/custom-rules.js"; import { getPlannerNotes } from "../../utils/get-notes.js"; import { PLANNER_NOTES_PROMPT, SYSTEM_PROMPT } from "./prompt.js"; import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants"; +import { filterMessagesWithoutContent } from "../../../../utils/message/content.js"; function formatSystemPrompt(state: PlannerGraphState): string { // It's a followup if there's more than one human message. @@ -68,6 +68,14 @@ export async function generatePlan( }); } + const inputMessages = filterMessagesWithoutContent([ + ...state.messages, + ...(optionalToolMessage ? [optionalToolMessage] : []), + ]); + if (!inputMessages.length) { + throw new Error("No messages to process."); + } + const response = await modelWithTools .withConfig({ tags: ["nostream"] }) .invoke([ @@ -75,8 +83,7 @@ export async function generatePlan( role: "system", content: formatSystemPrompt(state), }, - ...filterHiddenMessages(state.messages), - ...(optionalToolMessage ? [optionalToolMessage] : []), + ...inputMessages, ]); const toolCall = response.tool_calls?.[0]; diff --git a/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts b/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts index f87a8f60..9178eb65 100644 --- a/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts +++ b/apps/open-swe/src/graphs/planner/nodes/proposed-plan.ts @@ -58,7 +58,7 @@ function createAcceptedPlanMessage(input: { const acceptedPlanMessage = new AIMessage({ id: `${DO_NOT_RENDER_ID_PREFIX}${uuidv4()}`, - content: "", + content: "Accepted plan", additional_kwargs: { hidden: true, customNodeEvents: [acceptedPlanEvent], diff --git a/apps/open-swe/src/graphs/planner/nodes/take-action.ts b/apps/open-swe/src/graphs/planner/nodes/take-action.ts index bfdda6e5..900634ce 100644 --- a/apps/open-swe/src/graphs/planner/nodes/take-action.ts +++ b/apps/open-swe/src/graphs/planner/nodes/take-action.ts @@ -116,6 +116,10 @@ export async function takeActions( result = toolResult.result; toolCallStatus = toolResult.status; } + + if (!result) { + result = toolCallStatus; + } } catch (e) { toolCallStatus = "error"; if ( diff --git a/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts b/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts index e1e15ff2..d32ef320 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/generate-message/index.ts @@ -34,6 +34,7 @@ import { formatCodeReviewPrompt, getCodeReviewFields, } from "../../../../utils/review.js"; +import { filterMessagesWithoutContent } from "../../../../utils/message/content.js"; const logger = createLogger(LogLevel.INFO, "GenerateMessageNode"); @@ -117,6 +118,14 @@ export async function generateAction( getPlansFromIssue(state, config), ]); + const inputMessages = filterMessagesWithoutContent([ + ...state.internalMessages, + ...missingMessages, + ]); + if (!inputMessages.length) { + throw new Error("No messages to process."); + } + const response = await modelWithTools.invoke([ { role: "system", @@ -125,8 +134,7 @@ export async function generateAction( taskPlan: latestTaskPlan ?? state.taskPlan, }), }, - ...state.internalMessages, - ...missingMessages, + ...inputMessages, ]); const hasToolCalls = !!response.tool_calls?.length; diff --git a/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts b/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts index 92b505e2..818c6286 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/progress-plan-step.ts @@ -29,6 +29,7 @@ import { } from "@open-swe/shared/open-swe/tools"; import { calculateConversationHistoryTokenCount, + getMessagesSinceLastSummary, MAX_INTERNAL_TOKENS, } from "../../../utils/tokens.js"; import { z } from "zod"; @@ -117,8 +118,14 @@ Once you've determined the status of the current task, call either the \`mark_ta const newMessages = [response, toolMessage]; + const newMessageList = [...state.internalMessages, ...newMessages]; + const wouldBeConversationHistoryToSummarize = + await getMessagesSinceLastSummary(newMessageList, { + excludeHiddenMessages: true, + excludeCountFromEnd: 20, + }); const totalInternalTokenCount = calculateConversationHistoryTokenCount( - state.internalMessages, + wouldBeConversationHistoryToSummarize, { // Retain the last 20 messages from state excludeHiddenMessages: true, @@ -135,12 +142,15 @@ Once you've determined the status of the current task, call either the \`mark_ta internalMessages: newMessages, }; + // Check if we have any messages to summarize, and if we're at or above the max token limit. if (totalInternalTokenCount >= MAX_INTERNAL_TOKENS) { logger.info( "Internal messages list is at or above the max token limit. Routing to summarize history step.", { totalInternalTokenCount, maxInternalTokenCount: MAX_INTERNAL_TOKENS, + wouldBeConversationHistoryToSummarizeLength: + wouldBeConversationHistoryToSummarize.length, }, ); return new Command({ diff --git a/apps/open-swe/src/graphs/programmer/nodes/summarize-history.ts b/apps/open-swe/src/graphs/programmer/nodes/summarize-history.ts index b069fcbf..6818c523 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/summarize-history.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/summarize-history.ts @@ -18,8 +18,6 @@ import { getMessageContentString } from "@open-swe/shared/messages"; import { getMessageString } from "../../../utils/message/content.js"; import { getActivePlanItems } from "@open-swe/shared/open-swe/tasks"; import { createConversationHistorySummaryToolFields } from "@open-swe/shared/open-swe/tools"; -import { z } from "zod"; -import { DO_NOT_RENDER_ID_PREFIX } from "@open-swe/shared/constants"; import { getUserRequest } from "../../../utils/user-request.js"; import { getMessagesSinceLastSummary } from "../../../utils/tokens.js"; @@ -91,45 +89,9 @@ const formatPrompt = (inputs: { ); }; -/** - * Create an AI & tool message pair for the generated task summary. - * This is not included in the internal message state, but is exposed to - * users so they can see the summary of the actions that were taken. - */ -function createUserFacingConversationSummaryMessages( - summary: string, -): BaseMessage[] { - const conversationSummaryTool = createConversationHistorySummaryToolFields(); - const conversationSummaryToolCallArgs: z.infer< - typeof conversationSummaryTool.schema - > = { - conversation_history_summary: summary, - }; - const conversationSummaryToolCallId = uuidv4(); - const conversationSummaryPublicMessages = [ - new AIMessage({ - id: uuidv4(), - content: "", - tool_calls: [ - { - id: conversationSummaryToolCallId, - name: conversationSummaryTool.name, - args: conversationSummaryToolCallArgs, - }, - ], - }), - new ToolMessage({ - id: `${DO_NOT_RENDER_ID_PREFIX}${uuidv4()}`, - tool_call_id: conversationSummaryToolCallId, - content: "", - }), - ]; - - return conversationSummaryPublicMessages; -} - -function createInternalSummaryMessages(summary: string): BaseMessage[] { - const dummySummarizeHistoryToolName = "summarize_conversation_history"; +function createSummaryMessages(summary: string): BaseMessage[] { + const dummySummarizeHistoryToolName = + createConversationHistorySummaryToolFields().name; const dummySummarizeHistoryToolCallId = uuidv4(); return [ new AIMessage({ @@ -169,7 +131,7 @@ export async function summarizeHistory( const userRequest = getUserRequest(state.messages); const plan = getActivePlanItems(state.taskPlan); - const conversationHistoryToSummarize = getMessagesSinceLastSummary( + const conversationHistoryToSummarize = await getMessagesSinceLastSummary( state.internalMessages, { excludeHiddenMessages: true, @@ -193,14 +155,13 @@ export async function summarizeHistory( ]); const summaryString = getMessageContentString(response.content); - const taskSummaryMessages = - createUserFacingConversationSummaryMessages(summaryString); + const summaryMessages = createSummaryMessages(summaryString); const newInternalMessages = [ ...conversationHistoryToSummarize.map( (m) => new RemoveMessage({ id: m.id ?? "" }), ), - ...createInternalSummaryMessages(summaryString), + ...summaryMessages, ]; logger.info( @@ -208,7 +169,7 @@ export async function summarizeHistory( ); return { - messages: taskSummaryMessages, + messages: summaryMessages, internalMessages: newInternalMessages, }; } diff --git a/apps/open-swe/src/graphs/programmer/nodes/take-action.ts b/apps/open-swe/src/graphs/programmer/nodes/take-action.ts index 7dc8ce01..6308ada9 100644 --- a/apps/open-swe/src/graphs/programmer/nodes/take-action.ts +++ b/apps/open-swe/src/graphs/programmer/nodes/take-action.ts @@ -106,6 +106,10 @@ export async function takeAction( result = toolResult.result; toolCallStatus = toolResult.status; } + + if (!result) { + result = toolCallStatus; + } } catch (e) { toolCallStatus = "error"; if ( diff --git a/apps/open-swe/src/graphs/reviewer/nodes/initialize-state.ts b/apps/open-swe/src/graphs/reviewer/nodes/initialize-state.ts index d9a8ad13..be580ccc 100644 --- a/apps/open-swe/src/graphs/reviewer/nodes/initialize-state.ts +++ b/apps/open-swe/src/graphs/reviewer/nodes/initialize-state.ts @@ -37,7 +37,7 @@ function createReviewStartedMessage() { new ToolMessage({ id: uuidv4(), tool_call_id: toolCallId, - content: "", + content: "Review started", additional_kwargs: { hidden: true, }, diff --git a/apps/open-swe/src/graphs/reviewer/nodes/take-review-action.ts b/apps/open-swe/src/graphs/reviewer/nodes/take-review-action.ts index 534cff83..697eacde 100644 --- a/apps/open-swe/src/graphs/reviewer/nodes/take-review-action.ts +++ b/apps/open-swe/src/graphs/reviewer/nodes/take-review-action.ts @@ -96,6 +96,9 @@ export async function takeReviewerActions( status: "success" | "error"; }; result = toolResult.result; + if (!result) { + result = toolResult.status; + } toolCallStatus = toolResult.status; } catch (e) { toolCallStatus = "error"; diff --git a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts index b8f27997..424a9625 100644 --- a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts +++ b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts @@ -71,7 +71,7 @@ export async function initializeSandbox( const createEventsMessage = () => [ new AIMessage({ id: `${DO_NOT_RENDER_ID_PREFIX}${uuidv4()}`, - content: "", + content: "Initialize sandbox", additional_kwargs: { hidden: true, customNodeEvents: events, diff --git a/apps/open-swe/src/tools/shell.ts b/apps/open-swe/src/tools/shell.ts index 7c467d6c..efc18fa2 100644 --- a/apps/open-swe/src/tools/shell.ts +++ b/apps/open-swe/src/tools/shell.ts @@ -34,7 +34,7 @@ export function createShellTool( } return { - result: response.result, + result: response.result ?? `exit code: ${response.exitCode}`, status: "success", }; } catch (e) { diff --git a/apps/open-swe/src/utils/message/content.ts b/apps/open-swe/src/utils/message/content.ts index 65c1b1c9..8d6c884b 100644 --- a/apps/open-swe/src/utils/message/content.ts +++ b/apps/open-swe/src/utils/message/content.ts @@ -59,3 +59,20 @@ export function getMessageString(message: BaseMessage): string { return getUnknownMessageString(message); } + +export function filterMessagesWithoutContent( + messages: BaseMessage[], + filterHidden = true, +): BaseMessage[] { + return messages.filter((m) => { + if (filterHidden && m.additional_kwargs?.hidden) { + return false; + } + const messageContentStr = getMessageContentString(m.content); + if (!isAIMessage(m)) { + return !!messageContentStr; + } + const toolCallsCount = m.tool_calls?.length || 0; + return !!messageContentStr || toolCallsCount > 0; + }); +} diff --git a/apps/open-swe/src/utils/tokens.ts b/apps/open-swe/src/utils/tokens.ts index 38994f04..9d7a6992 100644 --- a/apps/open-swe/src/utils/tokens.ts +++ b/apps/open-swe/src/utils/tokens.ts @@ -5,6 +5,7 @@ import { isToolMessage, } from "@langchain/core/messages"; import { getMessageContentString } from "@open-swe/shared/messages"; +import { traceable } from "langsmith/traceable"; // After 60k tokens, summarize the conversation history. export const MAX_INTERNAL_TOKENS = 60_000; @@ -17,39 +18,121 @@ export function calculateConversationHistoryTokenCount( }, ) { let totalChars = 0; - const messagesToCount = options?.excludeCountFromEnd - ? messages.slice(0, -options.excludeCountFromEnd) - : messages; + let messagesToCount = messages; + + if (options?.excludeCountFromEnd && options.excludeCountFromEnd > 0) { + messagesToCount = getMessagesExcludingFromEnd( + messages, + options.excludeCountFromEnd, + ); + } messagesToCount.forEach((m) => { if (options?.excludeHiddenMessages && m.additional_kwargs?.hidden) { return; } if (isHumanMessage(m) || isToolMessage(m)) { const contentString = getMessageContentString(m.content); - totalChars += contentString.length; + // Divide each char by 4 as it's roughly one token per 4 characters. + totalChars += contentString.length / 4; } if (isAIMessage(m)) { const usageMetadata = m.usage_metadata; if (usageMetadata) { - // multiply by 4 here since we divide by 4 to estimate tokens. - totalChars += usageMetadata.total_tokens * 4; + totalChars += usageMetadata.output_tokens; } else { const contentString = getMessageContentString(m.content); - totalChars += contentString.length; + totalChars += contentString.length / 4; m.tool_calls?.forEach((tc) => { - totalChars += tc.name.length; - totalChars += JSON.stringify(tc.args).length; + const nameAndArgs = tc.name + JSON.stringify(tc.args); + totalChars += nameAndArgs.length / 4; }); } } }); - // Estimate 1 token for every 4 characters. - return Math.ceil(totalChars / 4); + return totalChars; } -export function getMessagesSinceLastSummary( +/** + * Helper function to exclude messages from the end while preserving AI/tool message pairs + */ +function getMessagesExcludingFromEnd( + messages: BaseMessage[], + excludeCount: number, +): BaseMessage[] { + if (excludeCount <= 0 || excludeCount >= messages.length) { + return excludeCount >= messages.length ? [] : messages; + } + + let endIndex = messages.length - excludeCount; + + // Check if we're breaking up an AI message with tool calls and its corresponding tool messages + // We need to look backwards from the cut point to ensure we don't separate AI/tool pairs + while (endIndex > 0 && endIndex < messages.length) { + const messageAtCutPoint = messages[endIndex - 1]; + + // If the message before the cut point is an AI message with tool calls, + // we need to check if there are corresponding tool messages after it + if ( + isAIMessage(messageAtCutPoint) && + (messageAtCutPoint as any).tool_calls && + (messageAtCutPoint as any).tool_calls.length > 0 + ) { + // Count how many tool messages follow this AI message + let toolMessageCount = 0; + for ( + let i = endIndex; + i < messages.length && isToolMessage(messages[i]); + i++ + ) { + toolMessageCount++; + } + + // If there are tool messages that would be cut off, move the cut point back + // to include the AI message and all its tool messages, or exclude them entirely + if (toolMessageCount > 0) { + // Move cut point back to exclude the AI message entirely (safer approach) + endIndex--; + continue; + } + } + + // If the message at the cut point is a tool message, check if it belongs to an AI message + if (isToolMessage(messages[endIndex])) { + // Look backwards to find the corresponding AI message + let aiMessageIndex = endIndex - 1; + while (aiMessageIndex >= 0 && isToolMessage(messages[aiMessageIndex])) { + aiMessageIndex--; + } + + // If we found an AI message with tool calls, include all related messages + if ( + aiMessageIndex >= 0 && + isAIMessage(messages[aiMessageIndex]) && + (messages[aiMessageIndex] as any).tool_calls && + (messages[aiMessageIndex] as any).tool_calls.length > 0 + ) { + // Move cut point to include the entire AI/tool group + let toolGroupEnd = endIndex; + while ( + toolGroupEnd < messages.length && + isToolMessage(messages[toolGroupEnd]) + ) { + toolGroupEnd++; + } + endIndex = toolGroupEnd; + break; + } + } + + break; + } + + return messages.slice(0, endIndex); +} + +export function getMessagesSinceLastSummaryFunc( messages: BaseMessage[], options?: { excludeHiddenMessages?: boolean; @@ -76,14 +159,18 @@ export function getMessagesSinceLastSummary( // Apply excludeCountFromEnd option if provided if (options?.excludeCountFromEnd && options.excludeCountFromEnd > 0) { - messagesAfterLastSummary = messagesAfterLastSummary.slice( - 0, - Math.max( - 0, - messagesAfterLastSummary.length - options.excludeCountFromEnd, - ), + messagesAfterLastSummary = getMessagesExcludingFromEnd( + messagesAfterLastSummary, + options.excludeCountFromEnd, ); } return messagesAfterLastSummary; } + +export const getMessagesSinceLastSummary = traceable( + getMessagesSinceLastSummaryFunc, + { + name: "get-messages-since-last-summary", + }, +); diff --git a/apps/web/src/components/gen-ui/task-summary.tsx b/apps/web/src/components/gen-ui/task-summary.tsx index 4eefe168..74983ce8 100644 --- a/apps/web/src/components/gen-ui/task-summary.tsx +++ b/apps/web/src/components/gen-ui/task-summary.tsx @@ -45,9 +45,9 @@ export function TaskSummary({ const getStatusText = () => { switch (status) { case "loading": - return "Preparing task summary..."; + return "Preparing action reflection..."; case "generating": - return "Generating summary..."; + return "Generating action reflection..."; case "done": if (completed === false) { return "Task not completed"; @@ -72,7 +72,7 @@ export function TaskSummary({ - {summaryText && status === "done" && ( + {summaryText && (
- {showSummary ? "Hide summary" : "Show summary"} + {showSummary ? "Hide reflection" : "Show reflection"} {showSummary && ( tr && tr.tool_call_id === conversationHistorySummaryToolCall.id, + ); return (
); diff --git a/packages/shared/src/open-swe/tools.ts b/packages/shared/src/open-swe/tools.ts index 98dc3b06..403bb59b 100644 --- a/packages/shared/src/open-swe/tools.ts +++ b/packages/shared/src/open-swe/tools.ts @@ -410,11 +410,11 @@ export function createWriteTechnicalNotesToolFields() { export function createConversationHistorySummaryToolFields() { const conversationHistorySummarySchema = z.object({ - conversation_history_summary: z.string(), + reasoning: z.string(), }); return { - name: "conversation_history_summary", + name: "summarize_conversation_history", description: "", schema: conversationHistorySummarySchema,