Merge branch 'main' into support-custom-baseUrl-for-google-ai-studio-gemini

This commit is contained in:
dqroid 2025-03-11 16:15:52 +08:00 • committed by GitHub
commit 8b0956666c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
263 changed files with 27671 additions and 8524 deletions

View file

@ -0,0 +1,5 @@
---
"roo-cline": patch
---
Update GitHub Actions workflow to automatically create and push git tags during release

View file

@ -0,0 +1,5 @@
---
"roo-cline": patch
---
App tab layout fixes

View file

@ -0,0 +1,5 @@
---
"roo-cline": patch
---
Fix usage tracking for SiliconFlow etc

View file

@ -6,22 +6,12 @@
2. Lint Rules:
- Never disable any lint rules without explicit user approval
- If a lint rule needs to be disabled, ask the user first and explain why
- Prefer fixing the underlying issue over disabling the lint rule
- Document any approved lint rule disabling with a comment explaining the reason
3. Logging Guidelines:
- Always instrument code changes using the logger exported from `src\utils\logging\index.ts`.
- This will facilitate efficient debugging without impacting production (as the logger no-ops outside of a test environment.)
- Logs can be found in `logs\app.log`
- Logfile is overwritten on each run to keep it to a manageable volume.
4. Styling Guidelines:
3. Styling Guidelines:
- Use Tailwind CSS classes instead of inline style objects for new markup
- VSCode CSS variables must be added to webview-ui/src/index.css before using them in Tailwind classes
- Example: `<div className="text-md text-vscode-descriptionForeground mb-2" />` instead of style objects
# Adding a New Setting
To add a new setting that persists its state, follow the steps in cline_docs/settings.md

1
.env.sample Normal file
View file

@ -0,0 +1 @@
POSTHOG_API_KEY=key-goes-here

View file

@ -19,5 +19,5 @@
"no-throw-literal": "warn",
"semi": "off"
},
"ignorePatterns": ["out", "dist", "**/*.d.ts"]
"ignorePatterns": ["out", "dist", "**/*.d.ts", "!roo-code.d.ts"]
}

View file

