Skip to content

feat(typescript-anthropic): add streaming support#20384

Merged
B-Step62 merged 2 commits intomlflow:masterfrom
rollyjoel:streaming
Feb 3, 2026
Merged

feat(typescript-anthropic): add streaming support#20384
B-Step62 merged 2 commits intomlflow:masterfrom
rollyjoel:streaming

Conversation

@rollyjoel
Copy link
Contributor

@rollyjoel rollyjoel commented Jan 27, 2026

Related Issues/PRs

#20382

What changes are proposed in this pull request?

Add tracing support for the messages.stream() method in the MLflow Anthropic TypeScript integration.

The stream method returns a MessageStream object that emits events as the response streams in. This implementation traces streaming by:

  • Wrapping the returned MessageStream object in a Proxy

  • Intercepting finalMessage() calls to create a span around the full streaming operation

  • Intercepting the async iterator (Symbol.asyncIterator) to trace for await loops

  • Extracting token usage from the final message after streaming completes

  • Setting appropriate span attributes (MESSAGE_FORMAT: 'anthropic', token usage, etc.)

  • libs/typescript/integrations/anthropic/src/index.ts - Main implementation

  • libs/typescript/integrations/anthropic/tests/index.test.ts - Tests

  • libs/typescript/integrations/anthropic/tests/mockAnthropicServer.ts - Mock SSE endpoint

How is this PR tested?

  • Existing unit/integration tests
  • New unit/integration tests
  • Manual tests
tracing-screenshot

Does this PR require documentation update?

  • No. You can skip the rest of this section.
  • Yes. I've updated:
    • Examples
    • API references
    • Instructions

Release Notes

Is this a user-facing change?

  • No. You can skip the rest of this section.
  • Yes. Give a description of this change to be included in the release notes for MLflow users.

Add support for @anthropic-ai/sdk's streaming responses

What component(s), interfaces, languages, and integrations does this PR affect?

Components

  • area/tracking: Tracking Service, tracking client APIs, autologging
  • area/models: MLmodel format, model serialization/deserialization, flavors
  • area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
  • area/scoring: MLflow Model server, model deployment tools, Spark UDFs
  • area/evaluation: MLflow model evaluation features, evaluation metrics, and evaluation workflows
  • area/gateway: MLflow AI Gateway client APIs, server, and third-party integrations
  • area/prompts: MLflow prompt engineering features, prompt templates, and prompt management
  • area/tracing: MLflow Tracing features, tracing APIs, and LLM tracing functionality
  • area/projects: MLproject format, project running backends
  • area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
  • area/build: Build and test infrastructure for MLflow
  • area/docs: MLflow documentation pages

How should the PR be classified in the release notes? Choose one:

  • rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
  • rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
  • rn/feature - A new user-facing feature worth mentioning in the release notes
  • rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
  • rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?
  • Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
    Bug fixes, doc updates and new features usually go into minor releases.
  • Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
    Bug fixes and doc updates usually go into patch releases.
  • Yes (this PR will be cherry-picked and included in the next patch release)
  • No (this PR will be included in the next minor release)

@rollyjoel rollyjoel marked this pull request as ready for review January 27, 2026 22:10
Copilot AI review requested due to automatic review settings January 27, 2026 22:10
@github-actions
Copy link
Contributor

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20384/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20384/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/20384/merge

@github-actions github-actions bot added area/tracing MLflow Tracing and its integrations rn/feature Mention under Features in Changelogs. labels Jan 27, 2026
@rollyjoel rollyjoel marked this pull request as draft January 27, 2026 22:11
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This pull request adds streaming support for the MLflow Anthropic TypeScript integration, addressing issue #20382. The implementation enables tracing for the messages.stream() method by wrapping the returned MessageStream object with a Proxy that intercepts both async iteration and finalMessage() calls.

Changes:

  • Added tracing support for Anthropic's messages.stream() method
  • Implemented Proxy-based wrapping for MessageStream with dual patterns: async iteration and finalMessage()
  • Added mock SSE endpoint and comprehensive test coverage for streaming scenarios

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File Description
libs/typescript/integrations/anthropic/src/index.ts Core streaming tracing implementation with Proxy-based MessageStream wrapping and async generator support
libs/typescript/integrations/anthropic/tests/index.test.ts Added three comprehensive tests for streaming: async iteration, finalMessage(), and parent span integration
libs/typescript/integrations/anthropic/tests/mockAnthropicServer.ts Added SSE streaming endpoint mock with proper event sequence simulation
docs/docs/genai/tracing/integrations/listing/anthropic.mdx Updated feature matrix to indicate streaming support

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@github-actions
Copy link
Contributor

github-actions bot commented Jan 29, 2026

Documentation preview for 58e9a55 is available at:

Changed Pages (1)

More info
  • Ignore this comment if this PR does not change the documentation.
  • The preview is updated when a new commit is pushed to this PR.
  • This comment was created by this workflow run.
  • The documentation was built by this workflow run.

@rollyjoel
Copy link
Contributor Author

Thank you to whomever approved the action runs! I've fixed the lint/prettier errors and would love another run!

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Collaborator

@B-Step62 B-Step62 left a comment

Choose a reason for hiding this comment

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

LGTM! Left a minor comment.

span.setAttribute(SpanAttributeKey.TOKEN_USAGE, usage);
}
} catch (error) {
console.debug('Error extracting token usage', error);
Copy link
Collaborator

Choose a reason for hiding this comment

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

Let's remove console.debug usage

Copy link
Contributor Author

Choose a reason for hiding this comment

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

seems to be consistent with how the other integrations handle failure to extract token messages. If you do want me to break the pattern, should I just catch (error) {}?

also, apologies I neglected to run lint 🤦 fixed inline with the other integrations and pushed up another commit (I squashed all the previous feedback commits)

Copy link
Collaborator

Choose a reason for hiding this comment

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

Oh my bad, I wasn't aware we use that extensibly in other flavors. Let's keep it as-is then:)

rollyjoel and others added 2 commits February 2, 2026 17:24
Add tracing support for the `messages.stream()` method in the MLflow
Anthropic TypeScript integration.

The `stream` method returns a `MessageStream` object that emits events
as the response streams in. This implementation traces streaming by:

- Wrapping the returned `MessageStream` object in a Proxy
- Intercepting `finalMessage()` calls to create a span around the full
  streaming operation
- Intercepting the async iterator (`Symbol.asyncIterator`) to trace
  `for await` loops
- Extracting token usage from the final message after streaming completes
- Setting appropriate span attributes (`MESSAGE_FORMAT: 'anthropic'`,
  token usage, etc.)

- `libs/typescript/integrations/anthropic/src/index.ts` - Main implementation
- `libs/typescript/integrations/anthropic/tests/index.test.ts` - Tests
- `libs/typescript/integrations/anthropic/tests/mockAnthropicServer.ts` - Mock SSE endpoint

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Joel Dodge <joeldodge@gmail.com>
Signed-off-by: Joel Dodge <joeldodge@gmail.com>
@B-Step62 B-Step62 enabled auto-merge February 3, 2026 09:56
@B-Step62 B-Step62 added this pull request to the merge queue Feb 3, 2026
Merged via the queue into mlflow:master with commit 55dd48e Feb 3, 2026
47 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/tracing MLflow Tracing and its integrations rn/feature Mention under Features in Changelogs. v3.9.1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants