feat(server): Add WithOAuth for composable OAuth middleware integration

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>

Author: tuannvm

docs/implementation.md

@@ -412,22 +412,124 @@ type TokenValidator interface {
 
 ## Phase 3: Simple API Implementation
 
-**Status:** ⏳ Not Started
+**Status:** ✅ Completed
 
-### Tasks
+**Started:** 2025-10-19
+**Completed:** 2025-10-19
 
-- [ ] Implement oauth.EnableOAuth() in ROOT package
-  - [ ] Call NewServer() with validation
-  - [ ] Apply middleware to mcpServer
-  - [ ] Register handlers on mux
-  - [ ] Set up HTTPContextFunc
-  - [ ] Auto-detect mode with validation
-- [ ] Test both native and proxy modes
-- [ ] Test error handling
+### Tasks Completed
+
+- [x] Implement `oauth.WithOAuth()` in ROOT package
+  - [x] Call NewServer() with validation
+  - [x] Apply middleware via server option
+  - [x] Register handlers on mux
+  - [x] Return mcpserver.ServerOption
+- [x] HTTPContextFunc already exists (CreateHTTPContextFunc)
+- [x] Test both native and proxy modes
+- [x] Test error handling
+- [x] Create simple example
+- [x] Update documentation
 
 ### Implementation Notes
 
-*TBD*
+**API Design Decision:**
+
+Following Gemini 2.5 Pro's recommendation, implemented **composable API** instead of monolithic `EnableOAuth()`.
+
+**Why:**
+- mcp-go v0.41.1 requires middleware at server **construction** (not after)
+- `server.NewMCPServer()` accepts options, not middleware methods
+- Composable API fits mcp-go patterns better
+
+**Implemented API:**
+
+```go
+// oauth.go
+func WithOAuth(mux *http.ServeMux, cfg *Config) (mcpserver.ServerOption, error)
+```
+
+**Usage Pattern (2 lines):**
+```go
+mux := http.NewServeMux()
+
+// Line 1: Get OAuth option
+oauthOption, _ := oauth.WithOAuth(mux, &oauth.Config{
+    Provider:  "hmac",
+    Audience:  "api://test",
+    JWTSecret: []byte("secret"),
+})
+
+// Line 2: Create server with OAuth
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+
+// Done! All tools are OAuth-protected
+```
+
+**What WithOAuth() Does:**
+1. Creates OAuth Server internally (`NewServer(cfg)`)
+2. Validates config (via `cfg.Validate()`)
+3. Registers HTTP handlers on mux
+4. Returns `server.WithToolHandlerMiddleware(middleware)`
+
+**Key Features:**
+- ✅ Server-wide middleware (all tools protected)
+- ✅ Composable with other `server.ServerOption`
+- ✅ Auto-detects mode (native vs proxy)
+- ✅ Validates config early (fail fast)
+- ✅ Compatible with mcp-go v0.41.1
+
+**Helper Function:**
+```go
+func CreateHTTPContextFunc() func(context.Context, *http.Request) context.Context
+```
+- Extracts Bearer token from HTTP headers
+- Adds to context via `WithOAuthToken()`
+- Use with `mcpserver.WithHTTPContextFunc()`
+
+**Files Created:**
+- `oauth.go` - Added `WithOAuth()` function
+- `examples/simple/main.go` - NEW: Simple API example
+- `phase3_test.go` - NEW: WithOAuth() tests
+- `examples/README.md` - Updated with comparison
+
+**Files Modified:**
+- `examples/embedded/main.go` - Moved from examples/embedded.go
+- `examples/README.md` - Added Simple vs Embedded comparison
+
+**Test Coverage:**
+- `TestWithOAuth` - 4 subtests
+  - BasicUsage_NativeMode
+  - ProxyMode
+  - InvalidConfig
+  - EndToEndWithHTTPContextFunc
+- `TestPhase3API` - 2 subtests
+  - TwoLineSetup
+  - ComposableWithOtherOptions
+
+**Build & Test Status:**
+- ✅ `go build ./...` - Success
+- ✅ `make test` - All tests passing
+- ✅ `examples/simple/main.go` - Compiles
+- ✅ `examples/embedded/main.go` - Compiles
+
+**Comparison to Original Plan:**
+
+Original plan called for `EnableOAuth(mcpServer, mux, cfg)` but this was impossible because:
+- mcp-go v0.41.1 requires middleware at server creation
+- Can't modify server after construction
+
+**New API is better:**
+- More composable (functional options pattern)
+- Idiomatic for mcp-go users
+- Same simplicity (2 lines vs 1 line)
+- More flexible (can combine with other options)
+
+**Key Achievements:**
+- ✅ 2-line OAuth setup (goal achieved)
+- ✅ Server-wide protection (all tools secured)
+- ✅ mcp-go v0.41.1 compatible
+- ✅ Composable design
+- ✅ Both examples working
 
 ---
 

examples/README.md

@@ -1,45 +1,99 @@
 # OAuth MCP Proxy Examples
 
-## Embedded Mode Example
+## 1. Simple API (Phase 3) - `simple/`
 
-`embedded.go` - Complete MCP server with OAuth authentication
+**Simplest OAuth setup using `WithOAuth()` API:**
 
-**Features:**
-- Full MCP server implementation (using mark3labs/mcp-go v0.41.1)
-- **Server-wide OAuth middleware** (WithToolHandlerMiddleware)
-- OAuth-protected MCP tool ("hello")
-- Context propagation (HTTP → MCP → OAuth → Tool)
-- HMAC token validation with caching
-- OAuth metadata endpoints
-- Auto-generates test token for easy testing
+```go
+mux := http.NewServeMux()
 
-**Run:**
-```bash
-go run examples/embedded.go
+// Line 1: Get OAuth option (registers HTTP handlers)
+oauthOption, _ := oauth.WithOAuth(mux, &oauth.Config{
+    Provider:  "hmac",
+    Audience:  "api://my-server",
+    JWTSecret: []byte("secret"),
+})
+
+// Line 2: Create MCP server with OAuth
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+
+// Add tools - automatically OAuth-protected!
+mcpServer.AddTool(tool, handler)
+
+// Setup MCP endpoint with token extraction
+streamable := server.NewStreamableHTTPServer(mcpServer,
+    server.WithHTTPContextFunc(oauth.CreateHTTPContextFunc()),
+)
 ```
 
+**Features:**
+- **2-line OAuth setup** (Phase 3 goal achieved!)
+- Server-wide middleware (all tools protected automatically)
+- Composable with other server options
+- Uses mcp-go v0.41.1 `WithToolHandlerMiddleware` pattern
+
+**Run:** `go run examples/simple/main.go`
+
 **Test:**
 ```bash
-# Server will print the test command on startup
-
-# Test MCP tool call with OAuth:
 curl -X POST http://localhost:8080/mcp \
   -H 'Authorization: Bearer <token-from-output>' \
   -H 'Content-Type: application/json' \
   -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"hello","arguments":{}}}'
 ```
 
-**What It Demonstrates (Phase 2 Features):**
+---
+
+## 2. Embedded Mode (Phase 2) - `embedded/`
+
+**Detailed implementation showing internal architecture:**
+
+`embedded/main.go` - Complete MCP server with OAuth authentication
+
+**Features:**
+- Full MCP server implementation (using mark3labs/mcp-go v0.41.1)
+- Shows internal `NewServer()` API
+- Demonstrates `Server.Middleware()` usage
+- Shows manual middleware application
+- provider/ package architecture visible
+- Context propagation explained
+- Instance-scoped state (no globals)
+- HMAC token validation with caching
+- OAuth metadata endpoints
+- Auto-generates test token
+
+**Run:** `go run examples/embedded/main.go`
+
+**What It Demonstrates (Phase 2 internals):**
 1. Creating OAuth server (`oauth.NewServer()`)
-2. **Server-wide middleware** (`WithToolHandlerMiddleware`) - mcp-go v0.41.1
-3. provider/ package isolation (HMACValidator from provider/)
-4. Context propagation (`ValidateToken(ctx, token)`)
-5. Instance-scoped state (Server.cache, no globals)
-6. OAuth context extraction from HTTP headers
-7. User context available in MCP tools
-8. Complete end-to-end OAuth flow
-
-**Endpoints:**
+2. Getting middleware (`server.Middleware()`)
+3. Server-wide middleware with `WithToolHandlerMiddleware`
+4. provider/ package isolation (HMACValidator from provider/)
+5. Context propagation (`ValidateToken(ctx, token)`)
+6. Instance-scoped cache (Server.cache, no globals)
+7. OAuth context extraction from HTTP headers
+8. User context available in MCP tools
+
+---
+
+## Comparison
+
+| Feature | `simple/` (Phase 3) | `embedded/` (Phase 2) |
+|---------|---------------------|----------------------|
+| **Setup Lines** | 2 lines | ~10 lines |
+| **API** | `WithOAuth()` convenience | `NewServer()` + manual setup |
+| **Recommended** | ✅ For production | For learning internals |
+| **Shows** | Simplest usage | Architecture details |
+
+**Recommendation:** Use `simple/` for real projects, read `embedded/` to understand how it works.
+
+---
+
+## OAuth Endpoints
+
+Both examples expose:
 - `POST /mcp` - MCP protocol endpoint (OAuth protected)
 - `GET /.well-known/oauth-authorization-server` - OAuth metadata
 - `GET /.well-known/oauth-protected-resource` - Resource metadata
+- `GET /.well-known/jwks.json` - JWKS keys (HMAC mode)
+- `GET /.well-known/openid-configuration` - OIDC discovery

examples/embedded.go renamed to examples/embedded/main.go

@@ -0,0 +1,107 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/golang-jwt/jwt/v5"
+	"github.com/mark3labs/mcp-go/mcp"
+	mcpserver "github.com/mark3labs/mcp-go/server"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+func main() {
+	log.Println("=== OAuth MCP Proxy - Simple API (Phase 3) ===")
+	log.Println()
+
+	// Setup HTTP mux
+	mux := http.NewServeMux()
+
+	// Line 1: Get OAuth server option (registers HTTP handlers)
+	oauthOption, err := oauth.WithOAuth(mux, &oauth.Config{
+		Provider:  "hmac",
+		Issuer:    "https://test.example.com",
+		Audience:  "api://simple-server",
+		JWTSecret: []byte("test-secret-key-must-be-32-bytes-long!"),
+	})
+	if err != nil {
+		log.Fatalf("WithOAuth failed: %v", err)
+	}
+
+	log.Println("✅ OAuth configured")
+	log.Println("   - HTTP handlers registered")
+	log.Println("   - Middleware ready")
+
+	// Line 2: Create MCP server with OAuth option
+	mcpServer := mcpserver.NewMCPServer("Simple OAuth Server", "1.0.0",
+		oauthOption, // OAuth middleware applied to ALL tools!
+	)
+
+	log.Println("✅ MCP server created with OAuth middleware")
+
+	// Add tools (automatically protected by OAuth)
+	mcpServer.AddTool(
+		mcp.Tool{
+			Name:        "hello",
+			Description: "Says hello to authenticated user",
+		},
+		func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			user, ok := oauth.GetUserFromContext(ctx)
+			if !ok {
+				return nil, fmt.Errorf("authentication required")
+			}
+
+			message := fmt.Sprintf("Hello, %s! (Subject: %s)", user.Username, user.Subject)
+			return mcp.NewToolResultText(message), nil
+		},
+	)
+
+	log.Println("✅ Tools added (automatically OAuth-protected)")
+
+	// Setup MCP endpoint with OAuth context extraction
+	streamableServer := mcpserver.NewStreamableHTTPServer(
+		mcpServer,
+		mcpserver.WithEndpointPath("/mcp"),
+		mcpserver.WithHTTPContextFunc(oauth.CreateHTTPContextFunc()),
+	)
+
+	mux.Handle("/mcp", streamableServer)
+
+	// Generate test token
+	testToken := generateTestToken(&oauth.Config{
+		Issuer:    "https://test.example.com",
+		Audience:  "api://simple-server",
+		JWTSecret: []byte("test-secret-key-must-be-32-bytes-long!"),
+	})
+
+	log.Println()
+	log.Println("📋 Test Command:")
+	log.Printf("curl -X POST http://localhost:8080/mcp \\\n")
+	log.Printf("  -H 'Authorization: Bearer %s' \\\n", testToken[:50]+"...")
+	log.Printf("  -H 'Content-Type: application/json' \\\n")
+	log.Printf("  -d '{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\",\"params\":{\"name\":\"hello\",\"arguments\":{}}}'\n")
+	log.Println()
+
+	log.Println("🚀 Server starting on http://localhost:8080")
+	if err := http.ListenAndServe(":8080", mux); err != nil {
+		log.Fatalf("Server failed: %v", err)
+	}
+}
+
+func generateTestToken(cfg *oauth.Config) string {
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"sub":                "test-user",
+		"email":              "test@example.com",
+		"preferred_username": "testuser",
+		"aud":                cfg.Audience,
+		"iss":                cfg.Issuer,
+		"exp":                time.Now().Add(time.Hour).Unix(),
+		"iat":                time.Now().Unix(),
+	})
+
+	tokenString, _ := token.SignedString(cfg.JWTSecret)
+	return tokenString
+}