Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
40 changes: 37 additions & 3 deletions packages/parser/src/protocols/morph-xml-protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,8 @@ export const morphXmlProtocol = (): ToolCallProtocol => ({
? Math.max(...toolNames.map(n => `<${n}>`.length))
: 0;
let buffer = "";
let currentToolCall: { name: string; content: string } | null = null;
let currentToolCall: { name: string; content: string; id: string } | null =
null;
let currentTextId: string | null = null;

const flushText = (
Expand Down Expand Up @@ -178,9 +179,16 @@ export const morphXmlProtocol = (): ToolCallProtocol => ({
// No additional fallback: RXML handles raw content for string fields

flushText(controller);

// Emit tool-input-end event
controller.enqueue({
type: "tool-input-end",
id: currentToolCall.id,
});

controller.enqueue({
type: "tool-call",
toolCallId: generateId(),
toolCallId: currentToolCall.id,
toolName: currentToolCall.name,
input: JSON.stringify(parsed),
});
Expand All @@ -200,6 +208,13 @@ export const morphXmlProtocol = (): ToolCallProtocol => ({
toolName: currentToolCall.name,
error,
});

// Emit tool-input-end event even on error
controller.enqueue({
type: "tool-input-end",
id: currentToolCall.id,
});

flushText(controller, originalCallText);
}
currentToolCall = null;
Expand Down Expand Up @@ -233,7 +248,20 @@ export const morphXmlProtocol = (): ToolCallProtocol => ({
buffer = buffer.substring(
earliestStartTagIndex + startTag.length
);
currentToolCall = { name: earliestToolName, content: "" };

const toolCallId = generateId();
currentToolCall = {
name: earliestToolName,
content: "",
id: toolCallId,
};

// Emit tool-input-start event
controller.enqueue({
type: "tool-input-start",
id: toolCallId,
toolName: earliestToolName,
});
} else {
// No start tag currently in buffer. Stream out as much as possible
// while keeping a small tail to catch a tag split across chunks.
Expand All @@ -253,6 +281,12 @@ export const morphXmlProtocol = (): ToolCallProtocol => ({
},
flush(controller) {
if (currentToolCall) {
// Emit tool-input-end for incomplete tool call
controller.enqueue({
type: "tool-input-end",
id: currentToolCall.id,
});

const unfinishedCall = `<${currentToolCall.name}>${buffer}`;
flushText(controller, unfinishedCall);
} else if (buffer) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,230 @@
import type { LanguageModelV2StreamPart } from "@ai-sdk/provider";
import { describe, expect, it, vi } from "vitest";

import { morphXmlProtocol } from "@/protocols/morph-xml-protocol";

vi.mock("@ai-sdk/provider-utils", () => ({
generateId: vi.fn(() => "mock-id"),
}));

function collect(stream: ReadableStream<LanguageModelV2StreamPart>) {
const out: LanguageModelV2StreamPart[] = [];
return (async () => {
for await (const c of stream) out.push(c);
return out;
})();
}

const tools = [
{
type: "function",
name: "get_weather",
description: "",
inputSchema: { type: "object" },
},
] as any;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using as any should be avoided as it bypasses TypeScript's type safety. Please cast to LanguageModelV2FunctionTool[] instead. You will need to add LanguageModelV2FunctionTool to the import from @ai-sdk/provider at the top of the file.

Suggested change
] as any;
] as LanguageModelV2FunctionTool[];

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.


describe("morphXmlProtocol tool-input events", () => {
it("emits tool-input-start and tool-input-end events for successful tool call", async () => {
const protocol = morphXmlProtocol();
const transformer = protocol.createStreamParser({ tools });
const rs = new ReadableStream<LanguageModelV2StreamPart>({
start(ctrl) {
ctrl.enqueue({ type: "text-delta", id: "1", delta: "prefix " });
ctrl.enqueue({ type: "text-delta", id: "1", delta: "<get_weather>" });
ctrl.enqueue({
type: "text-delta",
id: "1",
delta: "<location>NY</location>",
});
ctrl.enqueue({ type: "text-delta", id: "1", delta: "</get_weather>" });
ctrl.enqueue({ type: "text-delta", id: "1", delta: " suffix" });
ctrl.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
});
ctrl.close();
},
});

const out = await collect(rs.pipeThrough(transformer));

// Find tool-input events
const toolInputStart = out.find(c => c.type === "tool-input-start") as any;
const toolInputEnd = out.find(c => c.type === "tool-input-end") as any;
const toolCall = out.find(c => c.type === "tool-call") as any;
Comment on lines +54 to +56
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using as any should be avoided as it bypasses TypeScript's type safety. This is likely a temporary measure because the new stream part types (tool-input-start, tool-input-end) are not yet part of the LanguageModelV2StreamPart union type.

A better approach would be to define local interfaces for these new event types within the test file and cast to those specific types. This improves type safety and makes the test easier to understand.

For example:

interface ToolInputStartPart {
  type: 'tool-input-start';
  id: string;
  toolName: string;
}

// ... then later ...
const toolInputStart = out.find(
  c => c.type === 'tool-input-start'
) as ToolInputStartPart | undefined;

This pattern should be applied to all usages of as any in this file (e.g., lines 113-114, 155-156, and 221-223).

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

optional, skip


// Verify tool-input-start event
expect(toolInputStart).toBeTruthy();
expect(toolInputStart.type).toBe("tool-input-start");
expect(toolInputStart.id).toBe("mock-id");
expect(toolInputStart.toolName).toBe("get_weather");

// Verify tool-input-end event
expect(toolInputEnd).toBeTruthy();
expect(toolInputEnd.type).toBe("tool-input-end");
expect(toolInputEnd.id).toBe("mock-id");

// Verify tool-call event uses same ID
expect(toolCall).toBeTruthy();
expect(toolCall.toolCallId).toBe("mock-id");
expect(toolCall.toolName).toBe("get_weather");

// Verify event order: start -> end -> tool-call
const eventIndexes = {
start: out.findIndex(c => c.type === "tool-input-start"),
end: out.findIndex(c => c.type === "tool-input-end"),
call: out.findIndex(c => c.type === "tool-call"),
};
expect(eventIndexes.start).toBeLessThan(eventIndexes.end);
expect(eventIndexes.end).toBeLessThan(eventIndexes.call);
});

it("emits tool-input-start and tool-input-end events for failed tool call", async () => {
const onError = vi.fn();
const protocol = morphXmlProtocol();
const transformer = protocol.createStreamParser({
tools,
options: { onError },
});

const rs = new ReadableStream<LanguageModelV2StreamPart>({
start(ctrl) {
ctrl.enqueue({ type: "text-delta", id: "1", delta: "<get_weather>" });
ctrl.enqueue({
type: "text-delta",
id: "1",
delta: "<invalid>malformed xml",
});
ctrl.enqueue({ type: "text-delta", id: "1", delta: "</get_weather>" });
ctrl.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
});
ctrl.close();
},
});

const out = await collect(rs.pipeThrough(transformer));

// Find tool-input events
const toolInputStart = out.find(c => c.type === "tool-input-start") as any;
const toolInputEnd = out.find(c => c.type === "tool-input-end") as any;

// Verify both events are emitted even on error
expect(toolInputStart).toBeTruthy();
expect(toolInputStart.type).toBe("tool-input-start");
expect(toolInputStart.id).toBe("mock-id");
expect(toolInputStart.toolName).toBe("get_weather");

expect(toolInputEnd).toBeTruthy();
expect(toolInputEnd.type).toBe("tool-input-end");
expect(toolInputEnd.id).toBe("mock-id");

// Verify error callback was called
expect(onError).toHaveBeenCalled();
});

it("emits tool-input-start and tool-input-end for incomplete tool call at stream end", async () => {
const protocol = morphXmlProtocol();
const transformer = protocol.createStreamParser({ tools });

const rs = new ReadableStream<LanguageModelV2StreamPart>({
start(ctrl) {
ctrl.enqueue({ type: "text-delta", id: "1", delta: "<get_weather>" });
ctrl.enqueue({
type: "text-delta",
id: "1",
delta: "<location>NY</location>",
});
// Note: no closing tag - incomplete tool call
ctrl.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
});
ctrl.close();
},
});

const out = await collect(rs.pipeThrough(transformer));

// Find tool-input events
const toolInputStart = out.find(c => c.type === "tool-input-start") as any;
const toolInputEnd = out.find(c => c.type === "tool-input-end") as any;

// Verify both events are emitted even for incomplete calls
expect(toolInputStart).toBeTruthy();
expect(toolInputStart.type).toBe("tool-input-start");
expect(toolInputStart.id).toBe("mock-id");
expect(toolInputStart.toolName).toBe("get_weather");

expect(toolInputEnd).toBeTruthy();
expect(toolInputEnd.type).toBe("tool-input-end");
expect(toolInputEnd.id).toBe("mock-id");

// Verify the incomplete call content is emitted as text
const textParts = out
.filter(c => c.type === "text-delta")
.map((c: any) => c.delta);
const fullText = textParts.join("");
expect(fullText).toContain("<get_weather>");
expect(fullText).toContain("<location>NY</location>");
});

it("emits multiple paired tool-input events for multiple tool calls", async () => {
const protocol = morphXmlProtocol();
const transformer = protocol.createStreamParser({ tools });

const rs = new ReadableStream<LanguageModelV2StreamPart>({
start(ctrl) {
ctrl.enqueue({
type: "text-delta",
id: "1",
delta: "<get_weather><location>NY</location></get_weather>",
});
ctrl.enqueue({ type: "text-delta", id: "1", delta: " and " });
ctrl.enqueue({
type: "text-delta",
id: "1",
delta: "<get_weather><location>SF</location></get_weather>",
});
ctrl.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 },
});
ctrl.close();
},
});

const out = await collect(rs.pipeThrough(transformer));

// Find all tool-input events
const toolInputStarts = out.filter(c => c.type === "tool-input-start");
const toolInputEnds = out.filter(c => c.type === "tool-input-end");
const toolCalls = out.filter(c => c.type === "tool-call");

// Should have at least 1 tool call (implementation may coalesce)
expect(toolInputStarts.length).toBeGreaterThanOrEqual(1);
expect(toolInputEnds.length).toBeGreaterThanOrEqual(1);
expect(toolCalls.length).toBeGreaterThanOrEqual(1);
Comment on lines +211 to +213
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The assertion toBeGreaterThanOrEqual(1) is too lenient for this test case. Given the input stream contains two distinct tool calls, it would be more precise to assert that exactly two of each event type are found. This will make the test more robust against regressions where one of the tool calls might be missed during parsing.

Suggested change
expect(toolInputStarts.length).toBeGreaterThanOrEqual(1);
expect(toolInputEnds.length).toBeGreaterThanOrEqual(1);
expect(toolCalls.length).toBeGreaterThanOrEqual(1);
expect(toolInputStarts.length).toBe(2);
expect(toolInputEnds.length).toBe(2);
expect(toolCalls.length).toBe(2);

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@biubiupiu1 WDYT?


// Each start should have a corresponding end
expect(toolInputStarts.length).toBe(toolInputEnds.length);
expect(toolInputStarts.length).toBe(toolCalls.length);

// Verify IDs match between start, end, and tool-call events
for (let i = 0; i < toolInputStarts.length; i++) {
const start = toolInputStarts[i] as any;
const end = toolInputEnds[i] as any;
const call = toolCalls[i] as any;

expect(start.id).toBe(end.id);
expect(start.id).toBe(call.toolCallId);
expect(start.toolName).toBe("get_weather");
}
});
});
Loading