Skip to content

Commit

Permalink
More fixes to proxy (#1237)
Browse files Browse the repository at this point in the history 
<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> Improved proxy handling in Rust and TypeScript by refining error
messages, URL construction, and request routing logic.
> 
>   - **Rust Proxy Handling**:
> - Improved error message in `to_base64_with_inferred_mime_type()` in
`mod.rs` to include `media_url.url`.
> - Refactored URL construction in `fetch_with_proxy()` in `mod.rs` to
use `new_proxy_url` variable.
>   - **TypeScript Proxy Middleware**:
> - Added logic in `extension.ts` to handle image GET requests by
setting path to empty string.
>     - Removed trailing slashes from paths in `extension.ts`.
> - Ensured `baml-original-url` header is used to route requests in
`extension.ts`.
>   - **Miscellaneous**:
> - Disabled ESLint rule `@typescript-eslint/no-misused-promises` in
`extension.ts`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for e38f259. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->
  • Loading branch information
@aaronvg
aaronvg authored Dec 12, 2024
1 parent 859c699 commit 16054f5
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 41 deletions.
22 changes: 11 additions & 11 deletions engine/baml-runtime/src/internal/llm_client/traits/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -635,9 +635,10 @@ async fn to_base64_with_inferred_mime_type(
Ok((base64, mime_type))
} else {
Err(anyhow::anyhow!(
"Failed to fetch media: {}, {}",
"Failed to fetch media: {} {}, {}",
response.status(),
response.text().await.unwrap_or_default()
media_url.url,
response.text().await.unwrap_or_default(),
))
}
}
Expand All @@ -664,15 +665,14 @@ async fn fetch_with_proxy(
let client = reqwest::Client::new();

let request = if let Some(proxy) = proxy_url {
client
.get(format!(
"{}{}",
proxy,
url.parse::<url::Url>()
.map_err(|e| anyhow::anyhow!("Failed to parse URL: {}", e))?
.path()
))
.header("baml-original-url", url)
let new_proxy_url = format!(
"{}{}",
proxy,
url.parse::<url::Url>()
.map_err(|e| anyhow::anyhow!("Failed to parse URL: {}", e))?
.path()
);
client.get(new_proxy_url).header("baml-original-url", url)
} else {
client.get(url)
};
Expand Down
60 changes: 30 additions & 30 deletions typescript/vscode-ext/packages/vscode/src/extension.ts
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
/* eslint-disable @typescript-eslint/no-misused-promises */
import * as vscode from 'vscode'
import axios from 'axios'
import glooLens from './LanguageToBamlCodeLensProvider'
Expand Down Expand Up @@ -193,11 +194,21 @@ export function activate(context: vscode.ExtensionContext) {
createProxyMiddleware({
changeOrigin: true,
pathRewrite: (path, req) => {
// Ensure the URL does not end with a slash
console.log('reqmethod', req.method)

// Remove the path in the case of images. Since we request things differently for image GET requests, where we add the url to localhost:4500/actual-url.png
// to prevent caching issues with Rust reqwest.
// But for normal completion POST requests, we always call localhost:4500.
// The original url is always in baml-original-url header.

// Check for file extensions and set path to empty string.
if (/\.[a-zA-Z0-9]+$/.test(path) && req.method === 'GET') {
return ''
}
// Remove trailing slash
if (path.endsWith('/')) {
return path.slice(0, -1)
}
console.log('pathRewrite', path, req)
return path
},
router: (req) => {
Expand All @@ -209,11 +220,14 @@ export function activate(context: vscode.ExtensionContext) {
delete req.headers['origin']

// Ensure the URL does not end with a slash
console.log('originalUrl1', originalUrl)
if (originalUrl.endsWith('/')) {
originalUrl = originalUrl.slice(0, -1)
}
console.log('returning original url', originalUrl)
return new URL(originalUrl).origin
// return new URL(originalUrl).toString()

return originalUrl
} else {
console.log('baml-original-url header is missing or invalid')
throw new Error('baml-original-url header is missing or invalid')
Expand All @@ -222,33 +236,19 @@ export function activate(context: vscode.ExtensionContext) {
logger: console,
on: {
proxyReq: (proxyReq, req, res) => {
console.log('proxying request')

try {
const bamlOriginalUrl = req.headers['baml-original-url']
if (bamlOriginalUrl === undefined) {
return
}
const targetUrl = new URL(bamlOriginalUrl)
// proxyReq.path = targetUrl.pathname
// proxyReq.p
// It is very important that we ONLY resolve against API_KEY_INJECTION_ALLOWED
// by using the URL origin! (i.e. NOT using str.startsWith - the latter can still
// leak API keys to malicious subdomains e.g. https://api.openai.com.evil.com)
// const headers = API_KEY_INJECTION_ALLOWED[proxyOrigin]
// if (headers === undefined) {
// return
// }
// for (const [header, value] of Object.entries(headers)) {
// proxyReq.setHeader(header, value)
// }
// proxyReq.removeHeader('origin')
// proxyReq.setHeader('Origin', targetUrl.origin)
console.info('Proxying an LLM request (to bypass CORS)', { proxyReq, req, res })
} catch (err) {
// This is not console.warn because it's not important
console.log('baml-original-url is not parsable', err)
}
// const bamlOriginalUrl = req.headers['baml-original-url']
// if (typeof bamlOriginalUrl === 'string') {
// const targetUrl = new URL(bamlOriginalUrl)
// // Copy all original headers except those we want to modify/remove
// Object.entries(req.headers).forEach(([key, value]) => {
// if (key !== 'host' && key !== 'origin' && key !== 'baml-original-url') {
// proxyReq.setHeader(key, value)
// }
// })
// // Set the correct origin and host headers
// proxyReq.setHeader('origin', targetUrl.origin)
// proxyReq.setHeader('host', targetUrl.host)
// }
},
proxyRes: (proxyRes, req, res) => {
proxyRes.headers['Access-Control-Allow-Origin'] = '*'
Expand Down

0 comments on commit 16054f5

Please sign in to comment.