fix(client): correctly track input from server_tool_use input deltas

Author: dtmeadows

src/lib/BetaMessageStream.ts

@@ -10,6 +10,9 @@ import {
   type MessageCreateParamsBase as BetaMessageCreateParamsBase,
   type BetaTextBlock,
   type BetaTextCitation,
+  type BetaToolUseBlock,
+  type BetaServerToolUseBlock,
+  type BetaMCPToolUseBlock,
 } from '../resources/beta/messages/messages';
 import { Stream } from '../streaming';
 import { partialParse } from '../_vendor/partial-json-parser/parser';
@@ -39,6 +42,18 @@ type MessageStreamEventListeners<Event extends keyof MessageStreamEvents> = {
 
 const JSON_BUF_PROPERTY = '__json_buf';
 
+/**
+ * Types of content blocks that track tool input via input_json_delta events
+ */
+export type TracksToolInput = BetaToolUseBlock | BetaServerToolUseBlock | BetaMCPToolUseBlock;
+
+/**
+ * Type guard to check if a content block is one that tracks tool input
+ */
+function tracksToolInput(content: BetaContentBlock): content is TracksToolInput {
+  return content.type === 'tool_use' || content.type === 'server_tool_use' || content.type === 'mcp_tool_use';
+}
+
 export class BetaMessageStream implements AsyncIterable<BetaMessageStreamEvent> {
   messages: BetaMessageParam[] = [];
   receivedMessages: BetaMessage[] = [];
@@ -432,7 +447,7 @@ export class BetaMessageStream implements AsyncIterable<BetaMessageStreamEvent>
             break;
           }
           case 'input_json_delta': {
-            if ((content.type === 'tool_use' || content.type === 'mcp_tool_use') && content.input) {
+            if (tracksToolInput(content) && content.input) {
               this._emit('inputJson', event.delta.partial_json, content.input);
             }
             break;
@@ -571,7 +586,7 @@ export class BetaMessageStream implements AsyncIterable<BetaMessageStreamEvent>
             break;
           }
           case 'input_json_delta': {
-            if (snapshotContent?.type === 'tool_use' || snapshotContent?.type === 'mcp_tool_use') {
+            if (snapshotContent && tracksToolInput(snapshotContent)) {
               // we need to keep track of the raw JSON string as well so that we can
               // re-parse it for each delta, for now we just store it as an untyped
               // non-enumerable property on the snapshot

src/lib/MessageStream.ts

@@ -10,6 +10,8 @@ import {
   type MessageCreateParamsBase,
   type TextBlock,
   type TextCitation,
+  type ToolUseBlock,
+  type ServerToolUseBlock,
 } from '../resources/messages';
 import { Stream } from '../streaming';
 import { partialParse } from '../_vendor/partial-json-parser/parser';
@@ -39,6 +41,19 @@ type MessageStreamEventListeners<Event extends keyof MessageStreamEvents> = {
 
 const JSON_BUF_PROPERTY = '__json_buf';
 
+/**
+ * Types of content blocks that track tool input via input_json_delta events
+ */
+export type TracksToolInput = ToolUseBlock | ServerToolUseBlock;
+
+/**
+ * Type guard to check if a content block is one that tracks tool input
+ * by having an `input` property that is updated via input_json_delta events.
+ */
+function tracksToolInput(content: ContentBlock): content is TracksToolInput {
+  return content.type === 'tool_use' || content.type === 'server_tool_use';
+}
+
 export class MessageStream implements AsyncIterable<MessageStreamEvent> {
   messages: MessageParam[] = [];
   receivedMessages: Message[] = [];
@@ -432,7 +447,7 @@ export class MessageStream implements AsyncIterable<MessageStreamEvent> {
             break;
           }
           case 'input_json_delta': {
-            if (content.type === 'tool_use' && content.input) {
+            if (tracksToolInput(content) && content.input) {
               this._emit('inputJson', event.delta.partial_json, content.input);
             }
             break;
@@ -571,7 +586,7 @@ export class MessageStream implements AsyncIterable<MessageStreamEvent> {
             break;
           }
           case 'input_json_delta': {
-            if (snapshotContent?.type === 'tool_use') {
+            if (snapshotContent && tracksToolInput(snapshotContent)) {
               // we need to keep track of the raw JSON string as well so that we can
               // re-parse it for each delta, for now we just store it as an untyped
               // non-enumerable property on the snapshot

tests/lib/TracksToolInput.test.ts

@@ -0,0 +1,56 @@
+import { ContentBlock } from '../../src/resources/messages';
+
+import { TracksToolInput } from '@anthropic-ai/sdk/lib/MessageStream';
+import { TracksToolInput as BetaTracksToolInput } from '@anthropic-ai/sdk/lib/BetaMessageStream';
+import { BetaContentBlock } from '@anthropic-ai/sdk/resources/beta';
+
+/**
+ * This test ensures that our TracksToolInput type includes all content block types that have an input property.
+ * If any new content block types with input properties are added, they should be added to the TracksToolInput types.
+ */
+
+describe('TracksToolInput type', () => {
+  describe('Regular MessageStream', () => {
+    type ContentBlockWithInput = Extract<ContentBlock, { input: unknown }>;
+
+    it('TracksToolInput includes all content block types with input properties', () => {
+      // Define a type that extracts all ContentBlock types that have an input property
+
+      // TypeScript compile-time check: all ContentBlock types with input property
+      // should be assignable to TracksToolInput
+      type Test = ContentBlockWithInput extends TracksToolInput ? true : false;
+      const test: Test = true;
+      expect(test).toBe(true);
+    });
+
+    it('all TracksToolInput types should have an input property', () => {
+      // TypeScript compile-time check: all TracksToolInput types should have
+      // an input property (i.e., be assignable to ContentBlockWithInput)
+      type Test2 = TracksToolInput extends ContentBlockWithInput ? true : false;
+      const test2: Test2 = true;
+      expect(test2).toBe(true);
+    });
+  });
+
+  describe('Beta MessageStream', () => {
+    type BetaContentBlockWithInput = Extract<BetaContentBlock, { input: unknown }>;
+
+    it('TracksToolInput includes all content block types with input properties', () => {
+      // Define a type that extracts all BetaContentBlock types that have an input property
+
+      // TypeScript compile-time check: all BetaContentBlock types with input property
+      // should be assignable to BetaTracksToolInput
+      type Test = BetaContentBlockWithInput extends BetaTracksToolInput ? true : false;
+      const test: Test = true;
+      expect(test).toBe(true);
+    });
+
+    it('all BetaTracksToolInput types should have an input property', () => {
+      // TypeScript compile-time check: all BetaTracksToolInput types should have
+      // an input property (i.e., be assignable to BetaContentBlockWithInput)
+      type Test2 = BetaTracksToolInput extends BetaContentBlockWithInput ? true : false;
+      const test2: Test2 = true;
+      expect(test2).toBe(true);
+    });
+  });
+});