@ -1,37 +1,35 @@
<!-- **Note:** Consider creating PRs as a DRAFT. For early feedback and self-review. -->
## Context
## Description
<!-- Brief description of WHAT you’re doing and WHY. -->
## Type of change
## Implementation
<!-- Please ignore options that are not relevant -->
<!--
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update
Some description of HOW you achieved it. Perhaps give a high level description of the program flow. Did you need to refactor something? What tradeoffs did you take? Are there things in here which you’d particularly like people to pay close attention to?
## How Has This Been Tested?
-->
<!-- Please describe the tests that you ran to verify your changes -->
## Screenshots
## Checklist:
| before | after |
| ------ | ----- |
| | |
<!-- Go over all the following points, and put an `x` in all the boxes that apply -->
## How to Test
- [ ] My code follows the patterns of this project
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
<!--
## Additional context
A straightforward scenario of how to test your changes will help reviewers that are not familiar with the part of the code that you are changing but want to see it in action. This section can include a description or step-by-step instructions of how to get to the state of v2 that your change affects.
<!-- Add any other context or screenshots about the pull request here -->
A "How To Test" section can look something like this:
## Related Issues
- Sign in with a user with tracks
- Activate `show_awesome_cat_gifs` feature (add `?feature.show_awesome_cat_gifs=1` to your URL)
- You should see a GIF with cats dancing
<!-- List any related issues here. Use the GitHub issue linking syntax: #issue-number -->
-->
## Reviewers
## Get in Touch
<!-- @mention specific team members or individuals who should review this PR -->
<!-- We'd love to have a way to chat with you about your changes if necessary. If you're in the [Roo Code Discord](https://discord.gg/roocode), please share your handle here. -->

View file

@ -108,9 +108,11 @@ jobs:
with:
node-version: '18'
cache: 'npm'
- name: Create env.integration file
run: echo "OPENROUTER_API_KEY=${{ secrets.OPENROUTER_API_KEY }}" > .env.integration
- name: Install dependencies
run: npm run install:all
- name: Create env.integration file
working-directory: e2e
run: echo "OPENROUTER_API_KEY=${{ secrets.OPENROUTER_API_KEY }}" > .env.integration
- name: Run integration tests
run: xvfb-run -a npm run test:integration
working-directory: e2e
run: xvfb-run -a npm run ci

View file

@ -10,6 +10,8 @@ env:
jobs:
publish-extension:
runs-on: ubuntu-latest
permissions:
contents: write # Required for pushing tags.
if: >
( github.event_name == 'pull_request' &&
github.event.pull_request.base.ref == 'main' &&
@ -23,29 +25,40 @@ jobs:
- uses: actions/setup-node@v4
with:
node-version: 18
- run: |
git config user.name github-actions
git config user.email github-actions@github.com
- name: Install Dependencies
run: |
npm install -g vsce ovsx
npm install
cd webview-ui
npm install
cd ..
- name: Package and Publish Extension
env:
VSCE_PAT: ${{ secrets.VSCE_PAT }}
OVSX_PAT: ${{ secrets.OVSX_PAT }}
npm run install:all
- name: Create .env file
run: echo "POSTHOG_API_KEY=${{ secrets.POSTHOG_API_KEY }}" >> .env
- name: Package Extension
run: |
current_package_version=$(node -p "require('./package.json').version")
npm run vsix
package=$(unzip -l bin/roo-cline-${current_package_version}.vsix)
echo "$package"
echo "$package" | grep -q "dist/extension.js" || exit 1
echo "$package" | grep -q "extension/webview-ui/build/assets/index.js" || exit 1
echo "$package" | grep -q "extension/node_modules/@vscode/codicons/dist/codicon.ttf" || exit 1
echo "$package" | grep -q ".env" || exit 1
- name: Create and Push Git Tag
run: |
current_package_version=$(node -p "require('./package.json').version")
git tag -a "v${current_package_version}" -m "Release v${current_package_version}"
git push origin "v${current_package_version}"
echo "Successfully created and pushed git tag v${current_package_version}"
- name: Publish Extension
env:
VSCE_PAT: ${{ secrets.VSCE_PAT }}
OVSX_PAT: ${{ secrets.OVSX_PAT }}
run: |
current_package_version=$(node -p "require('./package.json').version")
npm run publish:marketplace
echo "Successfully published version $current_package_version to VS Code Marketplace"

1
.gitignore vendored
View file

@ -21,6 +21,7 @@ roo-cline-*.vsix
docs/_site/
# Dotenv
.env
.env.integration
#Local lint config

1
.rooignore Normal file
View file

@ -0,0 +1 @@
.env

View file

@ -4,6 +4,8 @@
.vscode/**
.vscode-test/**
out/**
out-integration/**
e2e/**
node_modules/**
src/**
.gitignore
@ -25,7 +27,6 @@ demo.gif
.roomodes
cline_docs/**
coverage/**
out-integration/**
# Ignore all webview-ui files except the build directory (https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/frameworks/hello-world-react-cra/.vscodeignore)
webview-ui/src/**
@ -47,3 +48,6 @@ webview-ui/node_modules/**
# Include icons
!assets/icons/**
# Include .env file for telemetry
!.env

View file

@ -1,43 +1,156 @@
# Roo Code Changelog
## [3.7.3]
## [3.8.4] - 2025-03-09
- Roll back multi-diff progress indicator temporarily to fix a double-confirmation in saving edits
- Add an option in the prompts tab to save tokens by disabling the ability to ask Roo to create/edit custom modes for you (thanks @hannesrudolph!)
## [3.8.3] - 2025-03-09
- Fix VS Code LM API model picker truncation issue
## [3.8.2] - 2025-03-08
- Create an auto-approval toggle for subtask creation and completion (thanks @shaybc!)
- Show a progress indicator when using the multi-diff editing strategy (thanks @qdaxb!)
- Add o3-mini support to the OpenAI-compatible provider (thanks @yt3trees!)
- Fix encoding issue where unreadable characters were sometimes getting added to the beginning of files
- Fix issue where settings dropdowns were getting truncated in some cases
## [3.8.1] - 2025-03-07
- Show the reserved output tokens in the context window visualization
- Improve the UI of the configuration profile dropdown (thanks @DeXtroTip!)
- Fix bug where custom temperature could not be unchecked (thanks @System233!)
- Fix bug where decimal prices could not be entered for OpenAI-compatible providers (thanks @System233!)
- Fix bug with enhance prompt on Sonnet 3.7 with a high thinking budget (thanks @moqimoqidea!)
- Fix bug with the context window management for thinking models (thanks @ReadyPlayerEmma!)
- Fix bug where checkpoints were no longer enabled by default
- Add extension and VSCode versions to telemetry
## [3.8.0] - 2025-03-07
- Add opt-in telemetry to help us improve Roo Code faster (thanks Cline!)
- Fix terminal overload / gray screen of death, and other terminal issues
- Add a new experimental diff editing strategy that applies multiple diff edits at once (thanks @qdaxb!)
- Add support for a .rooignore to prevent Roo Code from read/writing certain files, with a setting to also exclude them from search/lists (thanks Cline!)
- Update the new_task tool to return results to the parent task on completion, supporting better orchestration (thanks @shaybc!)
- Support running Roo in multiple editor windows simultaneously (thanks @samhvw8!)
- Make checkpoints asynchronous and exclude more files to speed them up
- Redesign the settings page to make it easier to navigate
- Add credential-based authentication for Vertex AI, enabling users to easily switch between Google Cloud accounts (thanks @eonghk!)
- Update the DeepSeek provider with the correct baseUrl and track caching correctly (thanks @olweraltuve!)
- Add a new “Human Relay” provider that allows you to manually copy information to a Web AI when needed, and then paste the AI's response back into Roo Code (thanks @NyxJae)!
- Add observability for OpenAI providers (thanks @refactorthis!)
- Support speculative decoding for LM Studio local models (thanks @adamwlarson!)
- Improve UI for mode/provider selectors in chat
- Improve styling of the task headers (thanks @monotykamary!)
- Improve context mention path handling on Windows (thanks @samhvw8!)
## [3.7.12] - 2025-03-03
- Expand max tokens of thinking models to 128k, and max thinking budget to over 100k (thanks @monotykamary!)
- Fix issue where keyboard mode switcher wasn't updating API profile (thanks @aheizi!)
- Use the count_tokens API in the Anthropic provider for more accurate context window management
- Default middle-out compression to on for OpenRouter
- Exclude MCP instructions from the prompt if the mode doesn't support MCP
- Add a checkbox to disable the browser tool
- Show a warning if checkpoints are taking too long to load
- Update the warning text for the VS LM API
- Correctly populate the default OpenRouter model on the welcome screen
## [3.7.11] - 2025-03-02
- Don't honor custom max tokens for non thinking models
- Include custom modes in mode switching keyboard shortcut
- Support read-only modes that can run commands
## [3.7.10] - 2025-03-01
- Add Gemini models on Vertex AI (thanks @ashktn!)
- Keyboard shortcuts to switch modes (thanks @aheizi!)
- Add support for Mermaid diagrams (thanks Cline!)
## [3.7.9] - 2025-03-01
- Delete task confirmation enhancements
- Smarter context window management
- Prettier thinking blocks
- Fix maxTokens defaults for Claude 3.7 Sonnet models
- Terminal output parsing improvements (thanks @KJ7LNW!)
- UI fix to dropdown hover colors (thanks @SamirSaji!)
- Add support for Claude Sonnet 3.7 thinking via Vertex AI (thanks @lupuletic!)
## [3.7.8] - 2025-02-27
- Add Vertex AI prompt caching support for Claude models (thanks @aitoroses and @lupuletic!)
- Add gpt-4.5-preview
- Add an advanced feature to customize the system prompt
## [3.7.7] - 2025-02-27
- Graduate checkpoints out of beta
- Fix enhance prompt button when using Thinking Sonnet
- Add tooltips to make what buttons do more obvious
## [3.7.6] - 2025-02-26
- Handle really long text better in the in the ChatRow similar to TaskHeader (thanks @joemanley201!)
- Support multiple files in drag-and-drop
- Truncate search_file output to avoid crashing the extension
- Better OpenRouter error handling (no more "Provider Error")
- Add slider to control max output tokens for thinking models
## [3.7.5] - 2025-02-26
- Fix context window truncation math (see [#1173](https://github.com/RooVetGit/Roo-Code/issues/1173))
- Fix various issues with the model picker (thanks @System233!)
- Fix model input / output cost parsing (thanks @System233!)
- Add drag-and-drop for files
- Enable the "Thinking Budget" slider for Claude 3.7 Sonnet on OpenRouter
## [3.7.4] - 2025-02-25
- Fix a bug that prevented the "Thinking" setting from properly updating when switching profiles.
## [3.7.3] - 2025-02-25
- Support for ["Thinking"](https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking) Sonnet 3.7 when using the Anthropic provider.
## [3.7.2]
## [3.7.2] - 2025-02-24
- Fix computer use and prompt caching for OpenRouter's `anthropic/claude-3.7-sonnet:beta` (thanks @cte!)
- Fix sliding window calculations for Sonnet 3.7 that were causing a context window overflow (thanks @cte!)
- Encourage diff editing more strongly in the system prompt (thanks @hannesrudolph!)
## [3.7.1]
## [3.7.1] - 2025-02-24
- Add AWS Bedrock support for Sonnet 3.7 and update some defaults to Sonnet 3.7 instead of 3.5
## [3.7.0]
## [3.7.0] - 2025-02-24
- Introducing Roo Code 3.7, with support for the new Claude Sonnet 3.7. Because who cares about skipping version numbers anymore? Thanks @lupuletic and @cte for the PRs!
## [3.3.26]
## [3.3.26] - 2025-02-27
- Adjust the default prompt for Debug mode to focus more on diagnosis and to require user confirmation before moving on to implementation
## [3.3.25]
## [3.3.25] - 2025-02-21
- Add a "Debug" mode that specializes in debugging tricky problems (thanks [Ted Werbel](https://x.com/tedx_ai/status/1891514191179309457) and [Carlos E. Perez](https://x.com/IntuitMachine/status/1891516362486337739)!)
- Add an experimental "Power Steering" option to significantly improve adherence to role definitions and custom instructions
## [3.3.24]
## [3.3.24] - 2025-02-20
- Fixed a bug with region selection preventing AWS Bedrock profiles from being saved (thanks @oprstchn!)
- Updated the price of gpt-4o (thanks @marvijo-code!)
## [3.3.23]
## [3.3.23] - 2025-02-20
- Handle errors more gracefully when reading custom instructions from files (thanks @joemanley201!)
- Bug fix to hitting "Done" on settings page with unsaved changes (thanks @System233!)
## [3.3.22]
## [3.3.22] - 2025-02-20
- Improve the Provider Settings configuration with clear Save buttons and warnings about unsaved changes (thanks @System233!)
- Correctly parse `<think>` reasoning tags from Ollama models (thanks @System233!)
@ -47,7 +160,7 @@
- Fix a bug where the .roomodes file was not automatically created when adding custom modes from the Prompts tab
- Allow setting a wildcard (`*`) to auto-approve all command execution (use with caution!)
## [3.3.21]
## [3.3.21] - 2025-02-17
- Fix input box revert issue and configuration loss during profile switch (thanks @System233!)
- Fix default preferred language for zh-cn and zh-tw (thanks @System233!)
@ -56,7 +169,7 @@
- Fix system prompt to make sure Roo knows about all available modes
- Enable streaming mode for OpenAI o1
## [3.3.20]
## [3.3.20] - 2025-02-14
- Support project-specific custom modes in a .roomodes file
- Add more Mistral models (thanks @d-oit and @bramburn!)
@ -64,7 +177,7 @@
- Add a setting to control the number of open editor tabs to tell the model about (665 is probably too many!)
- Fix race condition bug with entering API key on the welcome screen
## [3.3.19]
## [3.3.19] - 2025-02-12
- Fix a bug where aborting in the middle of file writes would not revert the write
- Honor the VS Code theme for dialog backgrounds
@ -72,7 +185,7 @@
- Add a help button that links to our new documentation site (which we would love help from the community to improve!)
- Switch checkpoints logic to use a shadow git repository to work around issues with hot reloads and polluting existing repositories (thanks Cline for the inspiration!)
## [3.3.18]
## [3.3.18] - 2025-02-11
- Add a per-API-configuration model temperature setting (thanks @joemanley201!)
- Add retries for fetching usage stats from OpenRouter (thanks @jcbdev!)
@ -83,18 +196,18 @@
- Fix logic error where automatic retries were waiting twice as long as intended
- Rework the checkpoints code to avoid conflicts with file locks on Windows (sorry for the hassle!)
## [3.3.17]
## [3.3.17] - 2025-02-09
- Fix the restore checkpoint popover
- Unset git config that was previously set incorrectly by the checkpoints feature
## [3.3.16]
## [3.3.16] - 2025-02-09
- Support Volcano Ark platform through the OpenAI-compatible provider
- Fix jumpiness while entering API config by updating on blur instead of input
- Add tooltips on checkpoint actions and fix an issue where checkpoints were overwriting existing git name/email settings - thanks for the feedback!
## [3.3.15]
## [3.3.15] - 2025-02-08
- Improvements to MCP initialization and server restarts (thanks @MuriloFP and @hannesrudolph!)
- Add a copy button to the recent tasks (thanks @hannesrudolph!)

37
PRIVACY.md Normal file
View file

@ -0,0 +1,37 @@
# Roo Code Privacy Policy
**Last Updated: March 7th, 2025**
Roo Code respects your privacy and is committed to transparency about how we handle your data. Below is a simple breakdown of where key pieces of data go—and, importantly, where they don’t.
### **Where Your Data Goes (And Where It Doesn’t)**
- **Code & Files**: Roo Code accesses files on your local machine when needed for AI-assisted features. When you send commands to Roo Code, relevant files may be transmitted to your chosen AI model provider (e.g., OpenAI, Anthropic, OpenRouter) to generate responses. We do not have access to this data, but AI providers may store it per their privacy policies.
- **Commands**: Any commands executed through Roo Code happen on your local environment. However, when you use AI-powered features, the relevant code and context from your commands may be transmitted to your chosen AI model provider (e.g., OpenAI, Anthropic, OpenRouter) to generate responses. We do not have access to or store this data, but AI providers may process it per their privacy policies.
- **Prompts & AI Requests**: When you use AI-powered features, your prompts and relevant project context are sent to your chosen AI model provider (e.g., OpenAI, Anthropic, OpenRouter) to generate responses. We do not store or process this data. These AI providers have their own privacy policies and may store data per their terms of service.
- **API Keys & Credentials**: If you enter an API key (e.g., to connect an AI model), it is stored locally on your device and never sent to us or any third party, except the provider you have chosen.
- **Telemetry (Usage Data)**: We only collect feature usage and error data if you explicitly opt-in. This telemetry is powered by PostHog and helps us understand feature usage to improve Roo Code. This includes your VS Code machine ID and feature usage patterns and exception reports. We do **not** collect personally identifiable information, your code, or AI prompts.
### **How We Use Your Data (If Collected)**
- If you opt-in to telemetry, we use it to understand feature usage and improve Roo Code.
- We do **not** sell or share your data.
- We do **not** train any models on your data.
### **Your Choices & Control**
- You can run models locally to prevent data being sent to third-parties.
- By default, telemetry collection is off and if you turn it on, you can opt out of telemetry at any time.
- You can delete Roo Code to stop all data collection.
### **Security & Updates**
We take reasonable measures to secure your data, but no system is 100% secure. If our privacy policy changes, we will notify you within the extension.
### **Contact Us**
For any privacy-related questions, reach out to us at support@roocode.com.
---
By using Roo Code, you agree to this Privacy Policy.

View file

@ -2,8 +2,8 @@
<h2>Join the Roo Code Community</h2>
<p>Connect with developers, contribute ideas, and stay ahead with the latest AI-powered coding tools.</p>
<a href="https://discord.gg/roocode" target="_blank"><img src="https://img.shields.io/badge/Join%20Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white" alt="Join Discord" height="60"></a>
<a href="https://www.reddit.com/r/RooCode/" target="_blank"><img src="https://img.shields.io/badge/Join%20Reddit-FF4500?style=for-the-badge&logo=reddit&logoColor=white" alt="Join Reddit" height="60"></a>
<a href="https://discord.gg/roocode" target="_blank"><img src="https://img.shields.io/badge/Join%20Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white" alt="Join Discord"></a>
<a href="https://www.reddit.com/r/RooCode/" target="_blank"><img src="https://img.shields.io/badge/Join%20Reddit-FF4500?style=for-the-badge&logo=reddit&logoColor=white" alt="Join Reddit"></a>
</div>
<br>
@ -34,15 +34,18 @@ Check out the [CHANGELOG](CHANGELOG.md) for detailed updates and fixes.
---
## New in 3.7: Claude 3.7 Sonnet Support 🚀
## 🎉 Roo Code 3.8 Released
We're excited to announce support for Anthropic's latest model, Claude 3.7 Sonnet! The model shows notable improvements in:
Roo Code 3.8 is out with performance boosts, new features, and bug fixes.
- Front-end development and full-stack updates
- Agentic workflows for multi-step processes
- More accurate math, coding, and instruction-following
Try it today in your provider of choice!
- Faster asynchronous checkpoints
- Support for .rooignore files
- Fixed terminal & gray screen issues
- Roo Code can run in multiple windows
- Experimental multi-diff editing strategy
- Subtask to parent task communication
- Updated DeepSeek provider
- New "Human Relay" provider
---
@ -112,31 +115,40 @@ Make Roo Code work your way with:
## Local Setup & Development
1. **Clone** the repo:
```bash
git clone https://github.com/RooVetGit/Roo-Code.git
```
```sh
git clone https://github.com/RooVetGit/Roo-Code.git
```
2. **Install dependencies**:
```bash
npm run install:all
```
3. **Build** the extension:
```bash
npm run build
```
- A `.vsix` file will appear in the `bin/` directory.
4. **Install** the `.vsix` manually if desired:
```bash
code --install-extension bin/roo-code-4.0.0.vsix
```
5. **Start the webview (Vite/React app with HMR)**:
```bash
npm run dev
```
6. **Debug**:
- Press `F5` (or **Run** → **Start Debugging**) in VSCode to open a new session with Roo Code loaded.
```sh
npm run install:all
```
3. **Start the webview (Vite/React app with HMR)**:
```sh
npm run dev
```
4. **Debug**:
Press `F5` (or **Run** → **Start Debugging**) in VSCode to open a new session with Roo Code loaded.
Changes to the webview will appear immediately. Changes to the core extension will require a restart of the extension host.
Alternatively you can build a .vsix and install it directly in VSCode:
```sh
npm run build
```
A `.vsix` file will appear in the `bin/` directory which can be installed with:
```sh
code --install-extension bin/roo-cline-<version>.vsix
```
We use [changesets](https://github.com/changesets/changesets) for versioning and publishing. Check our `CHANGELOG.md` for release notes.
---

View file

@ -6,7 +6,7 @@ import { defineConfig } from '@vscode/test-cli';
export default defineConfig({
label: 'integrationTest',
files: 'out-integration/test/**/*.test.js',
files: 'out/suite/**/*.test.js',
workspaceFolder: '.',
mocha: {
ui: 'tdd',

View file

@ -11,8 +11,8 @@ The integration tests use the `@vscode/test-electron` package to run tests in a
### Directory Structure
```
src/test/
├── runTest.ts # Main test runner
e2e/src/
├── runTest.ts # Main test runner
├── suite/
│ ├── index.ts # Test suite configuration
│ ├── modes.test.ts # Mode switching tests
@ -58,9 +58,9 @@ The following global objects are available in tests:
```typescript
declare global {
var api: ClineAPI
var api: RooCodeAPI
var provider: ClineProvider
var extension: vscode.Extension<ClineAPI>
var extension: vscode.Extension<RooCodeAPI>
var panel: vscode.WebviewPanel
}
```

2387
e2e/package-lock.json generated Normal file

File diff suppressed because it is too large Load diff

21
e2e/package.json Normal file
View file

@ -0,0 +1,21 @@
{
"name": "e2e",
"version": "0.1.0",
"private": true,
"scripts": {
"build": "cd .. && npm run build",
"compile": "tsc -p tsconfig.json",
"lint": "eslint src --ext ts",
"check-types": "tsc --noEmit",
"test": "npm run compile && npx dotenvx run -f .env.integration -- node ./out/runTest.js",
"ci": "npm run build && npm run test"
},
"dependencies": {},
"devDependencies": {
"@types/mocha": "^10.0.10",
"@vscode/test-cli": "^0.0.9",
"@vscode/test-electron": "^2.4.0",
"mocha": "^11.1.0",
"typescript": "^5.4.5"
}
}

View file

@ -1,14 +1,13 @@
import * as path from "path"
import Mocha from "mocha"
import { glob } from "glob"
import { ClineAPI } from "../../exports/cline"
import { ClineProvider } from "../../core/webview/ClineProvider"
import { RooCodeAPI, ClineProvider } from "../../../src/exports/roo-code"
import * as vscode from "vscode"
declare global {
var api: ClineAPI
var api: RooCodeAPI
var provider: ClineProvider
var extension: vscode.Extension<ClineAPI> | undefined
var extension: vscode.Extension<RooCodeAPI> | undefined
var panel: vscode.WebviewPanel | undefined
}

View file

@ -9,9 +9,8 @@
"strict": true,
"skipLibCheck": true,
"useUnknownInCatchVariables": false,
"rootDir": "src",
"outDir": "out-integration"
"outDir": "out"
},
"include": ["**/*.ts"],
"exclude": [".vscode-test", "benchmark", "dist", "**/node_modules/**", "out", "out-integration", "webview-ui"]
"include": ["src", "../src/exports/roo-code.d.ts"],
"exclude": [".vscode-test", "**/node_modules/**", "out"]
}

View file

@ -52,6 +52,7 @@ const copyWasmFiles = {
"java",
"php",
"swift",
"kotlin",
]
languages.forEach((lang) => {

View file

@ -30,9 +30,10 @@ module.exports = {
"^strip-ansi$": "<rootDir>/src/__mocks__/strip-ansi.js",
"^default-shell$": "<rootDir>/src/__mocks__/default-shell.js",
"^os-name$": "<rootDir>/src/__mocks__/os-name.js",
"^strip-bom$": "<rootDir>/src/__mocks__/strip-bom.js",
},
transformIgnorePatterns: [
"node_modules/(?!(@modelcontextprotocol|delay|p-wait-for|globby|serialize-error|strip-ansi|default-shell|os-name)/)",
"node_modules/(?!(@modelcontextprotocol|delay|p-wait-for|globby|serialize-error|strip-ansi|default-shell|os-name|strip-bom)/)",
],
roots: ["<rootDir>/src", "<rootDir>/webview-ui/src"],
modulePathIgnorePatterns: [".vscode-test"],

View file

@ -16,7 +16,9 @@
"src/activate/**",
"src/exports/**",
"src/extension.ts",
".vscode-test.mjs"
"e2e/.vscode-test.mjs",
"e2e/src/runTest.ts",
"e2e/src/suite/index.ts"
],
"workspaces": {
"webview-ui": {

992
package-lock.json generated

File diff suppressed because it is too large Load diff

View file

@ -1,9 +1,9 @@
{
"name": "roo-cline",
"displayName": "Roo Code (prev. Roo Cline)",
"description": "An AI-powered autonomous coding agent that lives in your editor.",
"description": "A whole dev team of AI agents in your editor.",
"publisher": "RooVeterinaryInc",
"version": "3.7.3",
"version": "3.8.4",
"icon": "assets/icons/rocket.png",
"galleryBanner": {
"color": "#617A91",
@ -128,31 +128,6 @@
"command": "roo-cline.addToContext",
"title": "Roo Code: Add To Context",
"category": "Roo Code"
},
{
"command": "roo-cline.terminalAddToContext",
"title": "Roo Code: Add Terminal Content to Context",
"category": "Terminal"
},
{
"command": "roo-cline.terminalFixCommand",
"title": "Roo Code: Fix This Command",
"category": "Terminal"
},
{
"command": "roo-cline.terminalExplainCommand",
"title": "Roo Code: Explain This Command",
"category": "Terminal"
},
{
"command": "roo-cline.terminalFixCommandInCurrentTask",
"title": "Roo Code: Fix This Command (Current Task)",
"category": "Terminal"
},
{
"command": "roo-cline.terminalExplainCommandInCurrentTask",
"title": "Roo Code: Explain This Command (Current Task)",
"category": "Terminal"
}
],
"menus": {
@ -178,28 +153,6 @@
"group": "Roo Code@4"
}
],
"terminal/context": [
{
"command": "roo-cline.terminalAddToContext",
"group": "Roo Code@1"
},
{
"command": "roo-cline.terminalFixCommand",
"group": "Roo Code@2"
},
{
"command": "roo-cline.terminalExplainCommand",
"group": "Roo Code@3"
},
{
"command": "roo-cline.terminalFixCommandInCurrentTask",
"group": "Roo Code@5"
},
{
"command": "roo-cline.terminalExplainCommandInCurrentTask",
"group": "Roo Code@6"
}
],
"view/title": [
{
"command": "roo-cline.plusButtonClicked",
@ -276,21 +229,26 @@
"scripts": {
"build": "npm run build:webview && npm run vsix",
"build:webview": "cd webview-ui && npm run build",
"changeset": "changeset",
"check-types": "tsc --noEmit && cd webview-ui && npm run check-types",
"compile": "tsc -p . --outDir out && node esbuild.js",
"compile:integration": "tsc -p tsconfig.integration.json",
"install:all": "npm install && cd webview-ui && npm install",
"knip": "knip --include files",
"lint": "eslint src --ext ts && npm run lint --prefix webview-ui",
"lint-local": "eslint -c .eslintrc.local.json src --ext ts && npm run lint --prefix webview-ui",
"lint-fix": "eslint src --ext ts --fix && npm run lint-fix --prefix webview-ui",
"lint-fix-local": "eslint -c .eslintrc.local.json src --ext ts --fix && npm run lint-fix --prefix webview-ui",
"install:all": "npm install npm-run-all && npm run install:_all",
"install:_all": "npm-run-all -p install-*",
"install-extension": "npm install",
"install-webview-ui": "cd webview-ui && npm install",
"install-e2e": "cd e2e && npm install",
"lint": "npm-run-all -p lint:*",
"lint:extension": "eslint src --ext ts",
"lint:webview-ui": "cd webview-ui && npm run lint",
"lint:e2e": "cd e2e && npm run lint",
"check-types": "npm-run-all -p check-types:*",
"check-types:extension": "tsc --noEmit",
"check-types:webview-ui": "cd webview-ui && npm run check-types",
"check-types:e2e": "cd e2e && npm run check-types",
"package": "npm run build:webview && npm run check-types && npm run lint && node esbuild.js --production",
"pretest": "npm run compile && npm run compile:integration",
"pretest": "npm run compile",
"dev": "cd webview-ui && npm run dev",
"test": "jest && cd webview-ui && npm run test",
"test:integration": "npm run build && npm run compile:integration && npx dotenvx run -f .env.integration -- node ./out-integration/test/runTest.js",
"test": "npm-run-all -p test:*",
"test:extension": "jest",
"test:webview": "cd webview-ui && npm run test",
"prepare": "husky",
"publish:marketplace": "vsce publish && ovsx publish",
"publish": "npm run build && changeset publish && npm install --package-lock-only",
@ -300,13 +258,16 @@
"watch": "npm-run-all -p watch:*",
"watch:esbuild": "node esbuild.js --watch",
"watch:tsc": "tsc --noEmit --watch --project tsconfig.json",
"watch-tests": "tsc -p . -w --outDir out"
"watch-tests": "tsc -p . -w --outDir out",
"changeset": "changeset",
"knip": "knip --include files"
},
"dependencies": {
"@anthropic-ai/bedrock-sdk": "^0.10.2",
"@anthropic-ai/sdk": "^0.37.0",
"@anthropic-ai/vertex-sdk": "^0.4.1",
"@anthropic-ai/vertex-sdk": "^0.7.0",
"@aws-sdk/client-bedrock-runtime": "^3.706.0",
"@google-cloud/vertexai": "^1.9.3",
"@google/generative-ai": "^0.18.0",
"@mistralai/mistralai": "^1.3.6",
"@modelcontextprotocol/sdk": "^1.0.1",
@ -329,12 +290,14 @@
"get-folder-size": "^5.0.0",
"globby": "^14.0.2",
"isbinaryfile": "^5.0.2",
"js-tiktoken": "^1.0.19",
"mammoth": "^1.8.0",
"monaco-vscode-textmate-theme-converter": "^0.1.7",
"openai": "^4.78.1",
"os-name": "^6.0.0",
"p-wait-for": "^5.0.2",
"pdf-parse": "^1.1.1",
"posthog-node": "^4.7.0",
"pretty-bytes": "^6.1.1",
"puppeteer-chromium-resolver": "^23.0.0",
"puppeteer-core": "^23.4.0",
@ -343,6 +306,7 @@
"sound-play": "^1.1.0",
"string-similarity": "^4.0.4",
"strip-ansi": "^7.1.0",
"strip-bom": "^5.0.0",
"tmp": "^0.2.3",
"tree-sitter-wasms": "^0.1.11",
"turndown": "^7.2.0",
@ -358,13 +322,10 @@
"@types/diff-match-patch": "^1.0.36",
"@types/glob": "^8.1.0",
"@types/jest": "^29.5.14",
"@types/mocha": "^10.0.10",
"@types/node": "20.x",
"@types/string-similarity": "^4.0.2",
"@typescript-eslint/eslint-plugin": "^7.14.1",
"@typescript-eslint/parser": "^7.11.0",
"@vscode/test-cli": "^0.0.9",
"@vscode/test-electron": "^2.4.0",
"esbuild": "^0.24.0",
"eslint": "^8.57.0",
"glob": "^11.0.1",
@ -374,7 +335,6 @@
"knip": "^5.44.4",
"lint-staged": "^15.2.11",
"mkdirp": "^3.0.1",
"mocha": "^11.1.0",
"npm-run-all": "^4.1.5",
"prettier": "^3.4.2",
"rimraf": "^6.0.1",

View file

@ -140,7 +140,6 @@ const mockFs = {
currentPath += "/" + parts[parts.length - 1]
mockDirectories.add(currentPath)
return Promise.resolve()
return Promise.resolve()
}),
access: jest.fn().mockImplementation(async (path: string) => {

View file

@ -15,3 +15,33 @@ jest.mock("../utils/logging", () => ({
}),
},
}))
// Add toPosix method to String prototype for all tests, mimicking src/utils/path.ts
// This is needed because the production code expects strings to have this method
// Note: In production, this is added via import in the entry point (extension.ts)
export {}
declare global {
interface String {
toPosix(): string
}
}
// Implementation that matches src/utils/path.ts
function toPosixPath(p: string) {
// Extended-Length Paths in Windows start with "\\?\" to allow longer paths
// and bypass usual parsing. If detected, we return the path unmodified.
const isExtendedLengthPath = p.startsWith("\\\\?\\")
if (isExtendedLengthPath) {
return p
}
return p.replace(/\\/g, "/")
}
if (!String.prototype.toPosix) {
String.prototype.toPosix = function (this: string): string {
return toPosixPath(this)
}
}

View file

@ -0,0 +1,13 @@
// Mock implementation of strip-bom
module.exports = function stripBom(string) {
if (typeof string !== "string") {
throw new TypeError("Expected a string")
}
// Removes UTF-8 BOM
if (string.charCodeAt(0) === 0xfeff) {
return string.slice(1)
}
return string
}

View file

@ -84,6 +84,12 @@ const vscode = {
this.uri = uri
}
},
RelativePattern: class {
constructor(base, pattern) {
this.base = base
this.pattern = pattern
}
},
}
module.exports = vscode

View file

@ -1,9 +1,11 @@
import * as vscode from "vscode"
import { ClineProvider } from "../core/webview/ClineProvider"
import { ClineAPI } from "./cline"
export function createClineAPI(outputChannel: vscode.OutputChannel, sidebarProvider: ClineProvider): ClineAPI {
const api: ClineAPI = {
import { ClineProvider } from "../core/webview/ClineProvider"
import { RooCodeAPI } from "../exports/roo-code"
export function createRooCodeAPI(outputChannel: vscode.OutputChannel, sidebarProvider: ClineProvider): RooCodeAPI {
return {
setCustomInstructions: async (value: string) => {
await sidebarProvider.updateCustomInstructions(value)
outputChannel.appendLine("Custom instructions set")
@ -15,7 +17,7 @@ export function createClineAPI(outputChannel: vscode.OutputChannel, sidebarProvi
startNewTask: async (task?: string, images?: string[]) => {
outputChannel.appendLine("Starting new task")
await sidebarProvider.clearTask()
await sidebarProvider.removeClineFromStack()
await sidebarProvider.postStateToWebview()
await sidebarProvider.postMessageToWebview({ type: "action", action: "chatButtonClicked" })
await sidebarProvider.postMessageToWebview({
@ -24,6 +26,7 @@ export function createClineAPI(outputChannel: vscode.OutputChannel, sidebarProvi
text: task,
images: images,
})
outputChannel.appendLine(
`Task started with message: ${task ? `"${task}"` : "undefined"} and ${images?.length || 0} image(s)`,
)
@ -33,6 +36,7 @@ export function createClineAPI(outputChannel: vscode.OutputChannel, sidebarProvi
outputChannel.appendLine(
`Sending message: ${message ? `"${message}"` : "undefined"} with ${images?.length || 0} image(s)`,
)
await sidebarProvider.postMessageToWebview({
type: "invoke",
invoke: "sendMessage",
@ -43,22 +47,14 @@ export function createClineAPI(outputChannel: vscode.OutputChannel, sidebarProvi
pressPrimaryButton: async () => {
outputChannel.appendLine("Pressing primary button")
await sidebarProvider.postMessageToWebview({
type: "invoke",
invoke: "primaryButtonClick",
})
await sidebarProvider.postMessageToWebview({ type: "invoke", invoke: "primaryButtonClick" })
},
pressSecondaryButton: async () => {
outputChannel.appendLine("Pressing secondary button")
await sidebarProvider.postMessageToWebview({
type: "invoke",
invoke: "secondaryButtonClick",
})
await sidebarProvider.postMessageToWebview({ type: "invoke", invoke: "secondaryButtonClick" })
},
sidebarProvider: sidebarProvider,
}
return api
}

View file

@ -0,0 +1,26 @@
// Callback mapping of human relay response.
const humanRelayCallbacks = new Map<string, (response: string | undefined) => void>()
/**
* Register a callback function for human relay response.
* @param requestId
* @param callback
*/
export const registerHumanRelayCallback = (requestId: string, callback: (response: string | undefined) => void) =>
humanRelayCallbacks.set(requestId, callback)
export const unregisterHumanRelayCallback = (requestId: string) => humanRelayCallbacks.delete(requestId)
export const handleHumanRelayResponse = (response: { requestId: string; text?: string; cancelled?: boolean }) => {
const callback = humanRelayCallbacks.get(response.requestId)
if (callback) {
if (response.cancelled) {
callback(undefined)
} else {
callback(response.text)
}
humanRelayCallbacks.delete(response.requestId)
}
}

View file

@ -1,4 +1,4 @@
export { handleUri } from "./handleUri"
export { registerCommands } from "./registerCommands"
export { registerCodeActions } from "./registerCodeActions"
export { registerTerminalActions } from "./registerTerminalActions"
export { createRooCodeAPI } from "./createRooCodeAPI"

View file

@ -3,6 +3,36 @@ import delay from "delay"
import { ClineProvider } from "../core/webview/ClineProvider"
import { registerHumanRelayCallback, unregisterHumanRelayCallback, handleHumanRelayResponse } from "./humanRelay"
// Store panel references in both modes
let sidebarPanel: vscode.WebviewView | undefined = undefined
let tabPanel: vscode.WebviewPanel | undefined = undefined
/**
* Get the currently active panel
* @returns WebviewPanel或WebviewView
*/
export function getPanel(): vscode.WebviewPanel | vscode.WebviewView | undefined {
return tabPanel || sidebarPanel
}
/**
* Set panel references
*/
export function setPanel(
newPanel: vscode.WebviewPanel | vscode.WebviewView | undefined,
type: "sidebar" | "tab",
): void {
if (type === "sidebar") {
sidebarPanel = newPanel as vscode.WebviewView
tabPanel = undefined
} else {
tabPanel = newPanel as vscode.WebviewPanel
sidebarPanel = undefined
}
}
export type RegisterCommandOptions = {
context: vscode.ExtensionContext
outputChannel: vscode.OutputChannel
@ -20,7 +50,7 @@ export const registerCommands = (options: RegisterCommandOptions) => {
const getCommandsMap = ({ context, outputChannel, provider }: RegisterCommandOptions) => {
return {
"roo-cline.plusButtonClicked": async () => {
await provider.clearTask()
await provider.removeClineFromStack()
await provider.postStateToWebview()
await provider.postMessageToWebview({ type: "action", action: "chatButtonClicked" })
},
@ -41,6 +71,20 @@ const getCommandsMap = ({ context, outputChannel, provider }: RegisterCommandOpt
"roo-cline.helpButtonClicked": () => {
vscode.env.openExternal(vscode.Uri.parse("https://docs.roocode.com"))
},
"roo-cline.showHumanRelayDialog": (params: { requestId: string; promptText: string }) => {
const panel = getPanel()
if (panel) {
panel?.webview.postMessage({
type: "showHumanRelayDialog",
requestId: params.requestId,
promptText: params.promptText,
})
}
},
"roo-cline.registerHumanRelayCallback": registerHumanRelayCallback,
"roo-cline.unregisterHumanRelayCallback": unregisterHumanRelayCallback,
"roo-cline.handleHumanRelayResponse": handleHumanRelayResponse,
}
}
@ -65,20 +109,28 @@ const openClineInNewTab = async ({ context, outputChannel }: Omit<RegisterComman
const targetCol = hasVisibleEditors ? Math.max(lastCol + 1, 1) : vscode.ViewColumn.Two
const panel = vscode.window.createWebviewPanel(ClineProvider.tabPanelId, "Roo Code", targetCol, {
const newPanel = vscode.window.createWebviewPanel(ClineProvider.tabPanelId, "Roo Code", targetCol, {
enableScripts: true,
retainContextWhenHidden: true,
localResourceRoots: [context.extensionUri],
})
// Save as tab type panel
setPanel(newPanel, "tab")
// TODO: use better svg icon with light and dark variants (see
// https://stackoverflow.com/questions/58365687/vscode-extension-iconpath).
panel.iconPath = {
newPanel.iconPath = {
light: vscode.Uri.joinPath(context.extensionUri, "assets", "icons", "rocket.png"),
dark: vscode.Uri.joinPath(context.extensionUri, "assets", "icons", "rocket.png"),
}
await tabProvider.resolveWebviewView(panel)
await tabProvider.resolveWebviewView(newPanel)
// Handle panel closing events
newPanel.onDidDispose(() => {
setPanel(undefined, "tab")
})
// Lock the editor group so clicking on files doesn't open them over the panel
await delay(100)

View file

@ -1,81 +0,0 @@
import * as vscode from "vscode"
import { ClineProvider } from "../core/webview/ClineProvider"
import { TerminalManager } from "../integrations/terminal/TerminalManager"
const TERMINAL_COMMAND_IDS = {
ADD_TO_CONTEXT: "roo-cline.terminalAddToContext",
FIX: "roo-cline.terminalFixCommand",
FIX_IN_CURRENT_TASK: "roo-cline.terminalFixCommandInCurrentTask",
EXPLAIN: "roo-cline.terminalExplainCommand",
EXPLAIN_IN_CURRENT_TASK: "roo-cline.terminalExplainCommandInCurrentTask",
} as const
export const registerTerminalActions = (context: vscode.ExtensionContext) => {
const terminalManager = new TerminalManager()
registerTerminalAction(context, terminalManager, TERMINAL_COMMAND_IDS.ADD_TO_CONTEXT, "TERMINAL_ADD_TO_CONTEXT")
registerTerminalActionPair(
context,
terminalManager,
TERMINAL_COMMAND_IDS.FIX,
"TERMINAL_FIX",
"What would you like Roo to fix?",
)
registerTerminalActionPair(
context,
terminalManager,
TERMINAL_COMMAND_IDS.EXPLAIN,
"TERMINAL_EXPLAIN",
"What would you like Roo to explain?",
)
}
const registerTerminalAction = (
context: vscode.ExtensionContext,
terminalManager: TerminalManager,
command: string,
promptType: "TERMINAL_ADD_TO_CONTEXT" | "TERMINAL_FIX" | "TERMINAL_EXPLAIN",
inputPrompt?: string,
) => {
context.subscriptions.push(
vscode.commands.registerCommand(command, async (args: any) => {
let content = args.selection
if (!content || content === "") {
content = await terminalManager.getTerminalContents(promptType === "TERMINAL_ADD_TO_CONTEXT" ? -1 : 1)
}
if (!content) {
vscode.window.showWarningMessage("No terminal content selected")
return
}
const params: Record<string, any> = {
terminalContent: content,
}
if (inputPrompt) {
params.userInput =
(await vscode.window.showInputBox({
prompt: inputPrompt,
})) ?? ""
}
await ClineProvider.handleTerminalAction(command, promptType, params)
}),
)
}
const registerTerminalActionPair = (
context: vscode.ExtensionContext,
terminalManager: TerminalManager,
baseCommand: string,
promptType: "TERMINAL_ADD_TO_CONTEXT" | "TERMINAL_FIX" | "TERMINAL_EXPLAIN",
inputPrompt?: string,
) => {
// Register new task version
registerTerminalAction(context, terminalManager, baseCommand, promptType, inputPrompt)
// Register current task version
registerTerminalAction(context, terminalManager, `${baseCommand}InCurrentTask`, promptType, inputPrompt)
}

View file

@ -0,0 +1,257 @@
// npx jest src/api/__tests__/index.test.ts
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta/messages/index.mjs"
import { getModelParams } from "../index"
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "../providers/constants"
describe("getModelParams", () => {
it("should return default values when no custom values are provided", () => {
const options = {}
const model = {
id: "test-model",
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
defaultMaxTokens: 1000,
defaultTemperature: 0.5,
})
expect(result).toEqual({
maxTokens: 1000,
thinking: undefined,
temperature: 0.5,
})
})
it("should use custom temperature from options when provided", () => {
const options = { modelTemperature: 0.7 }
const model = {
id: "test-model",
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
defaultMaxTokens: 1000,
defaultTemperature: 0.5,
})
expect(result).toEqual({
maxTokens: 1000,
thinking: undefined,
temperature: 0.7,
})
})
it("should use model maxTokens when available", () => {
const options = {}
const model = {
id: "test-model",
maxTokens: 2000,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
defaultMaxTokens: 1000,
})
expect(result).toEqual({
maxTokens: 2000,
thinking: undefined,
temperature: 0,
})
})
it("should handle thinking models correctly", () => {
const options = {}
const model = {
id: "test-model",
thinking: true,
maxTokens: 2000,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: 1600, // 80% of 2000
}
expect(result).toEqual({
maxTokens: 2000,
thinking: expectedThinking,
temperature: 1.0, // Thinking models require temperature 1.0.
})
})
it("should honor customMaxTokens for thinking models", () => {
const options = { modelMaxTokens: 3000 }
const model = {
id: "test-model",
thinking: true,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
defaultMaxTokens: 2000,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: 2400, // 80% of 3000
}
expect(result).toEqual({
maxTokens: 3000,
thinking: expectedThinking,
temperature: 1.0,
})
})
it("should honor customMaxThinkingTokens for thinking models", () => {
const options = { modelMaxThinkingTokens: 1500 }
const model = {
id: "test-model",
thinking: true,
maxTokens: 4000,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: 1500, // Using the custom value
}
expect(result).toEqual({
maxTokens: 4000,
thinking: expectedThinking,
temperature: 1.0,
})
})
it("should not honor customMaxThinkingTokens for non-thinking models", () => {
const options = { modelMaxThinkingTokens: 1500 }
const model = {
id: "test-model",
maxTokens: 4000,
contextWindow: 16000,
supportsPromptCache: true,
// Note: model.thinking is not set (so it's falsey).
}
const result = getModelParams({
options,
model,
})
expect(result).toEqual({
maxTokens: 4000,
thinking: undefined, // Should remain undefined despite customMaxThinkingTokens being set.
temperature: 0, // Using default temperature.
})
})
it("should clamp thinking budget to at least 1024 tokens", () => {
const options = { modelMaxThinkingTokens: 500 }
const model = {
id: "test-model",
thinking: true,
maxTokens: 2000,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: 1024, // Minimum is 1024
}
expect(result).toEqual({
maxTokens: 2000,
thinking: expectedThinking,
temperature: 1.0,
})
})
it("should clamp thinking budget to at most 80% of max tokens", () => {
const options = { modelMaxThinkingTokens: 5000 }
const model = {
id: "test-model",
thinking: true,
maxTokens: 4000,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: 3200, // 80% of 4000
}
expect(result).toEqual({
maxTokens: 4000,
thinking: expectedThinking,
temperature: 1.0,
})
})
it("should use ANTHROPIC_DEFAULT_MAX_TOKENS when no maxTokens is provided for thinking models", () => {
const options = {}
const model = {
id: "test-model",
thinking: true,
contextWindow: 16000,
supportsPromptCache: true,
}
const result = getModelParams({
options,
model,
})
const expectedThinking: BetaThinkingConfigParam = {
type: "enabled",
budget_tokens: Math.floor(ANTHROPIC_DEFAULT_MAX_TOKENS * 0.8),
}
expect(result).toEqual({
maxTokens: undefined,
thinking: expectedThinking,
temperature: 1.0,
})
})
})

View file

@ -1,6 +1,9 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta/messages/index.mjs"
import { ApiConfiguration, ModelInfo, ApiHandlerOptions } from "../shared/api"
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./providers/constants"
import { GlamaHandler } from "./providers/glama"
import { ApiConfiguration, ModelInfo } from "../shared/api"
import { AnthropicHandler } from "./providers/anthropic"
import { AwsBedrockHandler } from "./providers/bedrock"
import { OpenRouterHandler } from "./providers/openrouter"
@ -16,6 +19,7 @@ import { VsCodeLmHandler } from "./providers/vscode-lm"
import { ApiStream } from "./transform/stream"
import { UnboundHandler } from "./providers/unbound"
import { RequestyHandler } from "./providers/requesty"
import { HumanRelayHandler } from "./providers/human-relay"
export interface SingleCompletionHandler {
completePrompt(prompt: string): Promise<string>
@ -24,6 +28,16 @@ export interface SingleCompletionHandler {
export interface ApiHandler {
createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
getModel(): { id: string; info: ModelInfo }
/**
* Counts tokens for content blocks
* All providers extend BaseProvider which provides a default tiktoken implementation,
* but they can override this to use their native token counting endpoints
*
* @param content The content to count tokens for
* @returns A promise resolving to the token count
*/
countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number>
}
export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
@ -59,7 +73,47 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
return new UnboundHandler(options)
case "requesty":
return new RequestyHandler(options)
case "human-relay":
return new HumanRelayHandler(options)
default:
return new AnthropicHandler(options)
}
}
export function getModelParams({
options,
model,
defaultMaxTokens,
defaultTemperature = 0,
}: {
options: ApiHandlerOptions
model: ModelInfo
defaultMaxTokens?: number
defaultTemperature?: number
}) {
const {
modelMaxTokens: customMaxTokens,
modelMaxThinkingTokens: customMaxThinkingTokens,
modelTemperature: customTemperature,
} = options
let maxTokens = model.maxTokens ?? defaultMaxTokens
let thinking: BetaThinkingConfigParam | undefined = undefined
let temperature = customTemperature ?? defaultTemperature
if (model.thinking) {
// Only honor `customMaxTokens` for thinking models.
maxTokens = customMaxTokens ?? maxTokens
// Clamp the thinking budget to be at most 80% of max tokens and at
// least 1024 tokens.
const maxBudgetTokens = Math.floor((maxTokens || ANTHROPIC_DEFAULT_MAX_TOKENS) * 0.8)
const budgetTokens = Math.max(Math.min(customMaxThinkingTokens ?? maxBudgetTokens, maxBudgetTokens), 1024)
thinking = { type: "enabled", budget_tokens: budgetTokens }
// Anthropic "Thinking" models require a temperature of 1.0.
temperature = 1.0
}
return { maxTokens, thinking, temperature }
}

View file

@ -153,7 +153,7 @@ describe("AnthropicHandler", () => {
})
it("should handle API errors", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
mockCreate.mockRejectedValueOnce(new Error("Anthropic completion error: API Error"))
await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Anthropic completion error: API Error")
})
@ -194,5 +194,33 @@ describe("AnthropicHandler", () => {
expect(model.info.supportsImages).toBe(true)
expect(model.info.supportsPromptCache).toBe(true)
})
it("honors custom maxTokens for thinking models", () => {
const handler = new AnthropicHandler({
apiKey: "test-api-key",
apiModelId: "claude-3-7-sonnet-20250219:thinking",
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(32_768)
expect(result.thinking).toEqual({ type: "enabled", budget_tokens: 16_384 })
expect(result.temperature).toBe(1.0)
})
it("does not honor custom maxTokens for non-thinking models", () => {
const handler = new AnthropicHandler({
apiKey: "test-api-key",
apiModelId: "claude-3-7-sonnet-20250219",
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(16_384)
expect(result.thinking).toBeUndefined()
expect(result.temperature).toBe(0)
})
})
})

View file

@ -0,0 +1,75 @@
import { AwsBedrockHandler } from "../bedrock"
import { ApiHandlerOptions } from "../../../shared/api"
// Mock the AWS SDK
jest.mock("@aws-sdk/client-bedrock-runtime", () => {
const mockSend = jest.fn().mockImplementation(() => {
return Promise.resolve({
output: new TextEncoder().encode(JSON.stringify({ content: "Test response" })),
})
})
return {
BedrockRuntimeClient: jest.fn().mockImplementation(() => ({
send: mockSend,
config: {
region: "us-east-1",
},
})),
ConverseCommand: jest.fn(),
ConverseStreamCommand: jest.fn(),
}
})
describe("AwsBedrockHandler with custom ARN", () => {
const mockOptions: ApiHandlerOptions = {
apiModelId: "custom-arn",
awsCustomArn: "arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
awsRegion: "us-east-1",
}
it("should use the custom ARN as the model ID", async () => {
const handler = new AwsBedrockHandler(mockOptions)
const model = handler.getModel()
expect(model.id).toBe(mockOptions.awsCustomArn)
expect(model.info).toHaveProperty("maxTokens")
expect(model.info).toHaveProperty("contextWindow")
expect(model.info).toHaveProperty("supportsPromptCache")
})
it("should extract region from ARN and use it for client configuration", () => {
// Test with matching region
const handler1 = new AwsBedrockHandler(mockOptions)
expect((handler1 as any).client.config.region).toBe("us-east-1")
// Test with mismatched region
const mismatchOptions = {
...mockOptions,
awsRegion: "us-west-2",
}
const handler2 = new AwsBedrockHandler(mismatchOptions)
// Should use the ARN region, not the provided region
expect((handler2 as any).client.config.region).toBe("us-east-1")
})
it("should validate ARN format", async () => {
// Invalid ARN format
const invalidOptions = {
...mockOptions,
awsCustomArn: "invalid-arn-format",
}
const handler = new AwsBedrockHandler(invalidOptions)
// completePrompt should throw an error for invalid ARN
await expect(handler.completePrompt("test")).rejects.toThrow("Invalid ARN format")
})
it("should complete a prompt successfully with valid ARN", async () => {
const handler = new AwsBedrockHandler(mockOptions)
const response = await handler.completePrompt("test prompt")
expect(response).toBe("Test response")
})
})

View file

@ -315,5 +315,34 @@ describe("AwsBedrockHandler", () => {
expect(modelInfo.info.maxTokens).toBe(5000)
expect(modelInfo.info.contextWindow).toBe(128_000)
})
it("should use custom ARN when provided", () => {
const customArnHandler = new AwsBedrockHandler({
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
awsAccessKey: "test-access-key",
awsSecretKey: "test-secret-key",
awsRegion: "us-east-1",
awsCustomArn: "arn:aws:bedrock:us-east-1::foundation-model/custom-model",
})
const modelInfo = customArnHandler.getModel()
expect(modelInfo.id).toBe("arn:aws:bedrock:us-east-1::foundation-model/custom-model")
expect(modelInfo.info.maxTokens).toBe(4096)
expect(modelInfo.info.contextWindow).toBe(128_000)
expect(modelInfo.info.supportsPromptCache).toBe(false)
})
it("should use default model when custom-arn is selected but no ARN is provided", () => {
const customArnHandler = new AwsBedrockHandler({
apiModelId: "custom-arn",
awsAccessKey: "test-access-key",
awsSecretKey: "test-secret-key",
awsRegion: "us-east-1",
// No awsCustomArn provided
})
const modelInfo = customArnHandler.getModel()
// Should fall back to default model
expect(modelInfo.id).not.toBe("custom-arn")
expect(modelInfo.info).toBeDefined()
})
})
})

View file

@ -26,6 +26,10 @@ jest.mock("openai", () => {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
prompt_tokens_details: {
cache_miss_tokens: 8,
cached_tokens: 2,
},
},
}
}
@ -53,6 +57,10 @@ jest.mock("openai", () => {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
prompt_tokens_details: {
cache_miss_tokens: 8,
cached_tokens: 2,
},
},
}
},
@ -72,7 +80,7 @@ describe("DeepSeekHandler", () => {
mockOptions = {
deepSeekApiKey: "test-api-key",
apiModelId: "deepseek-chat",
deepSeekBaseUrl: "https://api.deepseek.com/v1",
deepSeekBaseUrl: "https://api.deepseek.com",
}
handler = new DeepSeekHandler(mockOptions)
mockCreate.mockClear()
@ -110,7 +118,7 @@ describe("DeepSeekHandler", () => {
// The base URL is passed to OpenAI client internally
expect(OpenAI).toHaveBeenCalledWith(
expect.objectContaining({
baseURL: "https://api.deepseek.com/v1",
baseURL: "https://api.deepseek.com",
}),
)
})
@ -149,7 +157,7 @@ describe("DeepSeekHandler", () => {
expect(model.info.maxTokens).toBe(8192)
expect(model.info.contextWindow).toBe(64_000)
expect(model.info.supportsImages).toBe(false)
expect(model.info.supportsPromptCache).toBe(false)
expect(model.info.supportsPromptCache).toBe(true) // Should be true now
})
it("should return provided model ID with default model info if model does not exist", () => {
@ -160,7 +168,12 @@ describe("DeepSeekHandler", () => {
const model = handlerWithInvalidModel.getModel()
expect(model.id).toBe("invalid-model") // Returns provided ID
expect(model.info).toBeDefined()
expect(model.info).toBe(handler.getModel().info) // But uses default model info
// With the current implementation, it's the same object reference when using default model info
expect(model.info).toBe(handler.getModel().info)
// Should have the same base properties
expect(model.info.contextWindow).toBe(handler.getModel().info.contextWindow)
// And should have supportsPromptCache set to true
expect(model.info.supportsPromptCache).toBe(true)
})
it("should return default model if no model ID is provided", () => {
@ -171,6 +184,13 @@ describe("DeepSeekHandler", () => {
const model = handlerWithoutModel.getModel()
expect(model.id).toBe(deepSeekDefaultModelId)
expect(model.info).toBeDefined()
expect(model.info.supportsPromptCache).toBe(true)
})
it("should include model parameters from getModelParams", () => {
const model = handler.getModel()
expect(model).toHaveProperty("temperature")
expect(model).toHaveProperty("maxTokens")
})
})
@ -213,5 +233,74 @@ describe("DeepSeekHandler", () => {
expect(usageChunks[0].inputTokens).toBe(10)
expect(usageChunks[0].outputTokens).toBe(5)
})
it("should include cache metrics in usage information", async () => {
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks.length).toBeGreaterThan(0)
expect(usageChunks[0].cacheWriteTokens).toBe(8)
expect(usageChunks[0].cacheReadTokens).toBe(2)
})
})
describe("processUsageMetrics", () => {
it("should correctly process usage metrics including cache information", () => {
// We need to access the protected method, so we'll create a test subclass
class TestDeepSeekHandler extends DeepSeekHandler {
public testProcessUsageMetrics(usage: any) {
return this.processUsageMetrics(usage)
}
}
const testHandler = new TestDeepSeekHandler(mockOptions)
const usage = {
prompt_tokens: 100,
completion_tokens: 50,
total_tokens: 150,
prompt_tokens_details: {
cache_miss_tokens: 80,
cached_tokens: 20,
},
}
const result = testHandler.testProcessUsageMetrics(usage)
expect(result.type).toBe("usage")
expect(result.inputTokens).toBe(100)
expect(result.outputTokens).toBe(50)
expect(result.cacheWriteTokens).toBe(80)
expect(result.cacheReadTokens).toBe(20)
})
it("should handle missing cache metrics gracefully", () => {
class TestDeepSeekHandler extends DeepSeekHandler {
public testProcessUsageMetrics(usage: any) {
return this.processUsageMetrics(usage)
}
}
const testHandler = new TestDeepSeekHandler(mockOptions)
const usage = {
prompt_tokens: 100,
completion_tokens: 50,
total_tokens: 150,
// No prompt_tokens_details
}
const result = testHandler.testProcessUsageMetrics(usage)
expect(result.type).toBe("usage")
expect(result.inputTokens).toBe(100)
expect(result.outputTokens).toBe(50)
expect(result.cacheWriteTokens).toBeUndefined()
expect(result.cacheReadTokens).toBeUndefined()
})
})
})

View file

@ -357,7 +357,7 @@ describe("OpenAiNativeHandler", () => {
const modelInfo = handler.getModel()
expect(modelInfo.id).toBe(mockOptions.apiModelId)
expect(modelInfo.info).toBeDefined()
expect(modelInfo.info.maxTokens).toBe(4096)
expect(modelInfo.info.maxTokens).toBe(16384)
expect(modelInfo.info.contextWindow).toBe(128_000)
})

View file

@ -0,0 +1,235 @@
import { OpenAiHandler } from "../openai"
import { ApiHandlerOptions } from "../../../shared/api"
import { Anthropic } from "@anthropic-ai/sdk"
// Mock OpenAI client with multiple chunks that contain usage data
const mockCreate = jest.fn()
jest.mock("openai", () => {
return {
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: mockCreate.mockImplementation(async (options) => {
if (!options.stream) {
return {
id: "test-completion",
choices: [
{
message: { role: "assistant", content: "Test response", refusal: null },
finish_reason: "stop",
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
},
}
}
// Return a stream with multiple chunks that include usage metrics
return {
[Symbol.asyncIterator]: async function* () {
// First chunk with partial usage
yield {
choices: [
{
delta: { content: "Test " },
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 2,
total_tokens: 12,
},
}
// Second chunk with updated usage
yield {
choices: [
{
delta: { content: "response" },
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 4,
total_tokens: 14,
},
}
// Final chunk with complete usage
yield {
choices: [
{
delta: {},
index: 0,
},
],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
},
}
},
}
}),
},
},
})),
}
})
describe("OpenAiHandler with usage tracking fix", () => {
let handler: OpenAiHandler
let mockOptions: ApiHandlerOptions
beforeEach(() => {
mockOptions = {
openAiApiKey: "test-api-key",
openAiModelId: "gpt-4",
openAiBaseUrl: "https://api.openai.com/v1",
}
handler = new OpenAiHandler(mockOptions)
mockCreate.mockClear()
})
describe("usage metrics with streaming", () => {
const systemPrompt = "You are a helpful assistant."
const messages: Anthropic.Messages.MessageParam[] = [
{
role: "user",
content: [
{
type: "text" as const,
text: "Hello!",
},
],
},
]
it("should only yield usage metrics once at the end of the stream", async () => {
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Check we have text chunks
const textChunks = chunks.filter((chunk) => chunk.type === "text")
expect(textChunks).toHaveLength(2)
expect(textChunks[0].text).toBe("Test ")
expect(textChunks[1].text).toBe("response")
// Check we only have one usage chunk and it's the last one
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks).toHaveLength(1)
expect(usageChunks[0]).toEqual({
type: "usage",
inputTokens: 10,
outputTokens: 5,
})
// Check the usage chunk is the last one reported from the API
const lastChunk = chunks[chunks.length - 1]
expect(lastChunk.type).toBe("usage")
expect(lastChunk.inputTokens).toBe(10)
expect(lastChunk.outputTokens).toBe(5)
})
it("should handle case where usage is only in the final chunk", async () => {
// Override the mock for this specific test
mockCreate.mockImplementationOnce(async (options) => {
if (!options.stream) {
return {
id: "test-completion",
choices: [{ message: { role: "assistant", content: "Test response" } }],
usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
}
}
return {
[Symbol.asyncIterator]: async function* () {
// First chunk with no usage
yield {
choices: [{ delta: { content: "Test " }, index: 0 }],
usage: null,
}
// Second chunk with no usage
yield {
choices: [{ delta: { content: "response" }, index: 0 }],
usage: null,
}
// Final chunk with usage data
yield {
choices: [{ delta: {}, index: 0 }],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
},
}
},
}
})
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Check usage metrics
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks).toHaveLength(1)
expect(usageChunks[0]).toEqual({
type: "usage",
inputTokens: 10,
outputTokens: 5,
})
})
it("should handle case where no usage is provided", async () => {
// Override the mock for this specific test
mockCreate.mockImplementationOnce(async (options) => {
if (!options.stream) {
return {
id: "test-completion",
choices: [{ message: { role: "assistant", content: "Test response" } }],
usage: null,
}
}
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [{ delta: { content: "Test response" }, index: 0 }],
usage: null,
}
yield {
choices: [{ delta: {}, index: 0 }],
usage: null,
}
},
}
})
const stream = handler.createMessage(systemPrompt, messages)
const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Check we don't have any usage chunks
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks).toHaveLength(0)
})
})
})

View file

@ -90,6 +90,20 @@ describe("OpenAiHandler", () => {
})
expect(handlerWithCustomUrl).toBeInstanceOf(OpenAiHandler)
})
it("should set default headers correctly", () => {
// Get the mock constructor from the jest mock system
const openAiMock = jest.requireMock("openai").default
expect(openAiMock).toHaveBeenCalledWith({
baseURL: expect.any(String),
apiKey: expect.any(String),
defaultHeaders: {
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
"X-Title": "Roo Code",
},
})
})
})
describe("createMessage", () => {

View file

@ -1,29 +1,30 @@
// npx jest src/api/providers/__tests__/openrouter.test.ts
import axios from "axios"
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { OpenRouterHandler } from "../openrouter"
import { ApiHandlerOptions, ModelInfo } from "../../../shared/api"
import OpenAI from "openai"
import axios from "axios"
import { Anthropic } from "@anthropic-ai/sdk"
// Mock dependencies
jest.mock("openai")
jest.mock("axios")
jest.mock("delay", () => jest.fn(() => Promise.resolve()))
const mockOpenRouterModelInfo: ModelInfo = {
maxTokens: 1000,
contextWindow: 2000,
supportsPromptCache: true,
inputPrice: 0.01,
outputPrice: 0.02,
}
describe("OpenRouterHandler", () => {
const mockOptions: ApiHandlerOptions = {
openRouterApiKey: "test-key",
openRouterModelId: "test-model",
openRouterModelInfo: {
name: "Test Model",
description: "Test Description",
maxTokens: 1000,
contextWindow: 2000,
supportsPromptCache: true,
inputPrice: 0.01,
outputPrice: 0.02,
} as ModelInfo,
openRouterModelInfo: mockOpenRouterModelInfo,
}
beforeEach(() => {
@ -50,6 +51,10 @@ describe("OpenRouterHandler", () => {
expect(result).toEqual({
id: mockOptions.openRouterModelId,
info: mockOptions.openRouterModelInfo,
maxTokens: 1000,
temperature: 0,
thinking: undefined,
topP: undefined,
})
})
@ -61,6 +66,38 @@ describe("OpenRouterHandler", () => {
expect(result.info.supportsPromptCache).toBe(true)
})
test("getModel honors custom maxTokens for thinking models", () => {
const handler = new OpenRouterHandler({
openRouterApiKey: "test-key",
openRouterModelId: "test-model",
openRouterModelInfo: {
...mockOpenRouterModelInfo,
maxTokens: 128_000,
thinking: true,
},
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(32_768)
expect(result.thinking).toEqual({ type: "enabled", budget_tokens: 16_384 })
expect(result.temperature).toBe(1.0)
})
test("getModel does not honor custom maxTokens for non-thinking models", () => {
const handler = new OpenRouterHandler({
...mockOptions,
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(1000)
expect(result.thinking).toBeUndefined()
expect(result.temperature).toBe(0)
})
test("createMessage generates correct stream chunks", async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockStream = {
@ -242,15 +279,7 @@ describe("OpenRouterHandler", () => {
test("completePrompt returns correct response", async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockResponse = {
choices: [
{
message: {
content: "test completion",
},
},
],
}
const mockResponse = { choices: [{ message: { content: "test completion" } }] }
const mockCreate = jest.fn().mockResolvedValue(mockResponse)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
@ -260,10 +289,13 @@ describe("OpenRouterHandler", () => {
const result = await handler.completePrompt("test prompt")
expect(result).toBe("test completion")
expect(mockCreate).toHaveBeenCalledWith({
model: mockOptions.openRouterModelId,
messages: [{ role: "user", content: "test prompt" }],
max_tokens: 1000,
thinking: undefined,
temperature: 0,
messages: [{ role: "user", content: "test prompt" }],
stream: false,
})
})
@ -292,8 +324,6 @@ describe("OpenRouterHandler", () => {
completions: { create: mockCreate },
} as any
await expect(handler.completePrompt("test prompt")).rejects.toThrow(
"OpenRouter completion error: Unexpected error",
)
await expect(handler.completePrompt("test prompt")).rejects.toThrow("Unexpected error")
})
})

View file

@ -22,8 +22,10 @@ describe("RequestyHandler", () => {
contextWindow: 4000,
supportsPromptCache: false,
supportsImages: true,
inputPrice: 0,
outputPrice: 0,
inputPrice: 1,
outputPrice: 10,
cacheReadsPrice: 0.1,
cacheWritesPrice: 1.5,
},
openAiStreamingEnabled: true,
includeMaxTokens: true, // Add this to match the implementation
@ -83,8 +85,12 @@ describe("RequestyHandler", () => {
yield {
choices: [{ delta: { content: " world" } }],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
prompt_tokens: 30,
completion_tokens: 10,
prompt_tokens_details: {
cached_tokens: 15,
caching_tokens: 5,
},
},
}
},
@ -105,10 +111,11 @@ describe("RequestyHandler", () => {
{ type: "text", text: " world" },
{
type: "usage",
inputTokens: 10,
outputTokens: 5,
cacheWriteTokens: undefined,
cacheReadTokens: undefined,
inputTokens: 30,
outputTokens: 10,
cacheWriteTokens: 5,
cacheReadTokens: 15,
totalCost: 0.000119, // (10 * 1 / 1,000,000) + (5 * 1.5 / 1,000,000) + (15 * 0.1 / 1,000,000) + (10 * 10 / 1,000,000)
},
])
@ -182,6 +189,9 @@ describe("RequestyHandler", () => {
type: "usage",
inputTokens: 10,
outputTokens: 5,
cacheWriteTokens: 0,
cacheReadTokens: 0,
totalCost: 0.00006, // (10 * 1 / 1,000,000) + (5 * 10 / 1,000,000)
},
])

View file

@ -192,6 +192,11 @@ describe("UnboundHandler", () => {
temperature: 0,
max_tokens: 8192,
}),
expect.objectContaining({
headers: expect.objectContaining({
"X-Unbound-Metadata": expect.stringContaining("roo-code"),
}),
}),
)
})
@ -233,6 +238,11 @@ describe("UnboundHandler", () => {
messages: [{ role: "user", content: "Test prompt" }],
temperature: 0,
}),
expect.objectContaining({
headers: expect.objectContaining({
"X-Unbound-Metadata": expect.stringContaining("roo-code"),
}),
}),
)
expect(mockCreate.mock.calls[0][0]).not.toHaveProperty("max_tokens")
})

View file

@ -2,8 +2,11 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta"
import { VertexHandler } from "../vertex"
import { ApiStreamChunk } from "../../transform/stream"
import { VertexAI } from "@google-cloud/vertexai"
// Mock Vertex SDK
jest.mock("@anthropic-ai/vertex-sdk", () => ({
@ -47,24 +50,100 @@ jest.mock("@anthropic-ai/vertex-sdk", () => ({
})),
}))
// Mock Vertex Gemini SDK
jest.mock("@google-cloud/vertexai", () => {
const mockGenerateContentStream = jest.fn().mockImplementation(() => {
return {
stream: {
async *[Symbol.asyncIterator]() {
yield {
candidates: [
{
content: {
parts: [{ text: "Test Gemini response" }],
},
},
],
}
},
},
response: {
usageMetadata: {
promptTokenCount: 5,
candidatesTokenCount: 10,
},
},
}
})
const mockGenerateContent = jest.fn().mockResolvedValue({
response: {
candidates: [
{
content: {
parts: [{ text: "Test Gemini response" }],
},
},
],
},
})
const mockGenerativeModel = jest.fn().mockImplementation(() => {
return {
generateContentStream: mockGenerateContentStream,
generateContent: mockGenerateContent,
}
})
return {
VertexAI: jest.fn().mockImplementation(() => {
return {
getGenerativeModel: mockGenerativeModel,
}
}),
GenerativeModel: mockGenerativeModel,
}
})
describe("VertexHandler", () => {
let handler: VertexHandler
beforeEach(() => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
})
describe("constructor", () => {
it("should initialize with provided config", () => {
it("should initialize with provided config for Claude", () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
expect(AnthropicVertex).toHaveBeenCalledWith({
projectId: "test-project",
region: "us-central1",
})
})
it("should initialize with provided config for Gemini", () => {
handler = new VertexHandler({
apiModelId: "gemini-1.5-pro-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
expect(VertexAI).toHaveBeenCalledWith({
project: "test-project",
location: "us-central1",
})
})
it("should throw error for invalid model", () => {
expect(() => {
new VertexHandler({
apiModelId: "invalid-model",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
}).toThrow("Unknown model ID: invalid-model")
})
})
describe("createMessage", () => {
@ -81,7 +160,13 @@ describe("VertexHandler", () => {
const systemPrompt = "You are a helpful assistant"
it("should handle streaming responses correctly", async () => {
it("should handle streaming responses correctly for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "message_start",
@ -125,10 +210,10 @@ describe("VertexHandler", () => {
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks = []
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
@ -158,13 +243,85 @@ describe("VertexHandler", () => {
model: "claude-3-5-sonnet-v2@20241022",
max_tokens: 8192,
temperature: 0,
system: systemPrompt,
messages: mockMessages,
system: [
{
type: "text",
text: "You are a helpful assistant",
cache_control: { type: "ephemeral" },
},
],
messages: [
{
role: "user",
content: [
{
type: "text",
text: "Hello",
cache_control: { type: "ephemeral" },
},
],
},
{
role: "assistant",
content: "Hi there!",
},
],
stream: true,
})
})
it("should handle multiple content blocks with line breaks", async () => {
it("should handle streaming responses correctly for Gemini", async () => {
const mockGemini = require("@google-cloud/vertexai")
const mockGenerateContentStream = mockGemini.VertexAI().getGenerativeModel().generateContentStream
handler = new VertexHandler({
apiModelId: "gemini-1.5-pro-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
expect(chunks.length).toBe(2)
expect(chunks[0]).toEqual({
type: "text",
text: "Test Gemini response",
})
expect(chunks[1]).toEqual({
type: "usage",
inputTokens: 5,
outputTokens: 10,
})
expect(mockGenerateContentStream).toHaveBeenCalledWith({
contents: [
{
role: "user",
parts: [{ text: "Hello" }],
},
{
role: "model",
parts: [{ text: "Hi there!" }],
},
],
generationConfig: {
maxOutputTokens: 16384,
temperature: 0,
},
})
})
it("should handle multiple content blocks with line breaks for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "content_block_start",
@ -193,10 +350,10 @@ describe("VertexHandler", () => {
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks = []
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
@ -217,10 +374,16 @@ describe("VertexHandler", () => {
})
})
it("should handle API errors", async () => {
it("should handle API errors for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockError = new Error("Vertex API error")
const mockCreate = jest.fn().mockRejectedValue(mockError)
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
@ -230,46 +393,469 @@ describe("VertexHandler", () => {
}
}).rejects.toThrow("Vertex API error")
})
it("should handle prompt caching for supported models for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "message_start",
message: {
usage: {
input_tokens: 10,
output_tokens: 0,
cache_creation_input_tokens: 3,
cache_read_input_tokens: 2,
},
},
},
{
type: "content_block_start",
index: 0,
content_block: {
type: "text",
text: "Hello",
},
},
{
type: "content_block_delta",
delta: {
type: "text_delta",
text: " world!",
},
},
{
type: "message_delta",
usage: {
output_tokens: 5,
},
},
]
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk
}
},
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, [
{
role: "user",
content: "First message",
},
{
role: "assistant",
content: "Response",
},
{
role: "user",
content: "Second message",
},
])
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Verify usage information
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks).toHaveLength(2)
expect(usageChunks[0]).toEqual({
type: "usage",
inputTokens: 10,
outputTokens: 0,
cacheWriteTokens: 3,
cacheReadTokens: 2,
})
expect(usageChunks[1]).toEqual({
type: "usage",
inputTokens: 0,
outputTokens: 5,
})
// Verify text content
const textChunks = chunks.filter((chunk) => chunk.type === "text")
expect(textChunks).toHaveLength(2)
expect(textChunks[0].text).toBe("Hello")
expect(textChunks[1].text).toBe(" world!")
// Verify cache control was added correctly
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
system: [
{
type: "text",
text: "You are a helpful assistant",
cache_control: { type: "ephemeral" },
},
],
messages: [
expect.objectContaining({
role: "user",
content: [
{
type: "text",
text: "First message",
cache_control: { type: "ephemeral" },
},
],
}),
expect.objectContaining({
role: "assistant",
content: "Response",
}),
expect.objectContaining({
role: "user",
content: [
{
type: "text",
text: "Second message",
cache_control: { type: "ephemeral" },
},
],
}),
],
}),
)
})
it("should handle cache-related usage metrics for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "message_start",
message: {
usage: {
input_tokens: 10,
output_tokens: 0,
cache_creation_input_tokens: 5,
cache_read_input_tokens: 3,
},
},
},
{
type: "content_block_start",
index: 0,
content_block: {
type: "text",
text: "Hello",
},
},
]
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk
}
},
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Check for cache-related metrics in usage chunk
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
expect(usageChunks.length).toBeGreaterThan(0)
expect(usageChunks[0]).toHaveProperty("cacheWriteTokens", 5)
expect(usageChunks[0]).toHaveProperty("cacheReadTokens", 3)
})
})
describe("thinking functionality", () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
role: "user",
content: "Hello",
},
]
const systemPrompt = "You are a helpful assistant"
it("should handle thinking content blocks and deltas for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "message_start",
message: {
usage: {
input_tokens: 10,
output_tokens: 0,
},
},
},
{
type: "content_block_start",
index: 0,
content_block: {
type: "thinking",
thinking: "Let me think about this...",
},
},
{
type: "content_block_delta",
delta: {
type: "thinking_delta",
thinking: " I need to consider all options.",
},
},
{
type: "content_block_start",
index: 1,
content_block: {
type: "text",
text: "Here's my answer:",
},
},
]
// Setup async iterator for mock stream
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk
}
},
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
// Verify thinking content is processed correctly
const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
expect(reasoningChunks).toHaveLength(2)
expect(reasoningChunks[0].text).toBe("Let me think about this...")
expect(reasoningChunks[1].text).toBe(" I need to consider all options.")
// Verify text content is processed correctly
const textChunks = chunks.filter((chunk) => chunk.type === "text")
expect(textChunks).toHaveLength(2) // One for the text block, one for the newline
expect(textChunks[0].text).toBe("\n")
expect(textChunks[1].text).toBe("Here's my answer:")
})
it("should handle multiple thinking blocks with line breaks for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockStream = [
{
type: "content_block_start",
index: 0,
content_block: {
type: "thinking",
thinking: "First thinking block",
},
},
{
type: "content_block_start",
index: 1,
content_block: {
type: "thinking",
thinking: "Second thinking block",
},
},
]
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk
}
},
}
const mockCreate = jest.fn().mockResolvedValue(asyncIterator)
;(handler["anthropicClient"].messages as any).create = mockCreate
const stream = handler.createMessage(systemPrompt, mockMessages)
const chunks: ApiStreamChunk[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}
expect(chunks.length).toBe(3)
expect(chunks[0]).toEqual({
type: "reasoning",
text: "First thinking block",
})
expect(chunks[1]).toEqual({
type: "reasoning",
text: "\n",
})
expect(chunks[2]).toEqual({
type: "reasoning",
text: "Second thinking block",
})
})
})
describe("completePrompt", () => {
it("should complete prompt successfully", async () => {
it("should complete prompt successfully for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("Test response")
expect(handler["client"].messages.create).toHaveBeenCalledWith({
expect(handler["anthropicClient"].messages.create).toHaveBeenCalledWith({
model: "claude-3-5-sonnet-v2@20241022",
max_tokens: 8192,
temperature: 0,
messages: [{ role: "user", content: "Test prompt" }],
system: "",
messages: [
{
role: "user",
content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }],
},
],
stream: false,
})
})
it("should handle API errors", async () => {
it("should complete prompt successfully for Gemini", async () => {
const mockGemini = require("@google-cloud/vertexai")
const mockGenerateContent = mockGemini.VertexAI().getGenerativeModel().generateContent
handler = new VertexHandler({
apiModelId: "gemini-1.5-pro-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("Test Gemini response")
expect(mockGenerateContent).toHaveBeenCalled()
expect(mockGenerateContent).toHaveBeenCalledWith({
contents: [{ role: "user", parts: [{ text: "Test prompt" }] }],
generationConfig: {
temperature: 0,
},
})
})
it("should handle API errors for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockError = new Error("Vertex API error")
const mockCreate = jest.fn().mockRejectedValue(mockError)
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
"Vertex completion error: Vertex API error",
)
})
it("should handle non-text content", async () => {
it("should handle API errors for Gemini", async () => {
const mockGemini = require("@google-cloud/vertexai")
const mockGenerateContent = mockGemini.VertexAI().getGenerativeModel().generateContent
mockGenerateContent.mockRejectedValue(new Error("Vertex API error"))
handler = new VertexHandler({
apiModelId: "gemini-1.5-pro-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
await expect(handler.completePrompt("Test prompt")).rejects.toThrow(
"Vertex completion error: Vertex API error",
)
})
it("should handle non-text content for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockCreate = jest.fn().mockResolvedValue({
content: [{ type: "image" }],
})
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
})
it("should handle empty response", async () => {
it("should handle empty response for Claude", async () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const mockCreate = jest.fn().mockResolvedValue({
content: [{ type: "text", text: "" }],
})
;(handler["client"].messages as any).create = mockCreate
;(handler["anthropicClient"].messages as any).create = mockCreate
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
})
it("should handle empty response for Gemini", async () => {
const mockGemini = require("@google-cloud/vertexai")
const mockGenerateContent = mockGemini.VertexAI().getGenerativeModel().generateContent
mockGenerateContent.mockResolvedValue({
response: {
candidates: [
{
content: {
parts: [{ text: "" }],
},
},
],
},
})
handler = new VertexHandler({
apiModelId: "gemini-1.5-pro-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
@ -277,7 +863,13 @@ describe("VertexHandler", () => {
})
describe("getModel", () => {
it("should return correct model info", () => {
it("should return correct model info for Claude", () => {
handler = new VertexHandler({
apiModelId: "claude-3-5-sonnet-v2@20241022",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const modelInfo = handler.getModel()
expect(modelInfo.id).toBe("claude-3-5-sonnet-v2@20241022")
expect(modelInfo.info).toBeDefined()
@ -285,14 +877,151 @@ describe("VertexHandler", () => {
expect(modelInfo.info.contextWindow).toBe(200_000)
})
it("should return default model if invalid model specified", () => {
const invalidHandler = new VertexHandler({
apiModelId: "invalid-model",
it("should return correct model info for Gemini", () => {
handler = new VertexHandler({
apiModelId: "gemini-2.0-flash-001",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})
const modelInfo = invalidHandler.getModel()
expect(modelInfo.id).toBe("claude-3-7-sonnet@20250219") // Default model
const modelInfo = handler.getModel()
expect(modelInfo.id).toBe("gemini-2.0-flash-001")
expect(modelInfo.info).toBeDefined()
expect(modelInfo.info.maxTokens).toBe(8192)
expect(modelInfo.info.contextWindow).toBe(1048576)
})
it("honors custom maxTokens for thinking models", () => {
const handler = new VertexHandler({
apiKey: "test-api-key",
apiModelId: "claude-3-7-sonnet@20250219:thinking",
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(32_768)
expect(result.thinking).toEqual({ type: "enabled", budget_tokens: 16_384 })
expect(result.temperature).toBe(1.0)
})
it("does not honor custom maxTokens for non-thinking models", () => {
const handler = new VertexHandler({
apiKey: "test-api-key",
apiModelId: "claude-3-7-sonnet@20250219",
modelMaxTokens: 32_768,
modelMaxThinkingTokens: 16_384,
})
const result = handler.getModel()
expect(result.maxTokens).toBe(16_384)
expect(result.thinking).toBeUndefined()
expect(result.temperature).toBe(0)
})
})
describe("thinking model configuration", () => {
it("should configure thinking for models with :thinking suffix", () => {
const thinkingHandler = new VertexHandler({
apiModelId: "claude-3-7-sonnet@20250219:thinking",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
modelMaxTokens: 16384,
modelMaxThinkingTokens: 4096,
})
const modelInfo = thinkingHandler.getModel()
// Verify thinking configuration
expect(modelInfo.id).toBe("claude-3-7-sonnet@20250219")
expect(modelInfo.thinking).toBeDefined()
const thinkingConfig = modelInfo.thinking as { type: "enabled"; budget_tokens: number }
expect(thinkingConfig.type).toBe("enabled")
expect(thinkingConfig.budget_tokens).toBe(4096)
expect(modelInfo.temperature).toBe(1.0) // Thinking requires temperature 1.0
})
it("should calculate thinking budget correctly", () => {
// Test with explicit thinking budget
const handlerWithBudget = new VertexHandler({
apiModelId: "claude-3-7-sonnet@20250219:thinking",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
modelMaxTokens: 16384,
modelMaxThinkingTokens: 5000,
})
expect((handlerWithBudget.getModel().thinking as any).budget_tokens).toBe(5000)
// Test with default thinking budget (80% of max tokens)
const handlerWithDefaultBudget = new VertexHandler({
apiModelId: "claude-3-7-sonnet@20250219:thinking",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
modelMaxTokens: 10000,
})
expect((handlerWithDefaultBudget.getModel().thinking as any).budget_tokens).toBe(8000) // 80% of 10000
// Test with minimum thinking budget (should be at least 1024)
const handlerWithSmallMaxTokens = new VertexHandler({
apiModelId: "claude-3-7-sonnet@20250219:thinking",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
modelMaxTokens: 1000, // This would result in 800 tokens for thinking, but minimum is 1024
})
expect((handlerWithSmallMaxTokens.getModel().thinking as any).budget_tokens).toBe(1024)
})
it("should pass thinking configuration to API", async () => {
const thinkingHandler = new VertexHandler({
apiModelId: "claude-3-7-sonnet@20250219:thinking",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
modelMaxTokens: 16384,
modelMaxThinkingTokens: 4096,
})
const mockCreate = jest.fn().mockImplementation(async (options) => {
if (!options.stream) {
return {
id: "test-completion",
content: [{ type: "text", text: "Test response" }],
role: "assistant",
model: options.model,
usage: {
input_tokens: 10,
output_tokens: 5,
},
}
}
return {
async *[Symbol.asyncIterator]() {
yield {
type: "message_start",
message: {
usage: {
input_tokens: 10,
output_tokens: 5,
},
},
}
},
}
})
;(thinkingHandler["anthropicClient"].messages as any).create = mockCreate
await thinkingHandler
.createMessage("You are a helpful assistant", [{ role: "user", content: "Hello" }])
.next()
expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
thinking: { type: "enabled", budget_tokens: 4096 },
temperature: 1.0, // Thinking requires temperature 1.0
}),
)
})
})
})

View file

@ -1,7 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming"
import { CacheControlEphemeral } from "@anthropic-ai/sdk/resources"
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta"
import {
anthropicDefaultModelId,
AnthropicModelId,
@ -9,20 +8,18 @@ import {
ApiHandlerOptions,
ModelInfo,
} from "../../shared/api"
import { ApiHandler, SingleCompletionHandler } from "../index"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./constants"
import { SingleCompletionHandler, getModelParams } from "../index"
const ANTHROPIC_DEFAULT_TEMPERATURE = 0
const THINKING_MODELS = ["claude-3-7-sonnet-20250219"]
export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
export class AnthropicHandler extends BaseProvider implements SingleCompletionHandler {
private options: ApiHandlerOptions
private client: Anthropic
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = new Anthropic({
apiKey: this.options.apiKey,
baseURL: this.options.anthropicBaseUrl || undefined,
@ -32,18 +29,7 @@ export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
let stream: AnthropicStream<Anthropic.Messages.RawMessageStreamEvent>
const cacheControl: CacheControlEphemeral = { type: "ephemeral" }
const modelId = this.getModel().id
const maxTokens = this.getModel().info.maxTokens || 8192
let temperature = this.options.modelTemperature ?? ANTHROPIC_DEFAULT_TEMPERATURE
let thinking: BetaThinkingConfigParam | undefined = undefined
if (THINKING_MODELS.includes(modelId)) {
thinking = this.options.anthropicThinking
? { type: "enabled", budget_tokens: this.options.anthropicThinking }
: { type: "disabled" }
temperature = 1.0
}
let { id: modelId, maxTokens, thinking, temperature, virtualId } = this.getModel()
switch (modelId) {
case "claude-3-7-sonnet-20250219":
@ -66,7 +52,7 @@ export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
stream = await this.client.messages.create(
{
model: modelId,
max_tokens: maxTokens,
max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
temperature,
thinking,
// Setting cache breakpoint for system prompt so new tasks can reuse it.
@ -96,13 +82,24 @@ export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
// prompt caching: https://x.com/alexalbert__/status/1823751995901272068
// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
// https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393
const betas = []
// Check for the thinking-128k variant first
if (virtualId === "claude-3-7-sonnet-20250219:thinking") {
betas.push("output-128k-2025-02-19")
}
// Then check for models that support prompt caching
switch (modelId) {
case "claude-3-7-sonnet-20250219":
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307":
betas.push("prompt-caching-2024-07-31")
return {
headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
headers: { "anthropic-beta": betas.join(",") },
}
default:
return undefined
@ -114,8 +111,8 @@ export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
default: {
stream = (await this.client.messages.create({
model: modelId,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: this.options.modelTemperature ?? ANTHROPIC_DEFAULT_TEMPERATURE,
max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
temperature,
system: [{ text: systemPrompt, type: "text" }],
messages,
// tools,
@ -193,40 +190,73 @@ export class AnthropicHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: AnthropicModelId; info: ModelInfo } {
getModel() {
const modelId = this.options.apiModelId
let id = modelId && modelId in anthropicModels ? (modelId as AnthropicModelId) : anthropicDefaultModelId
const info: ModelInfo = anthropicModels[id]
if (modelId && modelId in anthropicModels) {
const id = modelId as AnthropicModelId
return { id, info: anthropicModels[id] }
// Track the original model ID for special variant handling
const virtualId = id
// The `:thinking` variant is a virtual identifier for the
// `claude-3-7-sonnet-20250219` model with a thinking budget.
// We can handle this more elegantly in the future.
if (id === "claude-3-7-sonnet-20250219:thinking") {
id = "claude-3-7-sonnet-20250219"
}
return { id: anthropicDefaultModelId, info: anthropicModels[anthropicDefaultModelId] }
return {
id,
info,
virtualId, // Include the original ID to use for header selection
...getModelParams({ options: this.options, model: info, defaultMaxTokens: ANTHROPIC_DEFAULT_MAX_TOKENS }),
}
}
async completePrompt(prompt: string): Promise<string> {
async completePrompt(prompt: string) {
let { id: modelId, temperature } = this.getModel()
const message = await this.client.messages.create({
model: modelId,
max_tokens: ANTHROPIC_DEFAULT_MAX_TOKENS,
thinking: undefined,
temperature,
messages: [{ role: "user", content: prompt }],
stream: false,
})
const content = message.content.find(({ type }) => type === "text")
return content?.type === "text" ? content.text : ""
}
/**
* Counts tokens for the given content using Anthropic's API
*
* @param content The content blocks to count tokens for
* @returns A promise resolving to the token count
*/
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
try {
const response = await this.client.messages.create({
model: this.getModel().id,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: this.options.modelTemperature ?? ANTHROPIC_DEFAULT_TEMPERATURE,
messages: [{ role: "user", content: prompt }],
stream: false,
// Use the current model
const actualModelId = this.getModel().id
const response = await this.client.messages.countTokens({
model: actualModelId,
messages: [
{
role: "user",
content: content,
},
],
})
const content = response.content[0]
if (content.type === "text") {
return content.text
}
return ""
return response.input_tokens
} catch (error) {
if (error instanceof Error) {
throw new Error(`Anthropic completion error: ${error.message}`)
}
// Log error but fallback to tiktoken estimation
console.warn("Anthropic token counting failed, using fallback", error)
throw error
// Use the base provider's implementation as fallback
return super.countTokens(content)
}
}
}

View file

@ -0,0 +1,64 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler } from ".."
import { ModelInfo } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { Tiktoken } from "js-tiktoken/lite"
import o200kBase from "js-tiktoken/ranks/o200k_base"
// Reuse the fudge factor used in the original code
const TOKEN_FUDGE_FACTOR = 1.5
/**
* Base class for API providers that implements common functionality
*/
export abstract class BaseProvider implements ApiHandler {
// Cache the Tiktoken encoder instance since it's stateless
private encoder: Tiktoken | null = null
abstract createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
abstract getModel(): { id: string; info: ModelInfo }
/**
* Default token counting implementation using tiktoken
* Providers can override this to use their native token counting endpoints
*
* Uses a cached Tiktoken encoder instance for performance since it's stateless.
* The encoder is created lazily on first use and reused for subsequent calls.
*
* @param content The content to count tokens for
* @returns A promise resolving to the token count
*/
async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
if (!content || content.length === 0) return 0
let totalTokens = 0
// Lazily create and cache the encoder if it doesn't exist
if (!this.encoder) {
this.encoder = new Tiktoken(o200kBase)
}
// Process each content block using the cached encoder
for (const block of content) {
if (block.type === "text") {
// Use tiktoken for text token counting
const text = block.text || ""
if (text.length > 0) {
const tokens = this.encoder.encode(text)
totalTokens += tokens.length
}
} else if (block.type === "image") {
// For images, calculate based on data size
const imageSource = block.source
if (imageSource && typeof imageSource === "object" && "data" in imageSource) {
const base64Data = imageSource.data as string
totalTokens += Math.ceil(Math.sqrt(base64Data.length))
} else {
totalTokens += 300 // Conservative estimate for unknown images
}
}
}
// Add a fudge factor to account for the fact that tiktoken is not always accurate
return Math.ceil(totalTokens * TOKEN_FUDGE_FACTOR)
}
}

View file

@ -6,10 +6,52 @@ import {
} from "@aws-sdk/client-bedrock-runtime"
import { fromIni } from "@aws-sdk/credential-providers"
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler, SingleCompletionHandler } from "../"
import { SingleCompletionHandler } from "../"
import { ApiHandlerOptions, BedrockModelId, ModelInfo, bedrockDefaultModelId, bedrockModels } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { convertToBedrockConverseMessages } from "../transform/bedrock-converse-format"
import { BaseProvider } from "./base-provider"
import { logger } from "../../utils/logging"
/**
* Validates an AWS Bedrock ARN format and optionally checks if the region in the ARN matches the provided region
* @param arn The ARN string to validate
* @param region Optional region to check against the ARN's region
* @returns An object with validation results: { isValid, arnRegion, errorMessage }
*/
function validateBedrockArn(arn: string, region?: string) {
// Validate ARN format
const arnRegex = /^arn:aws:bedrock:([^:]+):(\d+):(foundation-model|provisioned-model|default-prompt-router)\/(.+)$/
const match = arn.match(arnRegex)
if (!match) {
return {
isValid: false,
arnRegion: undefined,
errorMessage:
"Invalid ARN format. ARN should follow the pattern: arn:aws:bedrock:region:account-id:resource-type/resource-name",
}
}
// Extract region from ARN
const arnRegion = match[1]
// Check if region in ARN matches provided region (if specified)
if (region && arnRegion !== region) {
return {
isValid: true,
arnRegion,
errorMessage: `Warning: The region in your ARN (${arnRegion}) does not match your selected region (${region}). This may cause access issues. The provider will use the region from the ARN.`,
}
}
// ARN is valid and region matches (or no region was provided to check against)
return {
isValid: true,
arnRegion,
errorMessage: undefined,
}
}
const BEDROCK_DEFAULT_TEMPERATURE = 0.3
@ -46,15 +88,39 @@ export interface StreamEvent {
}
}
export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class AwsBedrockHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: BedrockRuntimeClient
constructor(options: ApiHandlerOptions) {
super()
this.options = options
// Extract region from custom ARN if provided
let region = this.options.awsRegion || "us-east-1"
// If using custom ARN, extract region from the ARN
if (this.options.awsCustomArn) {
const validation = validateBedrockArn(this.options.awsCustomArn, region)
if (validation.isValid && validation.arnRegion) {
// If there's a region mismatch warning, log it and use the ARN region
if (validation.errorMessage) {
logger.info(
`Region mismatch: Selected region is ${region}, but ARN region is ${validation.arnRegion}. Using ARN region.`,
{
ctx: "bedrock",
selectedRegion: region,
arnRegion: validation.arnRegion,
},
)
region = validation.arnRegion
}
}
}
const clientConfig: BedrockRuntimeClientConfig = {
region: this.options.awsRegion || "us-east-1",
region: region,
}
if (this.options.awsUseProfile && this.options.awsProfile) {
@ -74,12 +140,46 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
this.client = new BedrockRuntimeClient(clientConfig)
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const modelConfig = this.getModel()
// Handle cross-region inference
let modelId: string
if (this.options.awsUseCrossRegionInference) {
// For custom ARNs, use the ARN directly without modification
if (this.options.awsCustomArn) {
modelId = modelConfig.id
// Validate ARN format and check region match
const clientRegion = this.client.config.region as string
const validation = validateBedrockArn(modelId, clientRegion)
if (!validation.isValid) {
logger.error("Invalid ARN format", {
ctx: "bedrock",
modelId,
errorMessage: validation.errorMessage,
})
yield {
type: "text",
text: `Error: ${validation.errorMessage}`,
}
yield { type: "usage", inputTokens: 0, outputTokens: 0 }
throw new Error("Invalid ARN format")
}
// Extract region from ARN
const arnRegion = validation.arnRegion!
// Log warning if there's a region mismatch
if (validation.errorMessage) {
logger.warn(validation.errorMessage, {
ctx: "bedrock",
arnRegion,
clientRegion,
})
}
} else if (this.options.awsUseCrossRegionInference) {
let regionPrefix = (this.options.awsRegion || "").slice(0, 3)
switch (regionPrefix) {
case "us-":
@ -105,7 +205,7 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
messages: formattedMessages,
system: [{ text: systemPrompt }],
inferenceConfig: {
maxTokens: modelConfig.info.maxTokens || 5000,
maxTokens: modelConfig.info.maxTokens || 4096,
temperature: this.options.modelTemperature ?? BEDROCK_DEFAULT_TEMPERATURE,
topP: 0.1,
...(this.options.awsUsePromptCache
@ -119,6 +219,16 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
}
try {
// Log the payload for debugging custom ARN issues
if (this.options.awsCustomArn) {
logger.debug("Using custom ARN for Bedrock request", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
clientRegion: this.client.config.region,
payload: JSON.stringify(payload, null, 2),
})
}
const command = new ConverseStreamCommand(payload)
const response = await this.client.send(command)
@ -132,7 +242,11 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
try {
streamEvent = typeof chunk === "string" ? JSON.parse(chunk) : (chunk as unknown as StreamEvent)
} catch (e) {
console.error("Failed to parse stream event:", e)
logger.error("Failed to parse stream event", {
ctx: "bedrock",
error: e instanceof Error ? e : String(e),
chunk: typeof chunk === "string" ? chunk : "binary data",
})
continue
}
@ -175,39 +289,257 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
}
}
} catch (error: unknown) {
console.error("Bedrock Runtime API Error:", error)
// Only access stack if error is an Error object
logger.error("Bedrock Runtime API Error", {
ctx: "bedrock",
error: error instanceof Error ? error : String(error),
})
// Enhanced error handling for custom ARN issues
if (this.options.awsCustomArn) {
logger.error("Error occurred with custom ARN", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
})
// Check for common ARN-related errors
if (error instanceof Error) {
const errorMessage = error.message.toLowerCase()
// Access denied errors
if (
errorMessage.includes("access") &&
(errorMessage.includes("model") || errorMessage.includes("denied"))
) {
logger.error("Permissions issue with custom ARN", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
errorType: "access_denied",
clientRegion: this.client.config.region,
})
yield {
type: "text",
text: `Error: You don't have access to the model with the specified ARN. Please verify:
1. The ARN is correct and points to a valid model
2. Your AWS credentials have permission to access this model (check IAM policies)
3. The region in the ARN (${this.client.config.region}) matches the region where the model is deployed
4. If using a provisioned model, ensure it's active and not in a failed state
5. If using a custom model, ensure your account has been granted access to it`,
}
}
// Model not found errors
else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) {
logger.error("Invalid ARN or non-existent model", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
errorType: "not_found",
})
yield {
type: "text",
text: `Error: The specified ARN does not exist or is invalid. Please check:
1. The ARN format is correct (arn:aws:bedrock:region:account-id:resource-type/resource-name)
2. The model exists in the specified region
3. The account ID in the ARN is correct
4. The resource type is one of: foundation-model, provisioned-model, or default-prompt-router`,
}
}
// Throttling errors
else if (
errorMessage.includes("throttl") ||
errorMessage.includes("rate") ||
errorMessage.includes("limit")
) {
logger.error("Throttling or rate limit issue with Bedrock", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
errorType: "throttling",
})
yield {
type: "text",
text: `Error: Request was throttled or rate limited. Please try:
1. Reducing the frequency of requests
2. If using a provisioned model, check its throughput settings
3. Contact AWS support to request a quota increase if needed`,
}
}
// Other errors
else {
logger.error("Unspecified error with custom ARN", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
errorStack: error.stack,
errorMessage: error.message,
})
yield {
type: "text",
text: `Error with custom ARN: ${error.message}
Please check:
1. Your AWS credentials are valid and have the necessary permissions
2. The ARN format is correct
3. The region in the ARN matches the region where you're making the request`,
}
}
} else {
yield {
type: "text",
text: `Unknown error occurred with custom ARN. Please check your AWS credentials and ARN format.`,
}
}
} else {
// Standard error handling for non-ARN cases
if (error instanceof Error) {
logger.error("Standard Bedrock error", {
ctx: "bedrock",
errorStack: error.stack,
errorMessage: error.message,
})
yield {
type: "text",
text: `Error: ${error.message}`,
}
} else {
logger.error("Unknown Bedrock error", {
ctx: "bedrock",
error: String(error),
})
yield {
type: "text",
text: "An unknown error occurred",
}
}
}
// Always yield usage info
yield {
type: "usage",
inputTokens: 0,
outputTokens: 0,
}
// Re-throw the error
if (error instanceof Error) {
console.error("Error stack:", error.stack)
yield {
type: "text",
text: `Error: ${error.message}`,
}
yield {
type: "usage",
inputTokens: 0,
outputTokens: 0,
}
throw error
} else {
const unknownError = new Error("An unknown error occurred")
yield {
type: "text",
text: unknownError.message,
}
yield {
type: "usage",
inputTokens: 0,
outputTokens: 0,
}
throw unknownError
throw new Error("An unknown error occurred")
}
}
}
getModel(): { id: BedrockModelId | string; info: ModelInfo } {
override getModel(): { id: BedrockModelId | string; info: ModelInfo } {
// If custom ARN is provided, use it
if (this.options.awsCustomArn) {
// Custom ARNs should not be modified with region prefixes
// as they already contain the full resource path
// Check if the ARN contains information about the model type
// This helps set appropriate token limits for models behind prompt routers
const arnLower = this.options.awsCustomArn.toLowerCase()
// Determine model info based on ARN content
let modelInfo: ModelInfo
if (arnLower.includes("claude-3-7-sonnet") || arnLower.includes("claude-3.7-sonnet")) {
// Claude 3.7 Sonnet has 8192 tokens in Bedrock
modelInfo = {
maxTokens: 8192,
contextWindow: 200_000,
supportsPromptCache: false,
supportsImages: true,
supportsComputerUse: true,
}
} else if (arnLower.includes("claude-3-5-sonnet") || arnLower.includes("claude-3.5-sonnet")) {
// Claude 3.5 Sonnet has 8192 tokens in Bedrock
modelInfo = {
maxTokens: 8192,
contextWindow: 200_000,
supportsPromptCache: false,
supportsImages: true,
supportsComputerUse: true,
}
} else if (arnLower.includes("claude-3-opus") || arnLower.includes("claude-3.0-opus")) {
// Claude 3 Opus has 4096 tokens in Bedrock
modelInfo = {
maxTokens: 4096,
contextWindow: 200_000,
supportsPromptCache: false,
supportsImages: true,
}
} else if (arnLower.includes("claude-3-haiku") || arnLower.includes("claude-3.0-haiku")) {
// Claude 3 Haiku has 4096 tokens in Bedrock
modelInfo = {
maxTokens: 4096,
contextWindow: 200_000,
supportsPromptCache: false,
supportsImages: true,
}
} else if (arnLower.includes("claude-3-5-haiku") || arnLower.includes("claude-3.5-haiku")) {
// Claude 3.5 Haiku has 8192 tokens in Bedrock
modelInfo = {
maxTokens: 8192,
contextWindow: 200_000,
supportsPromptCache: false,
supportsImages: false,
}
} else if (arnLower.includes("claude")) {
// Generic Claude model with conservative token limit
modelInfo = {
maxTokens: 4096,
contextWindow: 128_000,
supportsPromptCache: false,
supportsImages: true,
}
} else if (arnLower.includes("llama3") || arnLower.includes("llama-3")) {
// Llama 3 models typically have 8192 tokens in Bedrock
modelInfo = {
maxTokens: 8192,
contextWindow: 128_000,
supportsPromptCache: false,
supportsImages: arnLower.includes("90b") || arnLower.includes("11b"),
}
} else if (arnLower.includes("nova-pro")) {
// Amazon Nova Pro
modelInfo = {
maxTokens: 5000,
contextWindow: 300_000,
supportsPromptCache: false,
supportsImages: true,
}
} else {
// Default for unknown models or prompt routers
modelInfo = {
maxTokens: 4096,
contextWindow: 128_000,
supportsPromptCache: false,
supportsImages: true,
}
}
// If modelMaxTokens is explicitly set in options, override the default
if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) {
modelInfo.maxTokens = this.options.modelMaxTokens
}
return {
id: this.options.awsCustomArn,
info: modelInfo,
}
}
const modelId = this.options.apiModelId
if (modelId) {
// Special case for custom ARN option
if (modelId === "custom-arn") {
// This should not happen as we should have awsCustomArn set
// but just in case, return a default model
return {
id: bedrockDefaultModelId,
info: bedrockModels[bedrockDefaultModelId],
}
}
// For tests, allow any model ID
if (process.env.NODE_ENV === "test") {
return {
@ -237,7 +569,43 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
// Handle cross-region inference
let modelId: string
if (this.options.awsUseCrossRegionInference) {
// For custom ARNs, use the ARN directly without modification
if (this.options.awsCustomArn) {
modelId = modelConfig.id
logger.debug("Using custom ARN in completePrompt", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
})
// Validate ARN format and check region match
const clientRegion = this.client.config.region as string
const validation = validateBedrockArn(modelId, clientRegion)
if (!validation.isValid) {
logger.error("Invalid ARN format in completePrompt", {
ctx: "bedrock",
modelId,
errorMessage: validation.errorMessage,
})
throw new Error(
validation.errorMessage ||
"Invalid ARN format. ARN should follow the pattern: arn:aws:bedrock:region:account-id:resource-type/resource-name",
)
}
// Extract region from ARN
const arnRegion = validation.arnRegion!
// Log warning if there's a region mismatch
if (validation.errorMessage) {
logger.warn(validation.errorMessage, {
ctx: "bedrock",
arnRegion,
clientRegion,
})
}
} else if (this.options.awsUseCrossRegionInference) {
let regionPrefix = (this.options.awsRegion || "").slice(0, 3)
switch (regionPrefix) {
case "us-":
@ -263,12 +631,21 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
},
]),
inferenceConfig: {
maxTokens: modelConfig.info.maxTokens || 5000,
maxTokens: modelConfig.info.maxTokens || 4096,
temperature: this.options.modelTemperature ?? BEDROCK_DEFAULT_TEMPERATURE,
topP: 0.1,
},
}
// Log the payload for debugging custom ARN issues
if (this.options.awsCustomArn) {
logger.debug("Bedrock completePrompt request details", {
ctx: "bedrock",
clientRegion: this.client.config.region,
payload: JSON.stringify(payload, null, 2),
})
}
const command = new ConverseCommand(payload)
const response = await this.client.send(command)
@ -280,11 +657,67 @@ export class AwsBedrockHandler implements ApiHandler, SingleCompletionHandler {
return output.content
}
} catch (parseError) {
console.error("Failed to parse Bedrock response:", parseError)
logger.error("Failed to parse Bedrock response", {
ctx: "bedrock",
error: parseError instanceof Error ? parseError : String(parseError),
})
}
}
return ""
} catch (error) {
// Enhanced error handling for custom ARN issues
if (this.options.awsCustomArn) {
logger.error("Error occurred with custom ARN in completePrompt", {
ctx: "bedrock",
customArn: this.options.awsCustomArn,
error: error instanceof Error ? error : String(error),
})
if (error instanceof Error) {
const errorMessage = error.message.toLowerCase()
// Access denied errors
if (
errorMessage.includes("access") &&
(errorMessage.includes("model") || errorMessage.includes("denied"))
) {
throw new Error(
`Bedrock custom ARN error: You don't have access to the model with the specified ARN. Please verify:
1. The ARN is correct and points to a valid model
2. Your AWS credentials have permission to access this model (check IAM policies)
3. The region in the ARN matches the region where the model is deployed
4. If using a provisioned model, ensure it's active and not in a failed state`,
)
}
// Model not found errors
else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) {
throw new Error(
`Bedrock custom ARN error: The specified ARN does not exist or is invalid. Please check:
1. The ARN format is correct (arn:aws:bedrock:region:account-id:resource-type/resource-name)
2. The model exists in the specified region
3. The account ID in the ARN is correct
4. The resource type is one of: foundation-model, provisioned-model, or default-prompt-router`,
)
}
// Throttling errors
else if (
errorMessage.includes("throttl") ||
errorMessage.includes("rate") ||
errorMessage.includes("limit")
) {
throw new Error(
`Bedrock custom ARN error: Request was throttled or rate limited. Please try:
1. Reducing the frequency of requests
2. If using a provisioned model, check its throughput settings
3. Contact AWS support to request a quota increase if needed`,
)
} else {
throw new Error(`Bedrock custom ARN error: ${error.message}`)
}
}
}
// Standard error handling
if (error instanceof Error) {
throw new Error(`Bedrock completion error: ${error.message}`)
}

