Skip to content

feat: add support for MCP Sampling as AI provider#863

Merged
Crunchyman-ralph merged 50 commits into
eyaltoledano:nextfrom
OrenMe:orenm/mcpsampling
Jul 9, 2025
Merged

feat: add support for MCP Sampling as AI provider#863
Crunchyman-ralph merged 50 commits into
eyaltoledano:nextfrom
OrenMe:orenm/mcpsampling

Conversation

@OrenMe

@OrenMe OrenMe commented Jun 24, 2025

Copy link
Copy Markdown
Contributor

Description

This PR introduces MCP Sampling as a new AI provider for Task Master, enabling the system to use MCP (Model Context Protocol) client capabilities as an AI provider without requiring API keys. The implementation includes a dynamic provider registry system that allows runtime registration of remote providers while maintaining backward compatibility with existing static providers.

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation update

Testing

  • I have tested this locally
  • All existing tests pass
  • I have added tests for new functionality

Changeset

  • I have created a changeset (or this change doesn't need one)

Additional Notes

Key Features

🔧 Provider Registry System

  • New Singleton Registry: Created index.js for dynamic provider registration
  • Backward Compatibility: Maintains existing static PROVIDERS object while adding registry fallback
  • Runtime Registration: Enables MCP server to register providers when sessions connect

🌐 MCP Remote Provider

  • Session-Based Authentication: No API key required - uses MCP client sampling capabilities
  • Partial AI Operations Support: Implements generateText() and generateObject(), no streamText() support
  • Automatic Registration: Registers with provider registry when MCP sessions with sampling capabilities connect
  • Timeout Configuration: 2-minute timeout for complex AI operations

📚 Model Support

  • New Model ID: mcp-sampling for general text generation
  • Role Support: Compatible with main, fallback, and research roles

🧪 Comprehensive Testing

  • Unit Tests: Complete test coverage for MCP remote provider
  • Validation Tests: Session capability validation and error handling
  • Integration Tests: Provider registry integration testing

📖 Documentation

  • MCP Provider Guide: Comprehensive documentation with setup instructions
  • Configuration Examples: Clear examples for different use cases
  • Troubleshooting: Common issues and solutions

Future enhancements:

  • Migrate all providers to new registry
  • Clean and simplify ai-services-unified and make all provider logic generic and delegated to provider factory instance
  • Support server hints for better MCP sampling LLM selection

Additional notes

Had to patch FastMCP library to support passing requestSampling additional params punkpeye/fastmcp#125

task_master_sampling_demo_.mp4

@changeset-bot

changeset-bot Bot commented Jun 24, 2025

Copy link
Copy Markdown

🦋 Changeset detected

Latest commit: a8609e8

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 1 package
Name Type
task-master-ai Minor

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

@Crunchyman-ralph Crunchyman-ralph left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Great PR, will review again more thoroughly.

Mind if I ask which MCP clients this sampling will work for ?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

similar to how we made a custom sdk package for claude code, think we should have something similar for the sampling.

@OrenMe

OrenMe commented Jun 24, 2025

Copy link
Copy Markdown
Contributor Author

Thanks, tried to keep up to the standards
You are correct to ask - for now the only one I know to support it and the full MCP spec is VSCode
Fully tested on it
As long as they implement the spec to the dot it should be compatible with any

I wrote it in the PR description but more in depth support for model selection can be extended
Also I would suggest we revamp the provider registration and clean the ai unified class and utils so they won't know about hardcoded specific of providers

Maybe I forgot it so will write here a sampling has timeout and tasks like analyze might take a while
I think it would be best to move this to config value

@Crunchyman-ralph

Copy link
Copy Markdown
Collaborator

@OrenMe feel free to refactor what you're talking about in a separate, lets keep this PR scoped to adding the "sampling" as a provider

As I mentioned, it would be nice to make the "sampling provider" a custom ai-sdk provider, that would work better with what we currently have (check out the claude code example)

I'll be monitoring this more closely, let me know if there is anything else I can do to help!

ben-vargas added a commit to ben-vargas/ai-claude-task-master that referenced this pull request Jun 26, 2025
- Add MCP Sampling provider with custom-sdk pattern implementation
- Support MCP preference-based model selection (hints and priorities)
- Add comprehensive test coverage (69 tests)
- Add usage documentation and implementation notes
- Refactor from PR eyaltoledano#863 monolithic approach to modular design
- Follow established claude-code provider patterns
@Ton-Git

Ton-Git commented Jul 3, 2025

Copy link
Copy Markdown

Hi @OrenMe, sure I'll do it as soon as I can, thank you!

@OrenMe

OrenMe commented Jul 3, 2025

Copy link
Copy Markdown
Contributor Author

strange I can't repro the error locally, this was the last error but after I pushed the gemini fix I can't repro locally, unless I revert it
can we try and run the test again, clean the cache maybe?

@OrenMe OrenMe requested a review from Crunchyman-ralph July 4, 2025 06:41
@Ton-Git

Ton-Git commented Jul 4, 2025

Copy link
Copy Markdown

Hi @Ton-Git any chance u can create a PR on another branch on top of mine so I can see the diff more clearly A bit hard with the patch file

Hi @OrenMe, I started fresh by forking your repo and I don't have those errors anymore, so no need for an additional PR. Can't wait to see this in the next release, thank you :)

