Add full KT integration to Sheets and Docs guardrails - #803
Closed
cjdaley wants to merge 6 commits into
Closed
Conversation
- Force stateless_http=True for single-user streamable-http deployments. Clients like Relevance AI open fresh connections per call and never reuse Mcp-Session-Id; stateful mode caused FastMCP to accumulate abandoned session objects and grow memory unboundedly. - Route logging.basicConfig to sys.stdout in main.py and fastmcp_server.py. Railway tags all stderr as severity=error regardless of level; stdout is classified correctly by Railway log collector. - Suppress authlib.jose DeprecationWarning pending joserfc migration. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ropic strict mode compatibility Optional[dict] generates an anyOf schema without additionalProperties:false, which Anthropic rejects when strict_mode is enabled on an agent. Accepting a JSON string instead produces a clean string schema and parses it before passing to the Drive API. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughThis PR implements input validation and quota tracking guardrails for Google Docs and Sheets operations, updates Google Drive to handle JSON-formatted file properties, and configures consistent logging and deployment infrastructure across the application. ChangesGoogle Workspace Guardrails and Infrastructure
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
gsheets/test_guardrails.py (2)
361-363: 💤 Low valueTest count is misleading.
The message claims "21 tests passed" but there are only 8 test functions. If "21" refers to individual assertions, consider clarifying or counting actual test functions for accuracy.
💡 Suggested fix
print("\n" + "=" * 60) - print("✓ All 21 tests passed!") + print("✓ All 8 test functions passed!") print("=" * 60)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gsheets/test_guardrails.py` around lines 361 - 363, The final test summary prints a hardcoded "21 tests passed" which is misleading; update the end-of-run messaging in the print block (the three print(...) lines) to report the actual number of tests run or assertions — e.g., compute a variable for the count (derived from the test function list or a counter incremented on each successful assertion) and interpolate that value into the summary string, or change the text to "All tests passed!" without a number to avoid incorrect counts.
277-289: ⚡ Quick winRedundant nested patch - decorator's mock is unused.
The
@patchdecorator on line 277 injectsmock_apias a parameter, but inside the function, line 283 creates anotherwith patch(...)that re-patches the same target. The decorator'smock_apiargument is never used, making the decorator unnecessary.♻️ Option 1: Keep decorator, remove inner context manager
`@patch`("sheets_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api) def test_quota_persistence(mock_api): """Test quota persistence in Knowledge Table.""" print("\nTest 8: Quota persistence...") # Sub-test 1: Load existing quota from KT - with patch("sheets_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api): - quota_state = load_quota_state_from_kt("test@gmail.com") + quota_state = load_quota_state_from_kt("test@gmail.com")♻️ Option 2: Remove decorator, keep inner context manager only
-@patch("sheets_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api) -def test_quota_persistence(mock_api): +def test_quota_persistence(): """Test quota persistence in Knowledge Table.""" print("\nTest 8: Quota persistence...") # Sub-test 1: Load existing quota from KT with patch("sheets_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api): quota_state = load_quota_state_from_kt("test@gmail.com")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gsheets/test_guardrails.py` around lines 277 - 289, The test_quota_persistence function has a redundant `@patch`("sheets_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api) decorator because the test body re-patches the same target with a with patch(...) context manager; remove one of them to avoid an unused mock parameter (either delete the decorator and keep the inner with patch, or remove the inner with patch and use the decorator-injected mock_api), and ensure load_quota_state_from_kt("test@gmail.com") is called while relevance_raw_api is patched so the mock_relevance_raw_api is used.gsheets/sheets_guardrails.py (1)
332-343: 💤 Low valueConsider checking API response for errors.
The
responsefromrelevance_raw_apiis captured but never inspected. If the Knowledge Table API returns an error payload without raising an exception, the failure would go undetected.💡 Optional: Add response validation
response = relevance_raw_api( endpoint="/knowledge/update", method="PATCH", body={ "knowledge_set": knowledge_set, "document_id": drive_doc_id, "fields": { "quota_used": new_drive_quota, "updated_at": datetime.now(timezone.utc).isoformat() } } ) + if response and response.get("status") == "error": + logger.warning(f"Drive quota update returned error: {response}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gsheets/sheets_guardrails.py` around lines 332 - 343, The call to relevance_raw_api (in the block using variables knowledge_set, drive_doc_id, new_drive_quota) assigns response but never checks it; update this part to validate the API result (e.g., check response status, an "error" or "success" field, or HTTP status code) and handle failures by logging the error with context (include knowledge_set and drive_doc_id) and raising/returning an exception so failures don't go silently unnoticed; ensure the validation path covers both non-2xx statuses and error payloads returned without exceptions.gdocs/docs_guardrails.py (1)
90-96: ⚡ Quick winType hints should use
Optional[int]for mypy strict compliance.Per coding guidelines requiring mypy strict mode compatibility, parameters that can be
Noneshould be typed asOptional[int]rather thanint = None.Proposed fix
def validate_modify_doc_text_input( user_google_email: str, document_id: str, text: str = "", - start_index: int = None, - end_index: int = None + start_index: Optional[int] = None, + end_index: Optional[int] = None ) -> Dict[str, Any]:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gdocs/docs_guardrails.py` around lines 90 - 96, The function validate_modify_doc_text_input declares start_index and end_index defaulting to None but types them as int; change their annotations to Optional[int] and add/import Optional from typing so the signature uses start_index: Optional[int] and end_index: Optional[int] to satisfy mypy strict mode (update the function definition and any related type hints/usages inside validate_modify_doc_text_input accordingly).gdocs/test_docs_guardrails.py (1)
236-246: 💤 Low valueTest may flake at minute boundaries.
There's a potential race condition: if the clock rolls over to the next minute between
get_docs_quota_key()(line 240) anddatetime.now()(line 241), the keys won't match and the test will fail intermittently.Proposed fix using time mocking
+from unittest.mock import patch +from datetime import datetime, timezone + def test_quota_key_generation(): """Test 4: Quota key generation""" print("\n=== Test 4: Quota Key Generation ===") - key = get_docs_quota_key() - now = datetime.now(timezone.utc) - expected_prefix = f"docs_quota_{now.strftime('%Y_%m_%d_%H_%M')}" + fixed_time = datetime(2026, 5, 20, 14, 32, 0, tzinfo=timezone.utc) + with patch("gdocs.docs_guardrails.datetime") as mock_dt: + mock_dt.now.return_value = fixed_time + mock_dt.side_effect = lambda *args, **kw: datetime(*args, **kw) + key = get_docs_quota_key() + expected_prefix = "docs_quota_2026_05_20_14_32" assert key.startswith("docs_quota_"), "Test 4.1 failed: key should start with 'docs_quota_'" assert len(key) == len(expected_prefix), "Test 4.1 failed: key format incorrect" + assert key == expected_prefix, f"Test 4.1 failed: expected {expected_prefix}, got {key}" print(f"✓ Test 4.1: Quota key generated correctly: {key}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gdocs/test_docs_guardrails.py` around lines 236 - 246, The test is flaky because it calls get_docs_quota_key() and datetime.now() separately, allowing the minute to roll over; fix by making the test deterministic: either freeze/mock the current time before calling get_docs_quota_key() (e.g., monkeypatch datetime.now or use freezegun) and then build expected_prefix from that same frozen time, or change the assertion to validate the key's prefix pattern/timestamp format by extracting the timestamp from key (using get_docs_quota_key and test_quota_key_generation identifiers) rather than comparing two independently sampled times.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@gdocs/docs_guardrails.py`:
- Around line 338-353: The metadata currently always reports
"quota_state_updated": True even when the exception handler (except Exception as
e) runs; fix this by introducing a local boolean (e.g., quota_update_succeeded =
True) before attempting the quota update and set it to False inside the except
block, then use that boolean for the metadata field and compute
"quota_remaining" only when quota_update_succeeded is True (otherwise set
quota_remaining to None or leave previous known value); update references to
quota_state_updated, quota_remaining, and any dependent values (api_calls_made,
DOCS_QUOTA_PER_MINUTE, quota_state.get("docs_quota_key")) so the returned dict
accurately reflects whether the quota update actually succeeded.
In `@gdocs/test_docs_guardrails.py`:
- Around line 18-31: The test imports are relative and will fail when running
from repo root; change the import statement to use the fully qualified module
path (e.g., import the symbols from gdocs.docs_guardrails instead of
docs_guardrails) so pytest can locate the module, and also update any
unittest.mock.patch decorators that target docs_guardrails (the ones around the
tests at the spots mentioned) to use the fully qualified target names (e.g.,
"gdocs.docs_guardrails.<function_or_attribute>") so the patching points to the
correct module; update references to functions/constants like
validate_create_doc_input, validate_modify_doc_text_input, get_docs_quota_key,
load_quota_state_from_kt, check_docs_quota, validate_write_and_update_quota,
MAX_DOC_TITLE_LENGTH, MAX_TEXT_SIZE_BYTES, DOCS_QUOTA_PER_MINUTE,
DOCS_WARNING_THRESHOLD, DOCS_EXHAUSTED_THRESHOLD and the patch decorators
referenced on the two lines to use gdocs.docs_guardrails.<name>.
- Around line 344-375: The test uses `@patch`("docs_guardrails.relevance_raw_api",
side_effect=mock_relevance_raw_api) on test_quota_persistence but only exercises
the patch for Test 8.1 via a nested with patch; Test 8.2 calls
validate_write_and_update_quota without any active patch so the internal quota
update raises NotImplementedError and is silently logged. Fix by applying the
same mock_relevance_raw_api for Test 8.2 as well—either remove the nested with
and use the decorator-provided mock_api across both subtests or wrap the
validate_write_and_update_quota call in with
patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)
so load_quota_state_from_kt and validate_write_and_update_quota run with the
mocked relevance_raw_api.
In `@gdrive/drive_tools.py`:
- Around line 1465-1468: Replace the current error-return behavior when parsing
properties: instead of returning a string that echoes the raw properties, parse
with json.loads, ensure the result is a dict (not list/str/etc.), and on failure
raise a ValueError with a generic message like "Invalid properties: must be a
JSON object" (do not include the input). Update the code path that sets
update_body["properties"] so it only assigns when json.loads succeeds and
returns a dict; otherwise raise ValueError to signal validation failure.
In `@gsheets/test_guardrails.py`:
- Around line 10-17: Update the import to use the repo-root-qualified module
name: replace the current "from sheets_guardrails import validate_append_input,
normalize_rows, check_quota_limits, load_quota_state_from_kt, get_quota_keys,
validate_write_and_update_quota" with an import from "gsheets.sheets_guardrails"
so the test can be run from the repository root and still resolve the module;
keep the same function names (validate_append_input, normalize_rows,
check_quota_limits, load_quota_state_from_kt, get_quota_keys,
validate_write_and_update_quota).
---
Nitpick comments:
In `@gdocs/docs_guardrails.py`:
- Around line 90-96: The function validate_modify_doc_text_input declares
start_index and end_index defaulting to None but types them as int; change their
annotations to Optional[int] and add/import Optional from typing so the
signature uses start_index: Optional[int] and end_index: Optional[int] to
satisfy mypy strict mode (update the function definition and any related type
hints/usages inside validate_modify_doc_text_input accordingly).
In `@gdocs/test_docs_guardrails.py`:
- Around line 236-246: The test is flaky because it calls get_docs_quota_key()
and datetime.now() separately, allowing the minute to roll over; fix by making
the test deterministic: either freeze/mock the current time before calling
get_docs_quota_key() (e.g., monkeypatch datetime.now or use freezegun) and then
build expected_prefix from that same frozen time, or change the assertion to
validate the key's prefix pattern/timestamp format by extracting the timestamp
from key (using get_docs_quota_key and test_quota_key_generation identifiers)
rather than comparing two independently sampled times.
In `@gsheets/sheets_guardrails.py`:
- Around line 332-343: The call to relevance_raw_api (in the block using
variables knowledge_set, drive_doc_id, new_drive_quota) assigns response but
never checks it; update this part to validate the API result (e.g., check
response status, an "error" or "success" field, or HTTP status code) and handle
failures by logging the error with context (include knowledge_set and
drive_doc_id) and raising/returning an exception so failures don't go silently
unnoticed; ensure the validation path covers both non-2xx statuses and error
payloads returned without exceptions.
In `@gsheets/test_guardrails.py`:
- Around line 361-363: The final test summary prints a hardcoded "21 tests
passed" which is misleading; update the end-of-run messaging in the print block
(the three print(...) lines) to report the actual number of tests run or
assertions — e.g., compute a variable for the count (derived from the test
function list or a counter incremented on each successful assertion) and
interpolate that value into the summary string, or change the text to "All tests
passed!" without a number to avoid incorrect counts.
- Around line 277-289: The test_quota_persistence function has a redundant
`@patch`("sheets_guardrails.relevance_raw_api",
side_effect=mock_relevance_raw_api) decorator because the test body re-patches
the same target with a with patch(...) context manager; remove one of them to
avoid an unused mock parameter (either delete the decorator and keep the inner
with patch, or remove the inner with patch and use the decorator-injected
mock_api), and ensure load_quota_state_from_kt("test@gmail.com") is called while
relevance_raw_api is patched so the mock_relevance_raw_api is used.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b7cb1916-36a5-46a4-bbb6-75a315669941
📒 Files selected for processing (8)
fastmcp_server.pygdocs/docs_guardrails.pygdocs/test_docs_guardrails.pygdrive/drive_tools.pygsheets/sheets_guardrails.pygsheets/test_guardrails.pymain.pyrailway.toml
Comment on lines
+338
to
+353
| except Exception as e: | ||
| logger.error(f"Failed to update Docs quota: {str(e)}") | ||
|
|
||
| return { | ||
| "status": "success", | ||
| "operation": operation_type, | ||
| "is_complete": True, | ||
| "quota_impact": api_calls_made, | ||
| "metadata": { | ||
| "api_calls_consumed": api_calls_made, | ||
| "quota_state_updated": True, | ||
| "quota_remaining": DOCS_QUOTA_PER_MINUTE - (quota_state["docs_quota_used"] + api_calls_made), | ||
| "quota_key": quota_state.get("docs_quota_key"), | ||
| "note": "Shared 300/min pool with Drive, Sheets, and other APIs" | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
quota_state_updated is always True even when the update fails.
When an exception occurs during quota update (lines 338-339), the function logs the error but still returns quota_state_updated: True in the metadata (line 348). This is misleading to callers who may rely on this flag to know if quota tracking is accurate.
Proposed fix
+ quota_update_succeeded = False
try:
quota_key = quota_state.get("docs_quota_key")
new_quota = quota_state.get("docs_quota_used", 0) + api_calls_made
if quota_state.get("docs_doc_id"):
# PATCH update existing Docs quota
logger.debug(f"Updating Docs quota for {user_google_email} (doc_id={quota_state.get('docs_doc_id')})")
response = relevance_raw_api(
endpoint="/knowledge/update",
method="PATCH",
body={
"knowledge_set": knowledge_set,
"document_id": quota_state.get("docs_doc_id"),
"fields": {
"quota_used": new_quota,
"updated_at": datetime.now(timezone.utc).isoformat()
}
}
)
else:
# POST create new Docs quota entry
logger.debug(f"Creating Docs quota entry for {user_google_email}")
response = relevance_raw_api(
endpoint="/knowledge/add",
method="POST",
body={
"knowledge_set": knowledge_set,
"fields": {
"key": quota_key,
"quota_used": new_quota,
"user_google_email": user_google_email,
"service": "docs",
"updated_at": datetime.now(timezone.utc).isoformat()
}
}
)
logger.info(f"Docs quota updated: {new_quota}/{DOCS_QUOTA_PER_MINUTE}")
+ quota_update_succeeded = True
except Exception as e:
logger.error(f"Failed to update Docs quota: {str(e)}")
return {
"status": "success",
"operation": operation_type,
"is_complete": True,
"quota_impact": api_calls_made,
"metadata": {
"api_calls_consumed": api_calls_made,
- "quota_state_updated": True,
+ "quota_state_updated": quota_update_succeeded,
"quota_remaining": DOCS_QUOTA_PER_MINUTE - (quota_state["docs_quota_used"] + api_calls_made),
"quota_key": quota_state.get("docs_quota_key"),
"note": "Shared 300/min pool with Drive, Sheets, and other APIs"
}
}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gdocs/docs_guardrails.py` around lines 338 - 353, The metadata currently
always reports "quota_state_updated": True even when the exception handler
(except Exception as e) runs; fix this by introducing a local boolean (e.g.,
quota_update_succeeded = True) before attempting the quota update and set it to
False inside the except block, then use that boolean for the metadata field and
compute "quota_remaining" only when quota_update_succeeded is True (otherwise
set quota_remaining to None or leave previous known value); update references to
quota_state_updated, quota_remaining, and any dependent values (api_calls_made,
DOCS_QUOTA_PER_MINUTE, quota_state.get("docs_quota_key")) so the returned dict
accurately reflects whether the quota update actually succeeded.
Comment on lines
+18
to
+31
| from docs_guardrails import ( | ||
| validate_create_doc_input, | ||
| validate_modify_doc_text_input, | ||
| validate_find_replace_input, | ||
| get_docs_quota_key, | ||
| load_quota_state_from_kt, | ||
| check_docs_quota, | ||
| validate_write_and_update_quota, | ||
| MAX_DOC_TITLE_LENGTH, | ||
| MAX_TEXT_SIZE_BYTES, | ||
| DOCS_QUOTA_PER_MINUTE, | ||
| DOCS_WARNING_THRESHOLD, | ||
| DOCS_EXHAUSTED_THRESHOLD, | ||
| ) |
Contributor
There was a problem hiding this comment.
Import path may fail when tests run from repository root.
The import from docs_guardrails import ... assumes the test runs from within the gdocs/ directory. Standard test runners like pytest execute from the repository root, which would cause ModuleNotFoundError. Consider using the fully qualified module path.
Proposed fix
-from docs_guardrails import (
+from gdocs.docs_guardrails import (
validate_create_doc_input,
validate_modify_doc_text_input,
validate_find_replace_input,
get_docs_quota_key,
load_quota_state_from_kt,
check_docs_quota,
validate_write_and_update_quota,
MAX_DOC_TITLE_LENGTH,
MAX_TEXT_SIZE_BYTES,
DOCS_QUOTA_PER_MINUTE,
DOCS_WARNING_THRESHOLD,
DOCS_EXHAUSTED_THRESHOLD,
)Also update the patch decorator on line 344:
-@patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)
+@patch("gdocs.docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)And line 350:
- with patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api):
+ with patch("gdocs.docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gdocs/test_docs_guardrails.py` around lines 18 - 31, The test imports are
relative and will fail when running from repo root; change the import statement
to use the fully qualified module path (e.g., import the symbols from
gdocs.docs_guardrails instead of docs_guardrails) so pytest can locate the
module, and also update any unittest.mock.patch decorators that target
docs_guardrails (the ones around the tests at the spots mentioned) to use the
fully qualified target names (e.g.,
"gdocs.docs_guardrails.<function_or_attribute>") so the patching points to the
correct module; update references to functions/constants like
validate_create_doc_input, validate_modify_doc_text_input, get_docs_quota_key,
load_quota_state_from_kt, check_docs_quota, validate_write_and_update_quota,
MAX_DOC_TITLE_LENGTH, MAX_TEXT_SIZE_BYTES, DOCS_QUOTA_PER_MINUTE,
DOCS_WARNING_THRESHOLD, DOCS_EXHAUSTED_THRESHOLD and the patch decorators
referenced on the two lines to use gdocs.docs_guardrails.<name>.
Comment on lines
+344
to
+375
| @patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api) | ||
| def test_quota_persistence(mock_api): | ||
| """Test 8: Quota persistence in Knowledge Table""" | ||
| print("\n=== Test 8: Quota Persistence ===") | ||
|
|
||
| # Test 8.1: Load quota from KT | ||
| with patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api): | ||
| quota_state = load_quota_state_from_kt("test@gmail.com") | ||
| assert quota_state["user_google_email"] == "test@gmail.com" | ||
| assert quota_state["docs_quota_key"].startswith("docs_quota_") | ||
| print("✓ Test 8.1: Quota state initialized from KT") | ||
|
|
||
| # Test 8.2: Update quota after successful write | ||
| quota_state = { | ||
| "docs_quota_used": 150, | ||
| "docs_doc_id": "docs_quota_id_123", | ||
| "docs_quota_key": "docs_quota_2026_05_20_14_32", | ||
| "user_google_email": "test@gmail.com" | ||
| } | ||
|
|
||
| write_result = "Successfully created document with 500 character(s)" | ||
| result = validate_write_and_update_quota( | ||
| write_result=write_result, | ||
| operation_type="create_doc", | ||
| quota_state=quota_state, | ||
| user_google_email="test@gmail.com" | ||
| ) | ||
|
|
||
| assert result["status"] == "success", f"Test 8.2 failed: should be success, got {result['status']}" | ||
| assert result["metadata"]["quota_state_updated"] == True, "Test 8.2 failed: quota should be updated" | ||
| assert result["quota_impact"] == 1, "Test 8.2 failed: all ops should use 1 API call (atomic)" | ||
| print("✓ Test 8.2: Quota updated in KT after successful write") |
Contributor
There was a problem hiding this comment.
Test 8.2 doesn't patch relevance_raw_api, so quota update fails silently.
The @patch decorator at line 344 injects mock_api but it's unused. Test 8.1 (line 350) creates a nested context manager patch, but Test 8.2 (lines 356-375) calls validate_write_and_update_quota without any active patch. This means the quota update inside that function will raise NotImplementedError, which gets caught and logged. The test passes only because the function returns success regardless of update failure (which is a separate issue flagged in the implementation).
Consider applying the patch consistently to fully test the KT integration path.
Proposed fix
-@patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)
-def test_quota_persistence(mock_api):
+def test_quota_persistence():
"""Test 8: Quota persistence in Knowledge Table"""
print("\n=== Test 8: Quota Persistence ===")
# Test 8.1: Load quota from KT
- with patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api):
+ with patch("gdocs.docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api):
quota_state = load_quota_state_from_kt("test@gmail.com")
assert quota_state["user_google_email"] == "test@gmail.com"
assert quota_state["docs_quota_key"].startswith("docs_quota_")
print("✓ Test 8.1: Quota state initialized from KT")
# Test 8.2: Update quota after successful write
- quota_state = {
- "docs_quota_used": 150,
- "docs_doc_id": "docs_quota_id_123",
- "docs_quota_key": "docs_quota_2026_05_20_14_32",
- "user_google_email": "test@gmail.com"
- }
-
- write_result = "Successfully created document with 500 character(s)"
- result = validate_write_and_update_quota(
- write_result=write_result,
- operation_type="create_doc",
- quota_state=quota_state,
- user_google_email="test@gmail.com"
- )
+ with patch("gdocs.docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api):
+ quota_state = {
+ "docs_quota_used": 150,
+ "docs_doc_id": "docs_quota_id_123",
+ "docs_quota_key": "docs_quota_2026_05_20_14_32",
+ "user_google_email": "test@gmail.com"
+ }
+
+ write_result = "Successfully created document with 500 character(s)"
+ result = validate_write_and_update_quota(
+ write_result=write_result,
+ operation_type="create_doc",
+ quota_state=quota_state,
+ user_google_email="test@gmail.com"
+ )
assert result["status"] == "success", f"Test 8.2 failed: should be success, got {result['status']}"
assert result["metadata"]["quota_state_updated"] == True, "Test 8.2 failed: quota should be updated"
assert result["quota_impact"] == 1, "Test 8.2 failed: all ops should use 1 API call (atomic)"
print("✓ Test 8.2: Quota updated in KT after successful write")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gdocs/test_docs_guardrails.py` around lines 344 - 375, The test uses
`@patch`("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)
on test_quota_persistence but only exercises the patch for Test 8.1 via a nested
with patch; Test 8.2 calls validate_write_and_update_quota without any active
patch so the internal quota update raises NotImplementedError and is silently
logged. Fix by applying the same mock_relevance_raw_api for Test 8.2 as
well—either remove the nested with and use the decorator-provided mock_api
across both subtests or wrap the validate_write_and_update_quota call in with
patch("docs_guardrails.relevance_raw_api", side_effect=mock_relevance_raw_api)
so load_quota_state_from_kt and validate_write_and_update_quota run with the
mocked relevance_raw_api.
Comment on lines
+1465
to
+1468
| try: | ||
| update_body["properties"] = json.loads(properties) | ||
| except (json.JSONDecodeError, ValueError): | ||
| return f"Error: properties must be a valid JSON string, e.g. '{{\"key\": \"value\"}}'. Received: {properties}" |
Contributor
There was a problem hiding this comment.
❓ Verification inconclusive
Script executed:
#!/bin/bash
# Verify the JSON-parse error path no longer returns raw payload text.
rg -n -C2 'json.loads\(properties\)|Received: \{properties\}|return f"Error: properties|raise ValueError\(' gdrive/drive_tools.pyRepository: taylorwilsdon/google_workspace_mcp
Repository: taylorwilsdon/google_workspace_mcp
Exit code: 0
stdout:
87- total_bytes += len(chunk)
88- if total_bytes > MAX_DOWNLOAD_BYTES:
89: raise ValueError(
90- f"Download from {redacted_url} exceeded {MAX_DOWNLOAD_BYTES} byte limit "
91- f"({total_bytes} bytes)"
--
1003- source_count = sum(1 for x in [content, file_path, file_url] if x is not None)
1004- if source_count == 0:
1005: raise ValueError(
1006- "You must provide one of: 'content', 'file_path', or 'file_url'."
1007- )
1008- if source_count > 1:
1009: raise ValueError("Provide only one of: 'content', 'file_path', or 'file_url'.")
1010-
1011- # Determine source MIME type
--
1016- source_mime_type = GOOGLE_DOCS_IMPORT_FORMATS[format_key]
1017- else:
1018: raise ValueError(
1019- f"Unsupported source_format: '{source_format}'. "
1020- f"Supported: {', '.join(ext.lstrip('.') for ext in GOOGLE_DOCS_IMPORT_FORMATS.keys())}"
--
1064- actual_path = file_path
1065- else:
1066: raise ValueError(
1067- f"file_path should be a local path or file:// URL, got: {file_path}"
1068- )
--
1072- raise FileNotFoundError(f"File not found: {actual_path}")
1073- if not path_obj.is_file():
1074: raise ValueError(f"Path is not a file: {actual_path}")
1075-
1076- file_data = await asyncio.to_thread(path_obj.read_bytes)
--
1088- parsed_url = urlparse(file_url)
1089- if parsed_url.scheme not in ("http", "https"):
1090: raise ValueError(f"file_url must be http:// or https://, got: {file_url}")
1091-
1092- # SSRF protection: block internal/private network URLs and validate redirects
--
1464- if properties is not None:
1465- try:
1466: update_body["properties"] = json.loads(properties)
1467- except (json.JSONDecodeError, ValueError):
1468: return f"Error: properties must be a valid JSON string, e.g. '{{\"key\": \"value\"}}'. Received: {properties}"
1469-
1470- async def _resolve_parent_arguments(parent_arg: Optional[str]) -> Optional[str]:
--
1695- valid_actions = ("grant", "grant_batch", "update", "revoke", "transfer_owner")
1696- if action not in valid_actions:
1697: raise ValueError(
1698- f"Invalid action '{action}'. Must be one of: {', '.join(valid_actions)}"
1699- )
--
1711-
1712- if share_type in ("user", "group") and not share_with:
1713: raise ValueError(f"share_with is required for share_type '{share_type}'")
1714- if share_type == "domain" and not share_with:
1715: raise ValueError(
1716- "share_with (domain name) is required for share_type 'domain'"
1717- )
--
1767- if action == "grant_batch":
1768- if not recipients:
1769: raise ValueError("recipients list is required for 'grant_batch' action")
1770-
1771- resolved_file_id, file_metadata = await resolve_drive_item(
--
1869- if action == "update":
1870- if not permission_id:
1871: raise ValueError("permission_id is required for 'update' action")
1872- if not role and not expiration_time:
1873: raise ValueError(
1874- "Must provide at least one of: role, expiration_time for 'update' action"
1875- )
--
1927- if action == "revoke":
1928- if not permission_id:
1929: raise ValueError("permission_id is required for 'revoke' action")
1930-
1931- resolved_file_id, file_metadata = await resolve_drive_item(
--
1955- # action == "transfer_owner"
1956- if not new_owner_email:
1957: raise ValueError("new_owner_email is required for 'transfer_owner' action")
1958-
1959- resolved_file_id, file_metadata = await resolve_drive_item(
--
2116- and copy_requires_writer_permission is None
2117- ):
2118: raise ValueError(
2119- "Must provide at least one of: link_sharing, writers_can_share, copy_requires_writer_permission"
2120- )
--
2122- valid_link_sharing = {"off", "reader", "commenter", "writer"}
2123- if link_sharing is not None and link_sharing not in valid_link_sharing:
2124: raise ValueError(
2125- f"Invalid link_sharing '{link_sharing}'. Must be one of: {', '.join(sorted(valid_link_sharing))}"
2126- )
Raise ValueError instead of returning error string; remove echo of raw properties payload.
Line 1468 returns an error message that echoes the full user input, which can leak sensitive values. Additionally, validation failure should raise an exception rather than returning a string like a success response. Parse the JSON, validate it's a dict, then raise ValueError with a generic message instead of returning the error string.
Suggested fix
if properties is not None:
try:
- update_body["properties"] = json.loads(properties)
+ parsed_properties = json.loads(properties)
+ if not isinstance(parsed_properties, dict):
+ raise ValueError("properties JSON must decode to an object")
+ update_body["properties"] = parsed_properties
except (json.JSONDecodeError, ValueError):
- return f"Error: properties must be a valid JSON string, e.g. '{{\"key\": \"value\"}}'. Received: {properties}"
+ raise ValueError(
+ "properties must be a valid JSON object string, e.g. '{\"key\": \"value\"}'."
+ ) from NonePer coding guidelines: "Never log or leak secrets, refresh tokens, or PII in exceptions or event streams."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| try: | |
| update_body["properties"] = json.loads(properties) | |
| except (json.JSONDecodeError, ValueError): | |
| return f"Error: properties must be a valid JSON string, e.g. '{{\"key\": \"value\"}}'. Received: {properties}" | |
| try: | |
| parsed_properties = json.loads(properties) | |
| if not isinstance(parsed_properties, dict): | |
| raise ValueError("properties JSON must decode to an object") | |
| update_body["properties"] = parsed_properties | |
| except (json.JSONDecodeError, ValueError): | |
| raise ValueError( | |
| "properties must be a valid JSON object string, e.g. '{\"key\": \"value\"}'." | |
| ) from None |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gdrive/drive_tools.py` around lines 1465 - 1468, Replace the current
error-return behavior when parsing properties: instead of returning a string
that echoes the raw properties, parse with json.loads, ensure the result is a
dict (not list/str/etc.), and on failure raise a ValueError with a generic
message like "Invalid properties: must be a JSON object" (do not include the
input). Update the code path that sets update_body["properties"] so it only
assigns when json.loads succeeds and returns a dict; otherwise raise ValueError
to signal validation failure.
Comment on lines
+10
to
+17
| from sheets_guardrails import ( | ||
| validate_append_input, | ||
| normalize_rows, | ||
| check_quota_limits, | ||
| load_quota_state_from_kt, | ||
| get_quota_keys, | ||
| validate_write_and_update_quota, | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how other test files in the repo handle imports
fd -e py 'test_' --exec head -30 {} \; 2>/dev/null | grep -E "^from|^import" | head -20Repository: taylorwilsdon/google_workspace_mcp
Length of output: 576
🏁 Script executed:
#!/bin/bash
# Verify the file exists and show the actual import statement
if [ -f "gsheets/test_guardrails.py" ]; then
echo "=== File exists. Lines 1-20 ==="
head -20 gsheets/test_guardrails.py
else
echo "File not found at gsheets/test_guardrails.py"
fi
# Check if sheets_guardrails module exists
echo ""
echo "=== Checking for sheets_guardrails module ==="
fd 'sheets_guardrails' -e pyRepository: taylorwilsdon/google_workspace_mcp
Length of output: 670
🏁 Script executed:
#!/bin/bash
# Check pytest configuration that might affect imports
echo "=== Checking pyproject.toml for pytest config ==="
grep -A 10 "\[tool.pytest" pyproject.toml 2>/dev/null || echo "No pytest config found"
echo ""
echo "=== Checking for conftest.py files ==="
find . -name "conftest.py" -type fRepository: taylorwilsdon/google_workspace_mcp
Length of output: 449
Fix import path to use fully-qualified module name.
The import from sheets_guardrails import ... uses a relative path that will fail when running pytest from the repository root. Pytest is configured with pythonpath = ["."], so imports must be qualified from the repo root. All other test files in the repository follow this convention.
🔧 Proposed fix
-from sheets_guardrails import (
+from gsheets.sheets_guardrails import (
validate_append_input,
normalize_rows,
check_quota_limits,
load_quota_state_from_kt,
get_quota_keys,
validate_write_and_update_quota,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| from sheets_guardrails import ( | |
| validate_append_input, | |
| normalize_rows, | |
| check_quota_limits, | |
| load_quota_state_from_kt, | |
| get_quota_keys, | |
| validate_write_and_update_quota, | |
| ) | |
| from gsheets.sheets_guardrails import ( | |
| validate_append_input, | |
| normalize_rows, | |
| check_quota_limits, | |
| load_quota_state_from_kt, | |
| get_quota_keys, | |
| validate_write_and_update_quota, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gsheets/test_guardrails.py` around lines 10 - 17, Update the import to use
the repo-root-qualified module name: replace the current "from sheets_guardrails
import validate_append_input, normalize_rows, check_quota_limits,
load_quota_state_from_kt, get_quota_keys, validate_write_and_update_quota" with
an import from "gsheets.sheets_guardrails" so the test can be run from the
repository root and still resolve the module; keep the same function names
(validate_append_input, normalize_rows, check_quota_limits,
load_quota_state_from_kt, get_quota_keys, validate_write_and_update_quota).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes
Test Results
Summary by CodeRabbit
Release Notes
New Features
Breaking Changes