View file

@ -0,0 +1,3 @@
export const ANTHROPIC_DEFAULT_MAX_TOKENS = 8192
export const DEEP_SEEK_DEFAULT_TEMPERATURE = 0.6

View file

@ -1,6 +1,7 @@
import { OpenAiHandler, OpenAiHandlerOptions } from "./openai"
import { ModelInfo } from "../../shared/api"
import { deepSeekModels, deepSeekDefaultModelId } from "../../shared/api"
import { deepSeekModels, deepSeekDefaultModelId, ModelInfo } from "../../shared/api"
import { ApiStreamUsageChunk } from "../transform/stream" // Import for type
import { getModelParams } from "../index"
export class DeepSeekHandler extends OpenAiHandler {
constructor(options: OpenAiHandlerOptions) {
@ -8,7 +9,7 @@ export class DeepSeekHandler extends OpenAiHandler {
...options,
openAiApiKey: options.deepSeekApiKey ?? "not-provided",
openAiModelId: options.apiModelId ?? deepSeekDefaultModelId,
openAiBaseUrl: options.deepSeekBaseUrl ?? "https://api.deepseek.com/v1",
openAiBaseUrl: options.deepSeekBaseUrl ?? "https://api.deepseek.com",
openAiStreamingEnabled: true,
includeMaxTokens: true,
})
@ -16,9 +17,23 @@ export class DeepSeekHandler extends OpenAiHandler {
override getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.apiModelId ?? deepSeekDefaultModelId
const info = deepSeekModels[modelId as keyof typeof deepSeekModels] || deepSeekModels[deepSeekDefaultModelId]
return {
id: modelId,
info: deepSeekModels[modelId as keyof typeof deepSeekModels] || deepSeekModels[deepSeekDefaultModelId],
info,
...getModelParams({ options: this.options, model: info }),
}
}
// Override to handle DeepSeek's usage metrics, including caching.
protected override processUsageMetrics(usage: any): ApiStreamUsageChunk {
return {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
outputTokens: usage?.completion_tokens || 0,
cacheWriteTokens: usage?.prompt_tokens_details?.cache_miss_tokens,
cacheReadTokens: usage?.prompt_tokens_details?.cached_tokens,
}
}
}

