Skip to content

Commit 4bb59d7

Browse files
chrishayukchrishayuk
committed
improved coverage
1 parent a92d1ab commit 4bb59d7

File tree

13 files changed

+667
-157
lines changed

13 files changed

+667
-157
lines changed

‎Makefile‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@ help:
1616
@echo " lint - Run code linters"
1717
@echo " format - Auto-format code"
1818
@echo " typecheck - Run type checking"
19-
@echo " check - Run all checks (lint, typecheck, test)"
19+
@echo " security - Run security checks"
20+
@echo " check - Run all checks (lint, typecheck, security, test)"
2021
@echo " run - Run the server"
2122
@echo " build - Build the project"
2223
@echo " publish - Build and publish to PyPI"
@@ -205,8 +206,19 @@ typecheck:
205206
echo "MyPy not found. Install with: pip install mypy"; \
206207
fi
207208

209+
# Security checks
210+
security:
211+
@echo "Running security checks..."
212+
@if command -v uv >/dev/null 2>&1; then \
213+
uv run bandit -r src -ll; \
214+
elif command -v bandit >/dev/null 2>&1; then \
215+
bandit -r src -ll; \
216+
else \
217+
echo "Bandit not found. Install with: pip install bandit"; \
218+
fi
219+
208220
# Run all checks
209-
check: lint typecheck test
221+
check: lint typecheck security test
210222
@echo "All checks completed."
211223

212224
# Show project info

‎examples/test_notion_oauth.py‎

