Skip to content

Commit c523088

Browse files
chrishayukchrishayuk
committed
fixed timeouts
1 parent 780aad3 commit c523088

File tree

13 files changed

+847
-142
lines changed

13 files changed

+847
-142
lines changed

‎examples/atlassian_sse.py‎

Lines changed: 337 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,337 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Atlassian SSE OAuth transport example with chuk-tool-processor.
4+
5+
This script demonstrates:
6+
1. OAuth 2.1 authentication flow with Atlassian MCP server
7+
2. SSE transport connection with OAuth token
8+
3. Proper initialization_timeout handling for slow servers
9+
10+
Usage:
11+
cd /Users/chrishay/chris-source/chuk-ai/chuk-tool-processor
12+
uv run python examples/atlassian_sse.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.setup_mcp_sse import setup_mcp_sse
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-atlassian-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 SSE transport."""
204+
print("\n[5/5] Testing with chuk-tool-processor...")
205+
print("="*70)
206+
207+
atlassian_url = "https://mcp.atlassian.com/v1/sse"
208+
209+
print("\n✓ Test: Connecting to Atlassian MCP server with OAuth via SSE")
210+
print(f" URL: {atlassian_url}")
211+
print(" Initializing connection (may take 60-120 seconds)...")
212+
213+
try:
214+
# Use setup_mcp_sse with OAuth token in headers
215+
servers = [
216+
{
217+
"name": "atlassian",
218+
"url": atlassian_url,
219+
"headers": {"Authorization": f"Bearer {access_token}"},
220+
}
221+
]
222+
223+
# Use 120s initialization timeout for potentially slow servers
224+
processor, stream_manager = await setup_mcp_sse(
225+
servers=servers,
226+
namespace="atlassian",
227+
connection_timeout=30.0,
228+
default_timeout=30.0,
229+
initialization_timeout=120.0, # Extended for slow servers
230+
)
231+
232+
print(" ✅ Connection successful!")
233+
234+
# Get tools
235+
print("\n✓ Fetching tools from Atlassian")
236+
tools = stream_manager.get_all_tools()
237+
print(f" Retrieved {len(tools)} tools")
238+
239+
if tools:
240+
print("\n Available tools:")
241+
for tool in tools[:5]:
242+
name = tool.get('name', 'unknown')
243+
desc = tool.get('description', 'No description')[:50]
244+
print(f" • {name}: {desc}")
245+
if len(tools) > 5:
246+
print(f" ... and {len(tools) - 5} more")
247+
248+
await stream_manager.close()
249+
return True
250+
251+
except asyncio.TimeoutError:
252+
print(" ❌ Connection timed out after 120s")
253+
print("\n Possible issues:")
254+
print(" • Token may be invalid or expired")
255+
print(" • Atlassian MCP server not responding")
256+
print(" • Network connectivity issue")
257+
return False
258+
259+
except Exception as e:
260+
print(f" ❌ Error: {e}")
261+
logger.exception("Detailed error:")
262+
return False
263+
264+
265+
async def main():
266+
"""Main OAuth flow and test."""
267+
print("""
268+
╔══════════════════════════════════════════════════════════════════════╗
269+
║ Atlassian SSE OAuth Test with chuk-tool-processor ║
270+
║ ║
271+
║ This script performs complete OAuth 2.1 flow and tests ║
272+
║ that OAuth tokens work correctly with SSE transport ║
273+
╚══════════════════════════════════════════════════════════════════════╝
274+
""")
275+
276+
try:
277+
atlassian_server = "https://mcp.atlassian.com"
278+
279+
# Step 1: Discover OAuth server
280+
metadata = await discover_oauth_metadata(atlassian_server)
281+
282+
# Step 2: Register client
283+
registration = await register_client(metadata['registration_endpoint'])
284+
client_id = registration['client_id']
285+
286+
# Step 3: Generate PKCE challenge
287+
code_verifier, code_challenge = generate_pkce_challenge()
288+
289+
# Step 4: Get authorization code
290+
auth_code = await get_authorization_code(
291+
metadata['authorization_endpoint'],
292+
client_id,
293+
code_challenge
294+
)
295+
296+
# Step 5: Exchange for access token
297+
tokens = await exchange_code_for_token(
298+
metadata['token_endpoint'],
299+
client_id,
300+
auth_code,
301+
code_verifier
302+
)
303+
access_token = tokens['access_token']
304+
305+
# Step 6: Test with chuk-tool-processor
306+
success = await test_with_chuk_tool_processor(access_token)
307+
308+
if success:
309+
print("\n" + "="*70)
310+
print("✅ SUCCESS! OAuth + SSE + initialization_timeout working correctly")
311+
print("="*70)
312+
print("\nKey points proven:")
313+
print(" ✓ Complete OAuth 2.1 flow (RFC 8414 + RFC 7591 + PKCE)")
314+
print(" ✓ OAuth token passed via headers to SSE transport")
315+
print(" ✓ Extended initialization_timeout (120s) allows slow servers to initialize")
316+
print(" ✓ Successfully connected to Atlassian MCP server")
317+
print(" ✓ Retrieved tools from Atlassian")
318+
return 0
319+
else:
320+
print("\n" + "="*70)
321+
print("⚠️ OAuth flow completed but connection test failed")
322+
print("="*70)
323+
print("\nCheck the error messages above for details.")
324+
return 1
325+
326+
except KeyboardInterrupt:
327+
print("\n\n⚠️ Interrupted by user")
328+
return 1
329+
except Exception as e:
330+
print(f"\n❌ Error: {e}")
331+
logger.exception("Detailed error:")
332+
return 1
333+
334+
335+
if __name__ == "__main__":
336+
exit_code = asyncio.run(main())
337+
sys.exit(exit_code)

0 commit comments

Comments
 (0)