View file

@ -1,22 +1,24 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { GoogleGenerativeAI } from "@google/generative-ai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { SingleCompletionHandler } from "../"
import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "../../shared/api"
import { convertAnthropicMessageToGemini } from "../transform/gemini-format"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
const GEMINI_DEFAULT_TEMPERATURE = 0
export class GeminiHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class GeminiHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: GoogleGenerativeAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = new GoogleGenerativeAI(options.geminiApiKey ?? "not-provided")
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const model = this.client.getGenerativeModel(
{
model: this.getModel().id,
@ -26,6 +28,7 @@ export class GeminiHandler implements ApiHandler, SingleCompletionHandler {
baseUrl: this.options.googleGeminiBaseUrl || undefined,
},
)
const result = await model.generateContentStream({
contents: messages.map(convertAnthropicMessageToGemini),
generationConfig: {
@ -49,7 +52,7 @@ export class GeminiHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: GeminiModelId; info: ModelInfo } {
override getModel(): { id: GeminiModelId; info: ModelInfo } {
const modelId = this.options.apiModelId
if (modelId && modelId in geminiModels) {
const id = modelId as GeminiModelId

View file

@ -1,25 +1,44 @@
import { Anthropic } from "@anthropic-ai/sdk"
import axios from "axios"
import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, glamaDefaultModelId, glamaDefaultModelInfo } from "../../shared/api"
import { parseApiPrice } from "../../utils/cost"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { SingleCompletionHandler } from "../"
import { BaseProvider } from "./base-provider"
const GLAMA_DEFAULT_TEMPERATURE = 0
export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class GlamaHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
const baseURL = "https://glama.ai/api/gateway/openai/v1"
const apiKey = this.options.glamaApiKey ?? "not-provided"
this.client = new OpenAI({ baseURL, apiKey })
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
private supportsTemperature(): boolean {
return !this.getModel().id.startsWith("openai/o3-mini")
}
override getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.glamaModelId
const modelInfo = this.options.glamaModelInfo
if (modelId && modelInfo) {
return { id: modelId, info: modelInfo }
}
return { id: glamaDefaultModelId, info: glamaDefaultModelInfo }
}
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
// Convert Anthropic messages to OpenAI format
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
@ -69,7 +88,7 @@ export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
let maxTokens: number | undefined
if (this.getModel().id.startsWith("anthropic/")) {
maxTokens = 8_192
maxTokens = this.getModel().info.maxTokens
}
const requestOptions: OpenAI.Chat.ChatCompletionCreateParams = {
@ -150,21 +169,6 @@ export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
}
}
private supportsTemperature(): boolean {
return !this.getModel().id.startsWith("openai/o3-mini")
}
getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.glamaModelId
const modelInfo = this.options.glamaModelInfo
if (modelId && modelInfo) {
return { id: modelId, info: modelInfo }
}
return { id: glamaDefaultModelId, info: glamaDefaultModelInfo }
}
async completePrompt(prompt: string): Promise<string> {
try {
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
@ -177,7 +181,7 @@ export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
}
if (this.getModel().id.startsWith("anthropic/")) {
requestOptions.max_tokens = 8192
requestOptions.max_tokens = this.getModel().info.maxTokens
}
const response = await this.client.chat.completions.create(requestOptions)
@ -190,3 +194,44 @@ export class GlamaHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
export async function getGlamaModels() {
const models: Record<string, ModelInfo> = {}
try {
const response = await axios.get("https://glama.ai/api/gateway/v1/models")
const rawModels = response.data
for (const rawModel of rawModels) {
const modelInfo: ModelInfo = {
maxTokens: rawModel.maxTokensOutput,
contextWindow: rawModel.maxTokensInput,
supportsImages: rawModel.capabilities?.includes("input:image"),
supportsComputerUse: rawModel.capabilities?.includes("computer_use"),
supportsPromptCache: rawModel.capabilities?.includes("caching"),
inputPrice: parseApiPrice(rawModel.pricePerToken?.input),
outputPrice: parseApiPrice(rawModel.pricePerToken?.output),
description: undefined,
cacheWritesPrice: parseApiPrice(rawModel.pricePerToken?.cacheWrite),
cacheReadsPrice: parseApiPrice(rawModel.pricePerToken?.cacheRead),
}
switch (rawModel.id) {
case rawModel.id.startsWith("anthropic/claude-3-7-sonnet"):
modelInfo.maxTokens = 16384
break
case rawModel.id.startsWith("anthropic/"):
modelInfo.maxTokens = 8192
break
default:
break
}
models[rawModel.id] = modelInfo
}
} catch (error) {
console.error(`Error fetching Glama models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`)
}
return models
}

View file

@ -0,0 +1,139 @@
// filepath: e:\Project\Roo-Code\src\api\providers\human-relay.ts
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandlerOptions, ModelInfo } from "../../shared/api"
import { ApiHandler, SingleCompletionHandler } from "../index"
import { ApiStream } from "../transform/stream"
import * as vscode from "vscode"
import { ExtensionMessage } from "../../shared/ExtensionMessage"
import { getPanel } from "../../activate/registerCommands" // Import the getPanel function
/**
* Human Relay API processor
* This processor does not directly call the API, but interacts with the model through human operations copy and paste.
*/
export class HumanRelayHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
constructor(options: ApiHandlerOptions) {
this.options = options
}
countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
return Promise.resolve(0)
}
/**
* Create a message processing flow, display a dialog box to request human assistance
* @param systemPrompt System prompt words
* @param messages Message list
*/
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
// Get the most recent user message
const latestMessage = messages[messages.length - 1]
if (!latestMessage) {
throw new Error("No message to relay")
}
// If it is the first message, splice the system prompt word with the user message
let promptText = ""
if (messages.length === 1) {
promptText = `${systemPrompt}\n\n${getMessageContent(latestMessage)}`
} else {
promptText = getMessageContent(latestMessage)
}
// Copy to clipboard
await vscode.env.clipboard.writeText(promptText)
// A dialog box pops up to request user action
const response = await showHumanRelayDialog(promptText)
if (!response) {
// The user canceled the operation
throw new Error("Human relay operation cancelled")
}
// Return to the user input reply
yield { type: "text", text: response }
}
/**
* Get model information
*/
getModel(): { id: string; info: ModelInfo } {
// Human relay does not depend on a specific model, here is a default configuration
return {
id: "human-relay",
info: {
maxTokens: 16384,
contextWindow: 100000,
supportsImages: true,
supportsPromptCache: false,
supportsComputerUse: true,
inputPrice: 0,
outputPrice: 0,
description: "Calling web-side AI model through human relay",
},
}
}
/**
* Implementation of a single prompt
* @param prompt Prompt content
*/
async completePrompt(prompt: string): Promise<string> {
// Copy to clipboard
await vscode.env.clipboard.writeText(prompt)
// A dialog box pops up to request user action
const response = await showHumanRelayDialog(prompt)
if (!response) {
throw new Error("Human relay operation cancelled")
}
return response
}
}
/**
* Extract text content from message object
* @param message
*/
function getMessageContent(message: Anthropic.Messages.MessageParam): string {
if (typeof message.content === "string") {
return message.content
} else if (Array.isArray(message.content)) {
return message.content
.filter((item) => item.type === "text")
.map((item) => (item.type === "text" ? item.text : ""))
.join("\n")
}
return ""
}
/**
* Displays the human relay dialog and waits for user response.
* @param promptText The prompt text that needs to be copied.
* @returns The user's input response or undefined (if canceled).
*/
async function showHumanRelayDialog(promptText: string): Promise<string | undefined> {
return new Promise<string | undefined>((resolve) => {
// Create a unique request ID
const requestId = Date.now().toString()
// Register a global callback function
vscode.commands.executeCommand(
"roo-cline.registerHumanRelayCallback",
requestId,
(response: string | undefined) => {
resolve(response)
},
)
// Open the dialog box directly using the current panel
vscode.commands.executeCommand("roo-cline.showHumanRelayDialog", {
requestId,
promptText,
})
})
}

View file

