Skip to content

🪢 fix: Flatten Text-Only Assistant Content To String #205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
misdirection wants to merge 2 commits into
base: main
Choose a base branch
from
Open

🪢 fix: Flatten Text-Only Assistant Content To String #205

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce5248c
🪢 fix: Flatten Text-Only Assistant Content To String
May 28, 2026
fba7d7f
🪢 fix: Preserve whitespace when flattening assistant text blocks
Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 25 additions & 18 deletions src/messages/format.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,28 @@ function getToolUseId(part: MessageContentComplex): string | undefined {
return part.tool_use_id;
}

/**
* Builds the `content` for an assistant message in a single pass.
*
* Text-only parts are joined with a newline into a string. The original text is
* not trimmed, so whitespace-sensitive replies (indented code, trailing
* newlines) replay verbatim. Parts with any non-text block keep their array
* shape.
*
* @param parts - The accumulated assistant content parts.
* @returns A string for text-only content, otherwise the content array.
*/
function buildAssistantContent(parts: MessageContentComplex[]): MessageContent {
const texts: string[] = [];
for (const part of parts) {
if (part.type !== ContentTypes.TEXT) {
return toLangChainContent(parts);
}
texts.push(getTextContent(part));
}
return texts.join('\n');
}

/**
* Helper function to format an assistant message
* @param message The message to format
Expand All @@ -384,7 +406,6 @@ function formatAssistantMessage(
const formattedMessages: Array<AIMessage | ToolMessage> = [];
let currentContent: MessageContentComplex[] = [];
let lastAIMessage: AIMessage | null = null;
let hasReasoning = false;
let pendingReasoningContent = '';
const emittedServerToolUseIds = new Set<string>();
const pendingServerToolUses = new Map<string, MessageContentComplex>();
Expand Down Expand Up @@ -550,7 +571,6 @@ function formatAssistantMessage(
part.type === ContentTypes.REASONING_CONTENT ||
part.type === 'redacted_thinking'
) {
hasReasoning = true;
pendingReasoningContent += extractReasoningContent(part);
continue;
} else if (
Expand All @@ -571,24 +591,11 @@ function formatAssistantMessage(
}
}

if (hasReasoning && currentContent.length > 0) {
let content = '';
for (const part of currentContent) {
if (part.type !== ContentTypes.TEXT) {
formattedMessages.push(
createAIMessage(toLangChainContent(currentContent))
);
return formattedMessages;
}
content += `${getTextContent(part)}\n`;
}
content = content.trim();

if (content) {
if (currentContent.length > 0) {
const content = buildAssistantContent(currentContent);
if (content !== '') {
formattedMessages.push(createAIMessage(content));
}
} else if (currentContent.length > 0) {
formattedMessages.push(createAIMessage(toLangChainContent(currentContent)));
}

return formattedMessages;
Expand Down
107 changes: 56 additions & 51 deletions src/messages/formatAgentMessages.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,12 +97,7 @@ describe('formatAgentMessages', () => {
expect(result.summary!.tokenCount).toBe(12);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
expect(result.messages[1]).toBeInstanceOf(HumanMessage);
expect(
(result.messages[0].content as MessageContentComplex[])[0]
).toMatchObject({
type: ContentTypes.TEXT,
text: 'Preserved tail',
});
expect(result.messages[0].content).toBe('Preserved tail');
expect(result.indexTokenCountMap?.[0]).toBeLessThan(18);
expect(result.indexTokenCountMap?.[0]).toBeGreaterThan(0);
expect(result.indexTokenCountMap?.[1]).toBe(4);
Expand Down Expand Up @@ -134,12 +129,7 @@ describe('formatAgentMessages', () => {
expect(result.summary!.text).toBe('Newest summary');
expect(result.summary!.tokenCount).toBe(9);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
expect(
(result.messages[0].content as MessageContentComplex[])[0]
).toMatchObject({
type: ContentTypes.TEXT,
text: 'Keep this part',
});
expect(result.messages[0].content).toBe('Keep this part');
});

it('should format messages with content arrays', () => {
Expand Down Expand Up @@ -228,12 +218,7 @@ describe('formatAgentMessages', () => {
result.messages.some((message) => message instanceof ToolMessage)
).toBe(false);
expect((result.messages[1] as AIMessage).tool_calls).toHaveLength(0);
expect(result.messages[1].content).toEqual([
{
type: ContentTypes.TEXT,
[ContentTypes.TEXT]: 'Philadelphia 76ers',
},
]);
expect(result.messages[1].content).toBe('Philadelphia 76ers');
});

it('preserves paused Anthropic server tool calls without creating ToolMessages', () => {
Expand Down Expand Up @@ -266,8 +251,9 @@ describe('formatAgentMessages', () => {

expect(result.messages).toHaveLength(1);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
expect(result.messages.some((message) => message instanceof ToolMessage))
.toBe(false);
expect(
result.messages.some((message) => message instanceof ToolMessage)
).toBe(false);
expect((result.messages[0] as AIMessage).tool_calls).toHaveLength(0);
expect(result.messages[0].content).toEqual([
{
Expand Down Expand Up @@ -369,6 +355,10 @@ describe('formatAgentMessages', () => {

expect(anthropicPayload.messages).toHaveLength(3);
for (const message of anthropicPayload.messages) {
if (typeof message.content === 'string') {
expect(message.content.trim().length).toBeGreaterThan(0);
continue;
}
expect(Array.isArray(message.content)).toBe(true);
const content = message.content as Array<{
text?: unknown;
Expand Down Expand Up @@ -444,7 +434,7 @@ describe('formatAgentMessages', () => {
expect(result.messages.length).toBeGreaterThan(0);
});

it('should handle multiple content parts in assistant messages', () => {
it('should join multiple text-only parts in assistant messages into a string', () => {
const payload = [
{
role: 'assistant',
Expand All @@ -457,7 +447,45 @@ describe('formatAgentMessages', () => {
const result = formatAgentMessages(payload);
expect(result.messages).toHaveLength(1);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
expect(result.messages[0].content).toHaveLength(2);
expect(typeof result.messages[0].content).toBe('string');
expect(result.messages[0].content).toBe('Part 1\nPart 2');
});

it('flattens single-text-part assistant content to a string on multi-turn replays', () => {
const payload: TPayload = [
{ role: 'user', content: 'Test1' },
{
role: 'assistant',
content: [
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Test received.' },
],
},
{ role: 'user', content: 'Test2' },
];
const result = formatAgentMessages(payload);
expect(result.messages).toHaveLength(3);
expect(result.messages[1]).toBeInstanceOf(AIMessage);
expect(typeof result.messages[1].content).toBe('string');
expect(result.messages[1].content).toBe('Test received.');
});

it('preserves whitespace-sensitive text when flattening assistant content', () => {
const preformatted = ' def foo():\n return 1\n';
const payload: TPayload = [
{ role: 'user', content: 'Show me indented code' },
{
role: 'assistant',
content: [
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: preformatted },
],
},
{ role: 'user', content: 'Thanks' },
];
const result = formatAgentMessages(payload);
expect(result.messages[1]).toBeInstanceOf(AIMessage);
expect(typeof result.messages[1].content).toBe('string');
// Leading indentation and the trailing newline must survive the flatten.
expect(result.messages[1].content).toBe(preformatted);
});

it('should heal invalid tool call structure by creating a preceding AIMessage', () => {
Expand Down Expand Up @@ -614,9 +642,7 @@ describe('formatAgentMessages', () => {
expect(result.messages[3].content).toBe('23.89°C');

// Check final AIMessage
expect(result.messages[4].content).toStrictEqual([
{ [ContentTypes.TEXT]: 'Here\'s your answer.', type: ContentTypes.TEXT },
]);
expect(result.messages[4].content).toBe('Here\'s your answer.');
});

it('should dynamically discover tools from tool_search output and keep their tool calls', () => {
Expand Down Expand Up @@ -1423,13 +1449,8 @@ describe('formatAgentMessages', () => {

expect(result.messages).toHaveLength(1);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
const content = result.messages[0].content;
expect(Array.isArray(content)).toBe(true);
expect(
(content as { type: string; text?: string }[]).every(
(p) => (p.text ?? '').trim() !== ''
)
).toBe(true);
expect(typeof result.messages[0].content).toBe('string');
expect(result.messages[0].content).toBe('Actual content here.');
});

it('should preserve whitespace-only text that has tool_call_ids (common Bedrock pattern)', () => {
Expand Down Expand Up @@ -1511,19 +1532,8 @@ describe('formatAgentMessages', () => {

expect(result.messages).toHaveLength(1);
expect(result.messages[0]).toBeInstanceOf(AIMessage);
expect(result.messages[0].content).toEqual([
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Hello there' },
{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Final answer' },
]);

const hasErrorContent =
Array.isArray(result.messages[0].content) &&
result.messages[0].content.some(
(item) =>
item.type === ContentTypes.ERROR ||
JSON.stringify(item).includes('An error occurred')
);
expect(hasErrorContent).toBe(false);
expect(result.messages[0].content).toBe('Hello there\nFinal answer');
expect(result.messages[0].content).not.toContain('An error occurred');
});
it('should handle indexTokenCountMap and return updated map', () => {
const payload = [
Expand Down Expand Up @@ -2093,12 +2103,7 @@ describe('formatAgentMessages', () => {
// The AIMessage should have no tool calls or an empty array
const toolCalls = (result.messages[0] as AIMessage).tool_calls;
expect(toolCalls === undefined || toolCalls.length === 0).toBe(true);
expect(result.messages[0].content).toStrictEqual([
{
type: ContentTypes.TEXT,
[ContentTypes.TEXT]: 'Here is the information.',
},
]);
expect(result.messages[0].content).toBe('Here is the information.');
});

it('should NOT skip tool calls with no name but valid output', () => {
Expand Down