Lines changed: 362 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,362 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Standalone example: OAuth authentication with Notion MCP server using chuk-tool-processor.
4+
5+
This script:
6+
1. Performs MCP OAuth flow with Notion (RFC 8414 + RFC 7591 + PKCE)
7+
2. Uses the OAuth token with HTTPStreamableTransport
8+
3. Proves OAuth headers are preserved and not overwritten
9+
10+
Usage:
11+
cd /Users/chrishay/chris-source/chuk-ai/chuk-tool-processor
12+
uv run python examples/test_notion_oauth.py
13+
"""
14+
15+
import asyncio
16+
import hashlib
17+
import json
18+
import logging
19+
import secrets
20+
import sys
21+
import webbrowser
22+
from base64 import urlsafe_b64encode
23+
from http.server import HTTPServer, BaseHTTPRequestHandler
24+
from urllib.parse import parse_qs, urlencode, urlparse
25+
26+
import httpx
27+
28+
from chuk_tool_processor.mcp.transport import HTTPStreamableTransport
29+
30+
# Set up logging - WARNING level to reduce noise
31+
logging.basicConfig(
32+
level=logging.WARNING,
33+
format="%(asctime)s - %(levelname)s - %(message)s"
34+
)
35+
36+
# Only show INFO for httpx requests to see OAuth flow
37+
logging.getLogger("httpx").setLevel(logging.INFO)
38+
39+
logger = logging.getLogger(__name__)
40+
41+
42+
class OAuthCallbackHandler(BaseHTTPRequestHandler):
43+
"""Handle OAuth callback."""
44+
45+
authorization_code = None
46+
47+
def do_GET(self):
48+
"""Handle the callback request."""
49+
query = parse_qs(urlparse(self.path).query)
50+
51+
if 'code' in query:
52+
OAuthCallbackHandler.authorization_code = query['code'][0]
53+
self.send_response(200)
54+
self.send_header('Content-type', 'text/html')
55+
self.end_headers()
56+
self.wfile.write("""
57+
<html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
58+
<h1 style="color: green;">Authentication Successful!</h1>
59+
<p>You can close this window and return to the terminal.</p>
60+
</body></html>
61+
""".encode('utf-8'))
62+
else:
63+
self.send_response(400)
64+
self.send_header('Content-type', 'text/html')
65+
self.end_headers()
66+
error = query.get('error', ['Unknown error'])[0]
67+
self.wfile.write(f"""
68+
<html><body style="font-family: sans-serif; text-align: center; padding: 50px;">
69+
<h1 style="color: red;">Authentication Failed</h1>
70+
<p>Error: {error}</p>
71+
</body></html>
72+
""".encode('utf-8'))
73+
74+
def log_message(self, format, *args):
75+
"""Suppress HTTP server logs."""
76+
pass
77+
78+
79+
async def discover_oauth_metadata(server_url: str) -> dict:
80+
"""Discover OAuth Authorization Server metadata (RFC 8414)."""
81+
print("\n[1/5] Discovering OAuth Authorization Server...")
82+
print(f" Server: {server_url}")
83+
84+
well_known_url = f"{server_url}/.well-known/oauth-authorization-server"
85+
86+
async with httpx.AsyncClient() as client:
87+
response = await client.get(well_known_url)
88+
response.raise_for_status()
89+
metadata = response.json()
90+
91+
print(f" ✓ Authorization endpoint: {metadata['authorization_endpoint']}")
92+
print(f" ✓ Token endpoint: {metadata['token_endpoint']}")
93+
94+
return metadata
95+
96+
97+
async def register_client(registration_endpoint: str) -> dict:
98+
"""Register OAuth client dynamically (RFC 7591)."""
99+
print("\n[2/5] Registering OAuth client...")
100+
101+
client_metadata = {
102+
"client_name": "chuk-tool-processor-test",
103+
"redirect_uris": ["http://127.0.0.1:8765/callback"],
104+
"grant_types": ["authorization_code"],
105+
"response_types": ["code"],
106+
"token_endpoint_auth_method": "none", # PKCE provides security
107+
}
108+
109+
async with httpx.AsyncClient() as client:
110+
response = await client.post(
111+
registration_endpoint,
112+
json=client_metadata,
113+
headers={"Content-Type": "application/json"}
114+
)
115+
response.raise_for_status()
116+
registration = response.json()
117+
118+
print(f" ✓ Client ID: {registration['client_id']}")
119+
120+
return registration
121+
122+
123+
def generate_pkce_challenge():
124+
"""Generate PKCE code verifier and challenge."""
125+
code_verifier = urlsafe_b64encode(secrets.token_bytes(32)).decode('utf-8').rstrip('=')
126+
code_challenge = urlsafe_b64encode(
127+
hashlib.sha256(code_verifier.encode('utf-8')).digest()
128+
).decode('utf-8').rstrip('=')
129+
return code_verifier, code_challenge
130+
131+
132+
async def get_authorization_code(auth_endpoint: str, client_id: str, code_challenge: str) -> str:
133+
"""Get authorization code via browser flow."""
134+
print("\n[3/5] Starting authorization flow...")
135+
136+
# Start local callback server
137+
server = HTTPServer(('127.0.0.1', 8765), OAuthCallbackHandler)
138+
139+
# Build authorization URL
140+
params = {
141+
"client_id": client_id,
142+
"response_type": "code",
143+
"redirect_uri": "http://127.0.0.1:8765/callback",
144+
"code_challenge": code_challenge,
145+
"code_challenge_method": "S256",
146+
"scope": "mcp",
147+
}
148+
auth_url = f"{auth_endpoint}?{urlencode(params)}"
149+
150+
print(f" Opening browser for authorization...")
151+
print(f" URL: {auth_url}")
152+
webbrowser.open(auth_url)
153+
154+
print("\n ⏳ Waiting for authorization...")
155+
print(" (Please complete the OAuth flow in your browser)")
156+
157+
# Wait for callback
158+
while OAuthCallbackHandler.authorization_code is None:
159+
server.handle_request()
160+
161+
code = OAuthCallbackHandler.authorization_code
162+
print(f" ✓ Received authorization code: {code[:20]}...")
163+
164+
return code
165+
166+
167+
async def exchange_code_for_token(
168+
token_endpoint: str,
169+
client_id: str,
170+
code: str,
171+
code_verifier: str
172+
) -> dict:
173+
"""Exchange authorization code for access token."""
174+
print("\n[4/5] Exchanging code for access token...")
175+
176+
token_data = {
177+
"grant_type": "authorization_code",
178+
"code": code,
179+
"redirect_uri": "http://127.0.0.1:8765/callback",
180+
"client_id": client_id,
181+
"code_verifier": code_verifier,
182+
}
183+
184+
async with httpx.AsyncClient() as client:
185+
response = await client.post(
186+
token_endpoint,
187+
data=token_data,
188+
headers={"Content-Type": "application/x-www-form-urlencoded"}
189+
)
190+
response.raise_for_status()
191+
tokens = response.json()
192+
193+
access_token = tokens.get('access_token', '')
194+
print(f" ✓ Access token: {access_token[:30]}...")
195+
print(f" ✓ Token type: {tokens.get('token_type')}")
196+
if 'expires_in' in tokens:
197+
print(f" ✓ Expires in: {tokens['expires_in']} seconds")
198+
199+
return tokens
200+
201+
202+
async def test_with_chuk_tool_processor(access_token: str):
203+
"""Test OAuth token with chuk-tool-processor HTTPStreamableTransport."""
204+
print("\n[5/5] Testing with chuk-tool-processor...")
205+
print("="*70)
206+
207+
notion_url = "https://mcp.notion.com/mcp"
208+
209+
# Test 1: Verify OAuth header is set correctly
210+
print("\n✓ Test 1: OAuth header configuration")
211+
transport = HTTPStreamableTransport(
212+
url=notion_url,
213+
headers={"Authorization": f"Bearer {access_token}"},
214+
connection_timeout=30.0,
215+
default_timeout=30.0
216+
)
217+
218+
headers = transport._get_headers()
219+
auth_header = headers.get('Authorization', '')
220+
print(f" Authorization: {auth_header[:40]}...")
221+
222+
# Test 2: Verify OAuth is NOT overwritten by api_key
223+
print("\n✓ Test 2: OAuth precedence over api_key")
224+
transport_with_key = HTTPStreamableTransport(
225+
url=notion_url,
226+
api_key="fake-key", # This should be ignored
227+
headers={"Authorization": f"Bearer {access_token}"},
228+
connection_timeout=30.0,
229+
default_timeout=30.0
230+
)
231+
232+
headers = transport_with_key._get_headers()
233+
if "fake-key" in headers.get('Authorization', ''):
234+
print(" ❌ FAILED: OAuth was overwritten by api_key!")
235+
return False
236+
else:
237+
print(" ✓ OAuth token preserved (api_key ignored)")
238+
239+
# Test 3: Actually connect to Notion
240+
print("\n✓ Test 3: Connecting to Notion MCP server")
241+
print(f" URL: {notion_url}")
242+
print(" Initializing connection (may take 30-60 seconds)...")
243+
244+
try:
245+
# Notion responses are slow - need longer timeout
246+
success = await asyncio.wait_for(
247+
transport.initialize(),
248+
timeout=120.0
249+
)
250+
251+
if success:
252+
print(" ✅ Connection successful!")
253+
254+
# Get tools
255+
print("\n✓ Test 4: Fetching tools from Notion")
256+
tools = await transport.get_tools()
257+
print(f" Retrieved {len(tools)} tools")
258+
259+
if tools:
260+
print("\n Available tools:")
261+
for tool in tools[:5]:
262+
name = tool.get('name', 'unknown')
263+
desc = tool.get('description', 'No description')[:50]
264+
print(f" • {name}: {desc}")
265+
if len(tools) > 5:
266+
print(f" ... and {len(tools) - 5} more")
267+
268+
await transport.close()
269+
return True
270+
else:
271+
print(" ❌ Connection failed")
272+
return False
273+
274+
except asyncio.TimeoutError:
275+
print(" ❌ Connection timed out after 60s")
276+
print("\n Possible issues:")
277+
print(" • Token may be invalid or expired")
278+
print(" • Notion MCP server not responding")
279+
print(" • Network connectivity issue")
280+
return False
281+
282+
except Exception as e:
283+
print(f" ❌ Error: {e}")
284+
logger.exception("Detailed error:")
285+
return False
286+
287+
288+
async def main():
289+
"""Main OAuth flow and test."""
290+
print("""
291+
╔══════════════════════════════════════════════════════════════════════╗
292+
║ Notion OAuth Test with chuk-tool-processor ║
293+
║ ║
294+
║ This script performs complete MCP OAuth flow and tests ║
295+
║ that OAuth tokens are correctly preserved in HTTPStreamableTransport ║
296+
╚══════════════════════════════════════════════════════════════════════╝
297+
""")
298+
299+
server_url = "https://mcp.notion.com"
300+
301+
try:
302+
# Step 1: Discover OAuth metadata
303+
metadata = await discover_oauth_metadata(server_url)
304+
305+
# Step 2: Register client
306+
registration = await register_client(metadata['registration_endpoint'])
307+
client_id = registration['client_id']
308+
309+
# Step 3: Generate PKCE challenge
310+
code_verifier, code_challenge = generate_pkce_challenge()
311+
312+
# Step 4: Get authorization code
313+
auth_code = await get_authorization_code(
314+
metadata['authorization_endpoint'],
315+
client_id,
316+
code_challenge
317+
)
318+
319+
# Step 5: Exchange for token
320+
tokens = await exchange_code_for_token(
321+
metadata['token_endpoint'],
322+
client_id,
323+
auth_code,
324+
code_verifier
325+
)
326+
327+
access_token = tokens['access_token']
328+
329+
# Step 6: Test with chuk-tool-processor
330+
success = await test_with_chuk_tool_processor(access_token)
331+
332+
if success:
333+
print("\n" + "="*70)
334+
print("✅ SUCCESS! OAuth is working correctly with Notion MCP server")
335+
print("="*70)
336+
print("\nKey points proven:")
337+
print(" ✓ Complete MCP OAuth flow (RFC 8414 + RFC 7591 + PKCE)")
338+
print(" ✓ OAuth token passed to HTTPStreamableTransport")
339+
print(" ✓ OAuth token preserved (not overwritten by api_key)")
340+
print(" ✓ Successfully connected to Notion MCP server")
341+
print(" ✓ Retrieved tools from Notion")
342+
return 0
343+
else:
344+
print("\n" + "="*70)
345+
print("⚠️ OAuth flow completed but connection test failed")
346+
print("="*70)
347+
print("\nThe fix is working (OAuth headers preserved),")
348+
print("but there may be other issues with the connection.")
349+
return 1
350+
351+
except KeyboardInterrupt:
352+
print("\n\n⚠️ Interrupted by user")
353+
return 1
354+
except Exception as e:
355+
print(f"\n❌ Error: {e}")
356+
logger.exception("Detailed error:")
357+
return 1
358+
359+
360+
if __name__ == "__main__":
361+
exit_code = asyncio.run(main())
362+
sys.exit(exit_code)

0 commit comments

Comments
 (0)