@ -1,17 +1,21 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import axios from "axios"
import { SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
const LMSTUDIO_DEFAULT_TEMPERATURE = 0
export class LmStudioHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class LmStudioHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = new OpenAI({
baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1",
@ -19,20 +23,31 @@ export class LmStudioHandler implements ApiHandler, SingleCompletionHandler {
})
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
]
try {
const stream = await this.client.chat.completions.create({
// Create params object with optional draft model
const params: any = {
model: this.getModel().id,
messages: openAiMessages,
temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
stream: true,
})
for await (const chunk of stream) {
}
// Add draft model if speculative decoding is enabled and a draft model is specified
if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
params.draft_model = this.options.lmStudioDraftModelId
}
const results = await this.client.chat.completions.create(params)
// Stream handling
// @ts-ignore
for await (const chunk of results) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
@ -49,7 +64,7 @@ export class LmStudioHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: string; info: ModelInfo } {
override getModel(): { id: string; info: ModelInfo } {
return {
id: this.options.lmStudioModelId || "",
info: openAiModelInfoSaneDefaults,
@ -58,12 +73,20 @@ export class LmStudioHandler implements ApiHandler, SingleCompletionHandler {
async completePrompt(prompt: string): Promise<string> {
try {
const response = await this.client.chat.completions.create({
// Create params object with optional draft model
const params: any = {
model: this.getModel().id,
messages: [{ role: "user", content: prompt }],
temperature: this.options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE,
stream: false,
})
}
// Add draft model if speculative decoding is enabled and a draft model is specified
if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) {
params.draft_model = this.options.lmStudioDraftModelId
}
const response = await this.client.chat.completions.create(params)
return response.choices[0]?.message.content || ""
} catch (error) {
throw new Error(
@ -72,3 +95,17 @@ export class LmStudioHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
export async function getLmStudioModels(baseUrl = "http://localhost:1234") {
try {
if (!URL.canParse(baseUrl)) {
return []
}
const response = await axios.get(`${baseUrl}/v1/models`)
const modelsArray = response.data?.data?.map((model: any) => model.id) || []
return [...new Set<string>(modelsArray)]
} catch (error) {
return []
}
}

View file

@ -1,6 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { Mistral } from "@mistralai/mistralai"
import { ApiHandler } from "../"
import { SingleCompletionHandler } from "../"
import {
ApiHandlerOptions,
mistralDefaultModelId,
@ -13,14 +13,16 @@ import {
} from "../../shared/api"
import { convertToMistralMessages } from "../transform/mistral-format"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
const MISTRAL_DEFAULT_TEMPERATURE = 0
export class MistralHandler implements ApiHandler {
private options: ApiHandlerOptions
export class MistralHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: Mistral
constructor(options: ApiHandlerOptions) {
super()
if (!options.mistralApiKey) {
throw new Error("Mistral API key is required")
}
@ -48,7 +50,7 @@ export class MistralHandler implements ApiHandler {
return "https://api.mistral.ai"
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const response = await this.client.chat.stream({
model: this.options.apiModelId || mistralDefaultModelId,
messages: [{ role: "system", content: systemPrompt }, ...convertToMistralMessages(messages)],
@ -81,7 +83,7 @@ export class MistralHandler implements ApiHandler {
}
}
getModel(): { id: MistralModelId; info: ModelInfo } {
override getModel(): { id: MistralModelId; info: ModelInfo } {
const modelId = this.options.apiModelId
if (modelId && modelId in mistralModels) {
const id = modelId as MistralModelId

View file

@ -1,20 +1,22 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import axios from "axios"
import { SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { convertToR1Format } from "../transform/r1-format"
import { ApiStream } from "../transform/stream"
import { DEEP_SEEK_DEFAULT_TEMPERATURE } from "./openai"
import { DEEP_SEEK_DEFAULT_TEMPERATURE } from "./constants"
import { XmlMatcher } from "../../utils/xml-matcher"
import { BaseProvider } from "./base-provider"
const OLLAMA_DEFAULT_TEMPERATURE = 0
export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class OllamaHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = new OpenAI({
baseURL: (this.options.ollamaBaseUrl || "http://localhost:11434") + "/v1",
@ -22,7 +24,7 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
})
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const modelId = this.getModel().id
const useR1Format = modelId.toLowerCase().includes("deepseek-r1")
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
@ -33,7 +35,7 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
messages: openAiMessages,
temperature: this.options.modelTemperature ?? OLLAMA_DEFAULT_TEMPERATURE,
temperature: this.options.modelTemperature ?? 0,
stream: true,
})
const matcher = new XmlMatcher(
@ -58,7 +60,7 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: string; info: ModelInfo } {
override getModel(): { id: string; info: ModelInfo } {
return {
id: this.options.ollamaModelId || "",
info: openAiModelInfoSaneDefaults,
@ -74,9 +76,7 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
messages: useR1Format
? convertToR1Format([{ role: "user", content: prompt }])
: [{ role: "user", content: prompt }],
temperature:
this.options.modelTemperature ??
(useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : OLLAMA_DEFAULT_TEMPERATURE),
temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
stream: false,
})
return response.choices[0]?.message.content || ""
@ -88,3 +88,17 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
export async function getOllamaModels(baseUrl = "http://localhost:11434") {
try {
if (!URL.canParse(baseUrl)) {
return []
}
const response = await axios.get(`${baseUrl}/api/tags`)
const modelsArray = response.data?.models?.map((model: any) => model.name) || []
return [...new Set<string>(modelsArray)]
} catch (error) {
return []
}
}

View file

@ -1,6 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { SingleCompletionHandler } from "../"
import {
ApiHandlerOptions,
ModelInfo,
@ -10,20 +10,22 @@ import {
} from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
const OPENAI_NATIVE_DEFAULT_TEMPERATURE = 0
export class OpenAiNativeHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class OpenAiNativeHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
const apiKey = this.options.openAiNativeApiKey ?? "not-provided"
this.client = new OpenAI({ apiKey })
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const modelId = this.getModel().id
if (modelId.startsWith("o1")) {
@ -133,7 +135,7 @@ export class OpenAiNativeHandler implements ApiHandler, SingleCompletionHandler
}
}
getModel(): { id: OpenAiNativeModelId; info: ModelInfo } {
override getModel(): { id: OpenAiNativeModelId; info: ModelInfo } {
const modelId = this.options.apiModelId
if (modelId && modelId in openAiNativeModels) {
const id = modelId as OpenAiNativeModelId

View file

@ -1,5 +1,6 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI, { AzureOpenAI } from "openai"
import axios from "axios"
import {
ApiHandlerOptions,
@ -7,24 +8,28 @@ import {
ModelInfo,
openAiModelInfoSaneDefaults,
} from "../../shared/api"
import { ApiHandler, SingleCompletionHandler } from "../index"
import { SingleCompletionHandler } from "../index"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { convertToR1Format } from "../transform/r1-format"
import { convertToSimpleMessages } from "../transform/simple-format"
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
import { BaseProvider } from "./base-provider"
export interface OpenAiHandlerOptions extends ApiHandlerOptions {
defaultHeaders?: Record<string, string>
const DEEP_SEEK_DEFAULT_TEMPERATURE = 0.6
export const defaultHeaders = {
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
"X-Title": "Roo Code",
}
export const DEEP_SEEK_DEFAULT_TEMPERATURE = 0.6
const OPENAI_DEFAULT_TEMPERATURE = 0
export interface OpenAiHandlerOptions extends ApiHandlerOptions {}
export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
export class OpenAiHandler extends BaseProvider implements SingleCompletionHandler {
protected options: OpenAiHandlerOptions
private client: OpenAI
constructor(options: OpenAiHandlerOptions) {
super()
this.options = options
const baseURL = this.options.openAiBaseUrl ?? "https://api.openai.com/v1"
@ -46,13 +51,14 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
baseURL,
apiKey,
apiVersion: this.options.azureApiVersion || azureOpenAiDefaultApiVersion,
defaultHeaders,
})
} else {
this.client = new OpenAI({ baseURL, apiKey, defaultHeaders: this.options.defaultHeaders })
this.client = new OpenAI({ baseURL, apiKey, defaultHeaders })
}
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const modelInfo = this.getModel().info
const modelUrl = this.options.openAiBaseUrl ?? ""
const modelId = this.options.openAiModelId ?? ""
@ -60,6 +66,11 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
const deepseekReasoner = modelId.includes("deepseek-reasoner")
const ark = modelUrl.includes(".volces.com")
if (modelId.startsWith("o3-mini")) {
yield* this.handleO3FamilyMessage(modelId, systemPrompt, messages)
return
}
if (this.options.openAiStreamingEnabled ?? true) {
const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = {
role: "system",
@ -77,9 +88,7 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
model: modelId,
temperature:
this.options.modelTemperature ??
(deepseekReasoner ? DEEP_SEEK_DEFAULT_TEMPERATURE : OPENAI_DEFAULT_TEMPERATURE),
temperature: this.options.modelTemperature ?? (deepseekReasoner ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
messages: convertedMessages,
stream: true as const,
stream_options: { include_usage: true },
@ -90,6 +99,8 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
const stream = await this.client.chat.completions.create(requestOptions)
let lastUsage
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta ?? {}
@ -107,9 +118,13 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
}
}
if (chunk.usage) {
yield this.processUsageMetrics(chunk.usage)
lastUsage = chunk.usage
}
}
if (lastUsage) {
yield this.processUsageMetrics(lastUsage, modelInfo)
}
} else {
// o1 for instance doesnt support streaming, non-1 temp, or system prompt
const systemMessage: OpenAI.Chat.ChatCompletionUserMessageParam = {
@ -130,11 +145,11 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
type: "text",
text: response.choices[0]?.message.content || "",
}
yield this.processUsageMetrics(response.usage)
yield this.processUsageMetrics(response.usage, modelInfo)
}
}
protected processUsageMetrics(usage: any): ApiStreamUsageChunk {
protected processUsageMetrics(usage: any, modelInfo?: ModelInfo): ApiStreamUsageChunk {
return {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
@ -142,7 +157,7 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: string; info: ModelInfo } {
override getModel(): { id: string; info: ModelInfo } {
return {
id: this.options.openAiModelId ?? "",
info: this.options.openAiCustomModelInfo ?? openAiModelInfoSaneDefaults,
@ -165,4 +180,91 @@ export class OpenAiHandler implements ApiHandler, SingleCompletionHandler {
throw error
}
}
private async *handleO3FamilyMessage(
modelId: string,
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
): ApiStream {
if (this.options.openAiStreamingEnabled ?? true) {
const stream = await this.client.chat.completions.create({
model: "o3-mini",
messages: [
{
role: "developer",
content: `Formatting re-enabled\n${systemPrompt}`,
},
...convertToOpenAiMessages(messages),
],
stream: true,
stream_options: { include_usage: true },
reasoning_effort: this.getModel().info.reasoningEffort,
})
yield* this.handleStreamResponse(stream)
} else {
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
model: modelId,
messages: [
{
role: "developer",
content: `Formatting re-enabled\n${systemPrompt}`,
},
...convertToOpenAiMessages(messages),
],
}
const response = await this.client.chat.completions.create(requestOptions)
yield {
type: "text",
text: response.choices[0]?.message.content || "",
}
yield this.processUsageMetrics(response.usage)
}
}
private async *handleStreamResponse(stream: AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>): ApiStream {
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
}
}
if (chunk.usage) {
yield {
type: "usage",
inputTokens: chunk.usage.prompt_tokens || 0,
outputTokens: chunk.usage.completion_tokens || 0,
}
}
}
}
}
export async function getOpenAiModels(baseUrl?: string, apiKey?: string) {
try {
if (!baseUrl) {
return []
}
if (!URL.canParse(baseUrl)) {
return []
}
const config: Record<string, any> = {}
if (apiKey) {
config["headers"] = { Authorization: `Bearer ${apiKey}` }
}
const response = await axios.get(`${baseUrl}/models`, config)
const modelsArray = response.data?.data?.map((model: any) => model.id) || []
return [...new Set<string>(modelsArray)]
} catch (error) {
return []
}
}

View file

@ -1,73 +1,67 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { BetaThinkingConfigParam } from "@anthropic-ai/sdk/resources/beta"
import axios from "axios"
import OpenAI from "openai"
import { ApiHandler } from "../"
import delay from "delay"
import { ApiHandlerOptions, ModelInfo, openRouterDefaultModelId, openRouterDefaultModelInfo } from "../../shared/api"
import { parseApiPrice } from "../../utils/cost"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStreamChunk, ApiStreamUsageChunk } from "../transform/stream"
import delay from "delay"
import { DEEP_SEEK_DEFAULT_TEMPERATURE } from "./openai"
import { convertToR1Format } from "../transform/r1-format"
const OPENROUTER_DEFAULT_TEMPERATURE = 0
import { DEEP_SEEK_DEFAULT_TEMPERATURE } from "./constants"
import { getModelParams, SingleCompletionHandler } from ".."
import { BaseProvider } from "./base-provider"
import { defaultHeaders } from "./openai"
// Add custom interface for OpenRouter params
// Add custom interface for OpenRouter params.
type OpenRouterChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParams & {
transforms?: string[]
include_reasoning?: boolean
thinking?: BetaThinkingConfigParam
}
// Add custom interface for OpenRouter usage chunk
// Add custom interface for OpenRouter usage chunk.
interface OpenRouterApiStreamUsageChunk extends ApiStreamUsageChunk {
fullResponseText: string
}
import { SingleCompletionHandler } from ".."
import { convertToR1Format } from "../transform/r1-format"
export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class OpenRouterHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
const baseURL = this.options.openRouterBaseUrl || "https://openrouter.ai/api/v1"
const apiKey = this.options.openRouterApiKey ?? "not-provided"
const defaultHeaders = {
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
"X-Title": "Roo Code",
}
this.client = new OpenAI({ baseURL, apiKey, defaultHeaders })
}
async *createMessage(
override async *createMessage(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
): AsyncGenerator<ApiStreamChunk> {
// Convert Anthropic messages to OpenAI format
let { id: modelId, maxTokens, thinking, temperature, topP } = this.getModel()
// Convert Anthropic messages to OpenAI format.
let openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
]
// DeepSeek highly recommends using user instead of system role.
if (modelId.startsWith("deepseek/deepseek-r1") || modelId === "perplexity/sonar-reasoning") {
openAiMessages = convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
}
// prompt caching: https://openrouter.ai/docs/prompt-caching
// this is specifically for claude models (some models may 'support prompt caching' automatically without this)
switch (this.getModel().id) {
case "anthropic/claude-3.7-sonnet":
case "anthropic/claude-3.5-sonnet":
case "anthropic/claude-3.5-sonnet:beta":
case "anthropic/claude-3.5-sonnet-20240620":
case "anthropic/claude-3.5-sonnet-20240620:beta":
case "anthropic/claude-3-5-haiku":
case "anthropic/claude-3-5-haiku:beta":
case "anthropic/claude-3-5-haiku-20241022":
case "anthropic/claude-3-5-haiku-20241022:beta":
case "anthropic/claude-3-haiku":
case "anthropic/claude-3-haiku:beta":
case "anthropic/claude-3-opus":
case "anthropic/claude-3-opus:beta":
switch (true) {
case modelId.startsWith("anthropic/"):
openAiMessages[0] = {
role: "system",
content: [
@ -103,57 +97,28 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
break
}
// Not sure how openrouter defaults max tokens when no value is provided, but the anthropic api requires this value and since they offer both 4096 and 8192 variants, we should ensure 8192.
// (models usually default to max tokens allowed)
let maxTokens: number | undefined
switch (this.getModel().id) {
case "anthropic/claude-3.7-sonnet":
case "anthropic/claude-3.5-sonnet":
case "anthropic/claude-3.5-sonnet:beta":
case "anthropic/claude-3.5-sonnet-20240620":
case "anthropic/claude-3.5-sonnet-20240620:beta":
case "anthropic/claude-3-5-haiku":
case "anthropic/claude-3-5-haiku:beta":
case "anthropic/claude-3-5-haiku-20241022":
case "anthropic/claude-3-5-haiku-20241022:beta":
maxTokens = 8_192
break
}
let defaultTemperature = OPENROUTER_DEFAULT_TEMPERATURE
let topP: number | undefined = undefined
// Handle models based on deepseek-r1
if (
this.getModel().id.startsWith("deepseek/deepseek-r1") ||
this.getModel().id === "perplexity/sonar-reasoning"
) {
// Recommended temperature for DeepSeek reasoning models
defaultTemperature = DEEP_SEEK_DEFAULT_TEMPERATURE
// DeepSeek highly recommends using user instead of system role
openAiMessages = convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
// Some provider support topP and 0.95 is value that Deepseek used in their benchmarks
topP = 0.95
}
// https://openrouter.ai/docs/transforms
let fullResponseText = ""
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
const completionParams: OpenRouterChatCompletionParams = {
model: modelId,
max_tokens: maxTokens,
temperature: this.options.modelTemperature ?? defaultTemperature,
temperature,
thinking, // OpenRouter is temporarily supporting this.
top_p: topP,
messages: openAiMessages,
stream: true,
include_reasoning: true,
// This way, the transforms field will only be included in the parameters when openRouterUseMiddleOutTransform is true.
...(this.options.openRouterUseMiddleOutTransform && { transforms: ["middle-out"] }),
} as OpenRouterChatCompletionParams)
...((this.options.openRouterUseMiddleOutTransform ?? true) && { transforms: ["middle-out"] }),
}
const stream = await this.client.chat.completions.create(completionParams)
let genId: string | undefined
for await (const chunk of stream as unknown as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>) {
// openrouter returns an error object instead of the openai sdk throwing an error
// OpenRouter returns an error object instead of the OpenAI SDK throwing an error.
if ("error" in chunk) {
const error = chunk.error as { message?: string; code?: number }
console.error(`OpenRouter API Error: ${error?.code} - ${error?.message}`)
@ -165,12 +130,14 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
}
const delta = chunk.choices[0]?.delta
if ("reasoning" in delta && delta.reasoning) {
yield {
type: "reasoning",
text: delta.reasoning,
} as ApiStreamChunk
}
if (delta?.content) {
fullResponseText += delta.content
yield {
@ -178,6 +145,7 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
text: delta.content,
} as ApiStreamChunk
}
// if (chunk.usage) {
// yield {
// type: "usage",
@ -187,10 +155,12 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
// }
}
// retry fetching generation details
// Retry fetching generation details.
let attempt = 0
while (attempt++ < 10) {
await delay(200) // FIXME: necessary delay to ensure generation endpoint is ready
try {
const response = await axios.get(`https://openrouter.ai/api/v1/generation?id=${genId}`, {
headers: {
@ -200,7 +170,7 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
})
const generation = response.data?.data
console.log("OpenRouter generation details:", response.data)
yield {
type: "usage",
// cacheWriteTokens: 0,
@ -211,6 +181,7 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
totalCost: generation?.total_cost || 0,
fullResponseText,
} as OpenRouterApiStreamUsageChunk
return
} catch (error) {
// ignore if fails
@ -218,36 +189,119 @@ export class OpenRouterHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
getModel(): { id: string; info: ModelInfo } {
override getModel() {
const modelId = this.options.openRouterModelId
const modelInfo = this.options.openRouterModelInfo
if (modelId && modelInfo) {
return { id: modelId, info: modelInfo }
let id = modelId ?? openRouterDefaultModelId
const info = modelInfo ?? openRouterDefaultModelInfo
const isDeepSeekR1 = id.startsWith("deepseek/deepseek-r1") || modelId === "perplexity/sonar-reasoning"
const defaultTemperature = isDeepSeekR1 ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0
const topP = isDeepSeekR1 ? 0.95 : undefined
return {
id,
info,
...getModelParams({ options: this.options, model: info, defaultTemperature }),
topP,
}
return { id: openRouterDefaultModelId, info: openRouterDefaultModelInfo }
}
async completePrompt(prompt: string): Promise<string> {
try {
const response = await this.client.chat.completions.create({
model: this.getModel().id,
messages: [{ role: "user", content: prompt }],
temperature: this.options.modelTemperature ?? OPENROUTER_DEFAULT_TEMPERATURE,
stream: false,
})
async completePrompt(prompt: string) {
let { id: modelId, maxTokens, thinking, temperature } = this.getModel()
if ("error" in response) {
const error = response.error as { message?: string; code?: number }
throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`)
}
const completion = response as OpenAI.Chat.ChatCompletion
return completion.choices[0]?.message?.content || ""
} catch (error) {
if (error instanceof Error) {
throw new Error(`OpenRouter completion error: ${error.message}`)
}
throw error
const completionParams: OpenRouterChatCompletionParams = {
model: modelId,
max_tokens: maxTokens,
thinking,
temperature,
messages: [{ role: "user", content: prompt }],
stream: false,
}
const response = await this.client.chat.completions.create(completionParams)
if ("error" in response) {
const error = response.error as { message?: string; code?: number }
throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`)
}
const completion = response as OpenAI.Chat.ChatCompletion
return completion.choices[0]?.message?.content || ""
}
}
export async function getOpenRouterModels() {
const models: Record<string, ModelInfo> = {}
try {
const response = await axios.get("https://openrouter.ai/api/v1/models")
const rawModels = response.data.data
for (const rawModel of rawModels) {
const modelInfo: ModelInfo = {
maxTokens: rawModel.top_provider?.max_completion_tokens,
contextWindow: rawModel.context_length,
supportsImages: rawModel.architecture?.modality?.includes("image"),
supportsPromptCache: false,
inputPrice: parseApiPrice(rawModel.pricing?.prompt),
outputPrice: parseApiPrice(rawModel.pricing?.completion),
description: rawModel.description,
thinking: rawModel.id === "anthropic/claude-3.7-sonnet:thinking",
}
// NOTE: this needs to be synced with api.ts/openrouter default model info.
switch (true) {
case rawModel.id.startsWith("anthropic/claude-3.7-sonnet"):
modelInfo.supportsComputerUse = true
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 3.75
modelInfo.cacheReadsPrice = 0.3
modelInfo.maxTokens = rawModel.id === "anthropic/claude-3.7-sonnet:thinking" ? 128_000 : 16_384
break
case rawModel.id.startsWith("anthropic/claude-3.5-sonnet-20240620"):
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 3.75
modelInfo.cacheReadsPrice = 0.3
modelInfo.maxTokens = 8192
break
case rawModel.id.startsWith("anthropic/claude-3.5-sonnet"):
modelInfo.supportsComputerUse = true
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 3.75
modelInfo.cacheReadsPrice = 0.3
modelInfo.maxTokens = 8192
break
case rawModel.id.startsWith("anthropic/claude-3-5-haiku"):
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 1.25
modelInfo.cacheReadsPrice = 0.1
modelInfo.maxTokens = 8192
break
case rawModel.id.startsWith("anthropic/claude-3-opus"):
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 18.75
modelInfo.cacheReadsPrice = 1.5
modelInfo.maxTokens = 8192
break
case rawModel.id.startsWith("anthropic/claude-3-haiku"):
default:
modelInfo.supportsPromptCache = true
modelInfo.cacheWritesPrice = 0.3
modelInfo.cacheReadsPrice = 0.03
modelInfo.maxTokens = 8192
break
}
models[rawModel.id] = modelInfo
}
} catch (error) {
console.error(
`Error fetching OpenRouter models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`,
)
}
return models
}

View file

@ -1,6 +1,20 @@
import { OpenAiHandler, OpenAiHandlerOptions } from "./openai"
import axios from "axios"
import { ModelInfo, requestyModelInfoSaneDefaults, requestyDefaultModelId } from "../../shared/api"
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
import { calculateApiCostOpenAI, parseApiPrice } from "../../utils/cost"
import { ApiStreamUsageChunk } from "../transform/stream"
import { OpenAiHandler, OpenAiHandlerOptions } from "./openai"
import OpenAI from "openai"
// Requesty usage includes an extra field for Anthropic use cases.
// Safely cast the prompt token details section to the appropriate structure.
interface RequestyUsage extends OpenAI.CompletionUsage {
prompt_tokens_details?: {
caching_tokens?: number
cached_tokens?: number
}
total_cost?: number
}
export class RequestyHandler extends OpenAiHandler {
constructor(options: OpenAiHandlerOptions) {
@ -13,10 +27,6 @@ export class RequestyHandler extends OpenAiHandler {
openAiModelId: options.requestyModelId ?? requestyDefaultModelId,
openAiBaseUrl: "https://router.requesty.ai/v1",
openAiCustomModelInfo: options.requestyModelInfo ?? requestyModelInfoSaneDefaults,
defaultHeaders: {
"HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline",
"X-Title": "Roo Code",
},
})
}
@ -28,13 +38,68 @@ export class RequestyHandler extends OpenAiHandler {
}
}
protected override processUsageMetrics(usage: any): ApiStreamUsageChunk {
protected override processUsageMetrics(usage: any, modelInfo?: ModelInfo): ApiStreamUsageChunk {
const requestyUsage = usage as RequestyUsage
const inputTokens = requestyUsage?.prompt_tokens || 0
const outputTokens = requestyUsage?.completion_tokens || 0
const cacheWriteTokens = requestyUsage?.prompt_tokens_details?.caching_tokens || 0
const cacheReadTokens = requestyUsage?.prompt_tokens_details?.cached_tokens || 0
const totalCost = modelInfo
? calculateApiCostOpenAI(modelInfo, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens)
: 0
return {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
outputTokens: usage?.completion_tokens || 0,
cacheWriteTokens: usage?.cache_creation_input_tokens,
cacheReadTokens: usage?.cache_read_input_tokens,
inputTokens: inputTokens,
outputTokens: outputTokens,
cacheWriteTokens: cacheWriteTokens,
cacheReadTokens: cacheReadTokens,
totalCost: totalCost,
}
}
}
export async function getRequestyModels() {
const models: Record<string, ModelInfo> = {}
try {
const response = await axios.get("https://router.requesty.ai/v1/models")
const rawModels = response.data.data
for (const rawModel of rawModels) {
// {
// id: "anthropic/claude-3-5-sonnet-20240620",
// object: "model",
// created: 1740552655,
// owned_by: "system",
// input_price: 0.0000028,
// caching_price: 0.00000375,
// cached_price: 3e-7,
// output_price: 0.000015,
// max_output_tokens: 8192,
// context_window: 200000,
// supports_caching: true,
// description:
// "Anthropic's previous most intelligent model. High level of intelligence and capability. Excells in coding.",
// }
const modelInfo: ModelInfo = {
maxTokens: rawModel.max_output_tokens,
contextWindow: rawModel.context_window,
supportsPromptCache: rawModel.supports_caching,
supportsImages: rawModel.supports_vision,
supportsComputerUse: rawModel.supports_computer_use,
inputPrice: parseApiPrice(rawModel.input_price),
outputPrice: parseApiPrice(rawModel.output_price),
description: rawModel.description,
cacheWritesPrice: parseApiPrice(rawModel.caching_price),
cacheReadsPrice: parseApiPrice(rawModel.cached_price),
}
models[rawModel.id] = modelInfo
}
} catch (error) {
console.error(`Error fetching Requesty models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`)
}
return models
}

View file

@ -1,27 +1,31 @@
import { Anthropic } from "@anthropic-ai/sdk"
import axios from "axios"
import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, unboundDefaultModelId, unboundDefaultModelInfo } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
import { SingleCompletionHandler } from "../"
import { BaseProvider } from "./base-provider"
interface UnboundUsage extends OpenAI.CompletionUsage {
cache_creation_input_tokens?: number
cache_read_input_tokens?: number
}
export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class UnboundHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
super()
this.options = options
const baseURL = "https://api.getunbound.ai/v1"
const apiKey = this.options.unboundApiKey ?? "not-provided"
this.client = new OpenAI({ baseURL, apiKey })
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
// Convert Anthropic messages to OpenAI format
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
@ -71,7 +75,7 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
let maxTokens: number | undefined
if (this.getModel().id.startsWith("anthropic/")) {
maxTokens = 8_192
maxTokens = this.getModel().info.maxTokens
}
const { data: completion, response } = await this.client.chat.completions
@ -129,7 +133,7 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
}
}
getModel(): { id: string; info: ModelInfo } {
override getModel(): { id: string; info: ModelInfo } {
const modelId = this.options.unboundModelId
const modelInfo = this.options.unboundModelInfo
if (modelId && modelInfo) {
@ -150,10 +154,21 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
}
if (this.getModel().id.startsWith("anthropic/")) {
requestOptions.max_tokens = 8192
requestOptions.max_tokens = this.getModel().info.maxTokens
}
const response = await this.client.chat.completions.create(requestOptions)
const response = await this.client.chat.completions.create(requestOptions, {
headers: {
"X-Unbound-Metadata": JSON.stringify({
labels: [
{
key: "app",
value: "roo-code",
},
],
}),
},
})
return response.choices[0]?.message.content || ""
} catch (error) {
if (error instanceof Error) {
@ -163,3 +178,46 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
export async function getUnboundModels() {
const models: Record<string, ModelInfo> = {}
try {
const response = await axios.get("https://api.getunbound.ai/models")
if (response.data) {
const rawModels: Record<string, any> = response.data
for (const [modelId, model] of Object.entries(rawModels)) {
const modelInfo: ModelInfo = {
maxTokens: model?.maxTokens ? parseInt(model.maxTokens) : undefined,
contextWindow: model?.contextWindow ? parseInt(model.contextWindow) : 0,
supportsImages: model?.supportsImages ?? false,
supportsPromptCache: model?.supportsPromptCaching ?? false,
supportsComputerUse: model?.supportsComputerUse ?? false,
inputPrice: model?.inputTokenPrice ? parseFloat(model.inputTokenPrice) : undefined,
outputPrice: model?.outputTokenPrice ? parseFloat(model.outputTokenPrice) : undefined,
cacheWritesPrice: model?.cacheWritePrice ? parseFloat(model.cacheWritePrice) : undefined,
cacheReadsPrice: model?.cacheReadPrice ? parseFloat(model.cacheReadPrice) : undefined,
}
switch (true) {
case modelId.startsWith("anthropic/claude-3-7-sonnet"):
modelInfo.maxTokens = 16384
break
case modelId.startsWith("anthropic/"):
modelInfo.maxTokens = 8192
break
default:
break
}
models[modelId] = modelInfo
}
}
} catch (error) {
console.error(`Error fetching Unbound models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`)
}
return models
}

View file

@ -1,54 +1,330 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
import { ApiHandler, SingleCompletionHandler } from "../"
import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming"
import { VertexAI } from "@google-cloud/vertexai"
import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { convertAnthropicMessageToVertexGemini } from "../transform/vertex-gemini-format"
import { BaseProvider } from "./base-provider"
import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "./constants"
import { getModelParams, SingleCompletionHandler } from "../"
import { GoogleAuth } from "google-auth-library"
// Types for Vertex SDK
/**
* Vertex API has specific limitations for prompt caching:
* 1. Maximum of 4 blocks can have cache_control
* 2. Only text blocks can be cached (images and other content types cannot)
* 3. Cache control can only be applied to user messages, not assistant messages
*
* Our caching strategy:
* - Cache the system prompt (1 block)
* - Cache the last text block of the second-to-last user message (1 block)
* - Cache the last text block of the last user message (1 block)
* This ensures we stay under the 4-block limit while maintaining effective caching
* for the most relevant context.
*/
interface VertexTextBlock {
type: "text"
text: string
cache_control?: { type: "ephemeral" }
}
interface VertexImageBlock {
type: "image"
source: {
type: "base64"
media_type: "image/jpeg" | "image/png" | "image/gif" | "image/webp"
data: string
}
}
type VertexContentBlock = VertexTextBlock | VertexImageBlock
interface VertexUsage {
input_tokens?: number
output_tokens?: number
cache_creation_input_tokens?: number
cache_read_input_tokens?: number
}
interface VertexMessage extends Omit<Anthropic.Messages.MessageParam, "content"> {
content: string | VertexContentBlock[]
}
interface VertexMessageCreateParams {
model: string
max_tokens: number
temperature: number
system: string | VertexTextBlock[]
messages: VertexMessage[]
stream: boolean
}
interface VertexMessageResponse {
content: Array<{ type: "text"; text: string }>
}
interface VertexMessageStreamEvent {
type: "message_start" | "message_delta" | "content_block_start" | "content_block_delta"
message?: {
usage: VertexUsage
}
usage?: {
output_tokens: number
}
content_block?:
| {
type: "text"
text: string
}
| {
type: "thinking"
thinking: string
}
index?: number
delta?:
| {
type: "text_delta"
text: string
}
| {
type: "thinking_delta"
thinking: string
}
}
// https://docs.anthropic.com/en/api/claude-on-vertex-ai
export class VertexHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
private client: AnthropicVertex
export class VertexHandler extends BaseProvider implements SingleCompletionHandler {
MODEL_CLAUDE = "claude"
MODEL_GEMINI = "gemini"
protected options: ApiHandlerOptions
private anthropicClient: AnthropicVertex
private geminiClient: VertexAI
private modelType: string
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = new AnthropicVertex({
projectId: this.options.vertexProjectId ?? "not-provided",
// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
region: this.options.vertexRegion ?? "us-east5",
})
if (this.options.apiModelId?.startsWith(this.MODEL_CLAUDE)) {
this.modelType = this.MODEL_CLAUDE
} else if (this.options.apiModelId?.startsWith(this.MODEL_GEMINI)) {
this.modelType = this.MODEL_GEMINI
} else {
throw new Error(`Unknown model ID: ${this.options.apiModelId}`)
}
if (this.options.vertexJsonCredentials) {
this.anthropicClient = new AnthropicVertex({
projectId: this.options.vertexProjectId ?? "not-provided",
// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
region: this.options.vertexRegion ?? "us-east5",
googleAuth: new GoogleAuth({
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
credentials: JSON.parse(this.options.vertexJsonCredentials),
}),
})
} else if (this.options.vertexKeyFile) {
this.anthropicClient = new AnthropicVertex({
projectId: this.options.vertexProjectId ?? "not-provided",
// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
region: this.options.vertexRegion ?? "us-east5",
googleAuth: new GoogleAuth({
scopes: ["https://www.googleapis.com/auth/cloud-platform"],
keyFile: this.options.vertexKeyFile,
}),
})
} else {
this.anthropicClient = new AnthropicVertex({
projectId: this.options.vertexProjectId ?? "not-provided",
// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
region: this.options.vertexRegion ?? "us-east5",
})
}
if (this.options.vertexJsonCredentials) {
this.geminiClient = new VertexAI({
project: this.options.vertexProjectId ?? "not-provided",
location: this.options.vertexRegion ?? "us-east5",
googleAuthOptions: {
credentials: JSON.parse(this.options.vertexJsonCredentials),
},
})
} else if (this.options.vertexKeyFile) {
this.geminiClient = new VertexAI({
project: this.options.vertexProjectId ?? "not-provided",
location: this.options.vertexRegion ?? "us-east5",
googleAuthOptions: {
keyFile: this.options.vertexKeyFile,
},
})
} else {
this.geminiClient = new VertexAI({
project: this.options.vertexProjectId ?? "not-provided",
location: this.options.vertexRegion ?? "us-east5",
})
}
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const stream = await this.client.messages.create({
private formatMessageForCache(message: Anthropic.Messages.MessageParam, shouldCache: boolean): VertexMessage {
// Assistant messages are kept as-is since they can't be cached
if (message.role === "assistant") {
return message as VertexMessage
}
// For string content, we convert to array format with optional cache control
if (typeof message.content === "string") {
return {
...message,
content: [
{
type: "text" as const,
text: message.content,
// For string content, we only have one block so it's always the last
...(shouldCache && { cache_control: { type: "ephemeral" } }),
},
],
}
}
// For array content, find the last text block index once before mapping
const lastTextBlockIndex = message.content.reduce(
(lastIndex, content, index) => (content.type === "text" ? index : lastIndex),
-1,
)
// Then use this pre-calculated index in the map function
return {
...message,
content: message.content.map((content, contentIndex) => {
// Images and other non-text content are passed through unchanged
if (content.type === "image") {
return content as VertexImageBlock
}
// Check if this is the last text block using our pre-calculated index
const isLastTextBlock = contentIndex === lastTextBlockIndex
return {
type: "text" as const,
text: (content as { text: string }).text,
...(shouldCache && isLastTextBlock && { cache_control: { type: "ephemeral" } }),
}
}),
}
}
private async *createGeminiMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const model = this.geminiClient.getGenerativeModel({
model: this.getModel().id,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: this.options.modelTemperature ?? 0,
system: systemPrompt,
messages,
stream: true,
systemInstruction: systemPrompt,
})
const result = await model.generateContentStream({
contents: messages.map(convertAnthropicMessageToVertexGemini),
generationConfig: {
maxOutputTokens: this.getModel().info.maxTokens,
temperature: this.options.modelTemperature ?? 0,
},
})
for await (const chunk of result.stream) {
if (chunk.candidates?.[0]?.content?.parts) {
for (const part of chunk.candidates[0].content.parts) {
if (part.text) {
yield {
type: "text",
text: part.text,
}
}
}
}
}
const response = await result.response
yield {
type: "usage",
inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
}
}
private async *createClaudeMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const model = this.getModel()
let { id, info, temperature, maxTokens, thinking } = model
const useCache = model.info.supportsPromptCache
// Find indices of user messages that we want to cache
// We only cache the last two user messages to stay within the 4-block limit
// (1 block for system + 1 block each for last two user messages = 3 total)
const userMsgIndices = useCache
? messages.reduce((acc, msg, i) => (msg.role === "user" ? [...acc, i] : acc), [] as number[])
: []
const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
// Create the stream with appropriate caching configuration
const params = {
model: id,
max_tokens: maxTokens,
temperature,
thinking,
// Cache the system prompt if caching is enabled
system: useCache
? [
{
text: systemPrompt,
type: "text" as const,
cache_control: { type: "ephemeral" },
},
]
: systemPrompt,
messages: messages.map((message, index) => {
// Only cache the last two user messages
const shouldCache = useCache && (index === lastUserMsgIndex || index === secondLastMsgUserIndex)
return this.formatMessageForCache(message, shouldCache)
}),
stream: true,
}
const stream = (await this.anthropicClient.messages.create(
params as Anthropic.Messages.MessageCreateParamsStreaming,
)) as unknown as AnthropicStream<VertexMessageStreamEvent>
// Process the stream chunks
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
const usage = chunk.message.usage
case "message_start": {
const usage = chunk.message!.usage
yield {
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
cacheWriteTokens: usage.cache_creation_input_tokens,
cacheReadTokens: usage.cache_read_input_tokens,
}
break
case "message_delta":
}
case "message_delta": {
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.output_tokens || 0,
outputTokens: chunk.usage!.output_tokens || 0,
}
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
if (chunk.index > 0) {
}
case "content_block_start": {
switch (chunk.content_block!.type) {
case "text": {
if (chunk.index! > 0) {
yield {
type: "text",
text: "\n",
@ -56,49 +332,104 @@ export class VertexHandler implements ApiHandler, SingleCompletionHandler {
}
yield {
type: "text",
text: chunk.content_block.text,
text: chunk.content_block!.text,
}
break
}
case "thinking": {
if (chunk.index! > 0) {
yield {
type: "reasoning",
text: "\n",
}
}
yield {
type: "reasoning",
text: (chunk.content_block as any).thinking,
}
break
}
}
break
case "content_block_delta":
switch (chunk.delta.type) {
case "text_delta":
}
case "content_block_delta": {
switch (chunk.delta!.type) {
case "text_delta": {
yield {
type: "text",
text: chunk.delta.text,
text: chunk.delta!.text,
}
break
}
case "thinking_delta": {
yield {
type: "reasoning",
text: (chunk.delta as any).thinking,
}
break
}
}
break
}
}
}
}
getModel(): { id: VertexModelId; info: ModelInfo } {
const modelId = this.options.apiModelId
if (modelId && modelId in vertexModels) {
const id = modelId as VertexModelId
return { id, info: vertexModels[id] }
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
switch (this.modelType) {
case this.MODEL_CLAUDE: {
yield* this.createClaudeMessage(systemPrompt, messages)
break
}
case this.MODEL_GEMINI: {
yield* this.createGeminiMessage(systemPrompt, messages)
break
}
default: {
throw new Error(`Invalid model type: ${this.modelType}`)
}
}
return { id: vertexDefaultModelId, info: vertexModels[vertexDefaultModelId] }
}
async completePrompt(prompt: string): Promise<string> {
getModel() {
const modelId = this.options.apiModelId
let id = modelId && modelId in vertexModels ? (modelId as VertexModelId) : vertexDefaultModelId
const info: ModelInfo = vertexModels[id]
// The `:thinking` variant is a virtual identifier for thinking-enabled
// models (similar to how it's handled in the Anthropic provider.)
if (id.endsWith(":thinking")) {
id = id.replace(":thinking", "") as VertexModelId
}
return {
id,
info,
...getModelParams({ options: this.options, model: info, defaultMaxTokens: ANTHROPIC_DEFAULT_MAX_TOKENS }),
}
}
private async completePromptGemini(prompt: string) {
try {
const response = await this.client.messages.create({
const model = this.geminiClient.getGenerativeModel({
model: this.getModel().id,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: this.options.modelTemperature ?? 0,
messages: [{ role: "user", content: prompt }],
stream: false,
})
const content = response.content[0]
if (content.type === "text") {
return content.text
}
return ""
const result = await model.generateContent({
contents: [{ role: "user", parts: [{ text: prompt }] }],
generationConfig: {
temperature: this.options.modelTemperature ?? 0,
},
})
let text = ""
result.response.candidates?.forEach((candidate) => {
candidate.content.parts.forEach((part) => {
text += part.text
})
})
return text
} catch (error) {
if (error instanceof Error) {
throw new Error(`Vertex completion error: ${error.message}`)
@ -106,4 +437,63 @@ export class VertexHandler implements ApiHandler, SingleCompletionHandler {
throw error
}
}
private async completePromptClaude(prompt: string) {
try {
let { id, info, temperature, maxTokens, thinking } = this.getModel()
const useCache = info.supportsPromptCache
const params: Anthropic.Messages.MessageCreateParamsNonStreaming = {
model: id,
max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
temperature,
thinking,
system: "", // No system prompt needed for single completions
messages: [
{
role: "user",
content: useCache
? [
{
type: "text" as const,
text: prompt,
cache_control: { type: "ephemeral" },
},
]
: prompt,
},
],
stream: false,
}
const response = (await this.anthropicClient.messages.create(params)) as unknown as VertexMessageResponse
const content = response.content[0]
if (content.type === "text") {
return content.text
}
return ""
} catch (error) {
if (error instanceof Error) {
throw new Error(`Vertex completion error: ${error.message}`)
}
throw error
}
}
async completePrompt(prompt: string) {
switch (this.modelType) {
case this.MODEL_CLAUDE: {
return this.completePromptClaude(prompt)
}
case this.MODEL_GEMINI: {
return this.completePromptGemini(prompt)
}
default: {
throw new Error(`Invalid model type: ${this.modelType}`)
}
}
}
}

View file

@ -1,17 +1,19 @@
import { Anthropic } from "@anthropic-ai/sdk"
import * as vscode from "vscode"
import { ApiHandler, SingleCompletionHandler } from "../"
import { calculateApiCost } from "../../utils/cost"
import { SingleCompletionHandler } from "../"
import { calculateApiCostAnthropic } from "../../utils/cost"
import { ApiStream } from "../transform/stream"
import { convertToVsCodeLmMessages } from "../transform/vscode-lm-format"
import { SELECTOR_SEPARATOR, stringifyVsCodeLmModelSelector } from "../../shared/vsCodeSelectorUtils"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { BaseProvider } from "./base-provider"
/**
* Handles interaction with VS Code's Language Model API for chat-based operations.
* This handler implements the ApiHandler interface to provide VS Code LM specific functionality.
* This handler extends BaseProvider to provide VS Code LM specific functionality.
*
* @implements {ApiHandler}
* @extends {BaseProvider}
*
* @remarks
* The handler manages a VS Code language model chat client and provides methods to:
@ -34,13 +36,14 @@ import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../..
* }
* ```
*/
export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: vscode.LanguageModelChat | null
private disposable: vscode.Disposable | null
private currentRequestCancellation: vscode.CancellationTokenSource | null
constructor(options: ApiHandlerOptions) {
super()
this.options = options
this.client = null
this.disposable = null
@ -144,7 +147,33 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
}
}
private async countTokens(text: string | vscode.LanguageModelChatMessage): Promise<number> {
/**
* Implements the ApiHandler countTokens interface method
* Provides token counting for Anthropic content blocks
*
* @param content The content blocks to count tokens for
* @returns A promise resolving to the token count
*/
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
// Convert Anthropic content blocks to a string for VSCode LM token counting
let textContent = ""
for (const block of content) {
if (block.type === "text") {
textContent += block.text || ""
} else if (block.type === "image") {
// VSCode LM doesn't support images directly, so we'll just use a placeholder
textContent += "[IMAGE]"
}
}
return this.internalCountTokens(textContent)
}
/**
* Private implementation of token counting used internally by VsCodeLmHandler
*/
private async internalCountTokens(text: string | vscode.LanguageModelChatMessage): Promise<number> {
// Check for required dependencies
if (!this.client) {
console.warn("Roo Code <Language Model API>: No client available for token counting")
@ -215,9 +244,9 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
systemPrompt: string,
vsCodeLmMessages: vscode.LanguageModelChatMessage[],
): Promise<number> {
const systemTokens: number = await this.countTokens(systemPrompt)
const systemTokens: number = await this.internalCountTokens(systemPrompt)
const messageTokens: number[] = await Promise.all(vsCodeLmMessages.map((msg) => this.countTokens(msg)))
const messageTokens: number[] = await Promise.all(vsCodeLmMessages.map((msg) => this.internalCountTokens(msg)))
return systemTokens + messageTokens.reduce((sum: number, tokens: number): number => sum + tokens, 0)
}
@ -318,7 +347,7 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
return content
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
override async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
// Ensure clean state before starting a new request
this.ensureCleanState()
const client: vscode.LanguageModelChat = await this.getClient()
@ -426,14 +455,14 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
}
// Count tokens in the accumulated text after stream completion
const totalOutputTokens: number = await this.countTokens(accumulatedText)
const totalOutputTokens: number = await this.internalCountTokens(accumulatedText)
// Report final usage after stream completion
yield {
type: "usage",
inputTokens: totalInputTokens,
outputTokens: totalOutputTokens,
totalCost: calculateApiCost(this.getModel().info, totalInputTokens, totalOutputTokens),
totalCost: calculateApiCostAnthropic(this.getModel().info, totalInputTokens, totalOutputTokens),
}
} catch (error: unknown) {
this.ensureCleanState()
@ -466,7 +495,7 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
}
// Return model information based on the current client state
getModel(): { id: string; info: ModelInfo } {
override getModel(): { id: string; info: ModelInfo } {
if (this.client) {
// Validate client properties
const requiredProps = {
@ -545,3 +574,15 @@ export class VsCodeLmHandler implements ApiHandler, SingleCompletionHandler {
}
}
}
export async function getVsCodeLmModels() {
try {
const models = await vscode.lm.selectChatModels({})
return models || []
} catch (error) {
console.error(
`Error fetching VS Code LM models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`,
)
return []
}
}

View file

@ -0,0 +1,338 @@
// npx jest src/api/transform/__tests__/vertex-gemini-format.test.ts
import { Anthropic } from "@anthropic-ai/sdk"
import { convertAnthropicMessageToVertexGemini } from "../vertex-gemini-format"
describe("convertAnthropicMessageToVertexGemini", () => {
it("should convert a simple text message", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: "Hello, world!",
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [{ text: "Hello, world!" }],
})
})
it("should convert assistant role to model role", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "assistant",
content: "I'm an assistant",
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "model",
parts: [{ text: "I'm an assistant" }],
})
})
it("should convert a message with text blocks", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{ type: "text", text: "First paragraph" },
{ type: "text", text: "Second paragraph" },
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [{ text: "First paragraph" }, { text: "Second paragraph" }],
})
})
it("should convert a message with an image", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{ type: "text", text: "Check out this image:" },
{
type: "image",
source: {
type: "base64",
media_type: "image/jpeg",
data: "base64encodeddata",
},
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [
{ text: "Check out this image:" },
{
inlineData: {
data: "base64encodeddata",
mimeType: "image/jpeg",
},
},
],
})
})
it("should throw an error for unsupported image source type", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "image",
source: {
type: "url", // Not supported
url: "https://example.com/image.jpg",
} as any,
},
],
}
expect(() => convertAnthropicMessageToVertexGemini(anthropicMessage)).toThrow("Unsupported image source type")
})
it("should convert a message with tool use", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "assistant",
content: [
{ type: "text", text: "Let me calculate that for you." },
{
type: "tool_use",
id: "calc-123",
name: "calculator",
input: { operation: "add", numbers: [2, 3] },
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "model",
parts: [
{ text: "Let me calculate that for you." },
{
functionCall: {
name: "calculator",
args: { operation: "add", numbers: [2, 3] },
},
},
],
})
})
it("should convert a message with tool result as string", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{ type: "text", text: "Here's the result:" },
{
type: "tool_result",
tool_use_id: "calculator-123",
content: "The result is 5",
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [
{ text: "Here's the result:" },
{
functionResponse: {
name: "calculator",
response: {
name: "calculator",
content: "The result is 5",
},
},
},
],
})
})
it("should handle empty tool result content", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "tool_result",
tool_use_id: "calculator-123",
content: null as any, // Empty content
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
// Should skip the empty tool result
expect(result).toEqual({
role: "user",
parts: [],
})
})
it("should convert a message with tool result as array with text only", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "tool_result",
tool_use_id: "search-123",
content: [
{ type: "text", text: "First result" },
{ type: "text", text: "Second result" },
],
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [
{
functionResponse: {
name: "search",
response: {
name: "search",
content: "First result\n\nSecond result",
},
},
},
],
})
})
it("should convert a message with tool result as array with text and images", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "tool_result",
tool_use_id: "search-123",
content: [
{ type: "text", text: "Search results:" },
{
type: "image",
source: {
type: "base64",
media_type: "image/png",
data: "image1data",
},
},
{
type: "image",
source: {
type: "base64",
media_type: "image/jpeg",
data: "image2data",
},
},
],
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [
{
functionResponse: {
name: "search",
response: {
name: "search",
content: "Search results:\n\n(See next part for image)",
},
},
},
{
inlineData: {
data: "image1data",
mimeType: "image/png",
},
},
{
inlineData: {
data: "image2data",
mimeType: "image/jpeg",
},
},
],
})
})
it("should convert a message with tool result containing only images", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "tool_result",
tool_use_id: "imagesearch-123",
content: [
{
type: "image",
source: {
type: "base64",
media_type: "image/png",
data: "onlyimagedata",
},
},
],
},
],
}
const result = convertAnthropicMessageToVertexGemini(anthropicMessage)
expect(result).toEqual({
role: "user",
parts: [
{
functionResponse: {
name: "imagesearch",
response: {
name: "imagesearch",
content: "\n\n(See next part for image)",
},
},
},
{
inlineData: {
data: "onlyimagedata",
mimeType: "image/png",
},
},
],
})
})
it("should throw an error for unsupported content block type", () => {
const anthropicMessage: Anthropic.Messages.MessageParam = {
role: "user",
content: [
{
type: "unknown_type", // Unsupported type
data: "some data",
} as any,
],
}
expect(() => convertAnthropicMessageToVertexGemini(anthropicMessage)).toThrow(
"Unsupported content block type: unknown_type",
)
})
})

View file

@ -0,0 +1,83 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { Content, FunctionCallPart, FunctionResponsePart, InlineDataPart, Part, TextPart } from "@google-cloud/vertexai"
function convertAnthropicContentToVertexGemini(content: Anthropic.Messages.MessageParam["content"]): Part[] {
if (typeof content === "string") {
return [{ text: content } as TextPart]
}
return content.flatMap((block) => {
switch (block.type) {
case "text":
return { text: block.text } as TextPart
case "image":
if (block.source.type !== "base64") {
throw new Error("Unsupported image source type")
}
return {
inlineData: {
data: block.source.data,
mimeType: block.source.media_type,
},
} as InlineDataPart
case "tool_use":
return {
functionCall: {
name: block.name,
args: block.input,
},
} as FunctionCallPart
case "tool_result":
const name = block.tool_use_id.split("-")[0]
if (!block.content) {
return []
}
if (typeof block.content === "string") {
return {
functionResponse: {
name,
response: {
name,
content: block.content,
},
},
} as FunctionResponsePart
} else {
// The only case when tool_result could be array is when the tool failed and we're providing ie user feedback potentially with images
const textParts = block.content.filter((part) => part.type === "text")
const imageParts = block.content.filter((part) => part.type === "image")
const text = textParts.length > 0 ? textParts.map((part) => part.text).join("\n\n") : ""
const imageText = imageParts.length > 0 ? "\n\n(See next part for image)" : ""
return [
{
functionResponse: {
name,
response: {
name,
content: text + imageText,
},
},
} as FunctionResponsePart,
...imageParts.map(
(part) =>
({
inlineData: {
data: part.source.data,
mimeType: part.source.media_type,
},
}) as InlineDataPart,
),
]
}
default:
throw new Error(`Unsupported content block type: ${(block as any).type}`)
}
})
}
export function convertAnthropicMessageToVertexGemini(message: Anthropic.Messages.MessageParam): Content {
return {
role: message.role === "assistant" ? "model" : "user",
parts: convertAnthropicContentToVertexGemini(message.content),
}
}

File diff suppressed because it is too large Load diff

View file

