Skip to content

Commit 9e63323

Browse files
vincentkocvincentkoc
committed
perf(anthropic): index active stream blocks
1 parent 00f8b10 commit 9e63323

4 files changed

Lines changed: 144 additions & 20 deletions

File tree

‎src/agents/anthropic-transport-stream.test.ts‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1070,6 +1070,56 @@ describe("anthropic transport stream", () => {
10701070
});
10711071
});
10721072

1073+
it("routes interleaved active content blocks by their event indexes", async () => {
1074+
guardedFetchMock.mockResolvedValueOnce(
1075+
createSseResponse([
1076+
{
1077+
type: "message_start",
1078+
message: { id: "msg_interleaved", usage: { input_tokens: 1, output_tokens: 0 } },
1079+
},
1080+
{
1081+
type: "content_block_start",
1082+
index: 0,
1083+
content_block: { type: "text", text: "" },
1084+
},
1085+
{
1086+
type: "content_block_start",
1087+
index: 1,
1088+
content_block: { type: "text", text: "" },
1089+
},
1090+
{
1091+
type: "content_block_delta",
1092+
index: 1,
1093+
delta: { type: "text_delta", text: "second" },
1094+
},
1095+
{
1096+
type: "content_block_delta",
1097+
index: 0,
1098+
delta: { type: "text_delta", text: "first" },
1099+
},
1100+
{ type: "content_block_stop", index: 1 },
1101+
{ type: "content_block_stop", index: 0 },
1102+
{
1103+
type: "message_delta",
1104+
delta: { stop_reason: "end_turn" },
1105+
usage: { input_tokens: 1, output_tokens: 2 },
1106+
},
1107+
{ type: "message_stop" },
1108+
]),
1109+
);
1110+
1111+
const result = await runTransportStream(
1112+
makeAnthropicTransportModel(),
1113+
{ messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext,
1114+
{ apiKey: "sk-ant-api" } as AnthropicStreamOptions,
1115+
);
1116+
1117+
expect(result.content).toEqual([
1118+
{ type: "text", text: "first" },
1119+
{ type: "text", text: "second" },
1120+
]);
1121+
});
1122+
10731123
it("preserves provider-seeded thinking signatures when no signature_delta follows", async () => {
10741124
guardedFetchMock.mockResolvedValueOnce(
10751125
createSseResponse([

‎src/agents/anthropic-transport-stream.ts‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1116,6 +1116,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
11161116
transportOptions.signal ? { signal: transportOptions.signal } : undefined,
11171117
);
11181118
const blocks = output.content;
1119+
const blockIndexes = new Map<number, number>();
11191120
const signatureDeltaIndexes = new Set<number>();
11201121
const allowReasoningContentReplay = supportsReasoningContentReplay(model);
11211122
const reasoningContentThinkingBlocks = new Map<number, number>();
@@ -1276,6 +1277,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
12761277
const block: TransportContentBlock = { type: "text", text, index };
12771278
output.content.push(block);
12781279
const contentIndex = output.content.length - 1;
1280+
blockIndexes.set(index, contentIndex);
12791281
eventSink.push({
12801282
type: "text_start",
12811283
contentIndex,
@@ -1303,6 +1305,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
13031305
};
13041306
output.content.push(block);
13051307
const contentIndex = output.content.length - 1;
1308+
blockIndexes.set(index, contentIndex);
13061309
eventSink.push({
13071310
type: "thinking_start",
13081311
contentIndex,
@@ -1327,6 +1330,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
13271330
index,
13281331
};
13291332
output.content.push(block);
1333+
blockIndexes.set(index, output.content.length - 1);
13301334
eventSink.push({
13311335
type: "thinking_start",
13321336
contentIndex: output.content.length - 1,
@@ -1352,6 +1356,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
13521356
index,
13531357
};
13541358
output.content.push(block);
1359+
blockIndexes.set(index, output.content.length - 1);
13551360
eventSink.push({
13561361
type: "toolcall_start",
13571362
contentIndex: output.content.length - 1,
@@ -1362,8 +1367,9 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
13621367
}
13631368
if (event.type === "content_block_delta") {
13641369
const delta = event.delta as Record<string, unknown> | undefined;
1365-
let index = blocks.findIndex((block) => block.index === event.index);
1366-
let block = blocks[index];
1370+
const eventIndex = typeof event.index === "number" ? event.index : undefined;
1371+
let index = eventIndex === undefined ? undefined : blockIndexes.get(eventIndex);
1372+
let block = index === undefined ? undefined : blocks[index];
13671373
if (allowReasoningContentReplay) {
13681374
const appendedThinking = appendReasoningContentThinkingDelta(
13691375
event.index,
@@ -1405,6 +1411,9 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
14051411
block = { type: "text", text: "", index: recoveredIndex };
14061412
output.content.push(block);
14071413
index = output.content.length - 1;
1414+
if (typeof event.index === "number") {
1415+
blockIndexes.set(event.index, index);
1416+
}
14081417
eventSink.push({
14091418
type: "text_start",
14101419
contentIndex: index,
@@ -1470,12 +1479,14 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
14701479
continue;
14711480
}
14721481
if (event.type === "content_block_stop") {
1473-
const index = blocks.findIndex((block) => block.index === event.index);
1474-
const block = blocks[index];
1475-
if (!block) {
1482+
const eventIndex = typeof event.index === "number" ? event.index : undefined;
1483+
const index = eventIndex === undefined ? undefined : blockIndexes.get(eventIndex);
1484+
const block = index === undefined ? undefined : blocks[index];
1485+
if (index === undefined || !block) {
14761486
finishReasoningContentSidecars(event.index);
14771487
continue;
14781488
}
1489+
blockIndexes.delete(eventIndex);
14791490
delete block.index;
14801491
if (block.type === "text") {
14811492
eventSink.push({

‎src/llm/providers/anthropic.test.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -462,6 +462,63 @@ describe("Anthropic provider", () => {
462462
]);
463463
});
464464

465+
it("routes interleaved active content blocks by their event indexes", async () => {
466+
const client = {
467+
messages: {
468+
create: vi.fn(() => ({
469+
asResponse: () =>
470+
Promise.resolve(
471+
createSseResponse([
472+
{
473+
type: "message_start",
474+
message: { id: "msg_interleaved", usage: { input_tokens: 1, output_tokens: 0 } },
475+
},
476+
{
477+
type: "content_block_start",
478+
index: 0,
479+
content_block: { type: "text" },
480+
},
481+
{
482+
type: "content_block_start",
483+
index: 1,
484+
content_block: { type: "text" },
485+
},
486+
{
487+
type: "content_block_delta",
488+
index: 1,
489+
delta: { type: "text_delta", text: "second" },
490+
},
491+
{
492+
type: "content_block_delta",
493+
index: 0,
494+
delta: { type: "text_delta", text: "first" },
495+
},
496+
{ type: "content_block_stop", index: 1 },
497+
{ type: "content_block_stop", index: 0 },
498+
{
499+
type: "message_delta",
500+
delta: { stop_reason: "end_turn" },
501+
usage: { input_tokens: 1, output_tokens: 2 },
502+
},
503+
{ type: "message_stop" },
504+
]),
505+
),
506+
})),
507+
},
508+
};
509+
510+
const result = await streamAnthropic(
511+
makeAnthropicModel(),
512+
{ messages: [{ role: "user", content: "hello", timestamp: 0 }] },
513+
{ apiKey: "sk-ant-provider", client: client as never },
514+
).result();
515+
516+
expect(result.content).toEqual([
517+
{ type: "text", text: "first" },
518+
{ type: "text", text: "second" },
519+
]);
520+
});
521+
465522
it("discards buffered Fable output when the stream fails before terminal status", async () => {
466523
const client = {
467524
messages: {

‎src/llm/providers/anthropic.ts‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
536536
index: number;
537537
};
538538
const blocks = output.content as Block[];
539+
const blockIndexes = new Map<number, number>();
539540

540541
for await (const event of iterateAnthropicEvents(
541542
response,
@@ -568,6 +569,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
568569
index: event.index,
569570
};
570571
output.content.push(block);
572+
blockIndexes.set(event.index, output.content.length - 1);
571573
eventSink.push({
572574
type: "text_start",
573575
contentIndex: output.content.length - 1,
@@ -581,6 +583,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
581583
index: event.index,
582584
};
583585
output.content.push(block);
586+
blockIndexes.set(event.index, output.content.length - 1);
584587
eventSink.push({
585588
type: "thinking_start",
586589
contentIndex: output.content.length - 1,
@@ -595,6 +598,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
595598
index: event.index,
596599
};
597600
output.content.push(block);
601+
blockIndexes.set(event.index, output.content.length - 1);
598602
eventSink.push({
599603
type: "thinking_start",
600604
contentIndex: output.content.length - 1,
@@ -612,6 +616,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
612616
index: event.index,
613617
};
614618
output.content.push(block);
619+
blockIndexes.set(event.index, output.content.length - 1);
615620
eventSink.push({
616621
type: "toolcall_start",
617622
contentIndex: output.content.length - 1,
@@ -620,9 +625,9 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
620625
}
621626
} else if (event.type === "content_block_delta") {
622627
if (event.delta.type === "text_delta") {
623-
const index = blocks.findIndex((b) => b.index === event.index);
624-
const block = blocks[index];
625-
if (block && block.type === "text") {
628+
const index = blockIndexes.get(event.index);
629+
const block = index === undefined ? undefined : blocks[index];
630+
if (index !== undefined && block?.type === "text") {
626631
block.text += event.delta.text;
627632
eventSink.push({
628633
type: "text_delta",
@@ -632,9 +637,9 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
632637
});
633638
}
634639
} else if (event.delta.type === "thinking_delta") {
635-
const index = blocks.findIndex((b) => b.index === event.index);
636-
const block = blocks[index];
637-
if (block && block.type === "thinking") {
640+
const index = blockIndexes.get(event.index);
641+
const block = index === undefined ? undefined : blocks[index];
642+
if (index !== undefined && block?.type === "thinking") {
638643
block.thinking += event.delta.thinking;
639644
eventSink.push({
640645
type: "thinking_delta",
@@ -644,9 +649,9 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
644649
});
645650
}
646651
} else if (event.delta.type === "input_json_delta") {
647-
const index = blocks.findIndex((b) => b.index === event.index);
648-
const block = blocks[index];
649-
if (block && block.type === "toolCall") {
652+
const index = blockIndexes.get(event.index);
653+
const block = index === undefined ? undefined : blocks[index];
654+
if (index !== undefined && block?.type === "toolCall") {
650655
block.partialJson += event.delta.partial_json;
651656
block.arguments = parseStreamingJson(block.partialJson);
652657
eventSink.push({
@@ -657,17 +662,18 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
657662
});
658663
}
659664
} else if (event.delta.type === "signature_delta") {
660-
const index = blocks.findIndex((b) => b.index === event.index);
661-
const block = blocks[index];
662-
if (block && block.type === "thinking") {
665+
const index = blockIndexes.get(event.index);
666+
const block = index === undefined ? undefined : blocks[index];
667+
if (index !== undefined && block?.type === "thinking") {
663668
block.thinkingSignature = block.thinkingSignature || "";
664669
block.thinkingSignature += event.delta.signature;
665670
}
666671
}
667672
} else if (event.type === "content_block_stop") {
668-
const index = blocks.findIndex((b) => b.index === event.index);
669-
const block = blocks[index];
670-
if (block) {
673+
const index = blockIndexes.get(event.index);
674+
const block = index === undefined ? undefined : blocks[index];
675+
if (index !== undefined && block) {
676+
blockIndexes.delete(event.index);
671677
delete (block as Partial<Block>).index;
672678
if (block.type === "text") {
673679
eventSink.push({

0 commit comments

Comments
 (0)