feat(chat-adapter): implement tool call argument streaming events and enhance tool call handling

Author: jk4235

packages/aipex-react/src/adapters/chat-adapter.test.ts

@@ -205,6 +205,73 @@ describe("ChatAdapter", () => {
  expect(adapter.getStatus()).toBe("streaming");
  });
 
+ it("should create a pending tool call on tool_call_args_streaming_start", () => {
+ adapter.processEvent({
+ type: "tool_call_args_streaming_start",
+ toolName: "search",
+ });
+
+ const messages = adapter.getMessages();
+ expect(messages).toHaveLength(1);
+ expect(messages[0]).toMatchObject({ role: "assistant" });
+
+ const toolPart = messages[0]?.parts.find((p) => p.type === "tool");
+ expect(toolPart).toMatchObject({
+ toolName: "search",
+ state: "pending",
+ input: {},
+ });
+ expect(adapter.getStatus()).toBe("streaming");
+ });
+
+ it("should update pending tool params on tool_call_args_streaming_complete", () => {
+ adapter.processEvent({
+ type: "tool_call_args_streaming_start",
+ toolName: "search",
+ });
+ adapter.processEvent({
+ type: "tool_call_args_streaming_complete",
+ toolName: "search",
+ params: { q: "ts" },
+ });
+
+ const toolPart = adapter
+ .getMessages()[0]
+ ?.parts.find((p) => p.type === "tool");
+ expect(toolPart).toMatchObject({
+ toolName: "search",
+ state: "pending",
+ input: { q: "ts" },
+ });
+ });
+
+ it("should not duplicate tool parts when tool_call_start follows tool args streaming events", () => {
+ adapter.processEvent({
+ type: "tool_call_args_streaming_start",
+ toolName: "search",
+ });
+ adapter.processEvent({
+ type: "tool_call_args_streaming_complete",
+ toolName: "search",
+ params: { q: "ts" },
+ });
+ adapter.processEvent({
+ type: "tool_call_start",
+ toolName: "search",
+ params: { q: "ts" },
+ });
+
+ const toolParts =
+ adapter.getMessages()[0]?.parts.filter((p) => p.type === "tool") ?? [];
+ expect(toolParts).toHaveLength(1);
+ expect(toolParts[0]).toMatchObject({
+ toolName: "search",
+ state: "executing",
+ input: { q: "ts" },
+ });
+ expect(adapter.getStatus()).toBe("executing_tools");
+ });
+
  it("should append to existing assistant message", () => {
  adapter.processEvent({ type: "content_delta", delta: "Hello" });
  adapter.processEvent({ type: "content_delta", delta: " world" });

packages/aipex-react/src/adapters/chat-adapter.ts

@@ -139,9 +139,22 @@ export class ChatAdapter {
  this.appendContentDelta(event.delta);
  break;
 
+ case "tool_call_args_streaming_start":
+ this.ensureAssistantMessage();
+ this.ensurePendingToolCall(event.toolName, {});
+ this.updateStatus("streaming");
+ break;
+
+ case "tool_call_args_streaming_complete":
+ this.ensureAssistantMessage();
+ this.ensurePendingToolCall(event.toolName, event.params);
+ break;
+
  case "tool_call_start":
  this.ensureAssistantMessage();
- this.addToolCall(event.toolName, event.params);
+ if (!this.startExistingToolCall(event.toolName, event.params)) {
+ this.addToolCall(event.toolName, event.params);
+ }
  this.updateStatus("executing_tools");
  break;
 
@@ -304,6 +317,81 @@ export class ChatAdapter {
  this.toolsAddedSinceLastText = true;
  }
 
+ private ensurePendingToolCall(toolName: string, params: unknown): void {
+ const existingCallId = this.findPendingToolCallId(toolName);
+ if (existingCallId) {
+ this.updateToolPart(existingCallId, (toolPart) => ({
+ ...toolPart,
+ toolName,
+ input: params,
+ }));
+ return;
+ }
+
+ const callId = this.queueToolCall(toolName);
+ this.updateCurrentAssistantMessage((message) => {
+ const parts = [...message.parts];
+
+ const toolPart: UIToolPart = {
+ type: "tool",
+ toolCallId: callId,
+ toolName,
+ input: params,
+ state: "pending",
+ };
+
+ parts.push(toolPart);
+
+ return { ...message, parts };
+ });
+
+ this.toolsAddedSinceLastText = true;
+ }
+
+ private startExistingToolCall(toolName: string, params: unknown): boolean {
+ const callId = this.findPendingToolCallId(toolName);
+ if (!callId) {
+ return false;
+ }
+
+ this.updateToolPart(callId, (toolPart) => ({
+ ...toolPart,
+ toolName,
+ input: params,
+ state: "executing",
+ }));
+ this.toolsAddedSinceLastText = true;
+ return true;
+ }
+
+ private findPendingToolCallId(toolName: string): string | undefined {
+ const queue = this.pendingToolCalls.get(toolName);
+ if (!queue || queue.length === 0) {
+ return undefined;
+ }
+
+ const currentId = this.state.currentAssistantMessageId;
+ if (!currentId) {
+ return undefined;
+ }
+ const message = this.state.messages.find((m) => m.id === currentId);
+ if (!message) {
+ return undefined;
+ }
+
+ for (const callId of queue) {
+ const toolPart = message.parts.find(
+ (part): part is UIToolPart =>
+ part.type === "tool" && part.toolCallId === callId,
+ );
+ if (toolPart?.state === "pending") {
+ return callId;
+ }
+ }
+
+ return undefined;
+ }
+
  private updateToolComplete(toolName: string, result: unknown): void {
  const callId = this.dequeueToolCall(toolName);
  if (!callId) {

packages/aipex-react/src/components/chatbot/tools.ts

@@ -18,7 +18,7 @@ ${typeof output === "string" ? output : JSON.stringify(output, null, 2)}
 export function mapToolState(state: UIToolPart["state"]): ToolComponentState {
  switch (state) {
  case "pending":
- return "input-available";
+ return "input-streaming";
  case "executing":
  return "executing";
  case "completed":

packages/core/src/agent/aipex.test.ts

@@ -555,6 +555,114 @@ describe("AIPex", () => {
  });
 
  describe("tools and errors", () => {
+ it("should emit tool_call_args_streaming_complete before tool_call_start", async () => {
+ vi.mocked(run).mockResolvedValue(
+ createMockRunResult({
+ finalOutput: "",
+ streamEvents: [
+ {
+ type: "run_item_stream_event",
+ name: "tool_called",
+ item: { rawItem: { name: "calculator", arguments: '{"a":1}' } },
+ },
+ ],
+ }),
+ );
+
+ const agent = AIPex.create({
+ instructions: "Tools",
+ model: mockModel,
+ });
+
+ const events: AgentEvent[] = [];
+ for await (const event of agent.chat("use tool")) {
+ events.push(event);
+ }
+
+ const argsCompleteIndex = events.findIndex(
+ (event) => event.type === "tool_call_args_streaming_complete",
+ );
+ const toolStartIndex = events.findIndex(
+ (event) => event.type === "tool_call_start",
+ );
+
+ expect(argsCompleteIndex).toBeGreaterThanOrEqual(0);
+ expect(toolStartIndex).toBeGreaterThanOrEqual(0);
+ expect(argsCompleteIndex).toBeLessThan(toolStartIndex);
+
+ const argsComplete = events[argsCompleteIndex];
+ const toolStart = events[toolStartIndex];
+
+ expect(argsComplete?.type).toBe("tool_call_args_streaming_complete");
+ if (argsComplete?.type === "tool_call_args_streaming_complete") {
+ expect(argsComplete.toolName).toBe("calculator");
+ expect(argsComplete.params).toEqual({ a: 1 });
+ }
+
+ expect(toolStart?.type).toBe("tool_call_start");
+ if (toolStart?.type === "tool_call_start") {
+ expect(toolStart.toolName).toBe("calculator");
+ expect(toolStart.params).toEqual({ a: 1 });
+ }
+ });
+
+ it("should emit tool_call_args_streaming_start when tool args are streamed by the model", async () => {
+ vi.mocked(run).mockResolvedValue(
+ createMockRunResult({
+ finalOutput: "",
+ streamEvents: [
+ {
+ type: "raw_model_stream_event",
+ data: {
+ type: "model",
+ event: {
+ choices: [
+ {
+ delta: {
+ tool_calls: [
+ {
+ index: 0,
+ id: "call_1",
+ function: {
+ name: "calculator",
+ arguments: '{"a":',
+ },
+ },
+ ],
+ },
+ },
+ ],
+ },
+ },
+ },
+ {
+ type: "run_item_stream_event",
+ name: "tool_called",
+ item: { rawItem: { name: "calculator", arguments: '{"a":1}' } },
+ },
+ ],
+ }),
+ );
+
+ const agent = AIPex.create({
+ instructions: "Tools",
+ model: mockModel,
+ });
+
+ const events: AgentEvent[] = [];
+ for await (const event of agent.chat("use tool")) {
+ events.push(event);
+ }
+
+ const argsStart = events.find(
+ (event) => event.type === "tool_call_args_streaming_start",
+ );
+ expect(argsStart).toBeDefined();
+ if (argsStart?.type === "tool_call_args_streaming_start") {
+ expect(argsStart.toolName).toBe("calculator");
+ }
+ });
+
  it("should emit tool lifecycle events", async () => {
  vi.mocked(run).mockResolvedValue(
  createMockRunResult({