@ -9,6 +9,9 @@ import * as vscode from "vscode"
import * as os from "os"
import * as path from "path"
// Mock RooIgnoreController
jest.mock("../ignore/RooIgnoreController")
// Mock all MCP-related modules
jest.mock(
"@modelcontextprotocol/sdk/types.js",
@ -237,6 +240,7 @@ describe("Cline", () => {
return [
{
id: "123",
number: 0,
ts: Date.now(),
task: "historical task",
tokensIn: 100,
@ -374,7 +378,7 @@ describe("Cline", () => {
expect(cline.diffEnabled).toBe(true)
expect(cline.diffStrategy).toBeDefined()
expect(getDiffStrategySpy).toHaveBeenCalledWith("claude-3-5-sonnet-20241022", 0.9, false)
expect(getDiffStrategySpy).toHaveBeenCalledWith("claude-3-5-sonnet-20241022", 0.9, false, false)
getDiffStrategySpy.mockRestore()
@ -395,7 +399,7 @@ describe("Cline", () => {
expect(cline.diffEnabled).toBe(true)
expect(cline.diffStrategy).toBeDefined()
expect(getDiffStrategySpy).toHaveBeenCalledWith("claude-3-5-sonnet-20241022", 1.0, false)
expect(getDiffStrategySpy).toHaveBeenCalledWith("claude-3-5-sonnet-20241022", 1.0, false, false)
getDiffStrategySpy.mockRestore()

View file

@ -0,0 +1,331 @@
import * as vscode from "vscode"
import { ContextProxy } from "../contextProxy"
import { logger } from "../../utils/logging"
import { GLOBAL_STATE_KEYS, SECRET_KEYS } from "../../shared/globalState"
// Mock shared/globalState
jest.mock("../../shared/globalState", () => ({
GLOBAL_STATE_KEYS: ["apiProvider", "apiModelId", "mode"],
SECRET_KEYS: ["apiKey", "openAiApiKey"],
}))
// Mock VSCode API
jest.mock("vscode", () => ({
Uri: {
file: jest.fn((path) => ({ path })),
},
ExtensionMode: {
Development: 1,
Production: 2,
Test: 3,
},
}))
describe("ContextProxy", () => {
let proxy: ContextProxy
let mockContext: any
let mockGlobalState: any
let mockSecrets: any
beforeEach(() => {
// Reset mocks
jest.clearAllMocks()
// Mock globalState
mockGlobalState = {
get: jest.fn(),
update: jest.fn().mockResolvedValue(undefined),
}
// Mock secrets
mockSecrets = {
get: jest.fn().mockResolvedValue("test-secret"),
store: jest.fn().mockResolvedValue(undefined),
delete: jest.fn().mockResolvedValue(undefined),
}
// Mock the extension context
mockContext = {
globalState: mockGlobalState,
secrets: mockSecrets,
extensionUri: { path: "/test/extension" },
extensionPath: "/test/extension",
globalStorageUri: { path: "/test/storage" },
logUri: { path: "/test/logs" },
extension: { packageJSON: { version: "1.0.0" } },
extensionMode: vscode.ExtensionMode.Development,
}
// Create proxy instance
proxy = new ContextProxy(mockContext)
})
describe("read-only pass-through properties", () => {
it("should return extension properties from the original context", () => {
expect(proxy.extensionUri).toBe(mockContext.extensionUri)
expect(proxy.extensionPath).toBe(mockContext.extensionPath)
expect(proxy.globalStorageUri).toBe(mockContext.globalStorageUri)
expect(proxy.logUri).toBe(mockContext.logUri)
expect(proxy.extension).toBe(mockContext.extension)
expect(proxy.extensionMode).toBe(mockContext.extensionMode)
})
})
describe("constructor", () => {
it("should initialize state cache with all global state keys", () => {
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length)
for (const key of GLOBAL_STATE_KEYS) {
expect(mockGlobalState.get).toHaveBeenCalledWith(key)
}
})
it("should initialize secret cache with all secret keys", () => {
expect(mockSecrets.get).toHaveBeenCalledTimes(SECRET_KEYS.length)
for (const key of SECRET_KEYS) {
expect(mockSecrets.get).toHaveBeenCalledWith(key)
}
})
})
describe("getGlobalState", () => {
it("should return value from cache when it exists", async () => {
// Manually set a value in the cache
await proxy.updateGlobalState("test-key", "cached-value")
// Should return the cached value
const result = proxy.getGlobalState("test-key")
expect(result).toBe("cached-value")
// Original context should be called once during updateGlobalState
expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length) // Only from initialization
})
it("should handle default values correctly", async () => {
// No value in cache
const result = proxy.getGlobalState("unknown-key", "default-value")
expect(result).toBe("default-value")
})
})
describe("updateGlobalState", () => {
it("should update state directly in original context", async () => {
await proxy.updateGlobalState("test-key", "new-value")
// Should have called original context
expect(mockGlobalState.update).toHaveBeenCalledWith("test-key", "new-value")
// Should have stored the value in cache
const storedValue = await proxy.getGlobalState("test-key")
expect(storedValue).toBe("new-value")
})
})
describe("getSecret", () => {
it("should return value from cache when it exists", async () => {
// Manually set a value in the cache
await proxy.storeSecret("api-key", "cached-secret")
// Should return the cached value
const result = proxy.getSecret("api-key")
expect(result).toBe("cached-secret")
})
})
describe("storeSecret", () => {
it("should store secret directly in original context", async () => {
await proxy.storeSecret("api-key", "new-secret")
// Should have called original context
expect(mockSecrets.store).toHaveBeenCalledWith("api-key", "new-secret")
// Should have stored the value in cache
const storedValue = await proxy.getSecret("api-key")
expect(storedValue).toBe("new-secret")
})
it("should handle undefined value for secret deletion", async () => {
await proxy.storeSecret("api-key", undefined)
// Should have called delete on original context
expect(mockSecrets.delete).toHaveBeenCalledWith("api-key")
// Should have stored undefined in cache
const storedValue = await proxy.getSecret("api-key")
expect(storedValue).toBeUndefined()
})
})
describe("setValue", () => {
it("should route secret keys to storeSecret", async () => {
// Spy on storeSecret
const storeSecretSpy = jest.spyOn(proxy, "storeSecret")
// Test with a known secret key
await proxy.setValue("openAiApiKey", "test-api-key")
// Should have called storeSecret
expect(storeSecretSpy).toHaveBeenCalledWith("openAiApiKey", "test-api-key")
// Should have stored the value in secret cache
const storedValue = proxy.getSecret("openAiApiKey")
expect(storedValue).toBe("test-api-key")
})
it("should route global state keys to updateGlobalState", async () => {
// Spy on updateGlobalState
const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
// Test with a known global state key
await proxy.setValue("apiModelId", "gpt-4")
// Should have called updateGlobalState
expect(updateGlobalStateSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
// Should have stored the value in state cache
const storedValue = proxy.getGlobalState("apiModelId")
expect(storedValue).toBe("gpt-4")
})
it("should handle unknown keys as global state with warning", async () => {
// Spy on the logger
const warnSpy = jest.spyOn(logger, "warn")
// Spy on updateGlobalState
const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
// Test with an unknown key
await proxy.setValue("unknownKey", "some-value")
// Should have logged a warning
expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("Unknown key: unknownKey"))
// Should have called updateGlobalState
expect(updateGlobalStateSpy).toHaveBeenCalledWith("unknownKey", "some-value")
// Should have stored the value in state cache
const storedValue = proxy.getGlobalState("unknownKey")
expect(storedValue).toBe("some-value")
})
})
describe("setValues", () => {
it("should process multiple values correctly", async () => {
// Spy on setValue
const setValueSpy = jest.spyOn(proxy, "setValue")
// Test with multiple values
await proxy.setValues({
apiModelId: "gpt-4",
apiProvider: "openai",
mode: "test-mode",
})
// Should have called setValue for each key
expect(setValueSpy).toHaveBeenCalledTimes(3)
expect(setValueSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
expect(setValueSpy).toHaveBeenCalledWith("apiProvider", "openai")
expect(setValueSpy).toHaveBeenCalledWith("mode", "test-mode")
// Should have stored all values in state cache
expect(proxy.getGlobalState("apiModelId")).toBe("gpt-4")
expect(proxy.getGlobalState("apiProvider")).toBe("openai")
expect(proxy.getGlobalState("mode")).toBe("test-mode")
})
it("should handle both secret and global state keys", async () => {
// Spy on storeSecret and updateGlobalState
const storeSecretSpy = jest.spyOn(proxy, "storeSecret")
const updateGlobalStateSpy = jest.spyOn(proxy, "updateGlobalState")
// Test with mixed keys
await proxy.setValues({
apiModelId: "gpt-4", // global state
openAiApiKey: "test-api-key", // secret
unknownKey: "some-value", // unknown
})
// Should have called appropriate methods
expect(storeSecretSpy).toHaveBeenCalledWith("openAiApiKey", "test-api-key")
expect(updateGlobalStateSpy).toHaveBeenCalledWith("apiModelId", "gpt-4")
expect(updateGlobalStateSpy).toHaveBeenCalledWith("unknownKey", "some-value")
// Should have stored values in appropriate caches
expect(proxy.getSecret("openAiApiKey")).toBe("test-api-key")
expect(proxy.getGlobalState("apiModelId")).toBe("gpt-4")
expect(proxy.getGlobalState("unknownKey")).toBe("some-value")
})
})
describe("resetAllState", () => {
it("should clear all in-memory caches", async () => {
// Setup initial state in caches
await proxy.setValues({
apiModelId: "gpt-4", // global state
openAiApiKey: "test-api-key", // secret
unknownKey: "some-value", // unknown
})
// Verify initial state
expect(proxy.getGlobalState("apiModelId")).toBe("gpt-4")
expect(proxy.getSecret("openAiApiKey")).toBe("test-api-key")
expect(proxy.getGlobalState("unknownKey")).toBe("some-value")
// Reset all state
await proxy.resetAllState()
// Caches should be reinitialized with values from the context
// Since our mock globalState.get returns undefined by default,
// the cache should now contain undefined values
expect(proxy.getGlobalState("apiModelId")).toBeUndefined()
expect(proxy.getGlobalState("unknownKey")).toBeUndefined()
})
it("should update all global state keys to undefined", async () => {
// Setup initial state
await proxy.updateGlobalState("apiModelId", "gpt-4")
await proxy.updateGlobalState("apiProvider", "openai")
// Reset all state
await proxy.resetAllState()
// Should have called update with undefined for each key
for (const key of GLOBAL_STATE_KEYS) {
expect(mockGlobalState.update).toHaveBeenCalledWith(key, undefined)
}
// Total calls should include initial setup + reset operations
const expectedUpdateCalls = 2 + GLOBAL_STATE_KEYS.length
expect(mockGlobalState.update).toHaveBeenCalledTimes(expectedUpdateCalls)
})
it("should delete all secrets", async () => {
// Setup initial secrets
await proxy.storeSecret("apiKey", "test-api-key")
await proxy.storeSecret("openAiApiKey", "test-openai-key")
// Reset all state
await proxy.resetAllState()
// Should have called delete for each key
for (const key of SECRET_KEYS) {
expect(mockSecrets.delete).toHaveBeenCalledWith(key)
}
// Total calls should equal the number of secret keys
expect(mockSecrets.delete).toHaveBeenCalledTimes(SECRET_KEYS.length)
})
it("should reinitialize caches after reset", async () => {
// Spy on initialization methods
const initStateCache = jest.spyOn(proxy as any, "initializeStateCache")
const initSecretCache = jest.spyOn(proxy as any, "initializeSecretCache")
// Reset all state
await proxy.resetAllState()
// Should reinitialize caches
expect(initStateCache).toHaveBeenCalledTimes(1)
expect(initSecretCache).toHaveBeenCalledTimes(1)
})
})
})

View file

@ -1,3 +1,5 @@
// npx jest src/core/config/__tests__/CustomModesManager.test.ts
import * as vscode from "vscode"
import * as path from "path"
import * as fs from "fs/promises"
@ -15,9 +17,10 @@ describe("CustomModesManager", () => {
let mockOnUpdate: jest.Mock
let mockWorkspaceFolders: { uri: { fsPath: string } }[]
const mockStoragePath = "/mock/settings"
// Use path.sep to ensure correct path separators for the current platform
const mockStoragePath = `${path.sep}mock${path.sep}settings`
const mockSettingsPath = path.join(mockStoragePath, "settings", "cline_custom_modes.json")
const mockRoomodes = "/mock/workspace/.roomodes"
const mockRoomodes = `${path.sep}mock${path.sep}workspace${path.sep}.roomodes`
beforeEach(() => {
mockOnUpdate = jest.fn()
@ -243,7 +246,15 @@ describe("CustomModesManager", () => {
await manager.updateCustomMode("project-mode", projectMode)
// Verify .roomodes was created with the project mode
expect(fs.writeFile).toHaveBeenCalledWith(mockRoomodes, expect.stringContaining("project-mode"), "utf-8")
expect(fs.writeFile).toHaveBeenCalledWith(
expect.any(String), // Don't check exact path as it may have different separators on different platforms
expect.stringContaining("project-mode"),
"utf-8",
)
// Verify the path is correct regardless of separators
const writeCall = (fs.writeFile as jest.Mock).mock.calls[0]
expect(path.normalize(writeCall[0])).toBe(path.normalize(mockRoomodes))
// Verify the content written to .roomodes
expect(roomodesContent).toEqual({

157
src/core/contextProxy.ts Normal file
View file

@ -0,0 +1,157 @@
import * as vscode from "vscode"
import { logger } from "../utils/logging"
import { GLOBAL_STATE_KEYS, SECRET_KEYS } from "../shared/globalState"
export class ContextProxy {
private readonly originalContext: vscode.ExtensionContext
private stateCache: Map<string, any>
private secretCache: Map<string, string | undefined>
constructor(context: vscode.ExtensionContext) {
// Initialize properties first
this.originalContext = context
this.stateCache = new Map()
this.secretCache = new Map()
// Initialize state cache with all defined global state keys
this.initializeStateCache()
// Initialize secret cache with all defined secret keys
this.initializeSecretCache()
logger.debug("ContextProxy created")
}
// Helper method to initialize state cache
private initializeStateCache(): void {
for (const key of GLOBAL_STATE_KEYS) {
try {
const value = this.originalContext.globalState.get(key)
this.stateCache.set(key, value)
} catch (error) {
logger.error(`Error loading global ${key}: ${error instanceof Error ? error.message : String(error)}`)
}
}
}
// Helper method to initialize secret cache
private initializeSecretCache(): void {
for (const key of SECRET_KEYS) {
// Get actual value and update cache when promise resolves
;(this.originalContext.secrets.get(key) as Promise<string | undefined>)
.then((value) => {
this.secretCache.set(key, value)
})
.catch((error: Error) => {
logger.error(`Error loading secret ${key}: ${error.message}`)
})
}
}
get extensionUri(): vscode.Uri {
return this.originalContext.extensionUri
}
get extensionPath(): string {
return this.originalContext.extensionPath
}
get globalStorageUri(): vscode.Uri {
return this.originalContext.globalStorageUri
}
get logUri(): vscode.Uri {
return this.originalContext.logUri
}
get extension(): vscode.Extension<any> | undefined {
return this.originalContext.extension
}
get extensionMode(): vscode.ExtensionMode {
return this.originalContext.extensionMode
}
getGlobalState<T>(key: string): T | undefined
getGlobalState<T>(key: string, defaultValue: T): T
getGlobalState<T>(key: string, defaultValue?: T): T | undefined {
const value = this.stateCache.get(key) as T | undefined
return value !== undefined ? value : (defaultValue as T | undefined)
}
updateGlobalState<T>(key: string, value: T): Thenable<void> {
this.stateCache.set(key, value)
return this.originalContext.globalState.update(key, value)
}
getSecret(key: string): string | undefined {
return this.secretCache.get(key)
}
storeSecret(key: string, value?: string): Thenable<void> {
// Update cache
this.secretCache.set(key, value)
// Write directly to context
if (value === undefined) {
return this.originalContext.secrets.delete(key)
} else {
return this.originalContext.secrets.store(key, value)
}
}
/**
* Set a value in either secrets or global state based on key type.
* If the key is in SECRET_KEYS, it will be stored as a secret.
* If the key is in GLOBAL_STATE_KEYS or unknown, it will be stored in global state.
* @param key The key to set
* @param value The value to set
* @returns A promise that resolves when the operation completes
*/
setValue(key: string, value: any): Thenable<void> {
if (SECRET_KEYS.includes(key as any)) {
return this.storeSecret(key, value)
}
if (GLOBAL_STATE_KEYS.includes(key as any)) {
return this.updateGlobalState(key, value)
}
logger.warn(`Unknown key: ${key}. Storing as global state.`)
return this.updateGlobalState(key, value)
}
/**
* Set multiple values at once. Each key will be routed to either
* secrets or global state based on its type.
* @param values An object containing key-value pairs to set
* @returns A promise that resolves when all operations complete
*/
async setValues(values: Record<string, any>): Promise<void[]> {
const promises: Thenable<void>[] = []
for (const [key, value] of Object.entries(values)) {
promises.push(this.setValue(key, value))
}
return Promise.all(promises)
}
/**
* Resets all global state, secrets, and in-memory caches.
* This clears all data from both the in-memory caches and the VSCode storage.
* @returns A promise that resolves when all reset operations are complete
*/
async resetAllState(): Promise<void> {
// Clear in-memory caches
this.stateCache.clear()
this.secretCache.clear()
// Reset all global state values to undefined
const stateResetPromises = GLOBAL_STATE_KEYS.map((key) =>
this.originalContext.globalState.update(key, undefined),
)
// Delete all secrets
const secretResetPromises = SECRET_KEYS.map((key) => this.originalContext.secrets.delete(key))
// Wait for all reset operations to complete
await Promise.all([...stateResetPromises, ...secretResetPromises])
this.initializeStateCache()
this.initializeSecretCache()
}
}

View file

@ -2,6 +2,7 @@ import type { DiffStrategy } from "./types"
import { UnifiedDiffStrategy } from "./strategies/unified"
import { SearchReplaceDiffStrategy } from "./strategies/search-replace"
import { NewUnifiedDiffStrategy } from "./strategies/new-unified"
import { MultiSearchReplaceDiffStrategy } from "./strategies/multi-search-replace"
/**
* Get the appropriate diff strategy for the given model
* @param model The name of the model being used (e.g., 'gpt-4', 'claude-3-opus')
@ -11,11 +12,17 @@ export function getDiffStrategy(
model: string,
fuzzyMatchThreshold?: number,
experimentalDiffStrategy: boolean = false,
multiSearchReplaceDiffStrategy: boolean = false,
): DiffStrategy {
if (experimentalDiffStrategy) {
return new NewUnifiedDiffStrategy(fuzzyMatchThreshold)
}
return new SearchReplaceDiffStrategy(fuzzyMatchThreshold)
if (multiSearchReplaceDiffStrategy) {
return new MultiSearchReplaceDiffStrategy(fuzzyMatchThreshold)
} else {
return new SearchReplaceDiffStrategy(fuzzyMatchThreshold)
}
}
export type { DiffStrategy }

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,390 @@
import { DiffStrategy, DiffResult } from "../types"
import { addLineNumbers, everyLineHasLineNumbers, stripLineNumbers } from "../../../integrations/misc/extract-text"
import { distance } from "fastest-levenshtein"
import { ToolProgressStatus } from "../../../shared/ExtensionMessage"
import { ToolUse } from "../../assistant-message"
const BUFFER_LINES = 40 // Number of extra context lines to show before and after matches
function getSimilarity(original: string, search: string): number {
if (search === "") {
return 1
}
// Normalize strings by removing extra whitespace but preserve case
const normalizeStr = (str: string) => str.replace(/\s+/g, " ").trim()
const normalizedOriginal = normalizeStr(original)
const normalizedSearch = normalizeStr(search)
if (normalizedOriginal === normalizedSearch) {
return 1
}
// Calculate Levenshtein distance using fastest-levenshtein's distance function
const dist = distance(normalizedOriginal, normalizedSearch)
// Calculate similarity ratio (0 to 1, where 1 is an exact match)
const maxLength = Math.max(normalizedOriginal.length, normalizedSearch.length)
return 1 - dist / maxLength
}
export class MultiSearchReplaceDiffStrategy implements DiffStrategy {
private fuzzyThreshold: number
private bufferLines: number
constructor(fuzzyThreshold?: number, bufferLines?: number) {
// Use provided threshold or default to exact matching (1.0)
// Note: fuzzyThreshold is inverted in UI (0% = 1.0, 10% = 0.9)
// so we use it directly here
this.fuzzyThreshold = fuzzyThreshold ?? 1.0
this.bufferLines = bufferLines ?? BUFFER_LINES
}
getToolDescription(args: { cwd: string; toolOptions?: { [key: string]: string } }): string {
return `## apply_diff
Description: Request to replace existing code using a search and replace block.
This tool allows for precise, surgical replaces to files by specifying exactly what content to search for and what to replace it with.
The tool will maintain proper indentation and formatting while making changes.
Only a single operation is allowed per tool use.
The SEARCH section must exactly match existing content including whitespace and indentation.
If you're not confident in the exact content to search for, use the read_file tool first to get the exact content.
When applying the diffs, be extra careful to remember to change any closing brackets or other syntax that may be affected by the diff farther down in the file.
ALWAYS make as many changes in a single 'apply_diff' request as possible using multiple SEARCH/REPLACE blocks
Parameters:
- path: (required) The path of the file to modify (relative to the current working directory ${args.cwd})
- diff: (required) The search/replace block defining the changes.
Diff format:
\`\`\`
<<<<<<< SEARCH
:start_line: (required) The line number of original content where the search block starts.
:end_line: (required) The line number of original content where the search block ends.
-------
[exact content to find including whitespace]
=======
[new content to replace with]
>>>>>>> REPLACE
\`\`\`
Example:
Original file:
\`\`\`
1 | def calculate_total(items):
2 | total = 0
3 | for item in items:
4 | total += item
5 | return total
\`\`\`
Search/Replace content:
\`\`\`
<<<<<<< SEARCH
:start_line:1
:end_line:5
-------
def calculate_total(items):
total = 0
for item in items:
total += item
return total
=======
def calculate_total(items):
"""Calculate total with 10% markup"""
return sum(item * 1.1 for item in items)
>>>>>>> REPLACE
\`\`\`
Search/Replace content with multi edits:
\`\`\`
<<<<<<< SEARCH
:start_line:1
:end_line:2
-------
def calculate_sum(items):
sum = 0
=======
def calculate_sum(items):
sum = 0
>>>>>>> REPLACE
<<<<<<< SEARCH
:start_line:4
:end_line:5
-------
total += item
return total
=======
sum += item
return sum
>>>>>>> REPLACE
\`\`\`
Usage:
<apply_diff>
<path>File path here</path>
<diff>
Your search/replace content here
You can use multi search/replace block in one diff block, but make sure to include the line numbers for each block.
Only use a single line of '=======' between search and replacement content, because multiple '=======' will corrupt the file.
</diff>
</apply_diff>`
}
async applyDiff(
originalContent: string,
diffContent: string,
_paramStartLine?: number,
_paramEndLine?: number,
): Promise<DiffResult> {
let matches = [
...diffContent.matchAll(
/<<<<<<< SEARCH\n(:start_line:\s*(\d+)\n){0,1}(:end_line:\s*(\d+)\n){0,1}(-------\n){0,1}([\s\S]*?)\n?=======\n([\s\S]*?)\n?>>>>>>> REPLACE/g,
),
]
if (matches.length === 0) {
return {
success: false,
error: `Invalid diff format - missing required sections\n\nDebug Info:\n- Expected Format: <<<<<<< SEARCH\\n:start_line: start line\\n:end_line: end line\\n-------\\n[search content]\\n=======\\n[replace content]\\n>>>>>>> REPLACE\n- Tip: Make sure to include start_line/end_line/SEARCH/REPLACE sections with correct markers`,
}
}
// Detect line ending from original content
const lineEnding = originalContent.includes("\r\n") ? "\r\n" : "\n"
let resultLines = originalContent.split(/\r?\n/)
let delta = 0
let diffResults: DiffResult[] = []
let appliedCount = 0
const replacements = matches
.map((match) => ({
startLine: Number(match[2] ?? 0),
endLine: Number(match[4] ?? resultLines.length),
searchContent: match[6],
replaceContent: match[7],
}))
.sort((a, b) => a.startLine - b.startLine)
for (let { searchContent, replaceContent, startLine, endLine } of replacements) {
startLine += startLine === 0 ? 0 : delta
endLine += delta
// Strip line numbers from search and replace content if every line starts with a line number
if (everyLineHasLineNumbers(searchContent) && everyLineHasLineNumbers(replaceContent)) {
searchContent = stripLineNumbers(searchContent)
replaceContent = stripLineNumbers(replaceContent)
}
// Split content into lines, handling both \n and \r\n
const searchLines = searchContent === "" ? [] : searchContent.split(/\r?\n/)
const replaceLines = replaceContent === "" ? [] : replaceContent.split(/\r?\n/)
// Validate that empty search requires start line
if (searchLines.length === 0 && !startLine) {
diffResults.push({
success: false,
error: `Empty search content requires start_line to be specified\n\nDebug Info:\n- Empty search content is only valid for insertions at a specific line\n- For insertions, specify the line number where content should be inserted`,
})
continue
}
// Validate that empty search requires same start and end line
if (searchLines.length === 0 && startLine && endLine && startLine !== endLine) {
diffResults.push({
success: false,
error: `Empty search content requires start_line and end_line to be the same (got ${startLine}-${endLine})\n\nDebug Info:\n- Empty search content is only valid for insertions at a specific line\n- For insertions, use the same line number for both start_line and end_line`,
})
continue
}
// Initialize search variables
let matchIndex = -1
let bestMatchScore = 0
let bestMatchContent = ""
const searchChunk = searchLines.join("\n")
// Determine search bounds
let searchStartIndex = 0
let searchEndIndex = resultLines.length
// Validate and handle line range if provided
if (startLine && endLine) {
// Convert to 0-based index
const exactStartIndex = startLine - 1
const exactEndIndex = endLine - 1
if (exactStartIndex < 0 || exactEndIndex > resultLines.length || exactStartIndex > exactEndIndex) {
diffResults.push({
success: false,
error: `Line range ${startLine}-${endLine} is invalid (file has ${resultLines.length} lines)\n\nDebug Info:\n- Requested Range: lines ${startLine}-${endLine}\n- File Bounds: lines 1-${resultLines.length}`,
})
continue
}
// Try exact match first
const originalChunk = resultLines.slice(exactStartIndex, exactEndIndex + 1).join("\n")
const similarity = getSimilarity(originalChunk, searchChunk)
if (similarity >= this.fuzzyThreshold) {
matchIndex = exactStartIndex
bestMatchScore = similarity
bestMatchContent = originalChunk
} else {
// Set bounds for buffered search
searchStartIndex = Math.max(0, startLine - (this.bufferLines + 1))
searchEndIndex = Math.min(resultLines.length, endLine + this.bufferLines)
}
}
// If no match found yet, try middle-out search within bounds
if (matchIndex === -1) {
const midPoint = Math.floor((searchStartIndex + searchEndIndex) / 2)
let leftIndex = midPoint
let rightIndex = midPoint + 1
// Search outward from the middle within bounds
while (leftIndex >= searchStartIndex || rightIndex <= searchEndIndex - searchLines.length) {
// Check left side if still in range
if (leftIndex >= searchStartIndex) {
const originalChunk = resultLines.slice(leftIndex, leftIndex + searchLines.length).join("\n")
const similarity = getSimilarity(originalChunk, searchChunk)
if (similarity > bestMatchScore) {
bestMatchScore = similarity
matchIndex = leftIndex
bestMatchContent = originalChunk
}
leftIndex--
}
// Check right side if still in range
if (rightIndex <= searchEndIndex - searchLines.length) {
const originalChunk = resultLines.slice(rightIndex, rightIndex + searchLines.length).join("\n")
const similarity = getSimilarity(originalChunk, searchChunk)
if (similarity > bestMatchScore) {
bestMatchScore = similarity
matchIndex = rightIndex
bestMatchContent = originalChunk
}
rightIndex++
}
}
}
// Require similarity to meet threshold
if (matchIndex === -1 || bestMatchScore < this.fuzzyThreshold) {
const searchChunk = searchLines.join("\n")
const originalContentSection =
startLine !== undefined && endLine !== undefined
? `\n\nOriginal Content:\n${addLineNumbers(
resultLines
.slice(
Math.max(0, startLine - 1 - this.bufferLines),
Math.min(resultLines.length, endLine + this.bufferLines),
)
.join("\n"),
Math.max(1, startLine - this.bufferLines),
)}`
: `\n\nOriginal Content:\n${addLineNumbers(resultLines.join("\n"))}`
const bestMatchSection = bestMatchContent
? `\n\nBest Match Found:\n${addLineNumbers(bestMatchContent, matchIndex + 1)}`
: `\n\nBest Match Found:\n(no match)`
const lineRange =
startLine || endLine
? ` at ${startLine ? `start: ${startLine}` : "start"} to ${endLine ? `end: ${endLine}` : "end"}`
: ""
diffResults.push({
success: false,
error: `No sufficiently similar match found${lineRange} (${Math.floor(bestMatchScore * 100)}% similar, needs ${Math.floor(this.fuzzyThreshold * 100)}%)\n\nDebug Info:\n- Similarity Score: ${Math.floor(bestMatchScore * 100)}%\n- Required Threshold: ${Math.floor(this.fuzzyThreshold * 100)}%\n- Search Range: ${startLine && endLine ? `lines ${startLine}-${endLine}` : "start to end"}\n- Tip: Use read_file to get the latest content of the file before attempting the diff again, as the file content may have changed\n\nSearch Content:\n${searchChunk}${bestMatchSection}${originalContentSection}`,
})
continue
}
// Get the matched lines from the original content
const matchedLines = resultLines.slice(matchIndex, matchIndex + searchLines.length)
// Get the exact indentation (preserving tabs/spaces) of each line
const originalIndents = matchedLines.map((line) => {
const match = line.match(/^[\t ]*/)
return match ? match[0] : ""
})
// Get the exact indentation of each line in the search block
const searchIndents = searchLines.map((line) => {
const match = line.match(/^[\t ]*/)
return match ? match[0] : ""
})
// Apply the replacement while preserving exact indentation
const indentedReplaceLines = replaceLines.map((line, i) => {
// Get the matched line's exact indentation
const matchedIndent = originalIndents[0] || ""
// Get the current line's indentation relative to the search content
const currentIndentMatch = line.match(/^[\t ]*/)
const currentIndent = currentIndentMatch ? currentIndentMatch[0] : ""
const searchBaseIndent = searchIndents[0] || ""
// Calculate the relative indentation level
const searchBaseLevel = searchBaseIndent.length
const currentLevel = currentIndent.length
const relativeLevel = currentLevel - searchBaseLevel
// If relative level is negative, remove indentation from matched indent
// If positive, add to matched indent
const finalIndent =
relativeLevel < 0
? matchedIndent.slice(0, Math.max(0, matchedIndent.length + relativeLevel))
: matchedIndent + currentIndent.slice(searchBaseLevel)
return finalIndent + line.trim()
})
// Construct the final content
const beforeMatch = resultLines.slice(0, matchIndex)
const afterMatch = resultLines.slice(matchIndex + searchLines.length)
resultLines = [...beforeMatch, ...indentedReplaceLines, ...afterMatch]
delta = delta - matchedLines.length + replaceLines.length
appliedCount++
}
const finalContent = resultLines.join(lineEnding)
if (appliedCount === 0) {
return {
success: false,
failParts: diffResults,
}
}
return {
success: true,
content: finalContent,
failParts: diffResults,
}
}
getProgressStatus(toolUse: ToolUse, result?: DiffResult): ToolProgressStatus {
const diffContent = toolUse.params.diff
if (diffContent) {
const icon = "diff-multiple"
const searchBlockCount = (diffContent.match(/SEARCH/g) || []).length
if (toolUse.partial) {
if (diffContent.length < 1000 || (diffContent.length / 50) % 10 === 0) {
return { icon, text: `${searchBlockCount}` }
}
} else if (result) {
if (result.failParts?.length) {
return {
icon,
text: `${searchBlockCount - result.failParts.length}/${searchBlockCount}`,
}
} else {
return { icon, text: `${searchBlockCount}` }
}
}
}
return {}
}
}

View file

@ -2,11 +2,14 @@
* Interface for implementing different diff strategies
*/
import { ToolProgressStatus } from "../../shared/ExtensionMessage"
import { ToolUse } from "../assistant-message"
export type DiffResult =
| { success: true; content: string }
| {
| { success: true; content: string; failParts?: DiffResult[] }
| ({
success: false
error: string
error?: string
details?: {
similarity?: number
threshold?: number
@ -14,7 +17,8 @@ export type DiffResult =
searchContent?: string
bestMatch?: string
}
}
failParts?: DiffResult[]
} & ({ error: string } | { failParts: DiffResult[] }))
export interface DiffStrategy {
/**
@ -33,4 +37,6 @@ export interface DiffStrategy {
* @returns A DiffResult object containing either the successful result or error details
*/
applyDiff(originalContent: string, diffContent: string, startLine?: number, endLine?: number): Promise<DiffResult>
getProgressStatus?(toolUse: ToolUse, result?: any): ToolProgressStatus
}

View file

@ -0,0 +1,201 @@
import path from "path"
import { fileExistsAtPath } from "../../utils/fs"
import fs from "fs/promises"
import ignore, { Ignore } from "ignore"
import * as vscode from "vscode"
export const LOCK_TEXT_SYMBOL = "\u{1F512}"
/**
* Controls LLM access to files by enforcing ignore patterns.
* Designed to be instantiated once in Cline.ts and passed to file manipulation services.
* Uses the 'ignore' library to support standard .gitignore syntax in .rooignore files.
*/
export class RooIgnoreController {
private cwd: string
private ignoreInstance: Ignore
private disposables: vscode.Disposable[] = []
rooIgnoreContent: string | undefined
constructor(cwd: string) {
this.cwd = cwd
this.ignoreInstance = ignore()
this.rooIgnoreContent = undefined
// Set up file watcher for .rooignore
this.setupFileWatcher()
}
/**
* Initialize the controller by loading custom patterns
* Must be called after construction and before using the controller
*/
async initialize(): Promise<void> {
await this.loadRooIgnore()
}
/**
* Set up the file watcher for .rooignore changes
*/
private setupFileWatcher(): void {
const rooignorePattern = new vscode.RelativePattern(this.cwd, ".rooignore")
const fileWatcher = vscode.workspace.createFileSystemWatcher(rooignorePattern)
// Watch for changes and updates
this.disposables.push(
fileWatcher.onDidChange(() => {
this.loadRooIgnore()
}),
fileWatcher.onDidCreate(() => {
this.loadRooIgnore()
}),
fileWatcher.onDidDelete(() => {
this.loadRooIgnore()
}),
)
// Add fileWatcher itself to disposables
this.disposables.push(fileWatcher)
}
/**
* Load custom patterns from .rooignore if it exists
*/
private async loadRooIgnore(): Promise<void> {
try {
// Reset ignore instance to prevent duplicate patterns
this.ignoreInstance = ignore()
const ignorePath = path.join(this.cwd, ".rooignore")
if (await fileExistsAtPath(ignorePath)) {
const content = await fs.readFile(ignorePath, "utf8")
this.rooIgnoreContent = content
this.ignoreInstance.add(content)
this.ignoreInstance.add(".rooignore")
} else {
this.rooIgnoreContent = undefined
}
} catch (error) {
// Should never happen: reading file failed even though it exists
console.error("Unexpected error loading .rooignore:", error)
}
}
/**
* Check if a file should be accessible to the LLM
* @param filePath - Path to check (relative to cwd)
* @returns true if file is accessible, false if ignored
*/
validateAccess(filePath: string): boolean {
// Always allow access if .rooignore does not exist
if (!this.rooIgnoreContent) {
return true
}
try {
// Normalize path to be relative to cwd and use forward slashes
const absolutePath = path.resolve(this.cwd, filePath)
const relativePath = path.relative(this.cwd, absolutePath).toPosix()
// Ignore expects paths to be path.relative()'d
return !this.ignoreInstance.ignores(relativePath)
} catch (error) {
// console.error(`Error validating access for ${filePath}:`, error)
// Ignore is designed to work with relative file paths, so will throw error for paths outside cwd. We are allowing access to all files outside cwd.
return true
}
}
/**
* Check if a terminal command should be allowed to execute based on file access patterns
* @param command - Terminal command to validate
* @returns path of file that is being accessed if it is being accessed, undefined if command is allowed
*/
validateCommand(command: string): string | undefined {
// Always allow if no .rooignore exists
if (!this.rooIgnoreContent) {
return undefined
}
// Split command into parts and get the base command
const parts = command.trim().split(/\s+/)
const baseCommand = parts[0].toLowerCase()
// Commands that read file contents
const fileReadingCommands = [
// Unix commands
"cat",
"less",
"more",
"head",
"tail",
"grep",
"awk",
"sed",
// PowerShell commands and aliases
"get-content",
"gc",
"type",
"select-string",
"sls",
]
if (fileReadingCommands.includes(baseCommand)) {
// Check each argument that could be a file path
for (let i = 1; i < parts.length; i++) {
const arg = parts[i]
// Skip command flags/options (both Unix and PowerShell style)
if (arg.startsWith("-") || arg.startsWith("/")) {
continue
}
// Ignore PowerShell parameter names
if (arg.includes(":")) {
continue
}
// Validate file access
if (!this.validateAccess(arg)) {
return arg
}
}
}
return undefined
}
/**
* Filter an array of paths, removing those that should be ignored
* @param paths - Array of paths to filter (relative to cwd)
* @returns Array of allowed paths
*/
filterPaths(paths: string[]): string[] {
try {
return paths
.map((p) => ({
path: p,
allowed: this.validateAccess(p),
}))
.filter((x) => x.allowed)
.map((x) => x.path)
} catch (error) {
console.error("Error filtering paths:", error)
return [] // Fail closed for security
}
}
/**
* Clean up resources when the controller is no longer needed
*/
dispose(): void {
this.disposables.forEach((d) => d.dispose())
this.disposables = []
}
/**
* Get formatted instructions about the .rooignore file for the LLM
* @returns Formatted instructions or undefined if .rooignore doesn't exist
*/
getInstructions(): string | undefined {
if (!this.rooIgnoreContent) {
return undefined
}
return `# .rooignore\n\n(The following is provided by a root-level .rooignore file where the user has specified files and directories that should not be accessed. When using list_files, you'll notice a ${LOCK_TEXT_SYMBOL} next to files that are blocked. Attempting to access the file's contents e.g. through read_file will result in an error.)\n\n${this.rooIgnoreContent}\n.rooignore`
}
}

View file

@ -0,0 +1,38 @@
export const LOCK_TEXT_SYMBOL = "\u{1F512}"
export class RooIgnoreController {
rooIgnoreContent: string | undefined = undefined
constructor(cwd: string) {
// No-op constructor
}
async initialize(): Promise<void> {
// No-op initialization
return Promise.resolve()
}
validateAccess(filePath: string): boolean {
// Default implementation: allow all access
return true
}
validateCommand(command: string): string | undefined {
// Default implementation: allow all commands
return undefined
}
filterPaths(paths: string[]): string[] {
// Default implementation: allow all paths
return paths
}
dispose(): void {
// No-op dispose
}
getInstructions(): string | undefined {
// Default implementation: no instructions
return undefined
}
}

View file

@ -0,0 +1,323 @@
// npx jest src/core/ignore/__tests__/RooIgnoreController.security.test.ts
import { RooIgnoreController } from "../RooIgnoreController"
import * as path from "path"
import * as fs from "fs/promises"
import { fileExistsAtPath } from "../../../utils/fs"
import * as vscode from "vscode"
// Mock dependencies
jest.mock("fs/promises")
jest.mock("../../../utils/fs")
jest.mock("vscode", () => {
const mockDisposable = { dispose: jest.fn() }
return {
workspace: {
createFileSystemWatcher: jest.fn(() => ({
onDidCreate: jest.fn(() => mockDisposable),
onDidChange: jest.fn(() => mockDisposable),
onDidDelete: jest.fn(() => mockDisposable),
dispose: jest.fn(),
})),
},
RelativePattern: jest.fn().mockImplementation((base, pattern) => ({
base,
pattern,
})),
}
})
describe("RooIgnoreController Security Tests", () => {
const TEST_CWD = "/test/path"
let controller: RooIgnoreController
let mockFileExists: jest.MockedFunction<typeof fileExistsAtPath>
let mockReadFile: jest.MockedFunction<typeof fs.readFile>
beforeEach(async () => {
// Reset mocks
jest.clearAllMocks()
// Setup mocks
mockFileExists = fileExistsAtPath as jest.MockedFunction<typeof fileExistsAtPath>
mockReadFile = fs.readFile as jest.MockedFunction<typeof fs.readFile>
// By default, setup .rooignore to exist with some patterns
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**\n*.log\nprivate/")
// Create and initialize controller
controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
})
describe("validateCommand security", () => {
/**
* Tests Unix file reading commands with various arguments
*/
it("should block Unix file reading commands accessing ignored files", () => {
// Test simple cat command
expect(controller.validateCommand("cat node_modules/package.json")).toBe("node_modules/package.json")
// Test with command options
expect(controller.validateCommand("cat -n .git/config")).toBe(".git/config")
// Directory paths don't match in the implementation since it checks for exact files
// Instead, use a file path
expect(controller.validateCommand("grep -r 'password' secrets/keys.json")).toBe("secrets/keys.json")
// Multiple files with flags - first match is returned
expect(controller.validateCommand("head -n 5 app.log secrets/keys.json")).toBe("app.log")
// Commands with pipes
expect(controller.validateCommand("cat secrets/creds.json | grep password")).toBe("secrets/creds.json")
// The implementation doesn't handle quoted paths as expected
// Let's test with simple paths instead
expect(controller.validateCommand("less private/notes.txt")).toBe("private/notes.txt")
expect(controller.validateCommand("more private/data.csv")).toBe("private/data.csv")
})
/**
* Tests PowerShell file reading commands
*/
it("should block PowerShell file reading commands accessing ignored files", () => {
// Simple Get-Content
expect(controller.validateCommand("Get-Content node_modules/package.json")).toBe(
"node_modules/package.json",
)
// With parameters
expect(controller.validateCommand("Get-Content -Path .git/config -Raw")).toBe(".git/config")
// With parameter aliases
expect(controller.validateCommand("gc secrets/keys.json")).toBe("secrets/keys.json")
// Select-String (grep equivalent)
expect(controller.validateCommand("Select-String -Pattern 'password' -Path private/config.json")).toBe(
"private/config.json",
)
expect(controller.validateCommand("sls 'api-key' app.log")).toBe("app.log")
// Parameter form with colons is skipped by the implementation - replace with standard form
expect(controller.validateCommand("Get-Content -Path node_modules/package.json")).toBe(
"node_modules/package.json",
)
})
/**
* Tests non-file reading commands
*/
it("should allow non-file reading commands", () => {
// Directory commands
expect(controller.validateCommand("ls -la node_modules")).toBeUndefined()
expect(controller.validateCommand("dir .git")).toBeUndefined()
expect(controller.validateCommand("cd secrets")).toBeUndefined()
// Other system commands
expect(controller.validateCommand("ps -ef | grep node")).toBeUndefined()
expect(controller.validateCommand("npm install")).toBeUndefined()
expect(controller.validateCommand("git status")).toBeUndefined()
})
/**
* Tests command handling with special characters and spaces
*/
it("should handle complex commands with special characters", () => {
// The implementation doesn't handle quoted paths as expected
// Testing with unquoted paths instead
expect(controller.validateCommand("cat private/file-simple.txt")).toBe("private/file-simple.txt")
expect(controller.validateCommand("grep pattern secrets/file-with-dashes.json")).toBe(
"secrets/file-with-dashes.json",
)
expect(controller.validateCommand("less private/file_with_underscores.md")).toBe(
"private/file_with_underscores.md",
)
// Special characters - using simple paths without escapes since the implementation doesn't handle escaped spaces as expected
expect(controller.validateCommand("cat private/file.txt")).toBe("private/file.txt")
})
})
describe("Path traversal protection", () => {
/**
* Tests protection against path traversal attacks
*/
it("should handle path traversal attempts", () => {
// Setup complex ignore pattern
mockReadFile.mockResolvedValue("secrets/**")
// Reinitialize controller
return controller.initialize().then(() => {
// Test simple path
expect(controller.validateAccess("secrets/keys.json")).toBe(false)
// Attempt simple path traversal
expect(controller.validateAccess("secrets/../secrets/keys.json")).toBe(false)
// More complex traversal
expect(controller.validateAccess("public/../secrets/keys.json")).toBe(false)
// Deep traversal
expect(controller.validateAccess("public/css/../../secrets/keys.json")).toBe(false)
// Traversal with normalized path
expect(controller.validateAccess(path.normalize("public/../secrets/keys.json"))).toBe(false)
// Allowed files shouldn't be affected by traversal protection
expect(controller.validateAccess("public/css/../../public/app.js")).toBe(true)
})
})
/**
* Tests absolute path handling
*/
it("should handle absolute paths correctly", () => {
// Absolute path to ignored file within cwd
const absolutePathToIgnored = path.join(TEST_CWD, "secrets/keys.json")
expect(controller.validateAccess(absolutePathToIgnored)).toBe(false)
// Absolute path to allowed file within cwd
const absolutePathToAllowed = path.join(TEST_CWD, "src/app.js")
expect(controller.validateAccess(absolutePathToAllowed)).toBe(true)
// Absolute path outside cwd should be allowed
expect(controller.validateAccess("/etc/hosts")).toBe(true)
expect(controller.validateAccess("/var/log/system.log")).toBe(true)
})
/**
* Tests that paths outside cwd are allowed
*/
it("should allow paths outside the current working directory", () => {
// Paths outside cwd should be allowed
expect(controller.validateAccess("../outside-project/file.txt")).toBe(true)
expect(controller.validateAccess("../../other-project/secrets/keys.json")).toBe(true)
// Edge case: path that would be ignored if inside cwd
expect(controller.validateAccess("/other/path/secrets/keys.json")).toBe(true)
})
})
describe("Comprehensive path handling", () => {
/**
* Tests combinations of paths and patterns
*/
it("should correctly apply complex patterns to various paths", async () => {
// Setup complex patterns - but without negation patterns since they're not reliably handled
mockReadFile.mockResolvedValue(`
# Node modules and logs
node_modules
*.log
# Version control
.git
.svn
# Secrets and config
config/secrets/**
**/*secret*
**/password*.*
# Build artifacts
dist/
build/
# Comments and empty lines should be ignored
`)
// Reinitialize controller
await controller.initialize()
// Test standard ignored paths
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
expect(controller.validateAccess("app.log")).toBe(false)
expect(controller.validateAccess(".git/config")).toBe(false)
// Test wildcards and double wildcards
expect(controller.validateAccess("config/secrets/api-keys.json")).toBe(false)
expect(controller.validateAccess("src/config/secret-keys.js")).toBe(false)
expect(controller.validateAccess("lib/utils/password-manager.ts")).toBe(false)
// Test build artifacts
expect(controller.validateAccess("dist/main.js")).toBe(false)
expect(controller.validateAccess("build/index.html")).toBe(false)
// Test paths that should be allowed
expect(controller.validateAccess("src/app.js")).toBe(true)
expect(controller.validateAccess("README.md")).toBe(true)
// Test allowed paths
expect(controller.validateAccess("src/app.js")).toBe(true)
expect(controller.validateAccess("README.md")).toBe(true)
})
/**
* Tests non-standard file paths
*/
it("should handle unusual file paths", () => {
expect(controller.validateAccess(".node_modules_temp/file.js")).toBe(true) // Doesn't match node_modules
expect(controller.validateAccess("node_modules.bak/file.js")).toBe(true) // Doesn't match node_modules
expect(controller.validateAccess("not_secrets/file.json")).toBe(true) // Doesn't match secrets
// Files with dots
expect(controller.validateAccess("src/file.with.multiple.dots.js")).toBe(true)
// Files with no extension
expect(controller.validateAccess("bin/executable")).toBe(true)
// Hidden files
expect(controller.validateAccess(".env")).toBe(true) // Not ignored by default
})
})
describe("filterPaths security", () => {
/**
* Tests filtering paths for security
*/
it("should correctly filter mixed paths", () => {
const paths = [
"src/app.js", // allowed
"node_modules/package.json", // ignored
"README.md", // allowed
"secrets/keys.json", // ignored
".git/config", // ignored
"app.log", // ignored
"test/test.js", // allowed
]
const filtered = controller.filterPaths(paths)
// Should only contain allowed paths
expect(filtered).toEqual(["src/app.js", "README.md", "test/test.js"])
// Length should match allowed files
expect(filtered.length).toBe(3)
})
/**
* Tests error handling in filterPaths
*/
it("should fail closed (securely) when errors occur", () => {
// Mock validateAccess to throw error
jest.spyOn(controller, "validateAccess").mockImplementation(() => {
throw new Error("Test error")
})
// Spy on console.error
const consoleSpy = jest.spyOn(console, "error").mockImplementation()
// Even with mix of allowed/ignored paths, should return empty array on error
const filtered = controller.filterPaths(["src/app.js", "node_modules/package.json"])
// Should fail closed (return empty array)
expect(filtered).toEqual([])
// Should log error
expect(consoleSpy).toHaveBeenCalledWith("Error filtering paths:", expect.any(Error))
// Clean up
consoleSpy.mockRestore()
})
})
})

View file

@ -0,0 +1,503 @@
// npx jest src/core/ignore/__tests__/RooIgnoreController.test.ts
import { RooIgnoreController, LOCK_TEXT_SYMBOL } from "../RooIgnoreController"
import * as vscode from "vscode"
import * as path from "path"
import * as fs from "fs/promises"
import { fileExistsAtPath } from "../../../utils/fs"
// Mock dependencies
jest.mock("fs/promises")
jest.mock("../../../utils/fs")
// Mock vscode
jest.mock("vscode", () => {
const mockDisposable = { dispose: jest.fn() }
const mockEventEmitter = {
event: jest.fn(),
fire: jest.fn(),
}
return {
workspace: {
createFileSystemWatcher: jest.fn(() => ({
onDidCreate: jest.fn(() => mockDisposable),
onDidChange: jest.fn(() => mockDisposable),
onDidDelete: jest.fn(() => mockDisposable),
dispose: jest.fn(),
})),
},
RelativePattern: jest.fn().mockImplementation((base, pattern) => ({
base,
pattern,
})),
EventEmitter: jest.fn().mockImplementation(() => mockEventEmitter),
Disposable: {
from: jest.fn(),
},
}
})
describe("RooIgnoreController", () => {
const TEST_CWD = "/test/path"
let controller: RooIgnoreController
let mockFileExists: jest.MockedFunction<typeof fileExistsAtPath>
let mockReadFile: jest.MockedFunction<typeof fs.readFile>
let mockWatcher: any
beforeEach(() => {
// Reset mocks
jest.clearAllMocks()
// Setup mock file watcher
mockWatcher = {
onDidCreate: jest.fn().mockReturnValue({ dispose: jest.fn() }),
onDidChange: jest.fn().mockReturnValue({ dispose: jest.fn() }),
onDidDelete: jest.fn().mockReturnValue({ dispose: jest.fn() }),
dispose: jest.fn(),
}
// @ts-expect-error - Mocking
vscode.workspace.createFileSystemWatcher.mockReturnValue(mockWatcher)
// Setup fs mocks
mockFileExists = fileExistsAtPath as jest.MockedFunction<typeof fileExistsAtPath>
mockReadFile = fs.readFile as jest.MockedFunction<typeof fs.readFile>
// Create controller
controller = new RooIgnoreController(TEST_CWD)
})
describe("initialization", () => {
/**
* Tests the controller initialization when .rooignore exists
*/
it("should load .rooignore patterns on initialization when file exists", async () => {
// Setup mocks to simulate existing .rooignore file
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets.json")
// Initialize controller
await controller.initialize()
// Verify file was checked and read
expect(mockFileExists).toHaveBeenCalledWith(path.join(TEST_CWD, ".rooignore"))
expect(mockReadFile).toHaveBeenCalledWith(path.join(TEST_CWD, ".rooignore"), "utf8")
// Verify content was stored
expect(controller.rooIgnoreContent).toBe("node_modules\n.git\nsecrets.json")
// Test that ignore patterns were applied
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
expect(controller.validateAccess("src/app.ts")).toBe(true)
expect(controller.validateAccess(".git/config")).toBe(false)
expect(controller.validateAccess("secrets.json")).toBe(false)
})
/**
* Tests the controller behavior when .rooignore doesn't exist
*/
it("should allow all access when .rooignore doesn't exist", async () => {
// Setup mocks to simulate missing .rooignore file
mockFileExists.mockResolvedValue(false)
// Initialize controller
await controller.initialize()
// Verify no content was stored
expect(controller.rooIgnoreContent).toBeUndefined()
// All files should be accessible
expect(controller.validateAccess("node_modules/package.json")).toBe(true)
expect(controller.validateAccess("secrets.json")).toBe(true)
})
/**
* Tests the file watcher setup
*/
it("should set up file watcher for .rooignore changes", async () => {
// Check that watcher was created with correct pattern
expect(vscode.workspace.createFileSystemWatcher).toHaveBeenCalledWith(
expect.objectContaining({
base: TEST_CWD,
pattern: ".rooignore",
}),
)
// Verify event handlers were registered
expect(mockWatcher.onDidCreate).toHaveBeenCalled()
expect(mockWatcher.onDidChange).toHaveBeenCalled()
expect(mockWatcher.onDidDelete).toHaveBeenCalled()
})
/**
* Tests error handling during initialization
*/
it("should handle errors when loading .rooignore", async () => {
// Setup mocks to simulate error
mockFileExists.mockResolvedValue(true)
mockReadFile.mockRejectedValue(new Error("Test file read error"))
// Spy on console.error
const consoleSpy = jest.spyOn(console, "error").mockImplementation()
// Initialize controller - shouldn't throw
await controller.initialize()
// Verify error was logged
expect(consoleSpy).toHaveBeenCalledWith("Unexpected error loading .rooignore:", expect.any(Error))
// Cleanup
consoleSpy.mockRestore()
})
})
describe("validateAccess", () => {
beforeEach(async () => {
// Setup .rooignore content
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**\n*.log")
await controller.initialize()
})
/**
* Tests basic path validation
*/
it("should correctly validate file access based on ignore patterns", () => {
// Test different path patterns
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
expect(controller.validateAccess("node_modules")).toBe(false)
expect(controller.validateAccess("src/node_modules/file.js")).toBe(false)
expect(controller.validateAccess(".git/HEAD")).toBe(false)
expect(controller.validateAccess("secrets/api-keys.json")).toBe(false)
expect(controller.validateAccess("logs/app.log")).toBe(false)
// These should be allowed
expect(controller.validateAccess("src/app.ts")).toBe(true)
expect(controller.validateAccess("package.json")).toBe(true)
expect(controller.validateAccess("secret-file.json")).toBe(true)
})
/**
* Tests handling of absolute paths
*/
it("should handle absolute paths correctly", () => {
// Test with absolute paths
const absolutePath = path.join(TEST_CWD, "node_modules/package.json")
expect(controller.validateAccess(absolutePath)).toBe(false)
const allowedAbsolutePath = path.join(TEST_CWD, "src/app.ts")
expect(controller.validateAccess(allowedAbsolutePath)).toBe(true)
})
/**
* Tests handling of paths outside cwd
*/
it("should allow access to paths outside cwd", () => {
// Path traversal outside cwd
expect(controller.validateAccess("../outside-project/file.txt")).toBe(true)
// Completely different path
expect(controller.validateAccess("/etc/hosts")).toBe(true)
})
/**
* Tests the default behavior when no .rooignore exists
*/
it("should allow all access when no .rooignore content", async () => {
// Create a new controller with no .rooignore
mockFileExists.mockResolvedValue(false)
const emptyController = new RooIgnoreController(TEST_CWD)
await emptyController.initialize()
// All paths should be allowed
expect(emptyController.validateAccess("node_modules/package.json")).toBe(true)
expect(emptyController.validateAccess("secrets/api-keys.json")).toBe(true)
expect(emptyController.validateAccess(".git/HEAD")).toBe(true)
})
})
describe("validateCommand", () => {
beforeEach(async () => {
// Setup .rooignore content
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**\n*.log")
await controller.initialize()
})
/**
* Tests validation of file reading commands
*/
it("should block file reading commands accessing ignored files", () => {
// Cat command accessing ignored file
expect(controller.validateCommand("cat node_modules/package.json")).toBe("node_modules/package.json")
// Grep command accessing ignored file
expect(controller.validateCommand("grep pattern .git/config")).toBe(".git/config")
// Commands accessing allowed files should return undefined
expect(controller.validateCommand("cat src/app.ts")).toBeUndefined()
expect(controller.validateCommand("less README.md")).toBeUndefined()
})
/**
* Tests commands with various arguments and flags
*/
it("should handle command arguments and flags correctly", () => {
// Command with flags
expect(controller.validateCommand("cat -n node_modules/package.json")).toBe("node_modules/package.json")
// Command with multiple files (only first ignored file is returned)
expect(controller.validateCommand("grep pattern src/app.ts node_modules/index.js")).toBe(
"node_modules/index.js",
)
// Command with PowerShell parameter style
expect(controller.validateCommand("Get-Content -Path secrets/api-keys.json")).toBe("secrets/api-keys.json")
// Arguments with colons are skipped due to the implementation
// Adjust test to match actual implementation which skips arguments with colons
expect(controller.validateCommand("Select-String -Path secrets/api-keys.json -Pattern key")).toBe(
"secrets/api-keys.json",
)
})
/**
* Tests validation of non-file-reading commands
*/
it("should allow non-file-reading commands", () => {
// Commands that don't access files directly
expect(controller.validateCommand("ls -la")).toBeUndefined()
expect(controller.validateCommand("echo 'Hello'")).toBeUndefined()
expect(controller.validateCommand("cd node_modules")).toBeUndefined()
expect(controller.validateCommand("npm install")).toBeUndefined()
})
/**
* Tests behavior when no .rooignore exists
*/
it("should allow all commands when no .rooignore exists", async () => {
// Create a new controller with no .rooignore
mockFileExists.mockResolvedValue(false)
const emptyController = new RooIgnoreController(TEST_CWD)
await emptyController.initialize()
// All commands should be allowed
expect(emptyController.validateCommand("cat node_modules/package.json")).toBeUndefined()
expect(emptyController.validateCommand("grep pattern .git/config")).toBeUndefined()
})
})
describe("filterPaths", () => {
beforeEach(async () => {
// Setup .rooignore content
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**\n*.log")
await controller.initialize()
})
/**
* Tests filtering an array of paths
*/
it("should filter out ignored paths from an array", () => {
const paths = [
"src/app.ts",
"node_modules/package.json",
"README.md",
".git/HEAD",
"secrets/keys.json",
"build/app.js",
"logs/error.log",
]
const filtered = controller.filterPaths(paths)
// Expected filtered result
expect(filtered).toEqual(["src/app.ts", "README.md", "build/app.js"])
// Length should be reduced
expect(filtered.length).toBe(3)
})
/**
* Tests error handling in filterPaths
*/
it("should handle errors in filterPaths and fail closed", () => {
// Mock validateAccess to throw an error
jest.spyOn(controller, "validateAccess").mockImplementation(() => {
throw new Error("Test error")
})
// Spy on console.error
const consoleSpy = jest.spyOn(console, "error").mockImplementation()
// Should return empty array on error (fail closed)
const result = controller.filterPaths(["file1.txt", "file2.txt"])
expect(result).toEqual([])
// Verify error was logged
expect(consoleSpy).toHaveBeenCalledWith("Error filtering paths:", expect.any(Error))
// Cleanup
consoleSpy.mockRestore()
})
/**
* Tests empty array handling
*/
it("should handle empty arrays", () => {
const result = controller.filterPaths([])
expect(result).toEqual([])
})
})
describe("getInstructions", () => {
/**
* Tests instructions generation with .rooignore
*/
it("should generate formatted instructions when .rooignore exists", async () => {
// Setup .rooignore content
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**")
await controller.initialize()
const instructions = controller.getInstructions()
// Verify instruction format
expect(instructions).toContain("# .rooignore")
expect(instructions).toContain(LOCK_TEXT_SYMBOL)
expect(instructions).toContain("node_modules")
expect(instructions).toContain(".git")
expect(instructions).toContain("secrets/**")
})
/**
* Tests behavior when no .rooignore exists
*/
it("should return undefined when no .rooignore exists", async () => {
// Setup no .rooignore
mockFileExists.mockResolvedValue(false)
await controller.initialize()
const instructions = controller.getInstructions()
expect(instructions).toBeUndefined()
})
})
describe("dispose", () => {
/**
* Tests proper cleanup of resources
*/
it("should dispose all registered disposables", () => {
// Create spy for dispose methods
const disposeSpy = jest.fn()
// Manually add disposables to test
controller["disposables"] = [{ dispose: disposeSpy }, { dispose: disposeSpy }, { dispose: disposeSpy }]
// Call dispose
controller.dispose()
// Verify all disposables were disposed
expect(disposeSpy).toHaveBeenCalledTimes(3)
// Verify disposables array was cleared
expect(controller["disposables"]).toEqual([])
})
})
describe("file watcher", () => {
/**
* Tests behavior when .rooignore is created
*/
it("should reload .rooignore when file is created", async () => {
// Setup initial state without .rooignore
mockFileExists.mockResolvedValue(false)
await controller.initialize()
// Verify initial state
expect(controller.rooIgnoreContent).toBeUndefined()
expect(controller.validateAccess("node_modules/package.json")).toBe(true)
// Setup for the test
mockFileExists.mockResolvedValue(false) // Initially no file exists
// Create and initialize controller with no .rooignore
controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Initial state check
expect(controller.rooIgnoreContent).toBeUndefined()
// Now simulate file creation
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules")
// Find and trigger the onCreate handler
const onCreateHandler = mockWatcher.onDidCreate.mock.calls[0][0]
// Force reload of .rooignore content manually
await controller.initialize()
// Now verify content was updated
expect(controller.rooIgnoreContent).toBe("node_modules")
// Verify access validation changed
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
})
/**
* Tests behavior when .rooignore is changed
*/
it("should reload .rooignore when file is changed", async () => {
// Setup initial state with .rooignore
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules")
await controller.initialize()
// Verify initial state
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
expect(controller.validateAccess(".git/config")).toBe(true)
// Simulate file change
mockReadFile.mockResolvedValue("node_modules\n.git")
// Instead of relying on the onChange handler, manually reload
// This is because the mock watcher doesn't actually trigger the reload in tests
await controller.initialize()
// Verify content was updated
expect(controller.rooIgnoreContent).toBe("node_modules\n.git")
// Verify access validation changed
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
expect(controller.validateAccess(".git/config")).toBe(false)
})
/**
* Tests behavior when .rooignore is deleted
*/
it("should reset when .rooignore is deleted", async () => {
// Setup initial state with .rooignore
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules")
await controller.initialize()
// Verify initial state
expect(controller.validateAccess("node_modules/package.json")).toBe(false)
// Simulate file deletion
mockFileExists.mockResolvedValue(false)
// Find and trigger the onDelete handler
const onDeleteHandler = mockWatcher.onDidDelete.mock.calls[0][0]
await onDeleteHandler()
// Verify content was reset
expect(controller.rooIgnoreContent).toBeUndefined()
// Verify access validation changed
expect(controller.validateAccess("node_modules/package.json")).toBe(true)
})
})
})

View file

@ -2,13 +2,13 @@ import * as vscode from "vscode"
import * as path from "path"
import { openFile } from "../../integrations/misc/open-file"
import { UrlContentFetcher } from "../../services/browser/UrlContentFetcher"
import { mentionRegexGlobal, formatGitSuggestion, type MentionSuggestion } from "../../shared/context-mentions"
import { mentionRegexGlobal } from "../../shared/context-mentions"
import fs from "fs/promises"
import { extractTextFromFile } from "../../integrations/misc/extract-text"
import { isBinaryFile } from "isbinaryfile"
import { diagnosticsToProblemsString } from "../../integrations/diagnostics"
import { getCommitInfo, getWorkingState } from "../../utils/git"
import { getLatestTerminalOutput } from "../../integrations/terminal/get-latest-output"
import { getLatestTerminalOutput } from "../../integrations/terminal/getLatestTerminalOutput"
export async function openMention(mention?: string): Promise<void> {
if (!mention) {
@ -198,9 +198,9 @@ async function getFileOrFolderContent(mentionPath: string, cwd: string): Promise
}
}
function getWorkspaceProblems(cwd: string): string {
async function getWorkspaceProblems(cwd: string): Promise<string> {
const diagnostics = vscode.languages.getDiagnostics()
const result = diagnosticsToProblemsString(
const result = await diagnosticsToProblemsString(
diagnostics,
[vscode.DiagnosticSeverity.Error, vscode.DiagnosticSeverity.Warning],
cwd,

View file

@ -3899,9 +3899,17 @@ USER'S CUSTOM INSTRUCTIONS
The following additional instructions are provided by the user, and should be followed to the best of your ability without interfering with the TOOL USE guidelines.
Mode-specific Instructions:
Depending on the user's request, you may need to do some information gathering (for example using read_file or search_files) to get more context about the task. You may also ask the user clarifying questions to get a better understanding of the task. Once you've gained more context about the user's request, you should create a detailed plan for how to accomplish the task. (You can write the plan to a markdown file if it seems appropriate.)
1. Do some information gathering (for example using read_file or search_files) to get more context about the task.
Then you might ask the user if they are pleased with this plan, or if they would like to make any changes. Think of this as a brainstorming session where you can discuss the task and plan the best way to accomplish it. Finally once it seems like you've reached a good plan, use the switch_mode tool to request that the user switch to another mode to implement the solution.
2. You should also ask the user clarifying questions to get a better understanding of the task.
3. Once you've gained more context about the user's request, you should create a detailed plan for how to accomplish the task. Include Mermaid diagrams if they help make your plan clearer.
4. Ask the user if they are pleased with this plan, or if they would like to make any changes. Think of this as a brainstorming session where you can discuss the task and plan the best way to accomplish it.
5. Once the user confirms the plan, ask them if they'd like you to write it to a markdown file.
6. Use the switch_mode tool to request that the user switch to another mode to implement the solution.
Rules:
# Rules from .clinerules-architect:
@ -4176,7 +4184,7 @@ USER'S CUSTOM INSTRUCTIONS
The following additional instructions are provided by the user, and should be followed to the best of your ability without interfering with the TOOL USE guidelines.
Mode-specific Instructions:
You can analyze code, explain concepts, and access external resources. Make sure to answer the user's questions and don't rush to switch to implementing code.
You can analyze code, explain concepts, and access external resources. Make sure to answer the user's questions and don't rush to switch to implementing code. Include Mermaid diagrams if they help make your response clearer.
Rules:
# Rules from .clinerules-ask:

View file

@ -0,0 +1,172 @@
import { SYSTEM_PROMPT } from "../system"
import { defaultModeSlug, modes } from "../../../shared/modes"
import * as vscode from "vscode"
import * as fs from "fs/promises"
// Mock the fs/promises module
jest.mock("fs/promises", () => ({
readFile: jest.fn(),
mkdir: jest.fn().mockResolvedValue(undefined),
access: jest.fn().mockResolvedValue(undefined),
}))
// Get the mocked fs module
const mockedFs = fs as jest.Mocked<typeof fs>
// Mock the fileExistsAtPath function
jest.mock("../../../utils/fs", () => ({
fileExistsAtPath: jest.fn().mockResolvedValue(true),
createDirectoriesForFile: jest.fn().mockResolvedValue([]),
}))
// Create a mock ExtensionContext with relative paths instead of absolute paths
const mockContext = {
extensionPath: "mock/extension/path",
globalStoragePath: "mock/storage/path",
storagePath: "mock/storage/path",
logPath: "mock/log/path",
subscriptions: [],
workspaceState: {
get: () => undefined,
update: () => Promise.resolve(),
},
globalState: {
get: () => undefined,
update: () => Promise.resolve(),
setKeysForSync: () => {},
},
extensionUri: { fsPath: "mock/extension/path" },
globalStorageUri: { fsPath: "mock/settings/path" },
asAbsolutePath: (relativePath: string) => `mock/extension/path/${relativePath}`,
extension: {
packageJSON: {
version: "1.0.0",
},
},
} as unknown as vscode.ExtensionContext
describe("File-Based Custom System Prompt", () => {
const experiments = {}
beforeEach(() => {
// Reset mocks before each test
jest.clearAllMocks()
// Default behavior: file doesn't exist
mockedFs.readFile.mockRejectedValue({ code: "ENOENT" })
})
it("should use default generation when no file-based system prompt is found", async () => {
const customModePrompts = {
[defaultModeSlug]: {
roleDefinition: "Test role definition",
},
}
const prompt = await SYSTEM_PROMPT(
mockContext,
"test/path", // Using a relative path without leading slash
false,
undefined,
undefined,
undefined,
defaultModeSlug,
customModePrompts,
undefined,
undefined,
undefined,
undefined,
experiments,
true,
)
// Should contain default sections
expect(prompt).toContain("TOOL USE")
expect(prompt).toContain("CAPABILITIES")
expect(prompt).toContain("MODES")
expect(prompt).toContain("Test role definition")
})
it("should use file-based custom system prompt when available", async () => {
// Mock the readFile to return content from a file
const fileCustomSystemPrompt = "Custom system prompt from file"
// When called with utf-8 encoding, return a string
mockedFs.readFile.mockImplementation((filePath, options) => {
if (filePath.toString().includes(`.roo/system-prompt-${defaultModeSlug}`) && options === "utf-8") {
return Promise.resolve(fileCustomSystemPrompt)
}
return Promise.reject({ code: "ENOENT" })
})
const prompt = await SYSTEM_PROMPT(
mockContext,
"test/path", // Using a relative path without leading slash
false,
undefined,
undefined,
undefined,
defaultModeSlug,
undefined,
undefined,
undefined,
undefined,
undefined,
experiments,
true,
)
// Should contain role definition and file-based system prompt
expect(prompt).toContain(modes[0].roleDefinition)
expect(prompt).toContain(fileCustomSystemPrompt)
// Should not contain any of the default sections
expect(prompt).not.toContain("TOOL USE")
expect(prompt).not.toContain("CAPABILITIES")
expect(prompt).not.toContain("MODES")
})
it("should combine file-based system prompt with role definition and custom instructions", async () => {
// Mock the readFile to return content from a file
const fileCustomSystemPrompt = "Custom system prompt from file"
mockedFs.readFile.mockImplementation((filePath, options) => {
if (filePath.toString().includes(`.roo/system-prompt-${defaultModeSlug}`) && options === "utf-8") {
return Promise.resolve(fileCustomSystemPrompt)
}
return Promise.reject({ code: "ENOENT" })
})
// Define custom role definition
const customRoleDefinition = "Custom role definition"
const customModePrompts = {
[defaultModeSlug]: {
roleDefinition: customRoleDefinition,
},
}
const prompt = await SYSTEM_PROMPT(
mockContext,
"test/path", // Using a relative path without leading slash
false,
undefined,
undefined,
undefined,
defaultModeSlug,
customModePrompts,
undefined,
undefined,
undefined,
undefined,
experiments,
true,
)
// Should contain custom role definition and file-based system prompt
expect(prompt).toContain(customRoleDefinition)
expect(prompt).toContain(fileCustomSystemPrompt)
// Should not contain any of the default sections
expect(prompt).not.toContain("TOOL USE")
expect(prompt).not.toContain("CAPABILITIES")
expect(prompt).not.toContain("MODES")
})
})

View file

@ -0,0 +1,242 @@
// npx jest src/core/prompts/__tests__/responses-rooignore.test.ts
import { formatResponse } from "../responses"
import { RooIgnoreController, LOCK_TEXT_SYMBOL } from "../../ignore/RooIgnoreController"
import * as path from "path"
import { fileExistsAtPath } from "../../../utils/fs"
import * as fs from "fs/promises"
// Mock dependencies
jest.mock("../../../utils/fs")
jest.mock("fs/promises")
jest.mock("vscode", () => {
const mockDisposable = { dispose: jest.fn() }
return {
workspace: {
createFileSystemWatcher: jest.fn(() => ({
onDidCreate: jest.fn(() => mockDisposable),
onDidChange: jest.fn(() => mockDisposable),
onDidDelete: jest.fn(() => mockDisposable),
dispose: jest.fn(),
})),
},
RelativePattern: jest.fn(),
}
})
describe("RooIgnore Response Formatting", () => {
const TEST_CWD = "/test/path"
let mockFileExists: jest.MockedFunction<typeof fileExistsAtPath>
let mockReadFile: jest.MockedFunction<typeof fs.readFile>
beforeEach(() => {
// Reset mocks
jest.clearAllMocks()
// Setup fs mocks
mockFileExists = fileExistsAtPath as jest.MockedFunction<typeof fileExistsAtPath>
mockReadFile = fs.readFile as jest.MockedFunction<typeof fs.readFile>
// Default mock implementations
mockFileExists.mockResolvedValue(true)
mockReadFile.mockResolvedValue("node_modules\n.git\nsecrets/**\n*.log")
})
describe("formatResponse.rooIgnoreError", () => {
/**
* Tests the error message format for ignored files
*/
it("should format error message for ignored files", () => {
const errorMessage = formatResponse.rooIgnoreError("secrets/api-keys.json")
// Verify error message format
expect(errorMessage).toContain("Access to secrets/api-keys.json is blocked by the .rooignore file settings")
expect(errorMessage).toContain("continue in the task without using this file")
expect(errorMessage).toContain("ask the user to update the .rooignore file")
})
/**
* Tests with different file paths
*/
it("should include the file path in the error message", () => {
const paths = ["node_modules/package.json", ".git/HEAD", "secrets/credentials.env", "logs/app.log"]
// Test each path
for (const testPath of paths) {
const errorMessage = formatResponse.rooIgnoreError(testPath)
expect(errorMessage).toContain(`Access to ${testPath} is blocked`)
}
})
})
describe("formatResponse.formatFilesList with RooIgnoreController", () => {
/**
* Tests file listing with rooignore controller
*/
it("should format files list with lock symbols for ignored files", async () => {
// Create controller
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Mock validateAccess to control which files are ignored
controller.validateAccess = jest.fn().mockImplementation((filePath: string) => {
// Only allow files not matching these patterns
return (
!filePath.includes("node_modules") && !filePath.includes(".git") && !filePath.includes("secrets/")
)
})
// Files list with mixed allowed/ignored files
const files = [
"src/app.ts", // allowed
"node_modules/package.json", // ignored
"README.md", // allowed
".git/HEAD", // ignored
"secrets/keys.json", // ignored
]
// Format with controller
const result = formatResponse.formatFilesList(TEST_CWD, files, false, controller as any, true)
// Should contain each file
expect(result).toContain("src/app.ts")
expect(result).toContain("README.md")
// Should contain lock symbols for ignored files - case insensitive check using regex
expect(result).toMatch(new RegExp(`${LOCK_TEXT_SYMBOL}.*node_modules/package.json`, "i"))
expect(result).toMatch(new RegExp(`${LOCK_TEXT_SYMBOL}.*\\.git/HEAD`, "i"))
expect(result).toMatch(new RegExp(`${LOCK_TEXT_SYMBOL}.*secrets/keys.json`, "i"))
// No lock symbols for allowed files
expect(result).not.toContain(`${LOCK_TEXT_SYMBOL} src/app.ts`)
expect(result).not.toContain(`${LOCK_TEXT_SYMBOL} README.md`)
})
/**
* Tests formatFilesList when showRooIgnoredFiles is set to false
*/
it("should hide ignored files when showRooIgnoredFiles is false", async () => {
// Create controller
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Mock validateAccess to control which files are ignored
controller.validateAccess = jest.fn().mockImplementation((filePath: string) => {
// Only allow files not matching these patterns
return (
!filePath.includes("node_modules") && !filePath.includes(".git") && !filePath.includes("secrets/")
)
})
// Files list with mixed allowed/ignored files
const files = [
"src/app.ts", // allowed
"node_modules/package.json", // ignored
"README.md", // allowed
".git/HEAD", // ignored
"secrets/keys.json", // ignored
]
// Format with controller and showRooIgnoredFiles = false
const result = formatResponse.formatFilesList(
TEST_CWD,
files,
false,
controller as any,
false, // showRooIgnoredFiles = false
)
// Should contain allowed files
expect(result).toContain("src/app.ts")
expect(result).toContain("README.md")
// Should NOT contain ignored files (even with lock symbols)
expect(result).not.toContain("node_modules/package.json")
expect(result).not.toContain(".git/HEAD")
expect(result).not.toContain("secrets/keys.json")
// Double-check with regex to ensure no form of these filenames appears
expect(result).not.toMatch(/node_modules\/package\.json/i)
expect(result).not.toMatch(/\.git\/HEAD/i)
expect(result).not.toMatch(/secrets\/keys\.json/i)
})
/**
* Tests formatFilesList handles truncation correctly with RooIgnoreController
*/
it("should handle truncation with RooIgnoreController", async () => {
// Create controller
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Format with controller and truncation flag
const result = formatResponse.formatFilesList(
TEST_CWD,
["file1.txt", "file2.txt"],
true, // didHitLimit = true
controller as any,
true,
)
// Should contain truncation message (case-insensitive check)
expect(result).toContain("File list truncated")
expect(result).toMatch(/use list_files on specific subdirectories/i)
})
/**
* Tests formatFilesList handles empty results
*/
it("should handle empty file list with RooIgnoreController", async () => {
// Create controller
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Format with empty files array
const result = formatResponse.formatFilesList(TEST_CWD, [], false, controller as any, true)
// Should show "No files found"
expect(result).toBe("No files found.")
})
})
describe("getInstructions", () => {
/**
* Tests the instructions format
*/
it("should format .rooignore instructions for the LLM", async () => {
// Create controller
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Get instructions
const instructions = controller.getInstructions()
// Verify format and content
expect(instructions).toContain("# .rooignore")
expect(instructions).toContain(LOCK_TEXT_SYMBOL)
expect(instructions).toContain("node_modules")
expect(instructions).toContain(".git")
expect(instructions).toContain("secrets/**")
expect(instructions).toContain("*.log")
// Should explain what the lock symbol means
expect(instructions).toContain("you'll notice a")
expect(instructions).toContain("next to files that are blocked")
})
/**
* Tests null/undefined case
*/
it("should return undefined when no .rooignore exists", async () => {
// Set up no .rooignore
mockFileExists.mockResolvedValue(false)
// Create controller without .rooignore
const controller = new RooIgnoreController(TEST_CWD)
await controller.initialize()
// Should return undefined
expect(controller.getInstructions()).toBeUndefined()
})
})
})

View file

@ -1,6 +1,7 @@
import { Anthropic } from "@anthropic-ai/sdk"
import * as path from "path"
import * as diff from "diff"
import { RooIgnoreController, LOCK_TEXT_SYMBOL } from "../ignore/RooIgnoreController"
export const formatResponse = {
toolDenied: () => `The user denied this operation.`,
@ -13,6 +14,9 @@ export const formatResponse = {
toolError: (error?: string) => `The tool execution failed with the following error:\n<error>\n${error}\n</error>`,
rooIgnoreError: (path: string) =>
`Access to ${path} is blocked by the .rooignore file settings. You must try to continue in the task without using this file, or ask the user to update the .rooignore file.`,
noToolsUsed: () =>
`[ERROR] You did not use a tool in your previous response! Please retry with a tool use.
@ -52,7 +56,13 @@ Otherwise, if you have not completed the task and do not need additional informa
return formatImagesIntoBlocks(images)
},
formatFilesList: (absolutePath: string, files: string[], didHitLimit: boolean): string => {
formatFilesList: (
absolutePath: string,
files: string[],
didHitLimit: boolean,
rooIgnoreController: RooIgnoreController | undefined,
showRooIgnoredFiles: boolean,
): string => {
const sorted = files
.map((file) => {
// convert absolute path to relative path
@ -80,14 +90,38 @@ Otherwise, if you have not completed the task and do not need additional informa
// the shorter one comes first
return aParts.length - bParts.length
})
let rooIgnoreParsed: string[] = sorted
if (rooIgnoreController) {
rooIgnoreParsed = []
for (const filePath of sorted) {
// path is relative to absolute path, not cwd
// validateAccess expects either path relative to cwd or absolute path
// otherwise, for validating against ignore patterns like "assets/icons", we would end up with just "icons", which would result in the path not being ignored.
const absoluteFilePath = path.resolve(absolutePath, filePath)
const isIgnored = !rooIgnoreController.validateAccess(absoluteFilePath)
if (isIgnored) {
// If file is ignored and we're not showing ignored files, skip it
if (!showRooIgnoredFiles) {
continue
}
// Otherwise, mark it with a lock symbol
rooIgnoreParsed.push(LOCK_TEXT_SYMBOL + " " + filePath)
} else {
rooIgnoreParsed.push(filePath)
}
}
}
if (didHitLimit) {
return `${sorted.join(
return `${rooIgnoreParsed.join(
"\n",
)}\n\n(File list truncated. Use list_files on specific subdirectories if you need to explore further.)`
} else if (sorted.length === 0 || (sorted.length === 1 && sorted[0] === "")) {
} else if (rooIgnoreParsed.length === 0 || (rooIgnoreParsed.length === 1 && rooIgnoreParsed[0] === "")) {
return "No files found."
} else {
return sorted.join("\n")
return rooIgnoreParsed.join("\n")
}
},

View file

@ -33,7 +33,7 @@ export async function addCustomInstructions(
globalCustomInstructions: string,
cwd: string,
mode: string,
options: { preferredLanguage?: string } = {},
options: { preferredLanguage?: string; rooIgnoreInstructions?: string } = {},
): Promise<string> {
const sections = []
@ -70,6 +70,10 @@ export async function addCustomInstructions(
rules.push(`# Rules from ${modeRuleFile}:\n${modeRuleContent}`)
}
if (options.rooIgnoreInstructions) {
rules.push(options.rooIgnoreInstructions)
}
// Add generic rules
const genericRuleContent = await loadRuleFiles(cwd)
if (genericRuleContent && genericRuleContent.trim()) {

View file

@ -0,0 +1,60 @@
import fs from "fs/promises"
import path from "path"
import { Mode } from "../../../shared/modes"
import { fileExistsAtPath } from "../../../utils/fs"
/**
* Safely reads a file, returning an empty string if the file doesn't exist
*/
async function safeReadFile(filePath: string): Promise<string> {
try {
const content = await fs.readFile(filePath, "utf-8")
// When reading with "utf-8" encoding, content should be a string
return content.trim()
} catch (err) {
const errorCode = (err as NodeJS.ErrnoException).code
if (!errorCode || !["ENOENT", "EISDIR"].includes(errorCode)) {
throw err
}
return ""
}
}
/**
* Get the path to a system prompt file for a specific mode
*/
export function getSystemPromptFilePath(cwd: string, mode: Mode): string {
return path.join(cwd, ".roo", `system-prompt-${mode}`)
}
/**
* Loads custom system prompt from a file at .roo/system-prompt-[mode slug]
* If the file doesn't exist, returns an empty string
*/
export async function loadSystemPromptFile(cwd: string, mode: Mode): Promise<string> {
const filePath = getSystemPromptFilePath(cwd, mode)
return safeReadFile(filePath)
}
/**
* Ensures the .roo directory exists, creating it if necessary
*/
export async function ensureRooDirectory(cwd: string): Promise<void> {
const rooDir = path.join(cwd, ".roo")
// Check if directory already exists
if (await fileExistsAtPath(rooDir)) {
return
}
// Create the directory
try {
await fs.mkdir(rooDir, { recursive: true })
} catch (err) {
// If directory already exists (race condition), ignore the error
const errorCode = (err as NodeJS.ErrnoException).code
if (errorCode !== "EEXIST") {
throw err
}
}
}

View file

@ -11,12 +11,19 @@ export async function getModesSection(context: vscode.ExtensionContext): Promise
// Get all modes with their overrides from extension state
const allModes = await getAllModesWithPrompts(context)
return `====
// Get enableCustomModeCreation setting from extension state
const shouldEnableCustomModeCreation = await context.globalState.get<boolean>("enableCustomModeCreation") ?? true
let modesContent = `====
MODES
- These are the currently available modes:
${allModes.map((mode: ModeConfig) => ` * "${mode.name}" mode (${mode.slug}) - ${mode.roleDefinition.split(".")[0]}`).join("\n")}
${allModes.map((mode: ModeConfig) => ` * "${mode.name}" mode (${mode.slug}) - ${mode.roleDefinition.split(".")[0]}`).join("\n")}`
// Only include custom modes documentation if the feature is enabled
if (shouldEnableCustomModeCreation) {
modesContent += `
- Custom modes can be configured in two ways:
1. Globally via '${customModesPath}' (created automatically on startup)
@ -56,4 +63,7 @@ Both files should follow this structure:
}
]
}`
}
return modesContent
}

View file

@ -7,6 +7,7 @@ import {
defaultModeSlug,
ModeConfig,
getModeBySlug,
getGroupName,
} from "../../shared/modes"
import { DiffStrategy } from "../diff/DiffStrategy"
import { McpHub } from "../../services/mcp/McpHub"
@ -23,8 +24,7 @@ import {
getModesSection,
addCustomInstructions,
} from "./sections"
import fs from "fs/promises"
import path from "path"
import { loadSystemPromptFile } from "./sections/custom-system-prompt"
async function generatePrompt(
context: vscode.ExtensionContext,
@ -41,6 +41,7 @@ async function generatePrompt(
diffEnabled?: boolean,
experiments?: Record<string, boolean>,
enableMcpServerCreation?: boolean,
rooIgnoreInstructions?: string,
): Promise<string> {
if (!context) {
throw new Error("Extension context is required for generating system prompt")
@ -49,15 +50,17 @@ async function generatePrompt(
// If diff is disabled, don't pass the diffStrategy
const effectiveDiffStrategy = diffEnabled ? diffStrategy : undefined
const [mcpServersSection, modesSection] = await Promise.all([
getMcpServersSection(mcpHub, effectiveDiffStrategy, enableMcpServerCreation),
getModesSection(context),
])
// Get the full mode config to ensure we have the role definition
const modeConfig = getModeBySlug(mode, customModeConfigs) || modes.find((m) => m.slug === mode) || modes[0]
const roleDefinition = promptComponent?.roleDefinition || modeConfig.roleDefinition
const [modesSection, mcpServersSection] = await Promise.all([
getModesSection(context),
modeConfig.groups.some((groupEntry) => getGroupName(groupEntry) === "mcp")
? getMcpServersSection(mcpHub, effectiveDiffStrategy, enableMcpServerCreation)
: Promise.resolve(""),
])
const basePrompt = `${roleDefinition}
${getSharedToolUseSection()}
@ -87,7 +90,7 @@ ${getSystemInfoSection(cwd, mode, customModeConfigs)}
${getObjectiveSection()}
${await addCustomInstructions(promptComponent?.customInstructions || modeConfig.customInstructions || "", globalCustomInstructions || "", cwd, mode, { preferredLanguage })}`
${await addCustomInstructions(promptComponent?.customInstructions || modeConfig.customInstructions || "", globalCustomInstructions || "", cwd, mode, { preferredLanguage, rooIgnoreInstructions })}`
return basePrompt
}
@ -107,6 +110,7 @@ export const SYSTEM_PROMPT = async (
diffEnabled?: boolean,
experiments?: Record<string, boolean>,
enableMcpServerCreation?: boolean,
rooIgnoreInstructions?: string,
): Promise<string> => {
if (!context) {
throw new Error("Extension context is required for generating system prompt")
@ -119,11 +123,25 @@ export const SYSTEM_PROMPT = async (
return undefined
}
// Try to load custom system prompt from file
const fileCustomSystemPrompt = await loadSystemPromptFile(cwd, mode)
// Check if it's a custom mode
const promptComponent = getPromptComponent(customModePrompts?.[mode])
// Get full mode config from custom modes or fall back to built-in modes
const currentMode = getModeBySlug(mode, customModes) || modes.find((m) => m.slug === mode) || modes[0]
// If a file-based custom system prompt exists, use it
if (fileCustomSystemPrompt) {
const roleDefinition = promptComponent?.roleDefinition || currentMode.roleDefinition
return `${roleDefinition}
${fileCustomSystemPrompt}
${await addCustomInstructions(promptComponent?.customInstructions || currentMode.customInstructions || "", globalCustomInstructions || "", cwd, mode, { preferredLanguage, rooIgnoreInstructions })}`
}
// If diff is disabled, don't pass the diffStrategy
const effectiveDiffStrategy = diffEnabled ? diffStrategy : undefined
@ -142,5 +160,6 @@ export const SYSTEM_PROMPT = async (
diffEnabled,
experiments,
enableMcpServerCreation,
rooIgnoreInstructions,
)
}

View file

@ -3,8 +3,39 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { ModelInfo } from "../../../shared/api"
import { truncateConversation, truncateConversationIfNeeded } from "../index"
import { ApiHandler } from "../../../api"
import { BaseProvider } from "../../../api/providers/base-provider"
import { TOKEN_BUFFER_PERCENTAGE } from "../index"
import { estimateTokenCount, truncateConversation, truncateConversationIfNeeded } from "../index"
// Create a mock ApiHandler for testing
class MockApiHandler extends BaseProvider {
createMessage(): any {
throw new Error("Method not implemented.")
}
getModel(): { id: string; info: ModelInfo } {
return {
id: "test-model",
info: {
contextWindow: 100000,
maxTokens: 50000,
supportsPromptCache: true,
supportsImages: false,
inputPrice: 0,
outputPrice: 0,
description: "Test model",
},
}
}
}
// Create a singleton instance for tests
const mockApiHandler = new MockApiHandler()
/**
* Tests for the truncateConversation function
*/
describe("truncateConversation", () => {
it("should retain the first message", () => {
const messages: Anthropic.Messages.MessageParam[] = [
@ -91,10 +122,102 @@ describe("truncateConversation", () => {
})
})
/**
* Tests for the estimateTokenCount function
*/
describe("estimateTokenCount", () => {
it("should return 0 for empty or undefined content", async () => {
expect(await estimateTokenCount([], mockApiHandler)).toBe(0)
// @ts-ignore - Testing with undefined
expect(await estimateTokenCount(undefined, mockApiHandler)).toBe(0)
})
it("should estimate tokens for text blocks", async () => {
const content: Array<Anthropic.Messages.ContentBlockParam> = [
{ type: "text", text: "This is a text block with 36 characters" },
]
// With tiktoken, the exact token count may differ from character-based estimation
// Instead of expecting an exact number, we verify it's a reasonable positive number
const result = await estimateTokenCount(content, mockApiHandler)
expect(result).toBeGreaterThan(0)
// We can also verify that longer text results in more tokens
const longerContent: Array<Anthropic.Messages.ContentBlockParam> = [
{
type: "text",
text: "This is a longer text block with significantly more characters to encode into tokens",
},
]
const longerResult = await estimateTokenCount(longerContent, mockApiHandler)
expect(longerResult).toBeGreaterThan(result)
})
it("should estimate tokens for image blocks based on data size", async () => {
// Small image
const smallImage: Array<Anthropic.Messages.ContentBlockParam> = [
{ type: "image", source: { type: "base64", media_type: "image/jpeg", data: "small_dummy_data" } },
]
// Larger image with more data
const largerImage: Array<Anthropic.Messages.ContentBlockParam> = [
{ type: "image", source: { type: "base64", media_type: "image/png", data: "X".repeat(1000) } },
]
// Verify the token count scales with the size of the image data
const smallImageTokens = await estimateTokenCount(smallImage, mockApiHandler)
const largerImageTokens = await estimateTokenCount(largerImage, mockApiHandler)
// Small image should have some tokens
expect(smallImageTokens).toBeGreaterThan(0)
// Larger image should have proportionally more tokens
expect(largerImageTokens).toBeGreaterThan(smallImageTokens)
// Verify the larger image calculation matches our formula including the 50% fudge factor
expect(largerImageTokens).toBe(48)
})
it("should estimate tokens for mixed content blocks", async () => {
const content: Array<Anthropic.Messages.ContentBlockParam> = [
{ type: "text", text: "A text block with 30 characters" },
{ type: "image", source: { type: "base64", media_type: "image/jpeg", data: "dummy_data" } },
{ type: "text", text: "Another text with 24 chars" },
]
// We know image tokens calculation should be consistent
const imageTokens = Math.ceil(Math.sqrt("dummy_data".length)) * 1.5
// With tiktoken, we can't predict exact text token counts,
// but we can verify the total is greater than just the image tokens
const result = await estimateTokenCount(content, mockApiHandler)
expect(result).toBeGreaterThan(imageTokens)
// Also test against a version with only the image to verify text adds tokens
const imageOnlyContent: Array<Anthropic.Messages.ContentBlockParam> = [
{ type: "image", source: { type: "base64", media_type: "image/jpeg", data: "dummy_data" } },
]
const imageOnlyResult = await estimateTokenCount(imageOnlyContent, mockApiHandler)
expect(result).toBeGreaterThan(imageOnlyResult)
})
it("should handle empty text blocks", async () => {
const content: Array<Anthropic.Messages.ContentBlockParam> = [{ type: "text", text: "" }]
expect(await estimateTokenCount(content, mockApiHandler)).toBe(0)
})
it("should handle plain string messages", async () => {
const content = "This is a plain text message"
expect(await estimateTokenCount([{ type: "text", text: content }], mockApiHandler)).toBeGreaterThan(0)
})
})
/**
* Tests for the truncateConversationIfNeeded function
*/
describe("truncateConversationIfNeeded", () => {
const createModelInfo = (contextWindow: number, supportsPromptCache: boolean, maxTokens?: number): ModelInfo => ({
const createModelInfo = (contextWindow: number, maxTokens?: number): ModelInfo => ({
contextWindow,
supportsPromptCache,
supportsPromptCache: true,
maxTokens,
})
@ -106,25 +229,325 @@ describe("truncateConversationIfNeeded", () => {
{ role: "user", content: "Fifth message" },
]
it("should not truncate if tokens are below threshold for prompt caching models", () => {
const modelInfo = createModelInfo(200000, true, 50000)
const totalTokens = 100000 // Below threshold
const result = truncateConversationIfNeeded(messages, totalTokens, modelInfo)
expect(result).toEqual(messages)
it("should not truncate if tokens are below max tokens threshold", async () => {
const modelInfo = createModelInfo(100000, 30000)
const maxTokens = 100000 - 30000 // 70000
const dynamicBuffer = modelInfo.contextWindow * TOKEN_BUFFER_PERCENTAGE // 10000
const totalTokens = 70000 - dynamicBuffer - 1 // Just below threshold - buffer
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
const result = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens,
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result).toEqual(messagesWithSmallContent) // No truncation occurs
})
it("should not truncate if tokens are below threshold for non-prompt caching models", () => {
const modelInfo = createModelInfo(200000, false)
const totalTokens = 100000 // Below threshold
const result = truncateConversationIfNeeded(messages, totalTokens, modelInfo)
expect(result).toEqual(messages)
it("should truncate if tokens are above max tokens threshold", async () => {
const modelInfo = createModelInfo(100000, 30000)
const maxTokens = 100000 - 30000 // 70000
const totalTokens = 70001 // Above threshold
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// When truncating, always uses 0.5 fraction
// With 4 messages after the first, 0.5 fraction means remove 2 messages
const expectedResult = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
const result = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens,
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result).toEqual(expectedResult)
})
it("should use 80% of context window as threshold if it's greater than (contextWindow - buffer)", () => {
const modelInfo = createModelInfo(50000, true) // Small context window
const totalTokens = 40001 // Above 80% threshold (40000)
const mockResult = [messages[0], messages[3], messages[4]]
const result = truncateConversationIfNeeded(messages, totalTokens, modelInfo)
expect(result).toEqual(mockResult)
it("should work with non-prompt caching models the same as prompt caching models", async () => {
// The implementation no longer differentiates between prompt caching and non-prompt caching models
const modelInfo1 = createModelInfo(100000, 30000)
const modelInfo2 = createModelInfo(100000, 30000)
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// Test below threshold
const belowThreshold = 69999
const result1 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: belowThreshold,
contextWindow: modelInfo1.contextWindow,
maxTokens: modelInfo1.maxTokens,
apiHandler: mockApiHandler,
})
const result2 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: belowThreshold,
contextWindow: modelInfo2.contextWindow,
maxTokens: modelInfo2.maxTokens,
apiHandler: mockApiHandler,
})
expect(result1).toEqual(result2)
// Test above threshold
const aboveThreshold = 70001
const result3 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: aboveThreshold,
contextWindow: modelInfo1.contextWindow,
maxTokens: modelInfo1.maxTokens,
apiHandler: mockApiHandler,
})
const result4 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: aboveThreshold,
contextWindow: modelInfo2.contextWindow,
maxTokens: modelInfo2.maxTokens,
apiHandler: mockApiHandler,
})
expect(result3).toEqual(result4)
})
it("should consider incoming content when deciding to truncate", async () => {
const modelInfo = createModelInfo(100000, 30000)
const maxTokens = 30000
const availableTokens = modelInfo.contextWindow - maxTokens
// Test case 1: Small content that won't push us over the threshold
const smallContent = [{ type: "text" as const, text: "Small content" }]
const smallContentTokens = await estimateTokenCount(smallContent, mockApiHandler)
const messagesWithSmallContent: Anthropic.Messages.MessageParam[] = [
...messages.slice(0, -1),
{ role: messages[messages.length - 1].role, content: smallContent },
]
// Set base tokens so total is well below threshold + buffer even with small content added
const dynamicBuffer = modelInfo.contextWindow * TOKEN_BUFFER_PERCENTAGE
const baseTokensForSmall = availableTokens - smallContentTokens - dynamicBuffer - 10
const resultWithSmall = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: baseTokensForSmall,
contextWindow: modelInfo.contextWindow,
maxTokens,
apiHandler: mockApiHandler,
})
expect(resultWithSmall).toEqual(messagesWithSmallContent) // No truncation
// Test case 2: Large content that will push us over the threshold
const largeContent = [
{
type: "text" as const,
text: "A very large incoming message that would consume a significant number of tokens and push us over the threshold",
},
]
const largeContentTokens = await estimateTokenCount(largeContent, mockApiHandler)
const messagesWithLargeContent: Anthropic.Messages.MessageParam[] = [
...messages.slice(0, -1),
{ role: messages[messages.length - 1].role, content: largeContent },
]
// Set base tokens so we're just below threshold without content, but over with content
const baseTokensForLarge = availableTokens - Math.floor(largeContentTokens / 2)
const resultWithLarge = await truncateConversationIfNeeded({
messages: messagesWithLargeContent,
totalTokens: baseTokensForLarge,
contextWindow: modelInfo.contextWindow,
maxTokens,
apiHandler: mockApiHandler,
})
expect(resultWithLarge).not.toEqual(messagesWithLargeContent) // Should truncate
// Test case 3: Very large content that will definitely exceed threshold
const veryLargeContent = [{ type: "text" as const, text: "X".repeat(1000) }]
const veryLargeContentTokens = await estimateTokenCount(veryLargeContent, mockApiHandler)
const messagesWithVeryLargeContent: Anthropic.Messages.MessageParam[] = [
...messages.slice(0, -1),
{ role: messages[messages.length - 1].role, content: veryLargeContent },
]
// Set base tokens so we're just below threshold without content
const baseTokensForVeryLarge = availableTokens - Math.floor(veryLargeContentTokens / 2)
const resultWithVeryLarge = await truncateConversationIfNeeded({
messages: messagesWithVeryLargeContent,
totalTokens: baseTokensForVeryLarge,
contextWindow: modelInfo.contextWindow,
maxTokens,
apiHandler: mockApiHandler,
})
expect(resultWithVeryLarge).not.toEqual(messagesWithVeryLargeContent) // Should truncate
})
it("should truncate if tokens are within TOKEN_BUFFER_PERCENTAGE of the threshold", async () => {
const modelInfo = createModelInfo(100000, 30000)
const maxTokens = 100000 - 30000 // 70000
const dynamicBuffer = modelInfo.contextWindow * TOKEN_BUFFER_PERCENTAGE // 10% of 100000 = 10000
const totalTokens = 70000 - dynamicBuffer + 1 // Just within the dynamic buffer of threshold (70000)
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// When truncating, always uses 0.5 fraction
// With 4 messages after the first, 0.5 fraction means remove 2 messages
const expectedResult = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
const result = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens,
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result).toEqual(expectedResult)
})
})
/**
* Tests for the getMaxTokens function (private but tested through truncateConversationIfNeeded)
*/
describe("getMaxTokens", () => {
// We'll test this indirectly through truncateConversationIfNeeded
const createModelInfo = (contextWindow: number, maxTokens?: number): ModelInfo => ({
contextWindow,
supportsPromptCache: true, // Not relevant for getMaxTokens
maxTokens,
})
// Reuse across tests for consistency
const messages: Anthropic.Messages.MessageParam[] = [
{ role: "user", content: "First message" },
{ role: "assistant", content: "Second message" },
{ role: "user", content: "Third message" },
{ role: "assistant", content: "Fourth message" },
{ role: "user", content: "Fifth message" },
]
it("should use maxTokens as buffer when specified", async () => {
const modelInfo = createModelInfo(100000, 50000)
// Max tokens = 100000 - 50000 = 50000
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// Account for the dynamic buffer which is 10% of context window (10,000 tokens)
// Below max tokens and buffer - no truncation
const result1 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 39999, // Well below threshold + dynamic buffer
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result1).toEqual(messagesWithSmallContent)
// Above max tokens - truncate
const result2 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 50001, // Above threshold
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result2).not.toEqual(messagesWithSmallContent)
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
})
it("should use 20% of context window as buffer when maxTokens is undefined", async () => {
const modelInfo = createModelInfo(100000, undefined)
// Max tokens = 100000 - (100000 * 0.2) = 80000
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// Account for the dynamic buffer which is 10% of context window (10,000 tokens)
// Below max tokens and buffer - no truncation
const result1 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 69999, // Well below threshold + dynamic buffer
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result1).toEqual(messagesWithSmallContent)
// Above max tokens - truncate
const result2 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 80001, // Above threshold
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result2).not.toEqual(messagesWithSmallContent)
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
})
it("should handle small context windows appropriately", async () => {
const modelInfo = createModelInfo(50000, 10000)
// Max tokens = 50000 - 10000 = 40000
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// Below max tokens and buffer - no truncation
const result1 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 34999, // Well below threshold + buffer
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result1).toEqual(messagesWithSmallContent)
// Above max tokens - truncate
const result2 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 40001, // Above threshold
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result2).not.toEqual(messagesWithSmallContent)
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
})
it("should handle large context windows appropriately", async () => {
const modelInfo = createModelInfo(200000, 30000)
// Max tokens = 200000 - 30000 = 170000
// Create messages with very small content in the last one to avoid token overflow
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
// Account for the dynamic buffer which is 10% of context window (20,000 tokens for this test)
// Below max tokens and buffer - no truncation
const result1 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 149999, // Well below threshold + dynamic buffer
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result1).toEqual(messagesWithSmallContent)
// Above max tokens - truncate
const result2 = await truncateConversationIfNeeded({
messages: messagesWithSmallContent,
totalTokens: 170001, // Above threshold
contextWindow: modelInfo.contextWindow,
maxTokens: modelInfo.maxTokens,
apiHandler: mockApiHandler,
})
expect(result2).not.toEqual(messagesWithSmallContent)
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
})
})

View file

@ -1,6 +1,25 @@
import { Anthropic } from "@anthropic-ai/sdk"
import { ApiHandler } from "../../api"
import { ModelInfo } from "../../shared/api"
/**
* Default percentage of the context window to use as a buffer when deciding when to truncate
*/
export const TOKEN_BUFFER_PERCENTAGE = 0.1
/**
* Counts tokens for user content using the provider's token counting implementation.
*
* @param {Array<Anthropic.Messages.ContentBlockParam>} content - The content to count tokens for
* @param {ApiHandler} apiHandler - The API handler to use for token counting
* @returns {Promise<number>} A promise resolving to the token count
*/
export async function estimateTokenCount(
content: Array<Anthropic.Messages.ContentBlockParam>,
apiHandler: ApiHandler,
): Promise<number> {
if (!content || content.length === 0) return 0
return apiHandler.countTokens(content)
}
/**
* Truncates a conversation by removing a fraction of the messages.
@ -26,77 +45,56 @@ export function truncateConversation(
}
/**
* Conditionally truncates the conversation messages if the total token count exceeds the model's limit.
*
* Depending on whether the model supports prompt caching, different maximum token thresholds
* and truncation fractions are used. If the current total tokens exceed the threshold,
* the conversation is truncated using the appropriate fraction.
* Conditionally truncates the conversation messages if the total token count
* exceeds the model's limit, considering the size of incoming content.
*
* @param {Anthropic.Messages.MessageParam[]} messages - The conversation messages.
* @param {number} totalTokens - The total number of tokens in the conversation.
* @param {ModelInfo} modelInfo - Model metadata including context window size and prompt cache support.
* @param {number} totalTokens - The total number of tokens in the conversation (excluding the last user message).
* @param {number} contextWindow - The context window size.
* @param {number} maxTokens - The maximum number of tokens allowed.
* @param {ApiHandler} apiHandler - The API handler to use for token counting.
* @returns {Anthropic.Messages.MessageParam[]} The original or truncated conversation messages.
*/
export function truncateConversationIfNeeded(
messages: Anthropic.Messages.MessageParam[],
totalTokens: number,
modelInfo: ModelInfo,
): Anthropic.Messages.MessageParam[] {
if (modelInfo.supportsPromptCache) {
return totalTokens < getMaxTokensForPromptCachingModels(modelInfo)
? messages
: truncateConversation(messages, getTruncFractionForPromptCachingModels(modelInfo))
} else {
return totalTokens < getMaxTokensForNonPromptCachingModels(modelInfo)
? messages
: truncateConversation(messages, getTruncFractionForNonPromptCachingModels(modelInfo))
}
type TruncateOptions = {
messages: Anthropic.Messages.MessageParam[]
totalTokens: number
contextWindow: number
maxTokens?: number
apiHandler: ApiHandler
}
/**
* Calculates the maximum allowed tokens for models that support prompt caching.
* Conditionally truncates the conversation messages if the total token count
* exceeds the model's limit, considering the size of incoming content.
*
* The maximum is computed as the greater of (contextWindow - buffer) and 80% of the contextWindow.
*
* @param {ModelInfo} modelInfo - The model information containing the context window size.
* @returns {number} The maximum number of tokens allowed for prompt caching models.
* @param {TruncateOptions} options - The options for truncation
* @returns {Promise<Anthropic.Messages.MessageParam[]>} The original or truncated conversation messages.
*/
function getMaxTokensForPromptCachingModels(modelInfo: ModelInfo): number {
// The buffer needs to be at least as large as `modelInfo.maxTokens`.
const buffer = modelInfo.maxTokens ? Math.max(40_000, modelInfo.maxTokens) : 40_000
return Math.max(modelInfo.contextWindow - buffer, modelInfo.contextWindow * 0.8)
}
export async function truncateConversationIfNeeded({
messages,
totalTokens,
contextWindow,
maxTokens,
apiHandler,
}: TruncateOptions): Promise<Anthropic.Messages.MessageParam[]> {
// Calculate the maximum tokens reserved for response
const reservedTokens = maxTokens || contextWindow * 0.2
/**
* Provides the fraction of messages to remove for models that support prompt caching.
*
* @param {ModelInfo} modelInfo - The model information (unused in current implementation).
* @returns {number} The truncation fraction for prompt caching models (fixed at 0.5).
*/
function getTruncFractionForPromptCachingModels(modelInfo: ModelInfo): number {
return 0.5
}
// Estimate tokens for the last message (which is always a user message)
const lastMessage = messages[messages.length - 1]
const lastMessageContent = lastMessage.content
const lastMessageTokens = Array.isArray(lastMessageContent)
? await estimateTokenCount(lastMessageContent, apiHandler)
: await estimateTokenCount([{ type: "text", text: lastMessageContent as string }], apiHandler)
/**
* Calculates the maximum allowed tokens for models that do not support prompt caching.
*
* The maximum is computed as the greater of (contextWindow - 40000) and 80% of the contextWindow.
*
* @param {ModelInfo} modelInfo - The model information containing the context window size.
* @returns {number} The maximum number of tokens allowed for non-prompt caching models.
*/
function getMaxTokensForNonPromptCachingModels(modelInfo: ModelInfo): number {
// The buffer needs to be at least as large as `modelInfo.maxTokens`.
const buffer = modelInfo.maxTokens ? Math.max(40_000, modelInfo.maxTokens) : 40_000
return Math.max(modelInfo.contextWindow - buffer, modelInfo.contextWindow * 0.8)
}
// Calculate total effective tokens (totalTokens never includes the last message)
const effectiveTokens = totalTokens + lastMessageTokens
/**
* Provides the fraction of messages to remove for models that do not support prompt caching.
*
* @param {ModelInfo} modelInfo - The model information.
* @returns {number} The truncation fraction for non-prompt caching models (fixed at 0.1).
*/
function getTruncFractionForNonPromptCachingModels(modelInfo: ModelInfo): number {
return Math.min(40_000 / modelInfo.contextWindow, 0.2)
// Calculate available tokens for conversation history
// Truncate if we're within TOKEN_BUFFER_PERCENTAGE of the context window
const allowedTokens = contextWindow * (1 - TOKEN_BUFFER_PERCENTAGE) - reservedTokens
// Determine if truncation is needed and apply if necessary
return effectiveTokens > allowedTokens ? truncateConversation(messages, 0.5) : messages
}

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -1,55 +1,51 @@
# Cline API
# Roo Code API
The Cline extension exposes an API that can be used by other extensions. To use this API in your extension:
The Roo Code extension exposes an API that can be used by other extensions. To use this API in your extension:
1. Copy `src/extension-api/cline.d.ts` to your extension's source directory.
2. Include `cline.d.ts` in your extension's compilation.
1. Copy `src/extension-api/roo-code.d.ts` to your extension's source directory.
2. Include `roo-code.d.ts` in your extension's compilation.
3. Get access to the API with the following code:
```ts
const clineExtension = vscode.extensions.getExtension<ClineAPI>("rooveterinaryinc.roo-cline")
```typescript
const extension = vscode.extensions.getExtension<RooCodeAPI>("rooveterinaryinc.roo-cline")
if (!clineExtension?.isActive) {
throw new Error("Cline extension is not activated")
}
if (!extension?.isActive) {
throw new Error("Extension is not activated")
}
const cline = clineExtension.exports
const api = extension.exports
if (cline) {
// Now you can use the API
if (!api) {
throw new Error("API is not available")
}
// Set custom instructions
await cline.setCustomInstructions("Talk like a pirate")
// Set custom instructions.
await api.setCustomInstructions("Talk like a pirate")
// Get custom instructions
const instructions = await cline.getCustomInstructions()
console.log("Current custom instructions:", instructions)
// Get custom instructions.
const instructions = await api.getCustomInstructions()
console.log("Current custom instructions:", instructions)
// Start a new task with an initial message
await cline.startNewTask("Hello, Cline! Let's make a new project...")
// Start a new task with an initial message.
await api.startNewTask("Hello, Roo Code API! Let's make a new project...")
// Start a new task with an initial message and images
await cline.startNewTask("Use this design language", ["data:image/webp;base64,..."])
// Start a new task with an initial message and images.
await api.startNewTask("Use this design language", ["data:image/webp;base64,..."])
// Send a message to the current task
await cline.sendMessage("Can you fix the @problems?")
// Send a message to the current task.
await api.sendMessage("Can you fix the @problems?")
// Simulate pressing the primary button in the chat interface (e.g. 'Save' or 'Proceed While Running')
await cline.pressPrimaryButton()
// Simulate pressing the primary button in the chat interface (e.g. 'Save' or 'Proceed While Running').
await api.pressPrimaryButton()
// Simulate pressing the secondary button in the chat interface (e.g. 'Reject')
await cline.pressSecondaryButton()
} else {
console.error("Cline API is not available")
}
```
// Simulate pressing the secondary button in the chat interface (e.g. 'Reject').
await api.pressSecondaryButton()
```
**Note:** To ensure that the `rooveterinaryinc.roo-cline` extension is activated before your extension, add it to the `extensionDependencies` in your `package.json`:
**NOTE:** To ensure that the `rooveterinaryinc.roo-cline` extension is activated before your extension, add it to the `extensionDependencies` in your `package.json`:
```json
"extensionDependencies": [
"rooveterinaryinc.roo-cline"
]
```
```json
"extensionDependencies": ["rooveterinaryinc.roo-cline"]
```
For detailed information on the available methods and their usage, refer to the `cline.d.ts` file.
For detailed information on the available methods and their usage, refer to the `roo-code.d.ts` file.

Some files were not shown because too many files have changed in this diff Show more