Merge pull request #402 from stainless-sdks/dmeadows/add-server-tool-use-input-delta

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,12 @@ type MessageStreamEventListeners<Event extends keyof MessageStreamEvents> = {
 
 const JSON_BUF_PROPERTY = '__json_buf';
 
+export type TracksToolInput = BetaToolUseBlock | BetaServerToolUseBlock | BetaMCPToolUseBlock;
+
+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 +441,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 +580,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,12 @@ type MessageStreamEventListeners<Event extends keyof MessageStreamEvents> = {
 
 const JSON_BUF_PROPERTY = '__json_buf';
 
+export type TracksToolInput = ToolUseBlock | ServerToolUseBlock;
+
+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 +440,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 +579,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,44 @@
+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', () => {
+      type Test = ContentBlockWithInput extends TracksToolInput ? true : false;
+      const test: Test = true;
+      expect(test).toBe(true);
+    });
+
+    it('all TracksToolInput types should have an input property', () => {
+      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', () => {
+      type Test = BetaContentBlockWithInput extends BetaTracksToolInput ? true : false;
+      const test: Test = true;
+      expect(test).toBe(true);
+    });
+
+    it('all BetaTracksToolInput types should have an input property', () => {
+      type Test2 = BetaTracksToolInput extends BetaContentBlockWithInput ? true : false;
+      const test2: Test2 = true;
+      expect(test2).toBe(true);
+    });
+  });
+});