Feat/generic sdk (#4)

* chore(docs): Add initial implementation and refactoring plan

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

* ci(oauth): add OAuth server with token validation and caching support

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

* chore: update imports to use mark3labs for OAuth setup

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

* refactor(core): extract cache and context utilities for SDK-agnostic support

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

* fix(advancd/main.go): Rename file

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

* feat(oauth): add token validation and error handling in WrapHandler

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

* style: Add explanatory comments about PORT and SERVER_URL usage

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

* chore(github workflow): update example build to find all main.go files

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

* refactor(advanced/main.go): Enhance TLS detection and Logging for server startup

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

---------

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

Author: tuannvm

.github/workflows/test.yml

@@ -48,8 +48,10 @@ jobs:
 
       - name: Build examples
         run: |
-          go build ./examples/simple/main.go
-          go build ./examples/advanced/main.go
+          for example in $(find examples -name main.go); do
+            echo "Building $example..."
+            go build "$example"
+          done
 
       - name: Run go vet
         run: go vet ./...

CLAUDE.md

@@ -6,6 +6,8 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 **oauth-mcp-proxy** is an OAuth 2.1 authentication library for Go MCP servers. It provides server-side OAuth integration with minimal code (3-line integration via `WithOAuth()`), supporting multiple providers (HMAC, Okta, Google, Azure AD).
 
+**Version**: v2.0.0 (Supports both `mark3labs/mcp-go` and official `modelcontextprotocol/go-sdk`)
+
 ## Build Commands
 
 ```bash
@@ -36,34 +38,78 @@ make vuln
 
 ## Architecture
 
+### Package Structure (v2.0.0)
+
+```
+oauth-mcp-proxy/
+├── [core package - SDK-agnostic]
+│   ├── oauth.go         - Server type, NewServer, ValidateTokenCached
+│   ├── config.go        - Configuration validation and provider setup
+│   ├── cache.go         - Token cache with 5-minute TTL
+│   ├── context.go       - Context utilities (WithOAuthToken, GetUserFromContext, etc.)
+│   ├── handlers.go      - OAuth HTTP endpoints (/.well-known/*, /oauth/*)
+│   ├── middleware.go    - CreateHTTPContextFunc for token extraction
+│   ├── logger.go        - Logger interface
+│   ├── metadata.go      - OAuth metadata structures
+│   └── provider/        - Token validators (HMAC, OIDC)
+│
+├── mark3labs/          - Adapter for mark3labs/mcp-go SDK
+│   ├── oauth.go        - WithOAuth → ServerOption
+│   └── middleware.go   - Middleware for mark3labs types
+│
+└── mcp/                - Adapter for official modelcontextprotocol/go-sdk
+    └── oauth.go        - WithOAuth → http.Handler
+```
+
 ### Core Components
 
-1. **oauth.go** - Main entry point, provides `WithOAuth()` function that creates OAuth server and returns MCP server option
+**Core Package** (SDK-agnostic):
+1. **oauth.go** - `Server` type, `NewServer()`, `ValidateTokenCached()` (used by adapters)
 2. **config.go** - Configuration validation and provider setup
-3. **middleware.go** - Token validation middleware with 5-minute caching
-4. **handlers.go** - OAuth HTTP endpoints (/.well-known/*, /oauth/*)
-5. **provider/provider.go** - Token validators (HMACValidator, OIDCValidator)
+3. **cache.go** - Token caching logic (`TokenCache`, `CachedToken`)
+4. **context.go** - Context utilities (`WithOAuthToken`, `GetOAuthToken`, `WithUser`, `GetUserFromContext`)
+5. **handlers.go** - OAuth HTTP endpoints
+6. **provider/provider.go** - Token validators (HMACValidator, OIDCValidator)
+
+**Adapters** (SDK-specific):
+- **mark3labs/** - Middleware adapter for `mark3labs/mcp-go`
+- **mcp/** - HTTP handler wrapper for official SDK
 
 ### Key Design Patterns
 
+- **OpenTelemetry Pattern**: Core logic is SDK-agnostic; adapters provide SDK-specific integration
 - **Instance-scoped**: Each `Server` instance has its own token cache and validator (no globals)
 - **Provider abstraction**: `TokenValidator` interface supports multiple OAuth providers
 - **Caching strategy**: Tokens cached for 5 minutes using SHA-256 hash as key
-- **Context propagation**: OAuth token extracted from HTTP header → stored in context → validated by middleware → user added to context
+- **Context propagation**: OAuth token extracted from HTTP header → stored in context → validated → user added to context
 
 ### Integration Flow
 
+**mark3labs SDK:**
 ```text
 1. HTTP request with "Authorization: Bearer <token>" header
 2. CreateHTTPContextFunc() extracts token → adds to context via WithOAuthToken()
-3. OAuth middleware (Server.Middleware()) validates token:
-   - Checks cache first (5-minute TTL)
+3. mark3labs middleware validates token:
+   - Calls Server.ValidateTokenCached() (checks cache first)
    - If not cached, validates via provider (HMAC or OIDC)
-   - Caches result
-4. Adds authenticated User to context via userContextKey
+   - Caches result (5-minute TTL)
+4. Adds authenticated User to context via WithUser()
 5. Tool handler accesses user via GetUserFromContext(ctx)
 ```
 
+**Official SDK:**
+```text
+1. HTTP request with "Authorization: Bearer <token>" header
+2. mcp adapter's HTTP handler intercepts request
+3. Validates token via Server.ValidateTokenCached():
+   - Checks cache first (5-minute TTL)
+   - If not cached, validates via provider
+   - Caches result
+4. Adds token and user to context (WithOAuthToken, WithUser)
+5. Passes request to official SDK's StreamableHTTPHandler
+6. Tool handler accesses user via GetUserFromContext(ctx)
+```
+
 ### Provider System
 
 - **HMAC**: Validates JWT tokens with shared secret (testing/dev)
@@ -96,3 +142,29 @@ go test -v -run TestName ./...
 3. **Logging**: Config.Logger is optional. If nil, uses default logger (log.Printf with level prefixes)
 4. **Modes**: Library supports "native" (token validation only) and "proxy" (OAuth flow proxy) modes
 5. **Security**: All redirect URIs validated, state parameters HMAC-signed, tokens never logged (only hash previews)
+6. **v2.0.0 Breaking Change**: `WithOAuth()` moved to adapter packages (`mark3labs.WithOAuth()` or `mcp.WithOAuth()`). See `MIGRATION-V2.md`.
+
+## Using the Library
+
+### With mark3labs/mcp-go
+```go
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+)
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{...})
+mcpServer := server.NewMCPServer("name", "1.0.0", oauthOption)
+```
+
+### With Official SDK
+```go
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+mcpServer := mcp.NewServer(&mcp.Implementation{...}, nil)
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{...}, mcpServer)
+http.ListenAndServe(":8080", handler)
+```

README.md

@@ -2,21 +2,37 @@
 
 OAuth 2.1 authentication library for Go MCP servers.
 
+**Supports both MCP SDKs:**
+- ✅ `mark3labs/mcp-go`
+- ✅ `modelcontextprotocol/go-sdk` (official)
+
 **One-time setup:** Configure provider + add `WithOAuth()` to your server.
 **Result:** All tools automatically protected with token validation and caching.
 
+### mark3labs/mcp-go
 ```go
-// Enable OAuth authentication
-_, oauthOption, _ := oauth.WithOAuth(mux, &oauth.Config{
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",
     Issuer:   "https://your-company.okta.com",
     Audience: "api://your-mcp-server",
 })
 
-// All tools now require authentication
 mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
 ```
 
+### Official SDK
+```go
+import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+
+mcpServer := mcp.NewServer(&mcp.Implementation{...}, nil)
+_, handler, _ := mcpoauth.WithOAuth(mux, cfg, mcpServer)
+http.ListenAndServe(":8080", handler)
+```
+
+> **📢 Migrating from v1.x?** See [MIGRATION-V2.md](./MIGRATION-V2.md) (2 line change, ~5 min)
+
 [![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/tuannvm/oauth-mcp-proxy/test.yml?branch=main&label=Tests&logo=github)](https://github.com/tuannvm/oauth-mcp-proxy/actions/workflows/test.yml)
 [![Go Version](https://img.shields.io/github/go-mod/go-version/tuannvm/oauth-mcp-proxy?logo=go)](https://github.com/tuannvm/oauth-mcp-proxy/blob/main/go.mod)
 [![Go Report Card](https://goreportcard.com/badge/github.com/tuannvm/oauth-mcp-proxy)](https://goreportcard.com/report/github.com/tuannvm/oauth-mcp-proxy)
@@ -28,6 +44,7 @@ mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
 
 ## Why Use This Library?
 
+- **Dual SDK support** - Works with both mark3labs and official SDKs
 - **Simple integration** - One `WithOAuth()` call protects all tools
 - **Zero per-tool config** - All tools automatically protected
 - **Fast token caching** - 5-min cache, <5ms validation
@@ -64,21 +81,26 @@ sequenceDiagram
 
 ## Quick Start
 
-### 1. Install
+### Using mark3labs/mcp-go
+
+#### 1. Install
 
 ```bash
 go get github.com/tuannvm/oauth-mcp-proxy
 ```
 
-### 2. Add to Your Server
+#### 2. Add to Your Server
 
 ```go
-import oauth "github.com/tuannvm/oauth-mcp-proxy"
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+)
 
 mux := http.NewServeMux()
 
 // Enable OAuth (one time setup)
-_, oauthOption, _ := oauth.WithOAuth(mux, &oauth.Config{
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",                    // or "hmac", "google", "azure"
     Issuer:   "https://your-company.okta.com",
     Audience: "api://your-mcp-server",
@@ -99,7 +121,7 @@ streamable := mcpserver.NewStreamableHTTPServer(
 mux.Handle("/mcp", streamable)
 ```
 
-### 3. Access Authenticated User
+#### 3. Access Authenticated User
 
 ```go
 func myHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
@@ -111,16 +133,71 @@ func myHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResul
 }
 ```
 
+---
+
+### Using Official SDK
+
+#### 1. Install
+
+```bash
+go get github.com/modelcontextprotocol/go-sdk
+go get github.com/tuannvm/oauth-mcp-proxy
+```
+
+#### 2. Add to Your Server
+
+```go
+import (
+    "github.com/modelcontextprotocol/go-sdk/mcp"
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+mux := http.NewServeMux()
+
+// Create MCP server
+mcpServer := mcp.NewServer(&mcp.Implementation{
+    Name:    "my-server",
+    Version: "1.0.0",
+}, nil)
+
+// Add tools
+mcp.AddTool(mcpServer, &mcp.Tool{
+    Name: "greet",
+    Description: "Greet user",
+}, func(ctx context.Context, req *mcp.CallToolRequest, params *struct{}) (*mcp.CallToolResult, any, error) {
+    user, _ := oauth.GetUserFromContext(ctx)
+    return &mcp.CallToolResult{
+        Content: []mcp.Content{
+            &mcp.TextContent{Text: "Hello, " + user.Username},
+        },
+    }, nil, nil
+})
+
+// Add OAuth protection
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider: "okta",
+    Issuer:   "https://your-company.okta.com",
+    Audience: "api://your-mcp-server",
+}, mcpServer)
+
+http.ListenAndServe(":8080", handler)
+```
+
 Your MCP server now requires OAuth authentication.
 
 ---
 
 ## Examples
 
-| Example | Description |
-|---------|-------------|
-| **[Simple](examples/simple/)** | Minimal setup - copy/paste ready |
-| **[Advanced](examples/advanced/)** | All features - ConfigBuilder, WrapHandler, LogStartup |
+See [examples/README.md](examples/README.md) for detailed setup guide including Okta configuration.
+
+| SDK | Example | Description |
+|-----|---------|-------------|
+| **mark3labs** | [Simple](examples/mark3labs/simple/) | Minimal setup - copy/paste ready |
+| **mark3labs** | [Advanced](examples/mark3labs/advanced/) | ConfigBuilder, multiple tools, logging |
+| **Official** | [Simple](examples/official/simple/) | Minimal setup - copy/paste ready |
+| **Official** | [Advanced](examples/official/advanced/) | ConfigBuilder, multiple tools, logging |
 
 ---
 

cache.go

@@ -0,0 +1,64 @@
+package oauth
+
+import (
+	"sync"
+	"time"
+
+	"github.com/tuannvm/oauth-mcp-proxy/provider"
+)
+
+// Re-export User from provider for backwards compatibility
+type User = provider.User
+
+// TokenCache stores validated tokens to avoid re-validation
+type TokenCache struct {
+	mu    sync.RWMutex
+	cache map[string]*CachedToken
+}
+
+// CachedToken represents a cached token validation result
+type CachedToken struct {
+	User      *User
+	ExpiresAt time.Time
+}
+
+// getCachedToken retrieves a cached token validation result
+func (tc *TokenCache) getCachedToken(tokenHash string) (*CachedToken, bool) {
+	tc.mu.RLock()
+
+	cached, exists := tc.cache[tokenHash]
+	if !exists {
+		tc.mu.RUnlock()
+		return nil, false
+	}
+
+	if time.Now().After(cached.ExpiresAt) {
+		tc.mu.RUnlock()
+		go tc.deleteExpiredToken(tokenHash)
+		return nil, false
+	}
+
+	tc.mu.RUnlock()
+	return cached, true
+}
+
+// deleteExpiredToken safely deletes an expired token from the cache
+func (tc *TokenCache) deleteExpiredToken(tokenHash string) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+
+	if cached, exists := tc.cache[tokenHash]; exists && time.Now().After(cached.ExpiresAt) {
+		delete(tc.cache, tokenHash)
+	}
+}
+
+// setCachedToken stores a token validation result
+func (tc *TokenCache) setCachedToken(tokenHash string, user *User, expiresAt time.Time) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+
+	tc.cache[tokenHash] = &CachedToken{
+		User:      user,
+		ExpiresAt: expiresAt,
+	}
+}