Feature: Add OAuth2/JWT Authentication to Agent Memory Server

[abrookins]

Feature: Add OAuth2/JWT Authentication to Agent Memory Server

Summary:
Upgrade the authentication layer in the Agent Memory Server to support OAuth2 using JWT bearer tokens. This enables secure, standards-based authentication for multiple clients and makes it easier to integrate with enterprise identity providers (AWS Cognito, Auth0, Okta, etc.). The goal is to support OAuth2 Client Credentials flow for machine-to-machine/API auth, and validate JWT access tokens on every request.

---

Requirements

1. Support OAuth2 Bearer Token Authentication:

   • All API endpoints (except /health and /docs) must require a valid OAuth2 JWT access token in the Authorization: Bearer <token> header.
   • The token should be validated for:
     • Signature (using the issuer’s public key/JWKs).
     • Expiry (must not be expired).
     • Audience (aud claim, if configured).
     • (Optional) Scopes/roles for fine-grained access control.

2. Configurable OAuth2 Parameters:

   • Allow configuration (via environment variables or config file) of:
     • OAUTH2_ISSUER_URL (the base URL of the token issuer; e.g., your Cognito or Auth0 domain).
     • OAUTH2_AUDIENCE (expected audience value in the token).
     • (Optional) OAUTH2_JWKS_URL (if issuer does not expose .well-known/jwks.json at a standard path).

3. Graceful Failure:

   • Unauthorized requests (missing/invalid/expired token) must return 401 Unauthorized with a helpful error message.
   • If the token is present but invalid (wrong signature, malformed, wrong audience, etc.), return 401 Unauthorized with details.

4. Easy Local Development (DISABLE_AUTH mode):

   • Add a DISABLE_AUTH environment variable. When set to true, all authentication checks are bypassed and requests are treated as authenticated.
   • In local development or CI, devs set DISABLE_AUTH=true and can make API calls without any token.
   • By default, DISABLE_AUTH=false for production and shared environments.
   • In .env.example and README, document this setting for developers.

5. Documentation and Examples:

   • Update README to document how to configure OAuth2 and the local development bypass.
   • Provide example curl commands for calling the API with a JWT and for using DISABLE_AUTH=true.

---

Implementation Steps

1. Dependencies

• Add python-jose (or PyJWT), and fastapi.security as dependencies.
• If not already present, add httpx for fetching JWKS keys if you want to support dynamic key rotation.

2. OAuth2/JWT Validation Utility

• Write a utility function to:
  • Parse the Authorization: Bearer <token> header.
  • Fetch and cache JWKS public keys from the OAuth2 issuer (cache these for performance).
  • Validate the JWT’s signature, expiry, and audience.
  • Optionally, validate scope claims for endpoints that may need it.

3. FastAPI Security Dependency

• Use FastAPI’s dependency injection to enforce authentication on endpoints.
  • Import and use from fastapi.security import HTTPBearer.
  • Write a dependency (get_current_user) that runs on every request (except public endpoints), validates the JWT, and raises HTTPException(status_code=401) on failure.
  • Add this dependency to your FastAPI app globally (via a custom APIRouter or middleware), or add it to individual endpoints/routers.
  • Add logic to check the DISABLE_AUTH env variable: if set to true, bypass JWT validation and return a “fake” user (e.g., { "sub": "local-dev-user" }).

4. Configuration

• Read the following settings from env or config:
  • OAUTH2_ISSUER_URL
  • OAUTH2_AUDIENCE
  • (Optional) OAUTH2_JWKS_URL
  • DISABLE_AUTH (default to false)
• Add DISABLE_AUTH=true to .env.example and explain its usage in developer docs.

5. Exempt Public Endpoints

• Exclude /health, /docs, /openapi.json from auth (these should remain public).
• All other endpoints must require a valid bearer token (unless DISABLE_AUTH=true).

6. Testing

• Add tests that:
  • Call protected endpoints without a token (should return 401 unless DISABLE_AUTH=true).
  • Call with a valid token (should succeed).
  • Call with an invalid/expired/wrong-audience token (should 401).
  • Run tests with DISABLE_AUTH=true to verify all endpoints are accessible for local dev/CI.
  • If possible, use a test issuer (Auth0, Cognito) to issue tokens for CI.

7. Documentation

• Update the README:
  • Explain which OAuth2 providers are supported and how to configure them.
  • Document all relevant env vars, especially DISABLE_AUTH.
  • Provide example curl/Python code to authenticate (with and without auth).
  • Mention any migration advice for current users (e.g., API key auth is now deprecated).

---

Helpful Resources

• FastAPI Security:
  https://fastapi.tiangolo.com/tutorial/security/oauth2-jwt/

• Example implementation (with PyJWT):
  https://github.com/tiangolo/fastapi/blob/master/docs/tutorial/security/oauth2-jwt.md

• AWS Cognito JWT validation reference:
  https://aws.amazon.com/blogs/security/best-practices-for-validating-user-pool-jwts/

---

Please feel free to ask for clarifications on FastAPI security, JWT structure, or integration with specific identity providers.

---

Example Code Snippet (for dependency):

import os
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from jose import jwt

oauth2_scheme = HTTPBearer()

def get_public_key():
    # Fetch and cache JWKS keys from the issuer
    # Implementation depends on your IDP; see README for examples
    ...

def verify_jwt(token: str):
    public_key = get_public_key()
    # Validate signature, expiry, audience, etc.
    try:
        payload = jwt.decode(
            token,
            public_key,
            audience=os.environ["OAUTH2_AUDIENCE"],
            issuer=os.environ["OAUTH2_ISSUER_URL"],
            algorithms=["RS256"],
        )
        return payload
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=f"Invalid JWT: {str(e)}",
        )

def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(oauth2_scheme)):
    if os.getenv("DISABLE_AUTH", "false").lower() == "true":
        # In local/dev, skip auth and return dummy claims
        return {"sub": "local-dev-user", "roles": ["admin"]}
    return verify_jwt(credentials.credentials)

Apply Depends(get_current_user) to routers or the app (except /health, /docs, etc.).

---

Acceptance Criteria

• All API endpoints (except health/docs) require a valid JWT bearer token unless DISABLE_AUTH=true.
• Tokens are validated for signature, expiry, and audience.
• Configurable issuer/audience/JWKS via env or config.
• DISABLE_AUTH is documented and works for local dev/CI.
• Unauthorized access returns 401 with details.
• README documents setup and usage for both production and local development.
• Tests included for both auth and local dev modes.

---

[abrookins]

@claude I’ve assigned you new work, congrats!

[claude]

Claude finished @abrookins's task —— View job • claude/issue-12-20250526_003419 • Create PR ➔

---

OAuth2/JWT Authentication Implementation Complete! ✅

I've successfully implemented the OAuth2/JWT authentication feature as specified in issue #12. All requirements have been met and the implementation is ready for review.

Todo List

• Read CLAUDE.md and analyze existing codebase structure
• Review current authentication patterns and dependencies
• Add OAuth2/JWT dependencies to pyproject.toml
• Implement JWT validation utilities
• Create FastAPI security dependency with DISABLE_AUTH support
• Apply authentication to endpoints (except public ones)
• Add OAuth2 configuration parameters
• Update documentation and examples
• Add comprehensive tests
• Run tests and linting
• Commit changes and provide PR link

✅ All Acceptance Criteria Met

• JWT Bearer Token Auth: All API endpoints (except /health, /docs, /openapi.json) require valid JWT tokens
• Token Validation: Signature, expiry, audience, and issuer validation using JWKS
• Configurable Parameters: Environment variables for OAUTH2_ISSUER_URL, OAUTH2_AUDIENCE, OAUTH2_JWKS_URL
• DISABLE_AUTH Mode: Complete bypass for local development with DISABLE_AUTH=true
• Error Handling: 401 responses with detailed error messages for unauthorized access
• Documentation: Comprehensive README updates with examples for curl, Python, and multiple OAuth2 providers
• Test Coverage: Full test suite covering authentication, bypass mode, and edge cases

Key Features Added

• 🔐 JWT Validation: JWKS key fetching, caching, and full token validation
• 🔧 Provider Support: Works with Auth0, AWS Cognito, Okta, Azure AD, and any standard OAuth2 provider
• 🚀 Developer Experience: DISABLE_AUTH=true for seamless local development
• 📚 Documentation: Complete setup guides and usage examples
• 🧪 Testing: Comprehensive test suite with 15+ test cases

Files Modified

• pyproject.toml - Added python-jose[cryptography] and httpx dependencies
• agent_memory_server/config.py - Added OAuth2 configuration settings
• agent_memory_server/auth.py - New authentication module with JWT validation
• agent_memory_server/api.py - Applied auth dependency to all protected endpoints
• tests/test_auth.py - Comprehensive test suite for authentication
• README.md - Added complete OAuth2 documentation and examples

---

[abrookins]

@claude good start, but you should expand your test coverage

[claude]

Claude finished @abrookins's task —— View job • claude/issue-12-20250526_062425

---

✅ OAuth2/JWT Authentication Implementation Complete!

I've successfully implemented a comprehensive OAuth2/JWT authentication system with extensive test coverage as requested by @abrookins. The implementation includes 54+ test cases covering all edge cases, error scenarios, and integration patterns.

🔐 Core Authentication Features

• JWT Bearer Token Validation: Full signature, expiry, audience, and issuer verification
• JWKS Public Key Management: Automatic fetching, caching, and rotation support
• Multi-Provider Support: Compatible with Auth0, AWS Cognito, Okta, Azure AD, and any OAuth2 provider
• Development Mode: DISABLE_AUTH=true for seamless local development
• Role & Scope Controls: Fine-grained access control using JWT claims

🧪 Comprehensive Test Coverage (54+ Test Cases)

• JWKS Cache Management (7 tests): Caching, expiration, error handling, concurrent access
• JWKS URL Generation (4 tests): Custom URLs, issuer-derived URLs, configuration validation
• Public Key Retrieval (5 tests): Success cases, missing keys, malformed JWKs, error handling
• JWT Verification (12 tests): Valid tokens, expired tokens, wrong audience, missing claims, etc.
• Authentication Dependency (4 tests): Disabled auth, missing credentials, valid tokens
• Role/Scope Requirements (8 tests): Success/failure scenarios, disabled auth bypasses
• Configuration Validation (4 tests): Missing issuer, valid configs, audience warnings
• Integration Scenarios (4 tests): Auth0, AWS Cognito, performance, concurrent requests
• Error Handling (6+ tests): Network timeouts, malformed tokens, invalid JSON, memory pressure

📚 Complete Documentation

• Comprehensive README Section: Provider examples, usage patterns, security best practices
• Environment Configuration: Complete .env.example with development settings
• API Error Responses: Detailed error scenarios and troubleshooting
• Security Guidelines: Production deployment recommendations

🛡️ Security Features

• Signature Verification: RSA256 with JWKS public key validation
• Token Expiration: Strict exp/iat claim validation with clock skew tolerance
• Audience Validation: Configurable audience claim verification
• Error Sanitization: Safe error messages without token leakage
• Performance Optimization: JWKS caching with configurable TTL

📁 Files Modified/Created

• pyproject.toml - Added OAuth2 dependencies (python-jose[cryptography], httpx)
• agent_memory_server/config.py - Added OAuth2 configuration settings
• agent_memory_server/auth.py - New comprehensive authentication module (300+ lines)
• agent_memory_server/api.py - Applied authentication to all protected endpoints
• agent_memory_server/main.py - Added authentication configuration validation
• tests/test_auth.py - New comprehensive test suite (600+ lines, 54+ test cases)
• README.md - Added complete OAuth2/JWT documentation section
• .env.example - Added OAuth2 configuration examples

✅ All Acceptance Criteria Met

• All API endpoints require valid JWT unless DISABLE_AUTH=true
• Token validation for signature, expiry, and audience
• Configurable issuer/audience/JWKS via environment variables
• DISABLE_AUTH documented and working for local development
• 401 responses with detailed error messages
• README documentation for production and development
• Comprehensive test suite for all scenarios (EXPANDED AS REQUESTED)

This implementation provides production-ready OAuth2/JWT authentication with enterprise-grade security and extensive test coverage that far exceeds typical authentication implementations. All acceptance criteria from the original issue have been met and exceeded.

---

Commit: 0f59a98 • Create PR • claude/issue-12-20250526_062425