feat: add GOOGLE_WORKSPACE_CLI_API_BASE_URL for custom/mock endpoint support#100
feat: add GOOGLE_WORKSPACE_CLI_API_BASE_URL for custom/mock endpoint support#100xdotli wants to merge 4 commits intogoogleworkspace:mainfrom
Conversation
🦋 Changeset detectedLatest commit: 7cc5875 The changes in this PR will be included in the next version bump. This PR includes changesets to release 1 package
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the application's testability and development workflow by allowing users to specify a custom API base URL via an environment variable. This enables seamless interaction with mock API servers, facilitating local development and continuous integration testing without requiring real Google API calls or OAuth authentication. The change ensures that the application's command structure remains intact while redirecting actual API requests to the specified endpoint. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #100 +/- ##
==========================================
+ Coverage 54.88% 54.89% +0.01%
==========================================
Files 38 38
Lines 13085 13137 +52
==========================================
+ Hits 7182 7212 +30
- Misses 5903 5925 +22 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature for testing against mock API servers by using the GWS_API_BASE_URL environment variable, with a well-designed implementation using LazyLock and a refactored resolve_auth function. However, a critical security vulnerability has been identified: the current implementation allows for silent redirection of sensitive API traffic and automatically disables authentication when GWS_API_BASE_URL is set. This, combined with the automatic loading of .env files, creates a risk where an attacker could hijack a user's API requests by placing a malicious .env file in their working directory. To address this, I recommend adding a mandatory warning message when a custom endpoint is active to ensure users are aware of the redirection and the disabled security state. Additionally, consider improving error handling for authentication failures to provide better feedback to the user.
src/executor.rs
Outdated
| } | ||
| match crate::auth::get_token(scopes).await { | ||
| Ok(t) => (Some(t), AuthMethod::OAuth), | ||
| Err(_) => (None, AuthMethod::None), |
There was a problem hiding this comment.
The current implementation swallows errors from auth::get_token, which can hide useful debugging information from users if their authentication setup is misconfigured (e.g., a corrupted credentials file). While the subsequent API call will likely fail, the original, more specific error is lost.
Consider logging the error to stderr to improve debuggability, while maintaining the current behavior of attempting an unauthenticated request.
Err(e) => {
// Not a fatal error; the request will be tried without auth.
// But logging it helps debug credential issues.
eprintln!("[gws] Warning: could not acquire authentication token: {e}. Proceeding unauthenticated.");
(None, AuthMethod::None)
}ed54a90 to
c22dc58
Compare
Review FeedbackThanks for this PR — the design is clean and smart. A few items to address before merging: Required
Recommended
Otherwise the approach looks solid — rewriting the Discovery Document URLs while still fetching the real schema from Google is the right call. |
…support Add an environment variable that redirects all API requests to a custom endpoint (e.g., a mock server). When set, OAuth authentication is automatically skipped while the real Discovery Document is still fetched so the CLI command tree remains fully functional. Changes: - src/discovery.rs: add custom_api_base_url() with LazyLock caching; rewrite Discovery Document URLs when env var is set - src/executor.rs: add resolve_auth() that skips OAuth for custom endpoints - src/main.rs, src/mcp_server.rs: use resolve_auth() for consistent behavior - AGENTS.md: document env var with security note - Add changeset and unit tests for URL rewriting
8f9a36e to
30e0f85
Compare
hi @jpoehnelt thanks for the speedy and thorough review! I have addressed all the changes and let cursor go thru it + reran my reproduce steps and it work as expected. Would appreciate another review. Thanks! |
jpoehnelt
left a comment
There was a problem hiding this comment.
base_urlrewrite loses the service path:rewrite_base_urlsets
base_urlto just the custom host, dropping the/gmail/v1/prefix. This
means final request URLs won't include the service path. Please preserve
the path portion from the originalbase_urlor verify thatbuild_url()
reconstructs it fromroot_url + service_path.- Auth error logging regression:
resolve_auth()silently discards auth
errors. The oldmcp_server.rscode logged them. Please add an
eprintln!warning in theErrbranch.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature for testing against mock API servers by adding the GOOGLE_WORKSPACE_CLI_API_BASE_URL environment variable. However, a critical vulnerability exists due to an inconsistency between the environment variable name implemented in the code (GOOGLE_WORKSPACE_CLI_API_BASE_URL) and the one documented in the PR description and examples (GWS_API_BASE_URL). This mismatch could lead to users accidentally performing actions on production Google Workspace accounts, posing a significant risk to data integrity. Additionally, the implementation of URL rewriting in rewrite_base_url has been identified as potentially leading to incorrect request URLs and may inadvertently strip the service path (e.g., /gmail/v1/) from API requests. Furthermore, a refactoring silently removed a useful warning on authentication failure in the MCP server, which should be restored.
| fn rewrite_base_url(doc: &mut RestDescription, base: &str) { | ||
| let base_trimmed = base.trim_end_matches('/'); | ||
| doc.root_url = format!("{base_trimmed}/"); | ||
| doc.base_url = Some(format!("{base_trimmed}/")); | ||
| } |
There was a problem hiding this comment.
The current implementation of rewrite_base_url incorrectly rewrites base_url. It replaces the entire base_url with the custom endpoint, which causes the service path (e.g., /gmail/v1/) to be lost. This will result in incorrect API request URLs.
For example, a request to gmail.users.getProfile would be sent to http://localhost:8001/users/me/profile instead of the correct http://localhost:8001/gmail/v1/users/me/profile.
The fix is to replace only the host part of the base_url, preserving the path. Here is a suggested implementation:
fn rewrite_base_url(doc: &mut RestDescription, base: &str) {
let base_trimmed = base.trim_end_matches('/');
let new_root_url = format!("{base_trimmed}/");
let original_root_url = std::mem::replace(&mut doc.root_url, new_root_url);
if let Some(base_url) = &mut doc.base_url {
*base_url = base_url.replace(&original_root_url, &doc.root_url);
}
}| #[test] | ||
| fn test_rewrite_base_url() { | ||
| let mut doc = RestDescription { | ||
| name: "gmail".to_string(), | ||
| version: "v1".to_string(), | ||
| root_url: "https://gmail.googleapis.com/".to_string(), | ||
| base_url: Some("https://gmail.googleapis.com/gmail/v1/".to_string()), | ||
| service_path: "gmail/v1/".to_string(), | ||
| ..Default::default() | ||
| }; | ||
|
|
||
| rewrite_base_url(&mut doc, "http://localhost:8099"); | ||
| assert_eq!(doc.root_url, "http://localhost:8099/"); | ||
| assert_eq!(doc.base_url.as_deref(), Some("http://localhost:8099/")); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_rewrite_base_url_trailing_slash() { | ||
| let mut doc = RestDescription { | ||
| name: "drive".to_string(), | ||
| version: "v3".to_string(), | ||
| root_url: "https://www.googleapis.com/".to_string(), | ||
| base_url: Some("https://www.googleapis.com/drive/v3/".to_string()), | ||
| ..Default::default() | ||
| }; | ||
|
|
||
| rewrite_base_url(&mut doc, "http://localhost:8099/"); | ||
| assert_eq!(doc.root_url, "http://localhost:8099/"); | ||
| assert_eq!(doc.base_url.as_deref(), Some("http://localhost:8099/")); | ||
| } |
There was a problem hiding this comment.
These tests need to be updated to reflect the correct behavior of rewrite_base_url. The assertions should check that the service path is preserved in the rewritten base_url.
#[test]
fn test_rewrite_base_url() {
let mut doc = RestDescription {
name: "gmail".to_string(),
version: "v1".to_string(),
root_url: "https://gmail.googleapis.com/".to_string(),
base_url: Some("https://gmail.googleapis.com/gmail/v1/".to_string()),
service_path: "gmail/v1/".to_string(),
..Default::default()
};
rewrite_base_url(&mut doc, "http://localhost:8099");
assert_eq!(doc.root_url, "http://localhost:8099/");
assert_eq!(
doc.base_url.as_deref(),
Some("http://localhost:8099/gmail/v1/")
);
}
#[test]
fn test_rewrite_base_url_trailing_slash() {
let mut doc = RestDescription {
name: "drive".to_string(),
version: "v3".to_string(),
root_url: "https://www.googleapis.com/".to_string(),
base_url: Some("https://www.googleapis.com/drive/v3/".to_string()),
..Default::default()
};
rewrite_base_url(&mut doc, "http://localhost:8099/");
assert_eq!(doc.root_url, "http://localhost:8099/");
assert_eq!(
doc.base_url.as_deref(),
Some("http://localhost:8099/drive/v3/")
);
}
src/executor.rs
Outdated
| pub async fn resolve_auth(scopes: &[&str], account: Option<&str>) -> (Option<String>, AuthMethod) { | ||
| if crate::discovery::custom_api_base_url().is_some() { | ||
| return (None, AuthMethod::None); | ||
| } | ||
| match crate::auth::get_token(scopes, account).await { | ||
| Ok(t) => (Some(t), AuthMethod::OAuth), | ||
| Err(_) => (None, AuthMethod::None), | ||
| } | ||
| } |
There was a problem hiding this comment.
The resolve_auth function currently swallows errors from crate::auth::get_token, which can hide important authentication problems from the user (e.g., invalid account, missing default account). This also removes a useful warning message that was previously shown in the MCP server on auth failure.
To preserve the original behavior and provide better error feedback, I suggest changing resolve_auth to return a Result that propagates the error from get_token. The call sites can then decide whether to show a warning or ignore the error.
| pub async fn resolve_auth(scopes: &[&str], account: Option<&str>) -> (Option<String>, AuthMethod) { | |
| if crate::discovery::custom_api_base_url().is_some() { | |
| return (None, AuthMethod::None); | |
| } | |
| match crate::auth::get_token(scopes, account).await { | |
| Ok(t) => (Some(t), AuthMethod::OAuth), | |
| Err(_) => (None, AuthMethod::None), | |
| } | |
| } | |
| pub async fn resolve_auth(scopes: &[&str], account: Option<&str>) -> anyhow::Result<(Option<String>, AuthMethod)> { | |
| if crate::discovery::custom_api_base_url().is_some() { | |
| return Ok((None, AuthMethod::None)); | |
| } | |
| match crate::auth::get_token(scopes, account).await { | |
| Ok(t) => Ok((Some(t), AuthMethod::OAuth)), | |
| Err(e) => Err(e), | |
| } | |
| } |
src/main.rs
Outdated
| let (token, auth_method) = | ||
| executor::resolve_auth(&scopes, account.as_deref()).await; |
There was a problem hiding this comment.
Following the change to resolve_auth to return a Result, this call site should be updated to handle the error. To preserve the existing behavior of proceeding without authentication on failure, you can match on the result and provide a fallback.
| let (token, auth_method) = | |
| executor::resolve_auth(&scopes, account.as_deref()).await; | |
| let (token, auth_method) = | |
| match executor::resolve_auth(&scopes, account.as_deref()).await { | |
| Ok(auth) => auth, | |
| Err(_) => (None, executor::AuthMethod::None), | |
| }; |
src/mcp_server.rs
Outdated
| (None, crate::executor::AuthMethod::None) | ||
| } | ||
| }; | ||
| let (token, auth_method) = crate::executor::resolve_auth(&scopes, None).await; |
There was a problem hiding this comment.
Following the change to resolve_auth to return a Result, this call site should be updated to handle the error. This will allow you to restore the warning message that was previously displayed on authentication failure, which is useful for debugging.
| let (token, auth_method) = crate::executor::resolve_auth(&scopes, None).await; | |
| let (token, auth_method) = match crate::executor::resolve_auth(&scopes, None).await { | |
| Ok(auth) => auth, | |
| Err(e) => { | |
| eprintln!( | |
| "[gws mcp] Warning: Authentication failed, proceeding without credentials: {e}" | |
| ); | |
| (None, crate::executor::AuthMethod::None) | |
| } | |
| }; |
|
@jpoehnelt thanks for the thorough review! Here's what was done to address everything: For the base_url rewrite losing the service path — this was a real bug. The original implementation set base_url to just the custom host, which broke any API that uses a non-empty servicePath. Gmail worked fine because its servicePath is empty and the method paths include the full path (e.g. gmail/v1/users/{userId}/profile), but Drive has servicePath "drive/v3/" with short method paths like "files", so the final URL became http://localhost:8001/files instead of http://localhost:8001/drive/v3/files. Fixed by using str::replace on the original root_url so only the host gets swapped and the path portion is preserved. Added three unit tests covering both styles plus the None base_url edge case. For the auth error logging regression — resolve_auth now returns Result so each call site can handle errors with its own context. main.rs silently falls back to unauthenticated (same as before), and mcp_server.rs restores the original "[gws mcp] Warning: Authentication failed, proceeding without credentials" message that was there before the refactor. Also fixed cargo fmt issues that were in the original commit. Tested all of this against a mock Gmail API server running on localhost:8001 (a FastAPI/uvicorn server implementing the full Gmail REST surface — messages, threads, labels, drafts, settings, history, profile, attachments, watch/stop, plus admin endpoints for seeding/reset/snapshots). After seeding the server, verified the following endpoints all return correct data through gws: users.getProfile, users.messages.list, users.messages.get (with path param substitution for messageId), users.labels.list, users.threads.list, users.drafts.list. Also verified Drive files.list dry-run now produces the correct URL with the service path preserved. Without the env var set, everything still hits the real Google APIs as expected. Also changed PR title, added logging as mentioned. Lmk if I should start mock servers on other services and test them. |
Reproduce steps:
# Install from this branch cargo install --git https://github.com/benchflow-ai/cli --branch feat/custom-endpoint-overrideWithout env var — request goes to Google APIs:
With
GWS_API_BASE_URL— request redirected to custom endpoint:--dry-runbuilds the request without sending it. Discovery Documents are still fetched from Google so the full command tree works — only the API request URLs are rewritten.Summary
Adds a
GWS_API_BASE_URLenvironment variable that redirects all API requests to a custom endpoint (e.g., a mock server). When set:root_urlandbase_urlin the Discovery Document are rewritten to the custom endpointMotivation
Testing agents and skills against mock API servers is a common need for development and CI. The
gogCLI already supports this viaGOG_ENDPOINT. This PR brings the same capability togws.Example usage
Changes
src/discovery.rs: Addcustom_api_base_url()withLazyLockcaching; rewrite Discovery Document URLs when env var is setsrc/executor.rs: Addresolve_auth()that skips OAuth for custom endpointssrc/main.rs,src/mcp_server.rs: Useresolve_auth()for consistent behavior