@OrenMe

OrenMe commented Jul 4, 2025

Copy link
Copy Markdown
Contributor Author

Thanks @Ton-Git for verifying
I was looking into this today as well and beside the json replace I didn't see anything else which seemed like a must as well
So in this case @Crunchyman-ralph what r the chances we try to add this? I saw there is already a PR for the new version so guess it might be too late but hopefully there's a chance 😁

ben-vargas and others added 5 commits July 7, 2025 15:47
) (eyaltoledano#920)

Gracefully handle SyntaxError thrown by @anthropic-ai/claude-code when the CLI truncates large JSON outputs (4–16 kB cut-offs).\n\nKey points:\n• Detect JSON parse error + existing buffered text in both doGenerate() and doStream() code paths.\n• Convert the failure into a recoverable 'truncated' finish state and push a provider-warning.\n• Allows Task Master to continue parsing long PRDs / expand-task operations instead of crashing.\n\nA patch changeset (.changeset/claude-code-json-truncation.md) is included for the next release.\n\nRef: eyaltoledano#913
Remove erroneous 'gemini auth login' command references and replace with correct 'gemini' command authentication flow. Update documentation to reflect proper OAuth setup process via the gemini CLI interactive interface.
…cation (eyaltoledano#932)

- Fixed authentication compatibility issues with Google auth
- Added support for 'api-key' auth type alongside 'gemini-api-key'
- Resolved "Unsupported authType: undefined" runtime errors
- Updated @google/gemini-cli-core dependency to 0.1.9
- Improved documentation and removed invalid auth references
- Maintained backward compatibility while enhancing type validation
@Crunchyman-ralph

Copy link
Copy Markdown
Collaborator

Seems to be working pretty well on cursor and CLI, falls back to the right fallback model.

@Crunchyman-ralph

Copy link
Copy Markdown
Collaborator

Testing on vscode, watch out for how you use this.logger inside MCP, this can crash the mcp server in stdio:
/Volumes/Workspace/workspace/contrib/task-master/claude-task-master/mcp-server/src/index.js

OrenMe added 6 commits July 8, 2025 00:02
Need to patch upstream fastmcp to allow easier access and bootstrap the TM mcp logger to use the fastmcp logger which today is only exposed in the tools handler

@Crunchyman-ralph Crunchyman-ralph left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

lgtm, though feels like you have changes from other PRs, maybe we should merge origin/next ?

Comment thread docs/README.md Outdated
Comment thread mcp-server/src/index.js
Comment thread package.json
@Crunchyman-ralph Crunchyman-ralph merged commit b530657 into eyaltoledano:next Jul 9, 2025
3 checks passed
@github-actions github-actions Bot mentioned this pull request Jul 12, 2025
stephanschielke pushed a commit to stephanschielke/cursor-task-master that referenced this pull request Aug 22, 2025
* feat: support MCP sampling

* support provider registry

* use standard config options for MCP provider

* update fastmcp to support passing params to requestSampling

* move key name definition to base provider

* moved check for required api key to provider class

* remove unused code

* more cleanup

* more cleanup

* refactor provider

* remove not needed files

* more cleanup

* more cleanup

* more cleanup

* update docs

* fix tests

* add tests

* format fix

* clean files

* merge fixes

* format fix

* feat: add support for MCP Sampling as AI provider

* initial mcp ai sdk

* fix references to old provider

* update models

* lint

* fix gemini-cli conflicts

* ran format

* Update src/provider-registry/index.js

Co-authored-by: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>

* fix circular dependency

Circular Dependency Issue ✅ FIXED
Root Cause: BaseAIProvider was importing from index.js, which includes commands.js and other modules that eventually import back to AI providers
Solution: Changed imports to use direct paths to avoid circular dependencies:
Updated base-provider.js to import log directly from utils.js
Updated gemini-cli.js to import log directly from utils.js
Result: Fixed 11 failing tests in mcp-provider.test.js

* fix gemini test

* fix(claude-code): recover from CLI JSON truncation bug (eyaltoledano#913) (eyaltoledano#920)

Gracefully handle SyntaxError thrown by @anthropic-ai/claude-code when the CLI truncates large JSON outputs (4–16 kB cut-offs).\n\nKey points:\n• Detect JSON parse error + existing buffered text in both doGenerate() and doStream() code paths.\n• Convert the failure into a recoverable 'truncated' finish state and push a provider-warning.\n• Allows Task Master to continue parsing long PRDs / expand-task operations instead of crashing.\n\nA patch changeset (.changeset/claude-code-json-truncation.md) is included for the next release.\n\nRef: eyaltoledano#913

* docs: fix gemini-cli authentication documentation (eyaltoledano#923)

Remove erroneous 'gemini auth login' command references and replace with correct 'gemini' command authentication flow. Update documentation to reflect proper OAuth setup process via the gemini CLI interactive interface.

* fix tests

* fix: update ai-sdk-provider-gemini-cli to 0.0.4 for improved authentication (eyaltoledano#932)

- Fixed authentication compatibility issues with Google auth
- Added support for 'api-key' auth type alongside 'gemini-api-key'
- Resolved "Unsupported authType: undefined" runtime errors
- Updated @google/gemini-cli-core dependency to 0.1.9
- Improved documentation and removed invalid auth references
- Maintained backward compatibility while enhancing type validation

* call logging directly

Need to patch upstream fastmcp to allow easier access and bootstrap the TM mcp logger to use the fastmcp logger which today is only exposed in the tools handler

* fix tests

* removing logs until we figure out how to pass mcp logger

* format

* fix tests

* format

* clean up

* cleanup

* readme fix

---------

Co-authored-by: Oren Melamed <oren.m@gloat.com>
Co-authored-by: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>
Co-authored-by: Ben Vargas <ben@vargas.com>
stephanschielke pushed a commit to stephanschielke/cursor-task-master that referenced this pull request Aug 22, 2025
* feat: support MCP sampling

* support provider registry

* use standard config options for MCP provider

* update fastmcp to support passing params to requestSampling

* move key name definition to base provider

* moved check for required api key to provider class

* remove unused code

* more cleanup

* more cleanup

* refactor provider

* remove not needed files

* more cleanup

* more cleanup

* more cleanup

* update docs

* fix tests

* add tests

* format fix

* clean files

* merge fixes

* format fix

* feat: add support for MCP Sampling as AI provider

* initial mcp ai sdk

* fix references to old provider

* update models

* lint

* fix gemini-cli conflicts

* ran format

* Update src/provider-registry/index.js

Co-authored-by: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>

* fix circular dependency

Circular Dependency Issue ✅ FIXED
Root Cause: BaseAIProvider was importing from index.js, which includes commands.js and other modules that eventually import back to AI providers
Solution: Changed imports to use direct paths to avoid circular dependencies:
Updated base-provider.js to import log directly from utils.js
Updated gemini-cli.js to import log directly from utils.js
Result: Fixed 11 failing tests in mcp-provider.test.js

* fix gemini test

* fix(claude-code): recover from CLI JSON truncation bug (eyaltoledano#913) (eyaltoledano#920)

Gracefully handle SyntaxError thrown by @anthropic-ai/claude-code when the CLI truncates large JSON outputs (4–16 kB cut-offs).\n\nKey points:\n• Detect JSON parse error + existing buffered text in both doGenerate() and doStream() code paths.\n• Convert the failure into a recoverable 'truncated' finish state and push a provider-warning.\n• Allows Task Master to continue parsing long PRDs / expand-task operations instead of crashing.\n\nA patch changeset (.changeset/claude-code-json-truncation.md) is included for the next release.\n\nRef: eyaltoledano#913

* docs: fix gemini-cli authentication documentation (eyaltoledano#923)

Remove erroneous 'gemini auth login' command references and replace with correct 'gemini' command authentication flow. Update documentation to reflect proper OAuth setup process via the gemini CLI interactive interface.

* fix tests

* fix: update ai-sdk-provider-gemini-cli to 0.0.4 for improved authentication (eyaltoledano#932)

- Fixed authentication compatibility issues with Google auth
- Added support for 'api-key' auth type alongside 'gemini-api-key'
- Resolved "Unsupported authType: undefined" runtime errors
- Updated @google/gemini-cli-core dependency to 0.1.9
- Improved documentation and removed invalid auth references
- Maintained backward compatibility while enhancing type validation

* call logging directly

Need to patch upstream fastmcp to allow easier access and bootstrap the TM mcp logger to use the fastmcp logger which today is only exposed in the tools handler

* fix tests

* removing logs until we figure out how to pass mcp logger

* format

* fix tests

* format

* clean up

* cleanup

* readme fix

---------

Co-authored-by: Oren Melamed <oren.m@gloat.com>
Co-authored-by: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>
Co-authored-by: Ben Vargas <ben@vargas.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants