From 3ae423251e410a75eb6d7d620f39a5b4629b5426 Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 14:37:30 -0700
Subject: [PATCH 1/9] chore(docs): Add initial implementation and refactoring
 plan

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 docs/generic-implementation.md | 777 +++++++++++++++++++++++++++++++++
 docs/generic-plan.md           | 242 ++++++++++
 go.mod                         |   2 +
 go.sum                         |  10 +-
 4 files changed, 1029 insertions(+), 2 deletions(-)
 create mode 100644 docs/generic-implementation.md
 create mode 100644 docs/generic-plan.md

diff --git a/docs/generic-implementation.md b/docs/generic-implementation.md
new file mode 100644
index 0000000..a080bed
--- /dev/null
+++ b/docs/generic-implementation.md
@@ -0,0 +1,777 @@
+# OpenTelemetry Pattern Implementation Checkpoints
+
+## Overview
+
+This document tracks the implementation progress for refactoring oauth-mcp-proxy to support both mark3labs/mcp-go and the official modelcontextprotocol/go-sdk.
+
+**Status Legend:**
+- ⬜ Not Started
+- 🟡 In Progress
+- ✅ Completed
+- ❌ Blocked
+
+---
+
+## Phase 0: Pre-Implementation Verification ✅
+
+**Goal**: Verify critical assumptions before starting implementation.
+
+### Checkpoint 0.1: Verify Official SDK Context Propagation ✅
+
+**Task**: Confirm that official SDK propagates HTTP request context to tool handlers.
+
+**Critical Question**: Does `mcp.NewStreamableHTTPHandler()` pass HTTP request context through to tool handlers?
+
+**Why This Matters**: Our entire OAuth integration relies on injecting user identity into request context and accessing it in tool handlers via `GetUserFromContext(ctx)`. If context doesn't propagate, we need a completely different approach.
+
+**Test Created**: `verify_context_test.go:TestOfficialSDKContextPropagation`
+
+**Result**: ✅ **VERIFIED - Context propagation works correctly**
+
+```
+=== RUN   TestOfficialSDKContextPropagation
+    verify_context_test.go:99: ✅ VERIFIED: Official SDK DOES propagate HTTP request context to tool handlers
+--- PASS: TestOfficialSDKContextPropagation (0.00s)
+```
+
+**Implications**:
+- Our planned wrapping approach will work
+- Tool handlers can access authenticated user via `GetUserFromContext(ctx)`
+- No need for alternative authentication mechanisms
+
+**Full Report**: See `docs/verification-results.md`
+
+---
+
+### Checkpoint 0.2: Define Core API Contract ✅
+
+**Task**: Specify exactly what the core package exposes to adapters.
+
+**Core API Contract**:
+
+**What Core Provides**:
+```go
+// Core server and lifecycle
+func NewServer(cfg *Config) (*Server, error)
+func (s *Server) RegisterHandlers(mux *http.ServeMux)
+
+// HTTP handler wrapping (SDK-agnostic)
+func (s *Server) WrapHandler(next http.Handler) http.Handler
+
+// NEW: Token validation for adapters to use
+func (s *Server) ValidateTokenCached(ctx context.Context, token string) (*User, error)
+
+// Context utilities
+func WithOAuthToken(ctx context.Context, token string) context.Context
+func GetOAuthToken(ctx context.Context) (string, bool)
+func WithUser(ctx context.Context, user *User) context.Context
+func GetUserFromContext(ctx context.Context) (*User, bool)
+```
+
+**What Gets REMOVED from Core** (moves to adapters):
+- ❌ `Server.Middleware()` - mark3labs specific
+- ❌ `Server.GetHTTPServerOptions()` - mark3labs specific
+
+**Adapter Responsibilities**:
+- mark3labs adapter: Implements middleware using mark3labs types
+- Official SDK adapter: Wraps StreamableHTTPHandler with OAuth validation
+
+**Details**: See `docs/verification-results.md#core-api-contract-definition`
+
+---
+
+## Phase 1: Core Package Extraction
+
+**Goal**: Extract SDK-agnostic OAuth logic into core package without breaking existing functionality.
+
+### Checkpoint 1.1: Create cache.go ⬜
+
+**Task**: Extract token cache logic from middleware.go into separate file.
+
+**Files to Create**:
+- `cache.go`
+
+**What to Extract from middleware.go**:
+- `TokenCache` struct
+- `CachedToken` struct
+- `getCachedToken()` method
+- `setCachedToken()` method
+- `deleteExpiredToken()` method
+
+**Verification**:
+```bash
+go build ./...
+go test ./... -v
+```
+
+**Expected Outcome**: Build succeeds, all tests pass.
+
+---
+
+### Checkpoint 1.2: Create context.go ⬜
+
+**Task**: Extract context-related functions into separate file.
+
+**Files to Create**:
+- `context.go`
+
+**What to Extract from middleware.go**:
+- `contextKey` type
+- `oauthTokenKey` constant
+- `userContextKey` constant
+- `WithOAuthToken()` function
+- `GetOAuthToken()` function
+- `GetUserFromContext()` function
+- `User` type alias
+
+**Verification**:
+```bash
+go build ./...
+go test ./... -v
+```
+
+**Expected Outcome**: Build succeeds, all tests pass.
+
+---
+
+### Checkpoint 1.3: Update imports in existing files ⬜
+
+**Task**: Update all internal imports to use new file structure.
+
+**Files to Update**:
+- `middleware.go` (remove extracted code, update imports)
+- `oauth.go` (update imports if needed)
+- All test files (update imports)
+
+**Verification**:
+```bash
+go build ./...
+go test ./... -v
+go mod tidy
+```
+
+**Expected Outcome**: No import errors, all tests pass.
+
+---
+
+### Checkpoint 1.4: Add ValidateTokenCached method to Server ⬜
+
+**Task**: Add new core method that adapters can use for token validation.
+
+**Files to Modify**:
+- `oauth.go` (add method to Server)
+
+**Implementation**:
+```go
+// ValidateTokenCached validates a token with caching support.
+// This is the core validation method that adapters can use.
+func (s *Server) ValidateTokenCached(ctx context.Context, token string) (*User, error) {
+    // Create token hash for caching
+    tokenHash := fmt.Sprintf("%x", sha256.Sum256([]byte(token)))
+
+    // Check cache first
+    if cached, exists := s.cache.getCachedToken(tokenHash); exists {
+        s.logger.Info("Using cached authentication (hash: %s...)", tokenHash[:16])
+        return cached.User, nil
+    }
+
+    // Log token hash for debugging
+    s.logger.Info("Validating token (hash: %s...)", tokenHash[:16])
+
+    // Validate token using configured provider
+    user, err := s.validator.ValidateToken(ctx, token)
+    if err != nil {
+        s.logger.Error("Token validation failed: %v", err)
+        return nil, fmt.Errorf("authentication failed: %w", err)
+    }
+
+    // Cache the validation result (expire in 5 minutes)
+    expiresAt := time.Now().Add(5 * time.Minute)
+    s.cache.setCachedToken(tokenHash, user, expiresAt)
+
+    s.logger.Info("Authenticated user %s (cached for 5 minutes)", user.Username)
+    return user, nil
+}
+```
+
+**Also Add**:
+```go
+// WithUser adds an authenticated user to context
+func WithUser(ctx context.Context, user *User) context.Context {
+    return context.WithValue(ctx, userContextKey, user)
+}
+```
+
+**Verification**:
+```bash
+go build ./...
+go test ./... -v
+```
+
+**Expected Outcome**: Build succeeds, new method available.
+
+---
+
+## Phase 2: Create mark3labs Adapter Package
+
+**Goal**: Move mark3labs-specific code into dedicated adapter package.
+
+### Checkpoint 2.1: Create mark3labs directory structure ⬜
+
+**Task**: Create new package directory for mark3labs adapter.
+
+**Directories to Create**:
+- `mark3labs/`
+
+**Files to Create**:
+- `mark3labs/oauth.go`
+- `mark3labs/middleware.go`
+
+**Verification**:
+```bash
+ls -la mark3labs/
+```
+
+**Expected Outcome**: Directory and files exist.
+
+---
+
+### Checkpoint 2.2: Implement mark3labs/oauth.go ⬜
+
+**Task**: Create WithOAuth function for mark3labs SDK.
+
+**Implementation**:
+```go
+package mark3labs
+
+import (
+    "net/http"
+
+    mcpserver "github.com/mark3labs/mcp-go/server"
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+// WithOAuth returns a server option that enables OAuth authentication
+// for mark3labs/mcp-go SDK.
+func WithOAuth(mux *http.ServeMux, cfg *oauth.Config) (*oauth.Server, mcpserver.ServerOption, error) {
+    oauthServer, err := oauth.NewServer(cfg)
+    if err != nil {
+        return nil, nil, err
+    }
+
+    oauthServer.RegisterHandlers(mux)
+
+    return oauthServer, mcpserver.WithToolHandlerMiddleware(NewMiddleware(oauthServer)), nil
+}
+```
+
+**Verification**:
+```bash
+cd mark3labs && go build .
+```
+
+**Expected Outcome**: Package builds successfully.
+
+---
+
+### Checkpoint 2.3: Implement mark3labs/middleware.go ⬜
+
+**Task**: Create middleware adapter for mark3labs SDK.
+
+**What to Implement**:
+- `NewMiddleware()` function that wraps `oauth.Server` and returns mark3labs-compatible middleware
+- Adapt mark3labs-specific types (ToolHandlerFunc, CallToolRequest, CallToolResult)
+
+**Key Code**:
+```go
+func NewMiddleware(s *oauth.Server) func(server.ToolHandlerFunc) server.ToolHandlerFunc {
+    return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
+        return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+            // Token extraction and validation
+            // Delegate to core oauth.Server logic
+            // Add user to context
+            return next(ctx, req)
+        }
+    }
+}
+```
+
+**Verification**:
+```bash
+cd mark3labs && go build .
+go test ./mark3labs/...
+```
+
+**Expected Outcome**: Package builds, basic tests pass.
+
+---
+
+### Checkpoint 2.4: Update examples to use mark3labs package ⬜
+
+**Task**: Update example code to import from mark3labs package.
+
+**Files to Update**:
+- `examples/simple/main.go`
+- `examples/advanced/main.go`
+
+**Changes**:
+```diff
+- import "github.com/tuannvm/oauth-mcp-proxy"
++ import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+- oauth.WithOAuth(mux, cfg)
++ mark3labs.WithOAuth(mux, cfg)
+```
+
+**Verification**:
+```bash
+cd examples/simple && go build .
+cd examples/advanced && go build .
+```
+
+**Expected Outcome**: Examples build and run successfully.
+
+---
+
+## Phase 3: Create Official SDK Adapter Package
+
+**Goal**: Add support for official modelcontextprotocol/go-sdk.
+
+### Checkpoint 3.1: Add official SDK dependency ⬜
+
+**Task**: Add official SDK to go.mod.
+
+**Commands**:
+```bash
+go get github.com/modelcontextprotocol/go-sdk
+go mod tidy
+```
+
+**Files Modified**:
+- `go.mod`
+- `go.sum`
+
+**Verification**:
+```bash
+go mod verify
+```
+
+**Expected Outcome**: Dependency added successfully.
+
+---
+
+### Checkpoint 3.2: Create mcp directory structure ⬜
+
+**Task**: Create new package directory for official SDK adapter.
+
+**Directories to Create**:
+- `mcp/`
+
+**Files to Create**:
+- `mcp/oauth.go`
+
+**Verification**:
+```bash
+ls -la mcp/
+```
+
+**Expected Outcome**: Directory and files exist.
+
+---
+
+### Checkpoint 3.3: Implement mcp/oauth.go ⬜
+
+**Task**: Create WithOAuth function for official SDK.
+
+**Implementation**:
+```go
+package mcp
+
+import (
+    "net/http"
+
+    "github.com/modelcontextprotocol/go-sdk/mcp"
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+// WithOAuth returns an OAuth-protected HTTP handler for the official
+// modelcontextprotocol/go-sdk.
+func WithOAuth(mux *http.ServeMux, cfg *oauth.Config, mcpServer *mcp.Server) (*oauth.Server, http.Handler, error) {
+    oauthServer, err := oauth.NewServer(cfg)
+    if err != nil {
+        return nil, nil, err
+    }
+
+    oauthServer.RegisterHandlers(mux)
+
+    // Create MCP HTTP handler
+    handler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server {
+        return mcpServer
+    }, nil)
+
+    // Wrap with OAuth validation
+    wrappedHandler := oauthServer.WrapHandler(handler)
+
+    return oauthServer, wrappedHandler, nil
+}
+```
+
+**Verification**:
+```bash
+cd mcp && go build .
+```
+
+**Expected Outcome**: Package builds successfully.
+
+---
+
+### Checkpoint 3.4: Create official SDK example ⬜
+
+**Task**: Create example demonstrating official SDK integration.
+
+**Files to Create**:
+- `examples/official/main.go`
+
+**Example Structure**:
+```go
+package main
+
+import (
+    "github.com/modelcontextprotocol/go-sdk/mcp"
+    mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+func main() {
+    // Create MCP server
+    mcpServer := mcp.NewServer(&mcp.Implementation{
+        Name:    "official-example",
+        Version: "1.0.0",
+    }, nil)
+
+    // Add OAuth
+    mux := http.NewServeMux()
+    _, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{...}, mcpServer)
+
+    http.ListenAndServe(":8080", handler)
+}
+```
+
+**Verification**:
+```bash
+cd examples/official && go build .
+./official
+```
+
+**Expected Outcome**: Example builds and runs.
+
+---
+
+## Phase 4: Testing and Validation
+
+**Goal**: Ensure both SDK integrations work correctly.
+
+### Checkpoint 4.1: Update existing tests ⬜
+
+**Task**: Update all tests to use new package structure.
+
+**Files to Update**:
+- `api_test.go`
+- `integration_test.go`
+- `middleware_compatibility_test.go`
+- `context_propagation_test.go`
+- All other test files
+
+**Changes**:
+- Update imports to use `mark3labs` package where needed
+- Verify core package tests still pass
+- Update test helpers if needed
+
+**Verification**:
+```bash
+go test ./... -v
+go test ./... -race
+go test ./... -cover
+```
+
+**Expected Outcome**: All tests pass with race detector.
+
+---
+
+### Checkpoint 4.2: Create mark3labs integration tests ⬜
+
+**Task**: Create comprehensive tests for mark3labs adapter.
+
+**Files to Create**:
+- `mark3labs/integration_test.go`
+
+**Test Coverage**:
+- WithOAuth function returns correct types
+- Middleware properly validates tokens
+- Context propagation works
+- Error cases handled correctly
+
+**Verification**:
+```bash
+go test ./mark3labs/... -v -cover
+```
+
+**Expected Outcome**: Tests pass, coverage > 80%.
+
+---
+
+### Checkpoint 4.3: Create official SDK integration tests ⬜
+
+**Task**: Create comprehensive tests for official SDK adapter.
+
+**Files to Create**:
+- `mcp/integration_test.go`
+
+**Test Coverage**:
+- WithOAuth function returns correct types
+- HTTP handler validates tokens
+- Official SDK server receives authenticated requests
+- Error cases handled correctly
+
+**Verification**:
+```bash
+go test ./mcp/... -v -cover
+```
+
+**Expected Outcome**: Tests pass, coverage > 80%.
+
+---
+
+### Checkpoint 4.4: Run full test suite ⬜
+
+**Task**: Verify all tests pass across all packages.
+
+**Commands**:
+```bash
+make test
+make test-coverage
+make lint
+```
+
+**Expected Results**:
+- All tests pass
+- No race conditions
+- Test coverage remains high (> 85%)
+- No linter errors
+
+**Verification**:
+```bash
+open coverage.html
+```
+
+**Expected Outcome**: Coverage report shows good coverage across all packages.
+
+---
+
+## Phase 5: Documentation Updates
+
+**Goal**: Update all documentation to reflect new package structure.
+
+### Checkpoint 5.1: Update README.md ⬜
+
+**Task**: Update main README with new package structure.
+
+**Changes Needed**:
+- Update installation instructions (show both packages)
+- Update quick start examples (mark3labs and official)
+- Add "Which SDK should I use?" section
+- Update all code examples
+- Add migration guide link
+
+**Sections to Update**:
+- Installation
+- Quick Start
+- Usage Examples
+- API Documentation
+
+**Verification**: Manual review for clarity and correctness.
+
+---
+
+### Checkpoint 5.2: Update CLAUDE.md ⬜
+
+**Task**: Update project overview for Claude Code.
+
+**Changes Needed**:
+- Update architecture section
+- Document both adapter packages
+- Update integration flow
+- Add notes about package structure
+
+**Verification**: Manual review for accuracy.
+
+---
+
+### Checkpoint 5.3: Create MIGRATION.md ⬜
+
+**Task**: Create migration guide for v1 to v2.
+
+**File to Create**:
+- `docs/MIGRATION.md`
+
+**Contents**:
+- What changed and why
+- Step-by-step migration for mark3labs users
+- Step-by-step migration to official SDK
+- Breaking changes list
+- Common issues and solutions
+
+**Verification**: Manual review by following guide.
+
+---
+
+### Checkpoint 5.4: Update examples README ⬜
+
+**Task**: Update examples documentation.
+
+**Files to Update**:
+- `examples/README.md` (if exists, or create)
+
+**Contents**:
+- List all examples
+- Describe which SDK each uses
+- Link to relevant documentation
+
+**Verification**: Manual review.
+
+---
+
+## Phase 6: Release Preparation
+
+**Goal**: Prepare for v2.0.0 release.
+
+### Checkpoint 6.1: Update version and changelog ⬜
+
+**Task**: Prepare release artifacts.
+
+**Files to Update/Create**:
+- `CHANGELOG.md` (document v2.0.0 changes)
+- Version tags in code
+
+**Contents**:
+- Breaking changes
+- New features (official SDK support)
+- Migration guide link
+
+**Verification**: Manual review.
+
+---
+
+### Checkpoint 6.2: Final validation ⬜
+
+**Task**: Complete final validation checklist.
+
+**Checklist**:
+- [ ] All tests pass (`make test`)
+- [ ] Linter passes (`make lint`)
+- [ ] Coverage acceptable (`make test-coverage`)
+- [ ] Examples build and run
+- [ ] Documentation complete
+- [ ] Migration guide tested
+- [ ] CHANGELOG updated
+- [ ] No TODO comments in code
+
+**Verification**:
+```bash
+make clean
+make test
+make lint
+make test-coverage
+cd examples/simple && go run main.go
+cd examples/advanced && go run main.go
+cd examples/official && go run main.go
+```
+
+**Expected Outcome**: Everything works.
+
+---
+
+### Checkpoint 6.3: Create release PR ⬜
+
+**Task**: Create pull request for v2.0.0.
+
+**PR Contents**:
+- Link to this implementation doc
+- Summary of changes
+- Migration guide
+- Breaking changes highlighted
+
+**Verification**: PR review and approval.
+
+---
+
+## Notes and Blockers
+
+### Open Issues
+- [ ] None currently
+
+### Decisions Made
+- ✅ Using OpenTelemetry pattern for package structure
+- ✅ Package names: `mark3labs` and `mcp`
+- ✅ Core logic stays in root package
+- ✅ Maintaining backward compatibility not feasible (breaking change)
+- ✅ Official SDK DOES propagate context (verified via test)
+- ✅ Core API contract defined (see Phase 0.2)
+- ✅ New `ValidateTokenCached()` method to be added for adapters
+
+### Dependencies
+- Official SDK version: v1.0.0 (added to go.mod)
+- mark3labs SDK version: v0.41.1 (existing)
+
+---
+
+## Progress Summary
+
+| Phase | Status | Completion |
+|-------|--------|------------|
+| Phase 0: Pre-Implementation Verification | ✅ Completed | 100% |
+| Phase 1: Core Package Extraction | ⬜ Not Started | 0% |
+| Phase 2: mark3labs Adapter | ⬜ Not Started | 0% |
+| Phase 3: Official SDK Adapter | ⬜ Not Started | 0% |
+| Phase 4: Testing & Validation | ⬜ Not Started | 0% |
+| Phase 5: Documentation | ⬜ Not Started | 0% |
+| Phase 6: Release Preparation | ⬜ Not Started | 0% |
+| **Overall** | **🟡 In Progress** | **14%** |
+
+---
+
+## Quick Reference Commands
+
+```bash
+# Build everything
+go build ./...
+
+# Run all tests
+go test ./... -v
+
+# Run tests with race detector
+go test ./... -race
+
+# Generate coverage report
+make test-coverage
+
+# Run linters
+make lint
+
+# Clean build artifacts
+make clean
+
+# Run specific package tests
+go test ./mark3labs/... -v
+go test ./mcp/... -v
+go test ./provider/... -v
+```
+
+---
+
+**Last Updated**: 2025-10-22
+**Verification Date**: 2025-10-22 (Phase 0 completed)
+**Implementation Start Date**: TBD
+**Target Completion Date**: TBD
+
+**Current Status**: Phase 0 (Verification) completed. Ready to begin implementation once approved.
diff --git a/docs/generic-plan.md b/docs/generic-plan.md
new file mode 100644
index 0000000..d1ef013
--- /dev/null
+++ b/docs/generic-plan.md
@@ -0,0 +1,242 @@
+# OpenTelemetry Pattern Refactoring Plan
+
+## Overview
+
+This document outlines the plan to refactor `oauth-mcp-proxy` to support both mark3labs/mcp-go and the official modelcontextprotocol/go-sdk using the OpenTelemetry pattern approach.
+
+## Current State
+
+The library currently supports only `github.com/mark3labs/mcp-go` (v0.41.1) with a single `WithOAuth()` function that returns `mcpserver.ServerOption`.
+
+## Proposed Structure
+
+Following the OpenTelemetry instrumentation pattern, we'll organize the codebase as:
+
+```
+oauth-mcp-proxy/
+├── [core package - SDK-agnostic]
+│   ├── server.go         (Server, NewServer, RegisterHandlers, WrapHandler)
+│   ├── config.go         (Config, validation)
+│   ├── cache.go          (TokenCache, token caching logic)
+│   ├── context.go        (WithOAuthToken, GetOAuthToken, GetUserFromContext)
+│   ├── handlers.go       (OAuth HTTP endpoints)
+│   ├── logger.go         (Logger interface, defaultLogger)
+│   ├── metadata.go       (OAuth metadata structures)
+│   └── provider/         (TokenValidator interface, HMAC/OIDC validators)
+│       ├── provider.go
+│       └── provider_test.go
+│
+├── mark3labs/           [SDK-specific adapter]
+│   ├── oauth.go         (WithOAuth → ServerOption)
+│   └── middleware.go    (Middleware adapter for mark3labs types)
+│
+└── mcp/                 [SDK-specific adapter]
+    └── oauth.go         (WithOAuth → http.Handler)
+```
+
+## Package Naming Convention
+
+Following OpenTelemetry's pattern:
+
+```
+github.com/tuannvm/oauth-mcp-proxy              (core, SDK-agnostic)
+github.com/tuannvm/oauth-mcp-proxy/mark3labs    (mark3labs/mcp-go adapter)
+github.com/tuannvm/oauth-mcp-proxy/mcp          (official SDK adapter)
+```
+
+## API Examples
+
+### mark3labs (existing SDK)
+
+```go
+import (
+    "github.com/mark3labs/mcp-go/server"
+    "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+)
+
+mux := http.NewServeMux()
+oauthServer, oauthOption, err := mark3labs.WithOAuth(mux, &oauth.Config{
+    Provider: "okta",
+    Issuer:   "https://company.okta.com",
+    Audience: "api://my-server",
+})
+
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+```
+
+### Official SDK (new support)
+
+```go
+import (
+    "github.com/modelcontextprotocol/go-sdk/mcp"
+    mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+mux := http.NewServeMux()
+mcpServer := mcp.NewServer(&mcp.Implementation{
+    Name:    "time-server",
+    Version: "1.0.0",
+}, nil)
+
+oauthServer, handler, err := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider: "okta",
+    Issuer:   "https://company.okta.com",
+    Audience: "api://my-server",
+}, mcpServer)
+
+http.ListenAndServe(":8080", handler)
+```
+
+## Pros
+
+### 1. Clean Separation of Concerns
+
+90% of the OAuth logic (validation, caching, config, providers) remains SDK-agnostic. Only adapters are SDK-specific.
+
+### 2. Easier Maintenance
+
+Bug fixes and new features in the core benefit both SDKs automatically. No need to duplicate logic.
+
+### 3. Clear API Contracts
+
+Users explicitly import the SDK-specific package they need. The import path makes intent clear:
+
+- `oauth-mcp-proxy/mark3labs` → I'm using mark3labs SDK
+- `oauth-mcp-proxy/mcp` → I'm using official SDK
+
+### 4. Discoverability
+
+Package structure clearly communicates "this library supports multiple SDKs" and makes it easy to find the right integration.
+
+### 5. Future Extensibility
+
+Adding support for SDK v3 or another MCP implementation = create new adapter package. Core remains unchanged.
+
+### 6. Follows Go Ecosystem Patterns
+
+Same approach used by:
+
+- OpenTelemetry (`go.opentelemetry.io/contrib/instrumentation/.../otel{package}`)
+- Sentry (`github.com/getsentry/sentry-go/{framework}`)
+- Datadog, NewRelic, and other observability libraries
+
+## Cons
+
+### 1. Import Path Changes (Breaking Change)
+
+Existing users must update imports:
+
+**Before:**
+
+```go
+import "github.com/tuannvm/oauth-mcp-proxy"
+oauth.WithOAuth(mux, cfg)
+```
+
+**After:**
+
+```go
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+mark3labs.WithOAuth(mux, cfg)
+```
+
+### 2. Requires Major Version Bump
+
+This is a breaking change requiring v1 → v2 semver bump.
+
+### 3. More Files
+
+Adds 2-3 new files vs current monolithic approach. Slightly more complex directory structure.
+
+### 4. Documentation Updates Required
+
+All examples, README.md, CLAUDE.md, and tutorials need updates to reflect new import paths.
+
+## Refactoring Complexity Assessment
+
+**Overall Complexity: MEDIUM**
+
+### Code Distribution
+
+| Component | Lines | Location | Effort |
+|-----------|-------|----------|--------|
+| Core OAuth logic | ~800 | Root package | Move/rename (2 hours) |
+| mark3labs adapter | ~40 | New: mark3labs/ | Extract (30 min) |
+| Official SDK adapter | ~30 | New: mcp/ | Write new (30 min) |
+| Tests | ~1000 | Update imports | Update (1 hour) |
+| Documentation | Multiple files | Update all | Update (1 hour) |
+
+### What Moves Where
+
+**Core Package (stays at root):**
+
+- ✅ `Server` type (remove SDK-specific methods)
+- ✅ `Config`, validation, providers
+- ✅ `TokenCache`, `CachedToken`
+- ✅ Context functions (`WithOAuthToken`, `GetOAuthToken`, `GetUserFromContext`)
+- ✅ HTTP handlers (OAuth endpoints)
+- ✅ `WrapHandler` (already SDK-agnostic)
+- ✅ Logger interface and implementation
+
+**mark3labs/ (extract ~40 lines):**
+
+- `WithOAuth()` → returns `(*oauth.Server, mcpserver.ServerOption, error)`
+- `Middleware()` → wraps mark3labs-specific types
+- `GetHTTPServerOptions()` → returns `[]mcpserver.StreamableHTTPOption`
+
+**mcp/ (write ~30 new lines):**
+
+- `WithOAuth()` → returns `(*oauth.Server, http.Handler, error)`
+- Handler wrapper using `mcp.NewStreamableHTTPHandler()`
+- Integration with official SDK's HTTP model
+
+### Migration Effort Breakdown
+
+| Phase | Tasks | Time Estimate |
+|-------|-------|---------------|
+| **Code Refactoring** | Extract core, create adapters, update imports | 2-3 hours |
+| **Testing** | Verify both integrations, update tests | 1-2 hours |
+| **Documentation** | Update README, examples, migration guide | 1 hour |
+| **Validation** | Run full test suite, manual testing | 30 min |
+
+**Total Estimated Effort: 1 day of focused work**
+
+## Migration Strategy
+
+### For Library Maintainers
+
+1. **Phase 1**: Core extraction (keep existing API working)
+2. **Phase 2**: Create mark3labs adapter
+3. **Phase 3**: Create mcp adapter
+4. **Phase 4**: Update tests
+5. **Phase 5**: Update documentation
+6. **Phase 6**: Release v2.0.0 with migration guide
+
+### For Library Users
+
+Users have two migration paths:
+
+**Option 1: Quick Update (mark3labs users)**
+
+```diff
+- import "github.com/tuannvm/oauth-mcp-proxy"
++ import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+- oauth.WithOAuth(mux, cfg)
++ mark3labs.WithOAuth(mux, cfg)
+```
+
+**Option 2: Migrate to Official SDK**
+Follow the official SDK migration guide in the new documentation.
+
+## Open Questions
+
+1. Should we maintain v1 branch for bug fixes during transition period?
+2. How long should we support v1 before deprecating?
+3. Should we add compatibility shims in v2 to ease migration?
+
+## References
+
+- OpenTelemetry Go Contrib: <https://github.com/open-telemetry/opentelemetry-go-contrib>
+- Sentry Go SDK: <https://github.com/getsentry/sentry-go>
+- Official MCP Go SDK: <https://github.com/modelcontextprotocol/go-sdk>
diff --git a/go.mod b/go.mod
index 8de6cf9..b9f8ed2 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/coreos/go-oidc/v3 v3.16.0
 	github.com/golang-jwt/jwt/v5 v5.3.0
 	github.com/mark3labs/mcp-go v0.41.1
+	github.com/modelcontextprotocol/go-sdk v1.0.0
 	golang.org/x/oauth2 v0.32.0
 )
 
@@ -13,6 +14,7 @@ require (
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/go-jose/go-jose/v4 v4.1.3 // indirect
+	github.com/google/jsonschema-go v0.3.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/invopop/jsonschema v0.13.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
diff --git a/go.sum b/go.sum
index f4e3114..eb5198f 100644
--- a/go.sum
+++ b/go.sum
@@ -12,8 +12,10 @@ github.com/go-jose/go-jose/v4 v4.1.3 h1:CVLmWDhDVRa6Mi/IgCgaopNosCaHz7zrMeF9MlZR
 github.com/go-jose/go-jose/v4 v4.1.3/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08=
 github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
 github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
-github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
+github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
@@ -27,6 +29,8 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/mark3labs/mcp-go v0.41.1 h1:w78eWfiQam2i8ICL7AL0WFiq7KHNJQ6UB53ZVtH4KGA=
 github.com/mark3labs/mcp-go v0.41.1/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g=
+github.com/modelcontextprotocol/go-sdk v1.0.0 h1:Z4MSjLi38bTgLrd/LjSmofqRqyBiVKRyQSJgw8q8V74=
+github.com/modelcontextprotocol/go-sdk v1.0.0/go.mod h1:nYtYQroQ2KQiM0/SbyEPUWQ6xs4B95gJjEalc9AQyOs=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
@@ -41,6 +45,8 @@ github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zI
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 golang.org/x/oauth2 v0.32.0 h1:jsCblLleRMDrxMN29H3z/k1KliIvpLgCkE6R8FXXNgY=
 golang.org/x/oauth2 v0.32.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
+golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo=
+golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

From bf15601b589a787f75c2824d253ad6d21bd9609b Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 14:57:48 -0700
Subject: [PATCH 2/9] ci(oauth): add OAuth server with token validation and
 caching support

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 cache.go                | 64 +++++++++++++++++++++++++++
 context.go              | 45 +++++++++++++++++++
 mark3labs/middleware.go | 40 +++++++++++++++++
 mark3labs/oauth.go      | 48 ++++++++++++++++++++
 mcp/oauth.go            | 79 +++++++++++++++++++++++++++++++++
 middleware.go           | 97 -----------------------------------------
 oauth.go                | 36 +++++++++++++++
 7 files changed, 312 insertions(+), 97 deletions(-)
 create mode 100644 cache.go
 create mode 100644 context.go
 create mode 100644 mark3labs/middleware.go
 create mode 100644 mark3labs/oauth.go
 create mode 100644 mcp/oauth.go

diff --git a/cache.go b/cache.go
new file mode 100644
index 0000000..32759c2
--- /dev/null
+++ b/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,
+	}
+}
diff --git a/context.go b/context.go
new file mode 100644
index 0000000..9cc72b0
--- /dev/null
+++ b/context.go
@@ -0,0 +1,45 @@
+package oauth
+
+import "context"
+
+// Context keys
+type contextKey string
+
+const (
+	oauthTokenKey  contextKey = "oauth_token"
+	userContextKey contextKey = "user"
+)
+
+// WithOAuthToken adds an OAuth token to the context
+func WithOAuthToken(ctx context.Context, token string) context.Context {
+	return context.WithValue(ctx, oauthTokenKey, token)
+}
+
+// GetOAuthToken extracts an OAuth token from the context
+func GetOAuthToken(ctx context.Context) (string, bool) {
+	token, ok := ctx.Value(oauthTokenKey).(string)
+	return token, ok
+}
+
+// WithUser adds an authenticated user to context
+func WithUser(ctx context.Context, user *User) context.Context {
+	return context.WithValue(ctx, userContextKey, user)
+}
+
+// GetUserFromContext extracts the authenticated user from context.
+// Returns the User and true if authentication succeeded, or nil and false otherwise.
+//
+// Example:
+//
+//	func toolHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+//	    user, ok := oauth.GetUserFromContext(ctx)
+//	    if !ok {
+//	        return nil, fmt.Errorf("authentication required")
+//	    }
+//	    // Use user.Subject, user.Email, user.Username
+//	    return mcp.NewToolResultText("Hello, " + user.Username), nil
+//	}
+func GetUserFromContext(ctx context.Context) (*User, bool) {
+	user, ok := ctx.Value(userContextKey).(*User)
+	return user, ok
+}
diff --git a/mark3labs/middleware.go b/mark3labs/middleware.go
new file mode 100644
index 0000000..8c763fc
--- /dev/null
+++ b/mark3labs/middleware.go
@@ -0,0 +1,40 @@
+package mark3labs
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+// NewMiddleware creates an authentication middleware for mark3labs/mcp-go SDK.
+// It validates OAuth tokens, caches results, and adds authenticated user to context.
+//
+// The middleware:
+//  1. Extracts OAuth token from context (set by CreateHTTPContextFunc)
+//  2. Validates token using Server.ValidateTokenCached (with 5-minute cache)
+//  3. Adds User to context via oauth.WithUser
+//  4. Passes request to tool handler with authenticated context
+//
+// Use oauth.GetUserFromContext(ctx) in tool handlers to access authenticated user.
+func NewMiddleware(s *oauth.Server) func(server.ToolHandlerFunc) server.ToolHandlerFunc {
+	return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
+		return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			tokenString, ok := oauth.GetOAuthToken(ctx)
+			if !ok {
+				return nil, fmt.Errorf("authentication required: missing OAuth token")
+			}
+
+			user, err := s.ValidateTokenCached(ctx, tokenString)
+			if err != nil {
+				return nil, err
+			}
+
+			ctx = oauth.WithUser(ctx, user)
+
+			return next(ctx, req)
+		}
+	}
+}
diff --git a/mark3labs/oauth.go b/mark3labs/oauth.go
new file mode 100644
index 0000000..f31969d
--- /dev/null
+++ b/mark3labs/oauth.go
@@ -0,0 +1,48 @@
+package mark3labs
+
+import (
+	"fmt"
+	"net/http"
+
+	mcpserver "github.com/mark3labs/mcp-go/server"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+// WithOAuth returns a server option that enables OAuth authentication
+// for mark3labs/mcp-go SDK.
+//
+// Usage:
+//
+//	import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+//
+//	mux := http.NewServeMux()
+//	oauthServer, oauthOption, err := mark3labs.WithOAuth(mux, &oauth.Config{
+//	    Provider: "okta",
+//	    Issuer:   "https://company.okta.com",
+//	    Audience: "api://my-server",
+//	})
+//	mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+//
+// This function:
+// - Creates OAuth server instance
+// - Registers OAuth HTTP endpoints on mux
+// - Returns server instance and middleware as server option
+//
+// The returned Server instance provides access to:
+// - WrapHandler() - Wrap HTTP handlers with OAuth token validation
+// - GetHTTPServerOptions() - Get StreamableHTTPServer options
+// - LogStartup() - Log OAuth endpoint information
+// - Discovery URL helpers (GetCallbackURL, GetMetadataURL, etc.)
+//
+// Note: You must also configure HTTPContextFunc to extract the OAuth token
+// from HTTP headers. Use GetHTTPServerOptions() or CreateHTTPContextFunc().
+func WithOAuth(mux *http.ServeMux, cfg *oauth.Config) (*oauth.Server, mcpserver.ServerOption, error) {
+	oauthServer, err := oauth.NewServer(cfg)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create OAuth server: %w", err)
+	}
+
+	oauthServer.RegisterHandlers(mux)
+
+	return oauthServer, mcpserver.WithToolHandlerMiddleware(NewMiddleware(oauthServer)), nil
+}
diff --git a/mcp/oauth.go b/mcp/oauth.go
new file mode 100644
index 0000000..e8b3f5f
--- /dev/null
+++ b/mcp/oauth.go
@@ -0,0 +1,79 @@
+package mcp
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+)
+
+// WithOAuth returns an OAuth-protected HTTP handler for the official
+// modelcontextprotocol/go-sdk.
+//
+// Usage:
+//
+//	import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+//
+//	mux := http.NewServeMux()
+//	mcpServer := mcp.NewServer(&mcp.Implementation{
+//	    Name:    "time-server",
+//	    Version: "1.0.0",
+//	}, nil)
+//
+//	oauthServer, handler, err := mcpoauth.WithOAuth(mux, &oauth.Config{
+//	    Provider: "okta",
+//	    Issuer:   "https://company.okta.com",
+//	    Audience: "api://my-server",
+//	}, mcpServer)
+//
+//	http.ListenAndServe(":8080", handler)
+//
+// This function:
+// - Creates OAuth server instance
+// - Registers OAuth HTTP endpoints on mux
+// - Wraps MCP StreamableHTTPHandler with OAuth token validation
+// - Returns OAuth server and protected HTTP handler
+//
+// The returned oauth.Server instance provides access to:
+// - LogStartup() - Log OAuth endpoint information
+// - Discovery URL helpers (GetCallbackURL, GetMetadataURL, etc.)
+//
+// The HTTP handler validates OAuth tokens before delegating to the MCP server.
+// Tool handlers can access the authenticated user via oauth.GetUserFromContext(ctx).
+func WithOAuth(mux *http.ServeMux, cfg *oauth.Config, mcpServer *mcp.Server) (*oauth.Server, http.Handler, error) {
+	oauthServer, err := oauth.NewServer(cfg)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create OAuth server: %w", err)
+	}
+
+	oauthServer.RegisterHandlers(mux)
+
+	mcpHandler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
+		return mcpServer
+	}, nil)
+
+	wrappedHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		authHeader := r.Header.Get("Authorization")
+		if authHeader == "" || len(authHeader) < 7 || authHeader[:7] != "Bearer " {
+			http.Error(w, "Missing or invalid Authorization header", http.StatusUnauthorized)
+			return
+		}
+
+		token := authHeader[7:]
+
+		user, err := oauthServer.ValidateTokenCached(r.Context(), token)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("Authentication failed: %v", err), http.StatusUnauthorized)
+			return
+		}
+
+		ctx := oauth.WithOAuthToken(r.Context(), token)
+		ctx = oauth.WithUser(ctx, user)
+		r = r.WithContext(ctx)
+
+		mcpHandler.ServeHTTP(w, r)
+	})
+
+	return oauthServer, wrappedHandler, nil
+}
diff --git a/middleware.go b/middleware.go
index d4bce2b..e82870d 100644
--- a/middleware.go
+++ b/middleware.go
@@ -7,7 +7,6 @@ import (
 	"log"
 	"net/http"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
@@ -15,84 +14,6 @@ import (
 	"github.com/tuannvm/oauth-mcp-proxy/provider"
 )
 
-// Re-export User from provider for backwards compatibility
-type User = provider.User
-
-// Context keys
-type contextKey string
-
-const (
-	oauthTokenKey  contextKey = "oauth_token"
-	userContextKey contextKey = "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
-}
-
-// WithOAuthToken adds an OAuth token to the context
-func WithOAuthToken(ctx context.Context, token string) context.Context {
-	return context.WithValue(ctx, oauthTokenKey, token)
-}
-
-// GetOAuthToken extracts an OAuth token from the context
-func GetOAuthToken(ctx context.Context) (string, bool) {
-	token, ok := ctx.Value(oauthTokenKey).(string)
-	return token, ok
-}
-
-// 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
-	}
-
-	// Check if token is expired
-	if time.Now().After(cached.ExpiresAt) {
-		tc.mu.RUnlock()
-		// Schedule expired token deletion in a separate operation
-		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()
-
-	// Double-check if token is still expired before deleting
-	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,
-	}
-}
-
 // Middleware returns an authentication middleware for MCP tools.
 // Validates OAuth tokens, caches results, and adds authenticated user to context.
 //
@@ -183,24 +104,6 @@ func OAuthMiddleware(validator provider.TokenValidator, enabled bool) func(serve
 
 // validateJWT is deprecated - use provider-based validation instead
 
-// GetUserFromContext extracts the authenticated user from context.
-// Returns the User and true if authentication succeeded, or nil and false otherwise.
-//
-// Example:
-//
-//	func toolHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-//	    user, ok := oauth.GetUserFromContext(ctx)
-//	    if !ok {
-//	        return nil, fmt.Errorf("authentication required")
-//	    }
-//	    // Use user.Subject, user.Email, user.Username
-//	    return mcp.NewToolResultText("Hello, " + user.Username), nil
-//	}
-func GetUserFromContext(ctx context.Context) (*User, bool) {
-	user, ok := ctx.Value(userContextKey).(*User)
-	return user, ok
-}
-
 // CreateHTTPContextFunc creates an HTTP context function that extracts OAuth tokens
 // from Authorization headers. Use with mcpserver.WithHTTPContextFunc() to enable
 // token extraction from HTTP requests.
diff --git a/oauth.go b/oauth.go
index 6e1fc0d..a94f014 100644
--- a/oauth.go
+++ b/oauth.go
@@ -1,9 +1,12 @@
 package oauth
 
 import (
+	"context"
+	"crypto/sha256"
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"time"
 
 	mcpserver "github.com/mark3labs/mcp-go/server"
 	"github.com/tuannvm/oauth-mcp-proxy/provider"
@@ -96,6 +99,39 @@ func (s *Server) RegisterHandlers(mux *http.ServeMux) {
 	mux.HandleFunc("/.well-known/openid-configuration", s.handler.HandleOIDCDiscovery)
 }
 
+// ValidateTokenCached validates a token with caching support.
+// This is the core validation method that SDK adapters can use.
+//
+// The method:
+//  1. Checks token cache (5-minute TTL)
+//  2. Validates token using configured provider if not cached
+//  3. Caches validation result for future requests
+//  4. Returns authenticated User or error
+//
+// This method is used internally by both WrapHandler and adapter middleware.
+func (s *Server) ValidateTokenCached(ctx context.Context, token string) (*User, error) {
+	tokenHash := fmt.Sprintf("%x", sha256.Sum256([]byte(token)))
+
+	if cached, exists := s.cache.getCachedToken(tokenHash); exists {
+		s.logger.Info("Using cached authentication (hash: %s...)", tokenHash[:16])
+		return cached.User, nil
+	}
+
+	s.logger.Info("Validating token (hash: %s...)", tokenHash[:16])
+
+	user, err := s.validator.ValidateToken(ctx, token)
+	if err != nil {
+		s.logger.Error("Token validation failed: %v", err)
+		return nil, fmt.Errorf("authentication failed: %w", err)
+	}
+
+	expiresAt := time.Now().Add(5 * time.Minute)
+	s.cache.setCachedToken(tokenHash, user, expiresAt)
+
+	s.logger.Info("Authenticated user %s (cached for 5 minutes)", user.Username)
+	return user, nil
+}
+
 // GetAuthorizationServerMetadataURL returns the OAuth 2.0 authorization server metadata URL
 func (s *Server) GetAuthorizationServerMetadataURL() string {
 	return fmt.Sprintf("%s/.well-known/oauth-authorization-server", s.config.ServerURL)

From ac44964878a82d5156964f24b80504aa309da1da Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 14:58:02 -0700
Subject: [PATCH 3/9] chore: update imports to use mark3labs for OAuth setup

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 examples/advanced/main.go | 3 ++-
 examples/simple/main.go   | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/advanced/main.go b/examples/advanced/main.go
index 903705e..d4551f0 100644
--- a/examples/advanced/main.go
+++ b/examples/advanced/main.go
@@ -10,6 +10,7 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 	mcpserver "github.com/mark3labs/mcp-go/server"
 	oauth "github.com/tuannvm/oauth-mcp-proxy"
+	"github.com/tuannvm/oauth-mcp-proxy/mark3labs"
 )
 
 func main() {
@@ -31,7 +32,7 @@ func main() {
 	mux := http.NewServeMux()
 
 	// Feature 2: WithOAuth returns Server instance for helper methods
-	oauthServer, oauthOption, err := oauth.WithOAuth(mux, cfg)
+	oauthServer, oauthOption, err := mark3labs.WithOAuth(mux, cfg)
 	if err != nil {
 		log.Fatalf("OAuth setup failed: %v", err)
 	}
diff --git a/examples/simple/main.go b/examples/simple/main.go
index c033c27..4e2a576 100644
--- a/examples/simple/main.go
+++ b/examples/simple/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 	mcpserver "github.com/mark3labs/mcp-go/server"
 	oauth "github.com/tuannvm/oauth-mcp-proxy"
+	"github.com/tuannvm/oauth-mcp-proxy/mark3labs"
 )
 
 func main() {
@@ -16,7 +17,7 @@ func main() {
 	mux := http.NewServeMux()
 
 	// 2. Enable OAuth authentication
-	_, oauthOption, err := oauth.WithOAuth(mux, &oauth.Config{
+	_, oauthOption, err := mark3labs.WithOAuth(mux, &oauth.Config{
 		Provider:  "okta", // or "hmac", "google", "azure"
 		Issuer:    "https://your-company.okta.com",
 		Audience:  "api://your-mcp-server",

From a41ba571c037e292f92522468afa04ec551c3006 Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 14:58:57 -0700
Subject: [PATCH 4/9] refactor(core): extract cache and context utilities for
 SDK-agnostic support

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 CLAUDE.md                      | 90 ++++++++++++++++++++++++++++++----
 README.md                      | 89 ++++++++++++++++++++++++++++++---
 docs/generic-implementation.md | 90 ++++++++++++++++++++++++----------
 docs/generic-plan.md           | 79 +++++++++++++++++++++++++++--
 4 files changed, 299 insertions(+), 49 deletions(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index 336a459..f4d3e98 100644
--- a/CLAUDE.md
+++ b/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)
+```
diff --git a/README.md b/README.md
index 3700d33..496f52e 100644
--- a/README.md
+++ b/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,6 +133,57 @@ 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.
 
 ---
diff --git a/docs/generic-implementation.md b/docs/generic-implementation.md
index a080bed..63cf10c 100644
--- a/docs/generic-implementation.md
+++ b/docs/generic-implementation.md
@@ -80,11 +80,11 @@ func GetUserFromContext(ctx context.Context) (*User, bool)
 
 ---
 
-## Phase 1: Core Package Extraction
+## Phase 1: Core Package Extraction ✅
 
 **Goal**: Extract SDK-agnostic OAuth logic into core package without breaking existing functionality.
 
-### Checkpoint 1.1: Create cache.go ⬜
+### Checkpoint 1.1: Create cache.go ✅
 
 **Task**: Extract token cache logic from middleware.go into separate file.
 
@@ -106,9 +106,11 @@ go test ./... -v
 
 **Expected Outcome**: Build succeeds, all tests pass.
 
+**Actual Outcome**: ✅ Completed. File created with 68 lines. All tests pass.
+
 ---
 
-### Checkpoint 1.2: Create context.go ⬜
+### Checkpoint 1.2: Create context.go ✅
 
 **Task**: Extract context-related functions into separate file.
 
@@ -132,9 +134,11 @@ go test ./... -v
 
 **Expected Outcome**: Build succeeds, all tests pass.
 
+**Actual Outcome**: ✅ Completed. File created with 46 lines including WithUser() function. All tests pass.
+
 ---
 
-### Checkpoint 1.3: Update imports in existing files ⬜
+### Checkpoint 1.3: Update imports in existing files ✅
 
 **Task**: Update all internal imports to use new file structure.
 
@@ -152,9 +156,11 @@ go mod tidy
 
 **Expected Outcome**: No import errors, all tests pass.
 
+**Actual Outcome**: ✅ Completed. Removed sync import, extracted code to cache.go and context.go. All tests pass.
+
 ---
 
-### Checkpoint 1.4: Add ValidateTokenCached method to Server ⬜
+### Checkpoint 1.4: Add ValidateTokenCached method to Server ✅
 
 **Task**: Add new core method that adapters can use for token validation.
 
@@ -210,13 +216,15 @@ go test ./... -v
 
 **Expected Outcome**: Build succeeds, new method available.
 
+**Actual Outcome**: ✅ Completed. Added ValidateTokenCached() and WithUser() to core. All tests pass.
+
 ---
 
-## Phase 2: Create mark3labs Adapter Package
+## Phase 2: Create mark3labs Adapter Package ✅
 
 **Goal**: Move mark3labs-specific code into dedicated adapter package.
 
-### Checkpoint 2.1: Create mark3labs directory structure ⬜
+### Checkpoint 2.1: Create mark3labs directory structure ✅
 
 **Task**: Create new package directory for mark3labs adapter.
 
@@ -234,9 +242,11 @@ ls -la mark3labs/
 
 **Expected Outcome**: Directory and files exist.
 
+**Actual Outcome**: ✅ Completed. Created mark3labs/ directory with oauth.go and middleware.go files.
+
 ---
 
-### Checkpoint 2.2: Implement mark3labs/oauth.go ⬜
+### Checkpoint 2.2: Implement mark3labs/oauth.go ✅
 
 **Task**: Create WithOAuth function for mark3labs SDK.
 
@@ -272,9 +282,11 @@ cd mark3labs && go build .
 
 **Expected Outcome**: Package builds successfully.
 
+**Actual Outcome**: ✅ Completed. Created mark3labs/oauth.go with 45 lines. Package builds successfully.
+
 ---
 
-### Checkpoint 2.3: Implement mark3labs/middleware.go ⬜
+### Checkpoint 2.3: Implement mark3labs/middleware.go ✅
 
 **Task**: Create middleware adapter for mark3labs SDK.
 
@@ -304,9 +316,11 @@ go test ./mark3labs/...
 
 **Expected Outcome**: Package builds, basic tests pass.
 
+**Actual Outcome**: ✅ Completed. Created mark3labs/middleware.go with 38 lines using ValidateTokenCached(). Package builds successfully.
+
 ---
 
-### Checkpoint 2.4: Update examples to use mark3labs package ⬜
+### Checkpoint 2.4: Update examples to use mark3labs package ✅
 
 **Task**: Update example code to import from mark3labs package.
 
@@ -331,13 +345,15 @@ cd examples/advanced && go build .
 
 **Expected Outcome**: Examples build and run successfully.
 
+**Actual Outcome**: ✅ Completed. Updated both examples to use mark3labs.WithOAuth(). Both examples build successfully.
+
 ---
 
-## Phase 3: Create Official SDK Adapter Package
+## Phase 3: Create Official SDK Adapter Package ✅
 
 **Goal**: Add support for official modelcontextprotocol/go-sdk.
 
-### Checkpoint 3.1: Add official SDK dependency ⬜
+### Checkpoint 3.1: Add official SDK dependency ✅
 
 **Task**: Add official SDK to go.mod.
 
@@ -358,9 +374,11 @@ go mod verify
 
 **Expected Outcome**: Dependency added successfully.
 
+**Actual Outcome**: ✅ Completed. Added github.com/modelcontextprotocol/go-sdk v1.0.0 to go.mod during Phase 0 verification.
+
 ---
 
-### Checkpoint 3.2: Create mcp directory structure ⬜
+### Checkpoint 3.2: Create mcp directory structure ✅
 
 **Task**: Create new package directory for official SDK adapter.
 
@@ -377,9 +395,11 @@ ls -la mcp/
 
 **Expected Outcome**: Directory and files exist.
 
+**Actual Outcome**: ✅ Completed. Created mcp/ directory with oauth.go file.
+
 ---
 
-### Checkpoint 3.3: Implement mcp/oauth.go ⬜
+### Checkpoint 3.3: Implement mcp/oauth.go ✅
 
 **Task**: Create WithOAuth function for official SDK.
 
@@ -423,12 +443,16 @@ cd mcp && go build .
 
 **Expected Outcome**: Package builds successfully.
 
+**Actual Outcome**: ✅ Completed. Created mcp/oauth.go with 76 lines. Uses custom HTTP handler wrapper instead of WrapHandler for more control. Package builds successfully.
+
 ---
 
-### Checkpoint 3.4: Create official SDK example ⬜
+### Checkpoint 3.4: Create official SDK example ⏭️
 
 **Task**: Create example demonstrating official SDK integration.
 
+**Status**: Skipped for initial release. Can be added later.
+
 **Files to Create**:
 - `examples/official/main.go`
 
@@ -467,11 +491,13 @@ cd examples/official && go build .
 
 ---
 
-## Phase 4: Testing and Validation
+## Phase 4: Testing and Validation ⚠️
 
 **Goal**: Ensure both SDK integrations work correctly.
 
-### Checkpoint 4.1: Update existing tests ⬜
+**Status**: Core tests passing. Adapter-specific tests pending.
+
+### Checkpoint 4.1: Update existing tests ✅
 
 **Task**: Update all tests to use new package structure.
 
@@ -496,12 +522,16 @@ go test ./... -cover
 
 **Expected Outcome**: All tests pass with race detector.
 
+**Actual Outcome**: ✅ Completed. All existing core tests pass without modification. verify_context_test.go validates official SDK context propagation.
+
 ---
 
 ### Checkpoint 4.2: Create mark3labs integration tests ⬜
 
 **Task**: Create comprehensive tests for mark3labs adapter.
 
+**Status**: Pending - to be added in follow-up PR.
+
 **Files to Create**:
 - `mark3labs/integration_test.go`
 
@@ -524,6 +554,8 @@ go test ./mark3labs/... -v -cover
 
 **Task**: Create comprehensive tests for official SDK adapter.
 
+**Status**: Pending - to be added in follow-up PR.
+
 **Files to Create**:
 - `mcp/integration_test.go`
 
@@ -542,7 +574,7 @@ go test ./mcp/... -v -cover
 
 ---
 
-### Checkpoint 4.4: Run full test suite ⬜
+### Checkpoint 4.4: Run full test suite ✅
 
 **Task**: Verify all tests pass across all packages.
 
@@ -566,12 +598,16 @@ open coverage.html
 
 **Expected Outcome**: Coverage report shows good coverage across all packages.
 
+**Actual Outcome**: ✅ Completed. All core tests pass. Build successful across all packages (core, mark3labs, mcp, provider, examples).
+
 ---
 
-## Phase 5: Documentation Updates
+## Phase 5: Documentation Updates ⬜
 
 **Goal**: Update all documentation to reflect new package structure.
 
+**Status**: Pending - README and migration guide updates needed.
+
 ### Checkpoint 5.1: Update README.md ⬜
 
 **Task**: Update main README with new package structure.
@@ -730,13 +766,13 @@ cd examples/official && go run main.go
 | Phase | Status | Completion |
 |-------|--------|------------|
 | Phase 0: Pre-Implementation Verification | ✅ Completed | 100% |
-| Phase 1: Core Package Extraction | ⬜ Not Started | 0% |
-| Phase 2: mark3labs Adapter | ⬜ Not Started | 0% |
-| Phase 3: Official SDK Adapter | ⬜ Not Started | 0% |
-| Phase 4: Testing & Validation | ⬜ Not Started | 0% |
+| Phase 1: Core Package Extraction | ✅ Completed | 100% |
+| Phase 2: mark3labs Adapter | ✅ Completed | 100% |
+| Phase 3: Official SDK Adapter | ✅ Completed | 100% |
+| Phase 4: Testing & Validation | ⚠️ Partial | 75% |
 | Phase 5: Documentation | ⬜ Not Started | 0% |
 | Phase 6: Release Preparation | ⬜ Not Started | 0% |
-| **Overall** | **🟡 In Progress** | **14%** |
+| **Overall** | **🟢 Implementation Complete** | **82%** |
 
 ---
 
@@ -771,7 +807,7 @@ go test ./provider/... -v
 
 **Last Updated**: 2025-10-22
 **Verification Date**: 2025-10-22 (Phase 0 completed)
-**Implementation Start Date**: TBD
-**Target Completion Date**: TBD
+**Implementation Start Date**: 2025-10-22
+**Implementation Completion Date**: 2025-10-22
 
-**Current Status**: Phase 0 (Verification) completed. Ready to begin implementation once approved.
+**Current Status**: ✅ Core implementation complete (Phases 0-3 + core testing). Documentation updates pending (Phase 5).
diff --git a/docs/generic-plan.md b/docs/generic-plan.md
index d1ef013..97847ef 100644
--- a/docs/generic-plan.md
+++ b/docs/generic-plan.md
@@ -1,12 +1,22 @@
 # OpenTelemetry Pattern Refactoring Plan
 
+## Implementation Status
+
+**Status**: ✅ **IMPLEMENTED** (2025-10-22)
+
+**Completion**: 82% (Core implementation complete, documentation pending)
+
+**See**: `docs/generic-implementation.md` for detailed checkpoint tracking.
+
+---
+
 ## Overview
 
 This document outlines the plan to refactor `oauth-mcp-proxy` to support both mark3labs/mcp-go and the official modelcontextprotocol/go-sdk using the OpenTelemetry pattern approach.
 
-## Current State
+## Original State
 
-The library currently supports only `github.com/mark3labs/mcp-go` (v0.41.1) with a single `WithOAuth()` function that returns `mcpserver.ServerOption`.
+The library originally supported only `github.com/mark3labs/mcp-go` (v0.41.1) with a single `WithOAuth()` function that returns `mcpserver.ServerOption`.
 
 ## Proposed Structure
 
@@ -229,11 +239,70 @@ Users have two migration paths:
 **Option 2: Migrate to Official SDK**
 Follow the official SDK migration guide in the new documentation.
 
+---
+
+## Implementation Results
+
+### What Was Implemented
+
+**Date**: 2025-10-22
+
+**Files Created:**
+- `cache.go` (68 lines) - Token cache logic
+- `context.go` (46 lines) - Context utilities (WithOAuthToken, GetOAuthToken, WithUser, GetUserFromContext)
+- `mark3labs/oauth.go` (45 lines) - mark3labs SDK adapter
+- `mark3labs/middleware.go` (38 lines) - mark3labs middleware implementation
+- `mcp/oauth.go` (76 lines) - Official SDK adapter
+- `verify_context_test.go` - Context propagation verification test
+
+**Files Modified:**
+- `oauth.go` - Added ValidateTokenCached() method
+- `middleware.go` - Removed extracted code to new files
+- `examples/simple/main.go` - Updated to use mark3labs package
+- `examples/advanced/main.go` - Updated to use mark3labs package
+- `go.mod` - Added official SDK v1.0.0
+
+**Verification:**
+- ✅ All existing tests pass
+- ✅ Both example apps build successfully
+- ✅ Official SDK context propagation verified
+- ✅ Core API contract implemented as designed
+
+### Implementation Time
+
+**Actual Time**: ~3 hours (vs 1 day estimated)
+
+Faster than estimated due to:
+- Clear verification phase eliminated uncertainty
+- Well-defined core API contract
+- Minimal changes needed to existing tests
+
+### Deviations from Plan
+
+1. **Checkpoint 3.4 Skipped**: Official SDK example not created (can be added later)
+2. **Checkpoint 4.2 & 4.3 Pending**: Adapter-specific integration tests deferred to follow-up PR
+3. **mcp/oauth.go**: Implemented custom HTTP handler wrapper instead of using WrapHandler for more explicit control
+
+### Outstanding Work
+
+- **Phase 5**: README.md updates (show both SDKs, migration guide)
+- **Phase 6**: Release preparation (CHANGELOG, version bump, PR)
+- **Future**: Comprehensive adapter integration tests
+
+---
+
 ## Open Questions
 
-1. Should we maintain v1 branch for bug fixes during transition period?
-2. How long should we support v1 before deprecating?
-3. Should we add compatibility shims in v2 to ease migration?
+### Answered (Based on Gemini 2.5 Pro Review)
+
+1. **Should we maintain v1 branch for bug fixes during transition period?**
+   - ✅ Yes. Create `v1` branch from last commit before refactor. Support critical security fixes for 3-6 months.
+
+2. **How long should we support v1 before deprecating?**
+   - ✅ 3-6 months for critical security fixes only. No new features.
+
+3. **Should we add compatibility shims in v2 to ease migration?**
+   - ❌ No. Major version bump is the time for clean break. Shims add complexity and confusion. Use clear migration guide instead.
 
 ## References
 

From dad9dc760ba0bc2ca331f1f1b34813e14f73fa1f Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 15:35:17 -0700
Subject: [PATCH 5/9] fix(advancd/main.go): Rename file

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 examples/README.md                        | 367 ++++++++++++++++++++++
 examples/{ => mark3labs}/advanced/main.go |   7 +-
 examples/{ => mark3labs}/simple/main.go   |  35 ++-
 examples/official/advanced/main.go        | 150 +++++++++
 examples/official/simple/main.go          |  91 ++++++
 5 files changed, 641 insertions(+), 9 deletions(-)
 create mode 100644 examples/README.md
 rename examples/{ => mark3labs}/advanced/main.go (92%)
 rename examples/{ => mark3labs}/simple/main.go (51%)
 create mode 100644 examples/official/advanced/main.go
 create mode 100644 examples/official/simple/main.go

diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 0000000..608446d
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,367 @@
+# OAuth MCP Proxy Examples
+
+This directory contains example MCP servers demonstrating OAuth integration with both supported SDKs.
+
+## Directory Structure
+
+```
+examples/
+├── mark3labs/              (mark3labs/mcp-go SDK examples)
+│   ├── simple/            - Basic OAuth integration
+│   └── advanced/          - ConfigBuilder, env vars, multiple tools
+│
+└── official/               (modelcontextprotocol/go-sdk examples)
+    ├── simple/            - Basic OAuth integration
+    └── advanced/          - Multiple tools, env vars, logging
+```
+
+## Examples Overview
+
+| SDK | Example | Tools | Provider | Features |
+|-----|---------|-------|----------|----------|
+| **mark3labs** | simple | 1 (greet) | Okta | Basic OAuth, env vars |
+| **mark3labs** | advanced | 3 (greet, echo, time) | Okta | ConfigBuilder, env vars, logging |
+| **official** | simple | 1 (greet) | Okta | Basic OAuth, env vars |
+| **official** | advanced | 3 (greet, whoami, server_time) | Okta | ConfigBuilder, env vars, logging |
+
+---
+
+## Quick Start
+
+### mark3labs SDK
+
+**Simple:**
+```bash
+cd examples/mark3labs/simple
+go run main.go
+```
+
+**Advanced:**
+```bash
+cd examples/mark3labs/advanced
+go run main.go
+```
+
+### Official SDK
+
+**Simple:**
+```bash
+cd examples/official/simple
+go run main.go
+```
+
+**Advanced:**
+```bash
+cd examples/official/advanced
+go run main.go
+```
+
+All examples start a server on `http://localhost:8080` with OAuth protection.
+
+---
+
+## Okta Setup
+
+All examples use **Okta** as the OAuth provider. Before running, you need to set up Okta:
+
+### 1. Create Okta Account
+
+Sign up at https://developer.okta.com (free developer account)
+
+### 2. Create API in Okta
+
+1. Go to **Security > API** in Okta Admin Console
+2. Click **Add Authorization Server** or use the default
+3. Note your **Issuer URI** (e.g., `https://dev-12345.okta.com`)
+4. Create an **Audience** identifier (e.g., `api://my-mcp-server`)
+
+### 3. Set Environment Variables
+
+```bash
+export OKTA_DOMAIN="dev-12345.okta.com"        # Your Okta domain
+export OKTA_AUDIENCE="api://my-mcp-server"     # Your API identifier
+export SERVER_URL="http://localhost:8080"      # Your server URL
+```
+
+### 4. Get a Test Token
+
+**Option A: Using Okta CLI**
+```bash
+# Install Okta CLI
+brew install --cask oktacli  # macOS
+
+# Login and get token
+okta login
+okta get token --audience api://my-mcp-server
+```
+
+**Option B: Using Okta Dashboard**
+1. Go to **Security > API > Authorization Servers**
+2. Click your authorization server
+3. Go to **Token Preview** tab
+4. Generate a token with your audience
+
+### 5. Test the Server
+
+```bash
+# Save your Okta token
+TOKEN="<your-okta-access-token>"
+
+# Test with curl
+curl -H "Authorization: Bearer $TOKEN" \
+     -H "Content-Type: application/json" \
+     -H "Accept: application/json, text/event-stream" \
+     -X POST \
+     http://localhost:8080 \
+     -d '{
+       "jsonrpc": "2.0",
+       "method": "tools/list",
+       "id": 1
+     }'
+```
+
+---
+
+## Configuration Options
+
+### Environment Variables
+
+All examples support these environment variables:
+
+```bash
+# Required Okta Configuration
+export OKTA_DOMAIN="dev-12345.okta.com"         # Your Okta domain
+export OKTA_AUDIENCE="api://my-mcp-server"      # Your API identifier
+
+# Server Configuration
+export SERVER_URL="http://localhost:8080"       # Your server URL
+export PORT="8080"                              # Server port (default: 8080)
+export MCP_HOST="localhost"                     # Server host (default: localhost)
+export MCP_PORT="8080"                          # Server port for ConfigBuilder
+
+# Optional: For HTTPS
+export HTTPS_CERT_FILE="/path/to/cert.pem"      # If set, enables HTTPS
+```
+
+### Using Other Providers
+
+To use Google or Azure AD instead of Okta, modify the config:
+
+**Google:**
+```go
+&oauth.Config{
+    Provider: "google",
+    Issuer:   "https://accounts.google.com",
+    Audience: "your-google-client-id.apps.googleusercontent.com",
+}
+```
+
+**Azure AD:**
+```go
+&oauth.Config{
+    Provider: "azure",
+    Issuer:   "https://login.microsoftonline.com/YOUR-TENANT-ID/v2.0",
+    Audience: "api://your-app-id",
+}
+```
+
+---
+
+## Example Comparison
+
+### mark3labs/simple
+
+**What it shows:**
+- Basic OAuth integration with `mark3labs.WithOAuth()`
+- Single tool with user context access
+- Okta provider configuration
+- Environment variable support
+
+**Use when:** You want the simplest possible OAuth setup with mark3labs SDK.
+
+### mark3labs/advanced
+
+**What it shows:**
+- `ConfigBuilder` for flexible configuration
+- Environment variable support (Okta domain, audience, server URL)
+- Multiple tools with different functionality
+- Custom logging
+- OAuth endpoint discovery logging
+- Production-ready patterns
+
+**Use when:** You need production-ready configuration with mark3labs SDK.
+
+### official/simple
+
+**What it shows:**
+- Basic OAuth integration with `mcpoauth.WithOAuth()`
+- Single tool with user context access
+- Official SDK tool definition patterns
+- Okta provider configuration
+- Environment variable support
+
+**Use when:** You want the simplest possible OAuth setup with official SDK.
+
+### official/advanced
+
+**What it shows:**
+- `ConfigBuilder` for flexible configuration
+- Multiple tools (greet, whoami, server_time)
+- Environment variable support (Okta domain, audience)
+- OAuth endpoint discovery logging
+- Production-ready patterns
+- Official SDK patterns
+
+**Use when:** You need production-ready configuration with official SDK.
+
+---
+
+## Code Patterns Comparison
+
+### mark3labs SDK
+
+**Setup:**
+```go
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+)
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{...})
+mcpServer := mcpserver.NewMCPServer("name", "1.0.0", oauthOption)
+```
+
+**Adding Tools:**
+```go
+mcpServer.AddTool(tool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+    user, _ := oauth.GetUserFromContext(ctx)
+    return mcp.NewToolResultText("Hello, " + user.Username), nil
+})
+```
+
+### Official SDK
+
+**Setup:**
+```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)
+```
+
+**Adding Tools:**
+```go
+mcp.AddTool(mcpServer, &mcp.Tool{...},
+    func(ctx context.Context, req *mcp.CallToolRequest, params *P) (*mcp.CallToolResult, any, error) {
+        user, _ := oauth.GetUserFromContext(ctx)
+        return &mcp.CallToolResult{
+            Content: []mcp.Content{&mcp.TextContent{Text: "Hello, " + user.Username}},
+        }, nil, nil
+    })
+```
+
+**Key Difference**: mark3labs uses ServerOption before server creation, official SDK wraps the server with http.Handler after creation.
+
+---
+
+## Accessing User Information
+
+All examples show how to access authenticated user information:
+
+```go
+user, ok := oauth.GetUserFromContext(ctx)
+if !ok {
+    return nil, fmt.Errorf("authentication required")
+}
+
+// Available fields:
+user.Subject   // OAuth "sub" claim (user ID)
+user.Username  // "preferred_username" or "sub"
+user.Email     // "email" claim
+```
+
+---
+
+## Common Issues
+
+### "authentication required: missing OAuth token"
+
+**Cause:** No Authorization header or invalid format.
+
+**Solution:**
+```bash
+# Make sure to include Bearer token
+curl -H "Authorization: Bearer YOUR_TOKEN" ...
+```
+
+### "authentication failed: token validation failed"
+
+**Cause:** Invalid token or wrong secret.
+
+**Solution:**
+- For HMAC: Ensure `HMAC_SECRET` matches the secret used to sign the token
+- For OIDC: Verify issuer, audience, and that the token is from the correct provider
+
+### "Accept must contain both 'application/json' and 'text/event-stream'"
+
+**Cause:** Missing Accept header (official SDK only).
+
+**Solution:**
+```bash
+curl -H "Accept: application/json, text/event-stream" ...
+```
+
+---
+
+## Building for Production
+
+### Dockerfile Example
+
+```dockerfile
+FROM golang:1.24 AS builder
+WORKDIR /app
+COPY . .
+RUN go build -o server ./examples/advanced
+
+FROM alpine:latest
+RUN apk --no-cache add ca-certificates
+WORKDIR /root/
+COPY --from=builder /app/server .
+
+# Set production environment variables
+ENV OAUTH_PROVIDER=okta
+ENV SERVER_URL=https://your-server.com
+
+CMD ["./server"]
+```
+
+### Production Checklist
+
+- [ ] Use OIDC provider (Okta/Google/Azure), not HMAC
+- [ ] Set `SERVER_URL` to your actual domain (HTTPS)
+- [ ] Store secrets in environment variables or secret manager
+- [ ] Enable HTTPS (use reverse proxy like nginx or Caddy)
+- [ ] Configure proper CORS if needed
+- [ ] Set up monitoring and logging
+- [ ] Review OAuth scopes and permissions
+
+---
+
+## Further Reading
+
+- **Migration Guide**: [../MIGRATION-V2.md](../MIGRATION-V2.md)
+- **Main README**: [../README.md](../README.md)
+- **Project Documentation**: [../CLAUDE.md](../CLAUDE.md)
+- **Implementation Details**: [../docs/generic-implementation.md](../docs/generic-implementation.md)
+
+---
+
+## Need Help?
+
+- **Issues**: https://github.com/tuannvm/oauth-mcp-proxy/issues
+- **Discussions**: https://github.com/tuannvm/oauth-mcp-proxy/discussions
+- **Documentation**: See files in `/docs` directory
diff --git a/examples/advanced/main.go b/examples/mark3labs/advanced/main.go
similarity index 92%
rename from examples/advanced/main.go
rename to examples/mark3labs/advanced/main.go
index d4551f0..35ef7f7 100644
--- a/examples/advanced/main.go
+++ b/examples/mark3labs/advanced/main.go
@@ -15,10 +15,13 @@ import (
 
 func main() {
 	// Feature 1: ConfigBuilder - Auto-generates ServerURL from host/port/TLS
+	// Set these environment variables:
+	//   export OKTA_DOMAIN="dev-12345.okta.com"
+	//   export OKTA_AUDIENCE="api://my-mcp-server"
 	cfg, err := oauth.NewConfigBuilder().
 		WithProvider("okta").
-		WithIssuer("https://your-company.okta.com").
-		WithAudience("api://your-mcp-server").
+		WithIssuer(fmt.Sprintf("https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com"))).
+		WithAudience(getEnv("OKTA_AUDIENCE", "api://my-mcp-server")).
 		WithHost(getEnv("MCP_HOST", "localhost")).
 		WithPort(getEnv("MCP_PORT", "8080")).
 		WithTLS(getEnv("HTTPS_CERT_FILE", "") != "").
diff --git a/examples/simple/main.go b/examples/mark3labs/simple/main.go
similarity index 51%
rename from examples/simple/main.go
rename to examples/mark3labs/simple/main.go
index 4e2a576..03e7e18 100644
--- a/examples/simple/main.go
+++ b/examples/mark3labs/simple/main.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"os"
 
 	"github.com/mark3labs/mcp-go/mcp"
 	mcpserver "github.com/mark3labs/mcp-go/server"
@@ -16,12 +17,16 @@ func main() {
 	// 1. Create HTTP multiplexer
 	mux := http.NewServeMux()
 
-	// 2. Enable OAuth authentication
+	// 2. Configure OAuth with Okta
+	// Set these environment variables:
+	//   export OKTA_DOMAIN="dev-12345.okta.com"              (your Okta domain)
+	//   export OKTA_AUDIENCE="api://my-mcp-server"          (your API identifier)
+	//   export SERVER_URL="https://mcp.example.com"         (your server URL)
 	_, oauthOption, err := mark3labs.WithOAuth(mux, &oauth.Config{
-		Provider:  "okta", // or "hmac", "google", "azure"
-		Issuer:    "https://your-company.okta.com",
-		Audience:  "api://your-mcp-server",
-		ServerURL: "https://your-server.com",
+		Provider:  "okta",
+		Issuer:    fmt.Sprintf("https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com")),
+		Audience:  getEnv("OKTA_AUDIENCE", "api://my-mcp-server"),
+		ServerURL: getEnv("SERVER_URL", "http://localhost:8080"),
 	})
 	if err != nil {
 		log.Fatalf("OAuth setup failed: %v", err)
@@ -55,8 +60,24 @@ func main() {
 	mux.Handle("/mcp", streamableServer)
 
 	// 6. Start server
-	log.Println("Starting MCP server on :8080")
-	if err := http.ListenAndServe(":8080", mux); err != nil {
+	port := getEnv("PORT", "8080")
+	log.Printf("Starting MCP server on :%s", port)
+	log.Printf("OAuth Provider: Okta")
+	log.Printf("Issuer: https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com"))
+	log.Printf("Audience: %s", getEnv("OKTA_AUDIENCE", "api://my-mcp-server"))
+	log.Println("\nMake sure to set your Okta environment variables:")
+	log.Println("  export OKTA_DOMAIN=dev-12345.okta.com")
+	log.Println("  export OKTA_AUDIENCE=api://my-mcp-server")
+	log.Println("  export SERVER_URL=http://localhost:8080")
+
+	if err := http.ListenAndServe(":"+port, mux); err != nil {
 		log.Fatalf("Server failed: %v", err)
 	}
 }
+
+func getEnv(key, defaultValue string) string {
+	if value := os.Getenv(key); value != "" {
+		return value
+	}
+	return defaultValue
+}
diff --git a/examples/official/advanced/main.go b/examples/official/advanced/main.go
new file mode 100644
index 0000000..7144d5b
--- /dev/null
+++ b/examples/official/advanced/main.go
@@ -0,0 +1,150 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+	mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+func main() {
+	// Feature 1: ConfigBuilder - Auto-generates ServerURL from host/port/TLS
+	// Set these environment variables:
+	//   export OKTA_DOMAIN="dev-12345.okta.com"
+	//   export OKTA_AUDIENCE="api://my-mcp-server"
+	cfg, err := oauth.NewConfigBuilder().
+		WithProvider("okta").
+		WithIssuer(fmt.Sprintf("https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com"))).
+		WithAudience(getEnv("OKTA_AUDIENCE", "api://my-mcp-server")).
+		WithHost(getEnv("MCP_HOST", "localhost")).
+		WithPort(getEnv("MCP_PORT", "8080")).
+		WithTLS(getEnv("HTTPS_CERT_FILE", "") != "").
+		Build()
+
+	if err != nil {
+		log.Fatalf("Config setup failed: %v", err)
+	}
+
+	// Feature 2: Create MCP server with multiple tools
+	mcpServer := mcp.NewServer(&mcp.Implementation{
+		Name:    "Advanced Official SDK Example",
+		Version: "1.0.0",
+	}, nil)
+
+	// Tool 1: Greet user
+	type GreetParams struct{}
+	mcp.AddTool(mcpServer, &mcp.Tool{
+		Name:        "greet",
+		Description: "Greets the authenticated user",
+	}, func(ctx context.Context, req *mcp.CallToolRequest, params *GreetParams) (*mcp.CallToolResult, any, error) {
+		user, ok := oauth.GetUserFromContext(ctx)
+		if !ok {
+			return nil, nil, fmt.Errorf("authentication required")
+		}
+
+		message := fmt.Sprintf("Hello, %s! Your email is: %s", user.Username, user.Email)
+		log.Printf("[greet] Called by user: %s", user.Username)
+
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{
+				&mcp.TextContent{Text: message},
+			},
+		}, nil, nil
+	})
+
+	// Tool 2: Get user info
+	type UserInfoParams struct{}
+	mcp.AddTool(mcpServer, &mcp.Tool{
+		Name:        "whoami",
+		Description: "Returns information about the authenticated user",
+	}, func(ctx context.Context, req *mcp.CallToolRequest, params *UserInfoParams) (*mcp.CallToolResult, any, error) {
+		user, ok := oauth.GetUserFromContext(ctx)
+		if !ok {
+			return nil, nil, fmt.Errorf("authentication required")
+		}
+
+		info := fmt.Sprintf(`User Information:
+- Subject: %s
+- Username: %s
+- Email: %s`, user.Subject, user.Username, user.Email)
+
+		log.Printf("[whoami] Called by user: %s", user.Username)
+
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{
+				&mcp.TextContent{Text: info},
+			},
+		}, nil, nil
+	})
+
+	// Tool 3: Server time
+	type TimeParams struct{}
+	mcp.AddTool(mcpServer, &mcp.Tool{
+		Name:        "server_time",
+		Description: "Returns the current server time",
+	}, func(ctx context.Context, req *mcp.CallToolRequest, params *TimeParams) (*mcp.CallToolResult, any, error) {
+		user, ok := oauth.GetUserFromContext(ctx)
+		if !ok {
+			return nil, nil, fmt.Errorf("authentication required")
+		}
+
+		now := time.Now().Format(time.RFC3339)
+		message := fmt.Sprintf("Server time: %s (requested by %s)", now, user.Username)
+
+		log.Printf("[server_time] Called by user: %s", user.Username)
+
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{
+				&mcp.TextContent{Text: message},
+			},
+		}, nil, nil
+	})
+
+	// Feature 3: Enable OAuth with official SDK
+	mux := http.NewServeMux()
+	oauthServer, handler, err := mcpoauth.WithOAuth(mux, cfg, mcpServer)
+	if err != nil {
+		log.Fatalf("OAuth setup failed: %v", err)
+	}
+
+	// Feature 4: LogStartup - Displays OAuth endpoint information
+	useTLS := getEnv("HTTPS_CERT_FILE", "") != ""
+	oauthServer.LogStartup(useTLS)
+
+	// Additional OAuth endpoints available via helper methods
+	log.Printf("\nOAuth Discovery URLs:")
+	log.Printf("  - Metadata: %s", oauthServer.GetAuthorizationServerMetadataURL())
+	log.Printf("  - OIDC Discovery: %s", oauthServer.GetOIDCDiscoveryURL())
+
+	// Feature 5: Server status
+	log.Printf("\n%s", oauthServer.GetStatusString(useTLS))
+	log.Printf("Tools registered: greet, whoami, server_time")
+
+	// Start server
+	port := getEnv("PORT", "8080")
+	addr := ":" + port
+
+	log.Printf("\nStarting MCP server on %s", addr)
+	log.Println("\nMake sure to set your Okta environment variables:")
+	log.Println("  export OKTA_DOMAIN=dev-12345.okta.com")
+	log.Println("  export OKTA_AUDIENCE=api://my-mcp-server")
+	log.Println("\nTo test, obtain an access token from Okta and use:")
+	log.Printf("  curl -H 'Authorization: Bearer <okta-token>' http://localhost:%s", port)
+
+	if err := http.ListenAndServe(addr, handler); err != nil {
+		log.Fatalf("Server failed: %v", err)
+	}
+}
+
+func getEnv(key, defaultValue string) string {
+	if value := os.Getenv(key); value != "" {
+		return value
+	}
+	return defaultValue
+}
diff --git a/examples/official/simple/main.go b/examples/official/simple/main.go
new file mode 100644
index 0000000..a3aca68
--- /dev/null
+++ b/examples/official/simple/main.go
@@ -0,0 +1,91 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	oauth "github.com/tuannvm/oauth-mcp-proxy"
+	mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+func main() {
+	// 1. Create HTTP multiplexer
+	mux := http.NewServeMux()
+
+	// 2. Create MCP server using official SDK
+	mcpServer := mcp.NewServer(&mcp.Implementation{
+		Name:    "Official SDK OAuth Example",
+		Version: "1.0.0",
+	}, nil)
+
+	// 3. Add a simple greeting tool (all tools will be OAuth-protected)
+	type GreetParams struct{}
+
+	mcp.AddTool(mcpServer, &mcp.Tool{
+		Name:        "greet",
+		Description: "Greets the authenticated user by their username",
+	}, func(ctx context.Context, req *mcp.CallToolRequest, params *GreetParams) (*mcp.CallToolResult, any, error) {
+		// Access the authenticated user from context
+		user, ok := oauth.GetUserFromContext(ctx)
+		if !ok {
+			return nil, nil, fmt.Errorf("authentication required")
+		}
+
+		message := fmt.Sprintf("Hello, %s! Your email is: %s", user.Username, user.Email)
+
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{
+				&mcp.TextContent{Text: message},
+			},
+		}, nil, nil
+	})
+
+	// 4. Configure OAuth with Okta
+	// Set these environment variables:
+	//   export OKTA_DOMAIN="dev-12345.okta.com"
+	//   export OKTA_AUDIENCE="api://my-mcp-server"
+	//   export SERVER_URL="http://localhost:8080"
+	cfg := &oauth.Config{
+		Provider:  "okta",
+		Issuer:    fmt.Sprintf("https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com")),
+		Audience:  getEnv("OKTA_AUDIENCE", "api://my-mcp-server"),
+		ServerURL: getEnv("SERVER_URL", "http://localhost:8080"),
+	}
+
+	// 5. Enable OAuth protection - returns an http.Handler
+	oauthServer, handler, err := mcpoauth.WithOAuth(mux, cfg, mcpServer)
+	if err != nil {
+		log.Fatalf("OAuth setup failed: %v", err)
+	}
+
+	// 6. Log OAuth information
+	oauthServer.LogStartup(false) // false = HTTP (set true if using HTTPS)
+
+	// 7. Start server
+	port := getEnv("PORT", "8080")
+	addr := ":" + port
+
+	log.Printf("Starting MCP server on %s", addr)
+	log.Printf("OAuth Provider: Okta")
+	log.Printf("Issuer: https://%s", getEnv("OKTA_DOMAIN", "dev-12345.okta.com"))
+	log.Printf("Audience: %s", getEnv("OKTA_AUDIENCE", "api://my-mcp-server"))
+	log.Println("\nMake sure to set your Okta environment variables:")
+	log.Println("  export OKTA_DOMAIN=dev-12345.okta.com")
+	log.Println("  export OKTA_AUDIENCE=api://my-mcp-server")
+	log.Println("  export SERVER_URL=http://localhost:8080")
+
+	if err := http.ListenAndServe(addr, handler); err != nil {
+		log.Fatalf("Server failed: %v", err)
+	}
+}
+
+func getEnv(key, defaultValue string) string {
+	if value := os.Getenv(key); value != "" {
+		return value
+	}
+	return defaultValue
+}

From cb5ae40b082fb2b9cf36ed4b9d60670631b36f62 Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 16:04:38 -0700
Subject: [PATCH 6/9] feat(oauth): add token validation and error handling in
 WrapHandler

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 oauth.go | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/oauth.go b/oauth.go
index a94f014..4d66b86 100644
--- a/oauth.go
+++ b/oauth.go
@@ -279,8 +279,27 @@ func (s *Server) WrapHandler(next http.Handler) http.Handler {
 			return
 		}
 
-		contextFunc := CreateHTTPContextFunc()
-		ctx := contextFunc(r.Context(), r)
+		token := authHeader[7:]
+
+		user, err := s.ValidateTokenCached(r.Context(), token)
+		if err != nil {
+			s.logger.Info("OAuth: Token validation failed: %v", err)
+
+			metadataURL := s.GetProtectedResourceMetadataURL()
+			w.Header().Add("WWW-Authenticate", `Bearer realm="OAuth", error="invalid_token", error_description="Authentication failed"`)
+			w.Header().Add("WWW-Authenticate", fmt.Sprintf(`resource_metadata="%s"`, metadataURL))
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusUnauthorized)
+
+			_ = json.NewEncoder(w).Encode(oauthErrorResponse{
+				Error:            "invalid_token",
+				ErrorDescription: "Authentication failed",
+			})
+			return
+		}
+
+		ctx := WithOAuthToken(r.Context(), token)
+		ctx = WithUser(ctx, user)
 		r = r.WithContext(ctx)
 
 		next.ServeHTTP(w, r)

From a114747b92219d99c030692d3d20029fa313e77b Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 16:06:53 -0700
Subject: [PATCH 7/9] style: Add explanatory comments about PORT and SERVER_URL
 usage

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 examples/mark3labs/simple/main.go  | 4 ++++
 examples/official/advanced/main.go | 2 +-
 examples/official/simple/main.go   | 4 ++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/examples/mark3labs/simple/main.go b/examples/mark3labs/simple/main.go
index 03e7e18..0a0ebd5 100644
--- a/examples/mark3labs/simple/main.go
+++ b/examples/mark3labs/simple/main.go
@@ -60,6 +60,10 @@ func main() {
 	mux.Handle("/mcp", streamableServer)
 
 	// 6. Start server
+	// Note: PORT is the local bind port. If you change SERVER_URL port
+	// (e.g., http://localhost:9000), also set PORT=9000 to match.
+	// For production with reverse proxy, PORT is your local port while
+	// SERVER_URL is the public URL (e.g., SERVER_URL=https://api.example.com, PORT=8080)
 	port := getEnv("PORT", "8080")
 	log.Printf("Starting MCP server on :%s", port)
 	log.Printf("OAuth Provider: Okta")
diff --git a/examples/official/advanced/main.go b/examples/official/advanced/main.go
index 7144d5b..1dee456 100644
--- a/examples/official/advanced/main.go
+++ b/examples/official/advanced/main.go
@@ -127,7 +127,7 @@ func main() {
 	log.Printf("Tools registered: greet, whoami, server_time")
 
 	// Start server
-	port := getEnv("PORT", "8080")
+	port := getEnv("MCP_PORT", "8080")
 	addr := ":" + port
 
 	log.Printf("\nStarting MCP server on %s", addr)
diff --git a/examples/official/simple/main.go b/examples/official/simple/main.go
index a3aca68..ea80496 100644
--- a/examples/official/simple/main.go
+++ b/examples/official/simple/main.go
@@ -66,6 +66,10 @@ func main() {
 	oauthServer.LogStartup(false) // false = HTTP (set true if using HTTPS)
 
 	// 7. Start server
+	// Note: PORT is the local bind port. If you change SERVER_URL port
+	// (e.g., http://localhost:9000), also set PORT=9000 to match.
+	// For production with reverse proxy, PORT is your local port while
+	// SERVER_URL is the public URL (e.g., SERVER_URL=https://api.example.com, PORT=8080)
 	port := getEnv("PORT", "8080")
 	addr := ":" + port
 

From 5437063d583a51b6b204ffe7907d7b245fbafe82 Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 18:28:46 -0700
Subject: [PATCH 8/9] chore(github workflow): update example build to find all
 main.go files

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 .github/workflows/test.yml     |   6 +-
 README.md                      |  12 +-
 docs/CLIENT-SETUP.md           |  46 ++++-
 docs/MIGRATION.md              | 368 ---------------------------------
 docs/TROUBLESHOOTING.md        |   2 +-
 docs/generic-implementation.md |  32 ++-
 docs/providers/AZURE.md        |   3 +
 docs/providers/GOOGLE.md       |   3 +
 docs/providers/HMAC.md         |  62 +++++-
 docs/providers/OKTA.md         |  53 ++++-
 examples/README.md             |   2 +-
 11 files changed, 186 insertions(+), 403 deletions(-)
 delete mode 100644 docs/MIGRATION.md

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index a675873..eb00c83 100644
--- a/.github/workflows/test.yml
+++ b/.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 ./...
diff --git a/README.md b/README.md
index 496f52e..8fa4112 100644
--- a/README.md
+++ b/README.md
@@ -190,10 +190,14 @@ 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 |
 
 ---
 
diff --git a/docs/CLIENT-SETUP.md b/docs/CLIENT-SETUP.md
index 80102ba..6a67f09 100644
--- a/docs/CLIENT-SETUP.md
+++ b/docs/CLIENT-SETUP.md
@@ -160,14 +160,30 @@ Tells clients this is an OAuth-protected resource.
 
 ### Native Mode
 
-**Server config:**
+**Server config (mark3labs SDK):**
 
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",
     Issuer:   "https://company.okta.com",
     Audience: "api://my-server",
 })
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+```
+
+**Server config (official SDK):**
+
+```go
+import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+
+mcpServer := mcp.NewServer(&mcp.Implementation{...}, nil)
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider: "okta",
+    Issuer:   "https://company.okta.com",
+    Audience: "api://my-server",
+}, mcpServer)
 ```
 
 **Client discovers:**
@@ -193,16 +209,34 @@ Client fetches metadata, sees Okta issuer, handles OAuth with Okta directly.
 
 ### Proxy Mode
 
-**Server config:**
+**Server config (mark3labs SDK):**
 
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider:     "okta",
     ClientID:     "...",
     ClientSecret: "...",
     ServerURL:    "https://your-server.com",
     RedirectURIs: "https://your-server.com/oauth/callback",
 })
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+```
+
+**Server config (official SDK):**
+
+```go
+import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+
+mcpServer := mcp.NewServer(&mcp.Implementation{...}, nil)
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider:     "okta",
+    ClientID:     "...",
+    ClientSecret: "...",
+    ServerURL:    "https://your-server.com",
+    RedirectURIs: "https://your-server.com/oauth/callback",
+}, mcpServer)
 ```
 
 **Client discovers:**
@@ -297,9 +331,9 @@ jq '.issuer' # Should be your provider (native) or your server (proxy)
 
 ```bash
 # For HMAC - generate test token
-# See examples/simple/main.go for token generation
+# See examples/mark3labs/simple/ or examples/official/simple/ for token generation
 
-# For OIDC - get token from provider
+# For OIDC - get token from provider (see examples/README.md for Okta setup)
 # Test with curl
 curl -X POST https://your-server.com/mcp \
   -H "Authorization: Bearer <token>" \
diff --git a/docs/MIGRATION.md b/docs/MIGRATION.md
deleted file mode 100644
index 16df539..0000000
--- a/docs/MIGRATION.md
+++ /dev/null
@@ -1,368 +0,0 @@
-# Migration Guide: mcp-trino → oauth-mcp-proxy
-
-This guide helps mcp-trino users migrate to the standalone oauth-mcp-proxy library.
-
----
-
-## Why Migrate?
-
-**Benefits:**
-
-- ✅ Latest OAuth improvements and security fixes
-- ✅ Reusable across any MCP server (not Trino-specific)
-- ✅ Better API (`WithOAuth()` vs manual setup)
-- ✅ Pluggable logging support
-- ✅ Active maintenance in dedicated repo
-- ✅ No Trino dependencies
-
-**Timeline:** mcp-trino will migrate to oauth-mcp-proxy in a future release.
-
----
-
-## Breaking Changes
-
-### Import Path
-
-**Before (mcp-trino):**
-
-```go
-import "github.com/tuannvm/mcp-trino/internal/oauth"
-```
-
-**After (oauth-mcp-proxy):**
-
-```go
-import oauth "github.com/tuannvm/oauth-mcp-proxy"
-```
-
-### API Changes
-
-| Old (mcp-trino) | New (oauth-mcp-proxy) | Notes |
-|---|---|---|
-| `oauth.SetupOAuth()` | `oauth.WithOAuth()` | New API is simpler |
-| `oauth.OAuthMiddleware()` | `oauth.WithOAuth()` | Returns server option |
-| `internal/oauth` package | Root `oauth` package | Now public API |
-
----
-
-## Migration Steps
-
-### Step 1: Add Dependency
-
-```bash
-go get github.com/tuannvm/oauth-mcp-proxy
-```
-
-### Step 2: Update Imports
-
-```diff
-- import "github.com/tuannvm/mcp-trino/internal/oauth"
-+ import oauth "github.com/tuannvm/oauth-mcp-proxy"
-```
-
-### Step 3: Migrate Configuration
-
-**Before (mcp-trino):**
-
-```go
-// Old internal API
-validator, err := oauth.SetupOAuth(&oauth.Config{
-    Provider: "okta",
-    Issuer:   "https://company.okta.com",
-    Audience: "api://trino-server",
-})
-
-middleware := oauth.OAuthMiddleware(validator, true)
-
-mcpServer := server.NewMCPServer("Trino", "1.0.0",
-    server.WithToolHandlerMiddleware(middleware),
-)
-```
-
-**After (oauth-mcp-proxy):**
-
-```go
-// New simple API
-mux := http.NewServeMux()
-
-_, oauthOption, err := oauth.WithOAuth(mux, &oauth.Config{
-    Provider: "okta",
-    Issuer:   "https://company.okta.com",
-    Audience: "api://trino-server",
-})
-
-mcpServer := server.NewMCPServer("Trino", "1.0.0", oauthOption)
-```
-
-**Differences:**
-
-- ✅ Simpler: 1 function call vs 3
-- ✅ `mux` passed to WithOAuth (auto-registers endpoints)
-- ✅ Returns `mcpserver.ServerOption` directly
-- ✅ No manual middleware wrapping needed
-
-### Step 4: Update HTTP Context Setup
-
-**Before (mcp-trino):**
-
-```go
-// Manual token extraction
-oauthContextFunc := func(ctx context.Context, r *http.Request) context.Context {
-    authHeader := r.Header.Get("Authorization")
-    if strings.HasPrefix(authHeader, "Bearer ") {
-        token := strings.TrimPrefix(authHeader, "Bearer ")
-        ctx = oauth.WithOAuthToken(ctx, token)
-    }
-    return ctx
-}
-
-streamableServer := mcpserver.NewStreamableHTTPServer(
-    mcpServer,
-    mcpserver.WithHTTPContextFunc(oauthContextFunc),
-)
-```
-
-**After (oauth-mcp-proxy):**
-
-```go
-// Use helper function
-streamableServer := mcpserver.NewStreamableHTTPServer(
-    mcpServer,
-    mcpserver.WithHTTPContextFunc(oauth.CreateHTTPContextFunc()),
-)
-```
-
-**Difference:** Helper function provided for convenience.
-
-### Step 5: Update User Context Access
-
-**Before & After (same):**
-
-```go
-func toolHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-    user, ok := oauth.GetUserFromContext(ctx)
-    if !ok {
-        return nil, fmt.Errorf("authentication required")
-    }
-    // Use user.Subject, user.Email, user.Username
-}
-```
-
-No changes needed! ✅
-
----
-
-## Complete Example
-
-### Before (mcp-trino internal OAuth)
-
-```go
-package main
-
-import (
-    "github.com/tuannvm/mcp-trino/internal/oauth"
-    "github.com/mark3labs/mcp-go/server"
-)
-
-func main() {
-    // Step 1: Setup OAuth
-    validator, err := oauth.SetupOAuth(&oauth.Config{
-        Provider: "okta",
-        Issuer:   "https://company.okta.com",
-        Audience: "api://trino",
-    })
-    if err != nil {
-        log.Fatal(err)
-    }
-
-    // Step 2: Create middleware
-    middleware := oauth.OAuthMiddleware(validator, true)
-
-    // Step 3: Create server with middleware
-    mcpServer := server.NewMCPServer("Trino", "1.0.0",
-        server.WithToolHandlerMiddleware(middleware),
-    )
-
-    // Step 4: Manual HTTP setup
-    mux := http.NewServeMux()
-    // ... register OAuth handlers manually ...
-
-    // Step 5: Create context func manually
-    contextFunc := func(ctx context.Context, r *http.Request) context.Context {
-        authHeader := r.Header.Get("Authorization")
-        if strings.HasPrefix(authHeader, "Bearer ") {
-            token := strings.TrimPrefix(authHeader, "Bearer ")
-            ctx = oauth.WithOAuthToken(ctx, token)
-        }
-        return ctx
-    }
-
-    streamable := server.NewStreamableHTTPServer(mcpServer,
-        server.WithHTTPContextFunc(contextFunc),
-    )
-    mux.Handle("/mcp", streamable)
-
-    http.ListenAndServeTLS(":443", "cert.pem", "key.pem", mux)
-}
-```
-
-### After (oauth-mcp-proxy)
-
-```go
-package main
-
-import (
-    oauth "github.com/tuannvm/oauth-mcp-proxy"
-    "github.com/mark3labs/mcp-go/server"
-)
-
-func main() {
-    mux := http.NewServeMux()
-
-    // Step 1: Enable OAuth (one call!)
-    _, oauthOption, err := oauth.WithOAuth(mux, &oauth.Config{
-        Provider: "okta",
-        Issuer:   "https://company.okta.com",
-        Audience: "api://trino",
-    })
-    if err != nil {
-        log.Fatal(err)
-    }
-
-    // Step 2: Create server with OAuth
-    mcpServer := server.NewMCPServer("Trino", "1.0.0", oauthOption)
-
-    // Step 3: Setup MCP endpoint (use helper)
-    streamable := server.NewStreamableHTTPServer(mcpServer,
-        server.WithHTTPContextFunc(oauth.CreateHTTPContextFunc()),
-    )
-    mux.Handle("/mcp", streamable)
-
-    http.ListenAndServeTLS(":443", "cert.pem", "key.pem", mux)
-}
-```
-
-**From ~40 lines → ~20 lines** ✅
-
----
-
-## Configuration Mapping
-
-| mcp-trino Config | oauth-mcp-proxy Config | Notes |
-|---|---|---|
-| `Provider` | `Provider` | Same |
-| `Issuer` | `Issuer` | Same |
-| `Audience` | `Audience` | Same |
-| `ClientID` | `ClientID` | Same |
-| `ClientSecret` | `ClientSecret` | Same |
-| `MCPHost + MCPPort` | `ServerURL` | Simplified to one field |
-| `RedirectURIs` | `RedirectURIs` | Same |
-| `JWTSecret` | `JWTSecret` | Same |
-| N/A | `Logger` | **New:** Pluggable logging |
-| N/A | `Mode` | **New:** Auto-detected |
-
----
-
-## New Features
-
-### Pluggable Logging
-
-```go
-oauth.WithOAuth(mux, &oauth.Config{
-    Provider: "okta",
-    Issuer:   "https://company.okta.com",
-    Audience: "api://trino",
-    Logger:   &myCustomLogger{},  // NEW!
-})
-```
-
-### Auto-Mode Detection
-
-```go
-// Native mode auto-detected (no ClientID)
-oauth.WithOAuth(mux, &oauth.Config{
-    Provider: "okta",
-    Issuer:   "...",
-    Audience: "...",
-})
-
-// Proxy mode auto-detected (has ClientID)
-oauth.WithOAuth(mux, &oauth.Config{
-    Provider: "okta",
-    ClientID: "...",  // Triggers proxy mode
-    ServerURL: "...",
-})
-```
-
-No need to set `Mode` explicitly unless you want to.
-
----
-
-## Testing Migration
-
-### 1. Keep Old Code Commented
-
-```go
-// Old mcp-trino OAuth
-// validator, err := trinoOAuth.SetupOAuth(...)
-
-// New oauth-mcp-proxy
-_, oauthOption, err := oauth.WithOAuth(mux, &oauth.Config{...})
-```
-
-### 2. Test Locally
-
-```bash
-go run main.go
-# Verify OAuth endpoints work
-curl http://localhost:8080/.well-known/oauth-authorization-server
-```
-
-### 3. Test Authentication
-
-Use same test tokens as before - token validation logic unchanged.
-
-### 4. Deploy & Monitor
-
-- Watch logs for OAuth errors
-- Verify users can authenticate
-- Check token caching works (look for cache hit logs)
-
----
-
-## Rollback Plan
-
-If issues occur:
-
-```go
-// Comment out new code
-// _, oauthOption, err := oauth.WithOAuth(...)
-
-// Uncomment old code
-validator, err := trinoOAuth.SetupOAuth(...)
-middleware := trinoOAuth.OAuthMiddleware(validator, true)
-mcpServer := server.NewMCPServer("Trino", "1.0.0",
-    server.WithToolHandlerMiddleware(middleware),
-)
-```
-
-Redeploy. OAuth logic is identical, just packaged differently.
-
----
-
-## Support
-
-Questions? Check:
-
-- [README.md](../README.md) - Quick start
-- [Provider Guides](./providers/) - Provider-specific setup
-- [SECURITY.md](./SECURITY.md) - Security best practices
-- [GitHub Issues](https://github.com/tuannvm/oauth-mcp-proxy/issues)
-
----
-
-## Timeline
-
-- **Now:** oauth-mcp-proxy v0.1.0 available
-- **Future:** mcp-trino will update to use oauth-mcp-proxy library
-- **Support:** Both approaches work, new approach recommended
diff --git a/docs/TROUBLESHOOTING.md b/docs/TROUBLESHOOTING.md
index eb65dd1..a5c2fe0 100644
--- a/docs/TROUBLESHOOTING.md
+++ b/docs/TROUBLESHOOTING.md
@@ -474,7 +474,7 @@ echo "<token>" | cut -d. -f2 | base64 -d 2>/dev/null | jq
 
 ```bash
 # Generate test token (HMAC)
-go run examples/simple/main.go
+go run examples/mark3labs/simple/ or examples/official/simple/
 # Copy token from output, test with curl
 
 # For OIDC providers, get token from provider:
diff --git a/docs/generic-implementation.md b/docs/generic-implementation.md
index 63cf10c..86eaee7 100644
--- a/docs/generic-implementation.md
+++ b/docs/generic-implementation.md
@@ -322,30 +322,38 @@ go test ./mark3labs/...
 
 ### Checkpoint 2.4: Update examples to use mark3labs package ✅
 
-**Task**: Update example code to import from mark3labs package.
+**Task**: Update example code to import from SDK-specific packages and create examples for both SDKs.
 
-**Files to Update**:
-- `examples/simple/main.go`
-- `examples/advanced/main.go`
+**Files Updated**:
+- `examples/mark3labs/simple/main.go`
+- `examples/mark3labs/advanced/main.go`
+- `examples/official/simple/main.go` (new)
+- `examples/official/advanced/main.go` (new)
 
 **Changes**:
 ```diff
+# mark3labs examples:
 - import "github.com/tuannvm/oauth-mcp-proxy"
 + import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
 
 - oauth.WithOAuth(mux, cfg)
 + mark3labs.WithOAuth(mux, cfg)
+
+# official SDK examples (new):
++ import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
++ mcpoauth.WithOAuth(mux, cfg, mcpServer)
 ```
 
 **Verification**:
 ```bash
-cd examples/simple && go build .
-cd examples/advanced && go build .
+for example in $(find examples -name main.go); do
+  echo "Building $example..." && go build "$example"
+done
 ```
 
-**Expected Outcome**: Examples build and run successfully.
+**Expected Outcome**: All 4 examples build and run successfully.
 
-**Actual Outcome**: ✅ Completed. Updated both examples to use mark3labs.WithOAuth(). Both examples build successfully.
+**Actual Outcome**: ✅ Completed. Created 4 examples (2 per SDK). All examples build successfully with Okta configuration.
 
 ---
 
@@ -718,9 +726,11 @@ make clean
 make test
 make lint
 make test-coverage
-cd examples/simple && go run main.go
-cd examples/advanced && go run main.go
-cd examples/official && go run main.go
+
+# Test all examples build
+for example in $(find examples -name main.go); do
+  echo "Building $example..." && go build "$example"
+done
 ```
 
 **Expected Outcome**: Everything works.
diff --git a/docs/providers/AZURE.md b/docs/providers/AZURE.md
index c52950c..01a5bbc 100644
--- a/docs/providers/AZURE.md
+++ b/docs/providers/AZURE.md
@@ -1,3 +1,6 @@
+> **📢 v2.0.0:** This guide shows examples for both `mark3labs/mcp-go` and official `modelcontextprotocol/go-sdk`.
+> See [examples/README.md](../../examples/README.md) for complete setup guide.
+
 # Azure AD Provider Guide
 
 ## Overview
diff --git a/docs/providers/GOOGLE.md b/docs/providers/GOOGLE.md
index 77dfe74..76cd4ba 100644
--- a/docs/providers/GOOGLE.md
+++ b/docs/providers/GOOGLE.md
@@ -1,3 +1,6 @@
+> **📢 v2.0.0:** This guide shows examples for both `mark3labs/mcp-go` and official `modelcontextprotocol/go-sdk`.
+> See [examples/README.md](../../examples/README.md) for complete setup guide.
+
 # Google Provider Guide
 
 ## Overview
diff --git a/docs/providers/HMAC.md b/docs/providers/HMAC.md
index bc1aa05..759f02e 100644
--- a/docs/providers/HMAC.md
+++ b/docs/providers/HMAC.md
@@ -1,5 +1,8 @@
 # HMAC Provider Guide
 
+> **📢 v2.0.0:** This guide shows examples for both `mark3labs/mcp-go` and official `modelcontextprotocol/go-sdk`.
+> See [MIGRATION-V2.md](../../MIGRATION-V2.md) for upgrade details.
+
 ## Overview
 
 HMAC provider uses shared secret JWT validation with HS256 algorithm. Best for testing, development, and service-to-service authentication.
@@ -23,12 +26,46 @@ HMAC provider uses shared secret JWT validation with HS256 algorithm. Best for t
 
 ## Configuration
 
+### Using mark3labs/mcp-go
+
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+)
+
+mux := http.NewServeMux()
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider:  "hmac",
     Audience:  "api://my-mcp-server",      // Your server's identifier
     JWTSecret: []byte("your-secret-key"),  // 32+ bytes recommended
 })
+
+mcpServer := server.NewMCPServer("My Server", "1.0.0", oauthOption)
+```
+
+### Using Official SDK
+
+```go
+import (
+    oauth "github.com/tuannvm/oauth-mcp-proxy"
+    mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+)
+
+mux := http.NewServeMux()
+mcpServer := mcp.NewServer(&mcp.Implementation{
+    Name:    "My Server",
+    Version: "1.0.0",
+}, nil)
+
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider:  "hmac",
+    Audience:  "api://my-mcp-server",      // Your server's identifier
+    JWTSecret: []byte("your-secret-key"),  // 32+ bytes recommended
+}, mcpServer)
+
+http.ListenAndServe(":8080", handler)
 ```
 
 ### Required Fields
@@ -92,11 +129,20 @@ if len(secret) < 32 {
     log.Fatal("JWT_SECRET must be at least 32 bytes")
 }
 
-oauth.WithOAuth(mux, &oauth.Config{
+// mark3labs SDK:
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider:  "hmac",
     Audience:  "api://my-server",
     JWTSecret: secret,
 })
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+
+// OR official SDK:
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider:  "hmac",
+    Audience:  "api://my-server",
+    JWTSecret: secret,
+}, mcpServer)
 ```
 
 ### Secret Strength
@@ -152,9 +198,17 @@ curl -X POST http://localhost:8080/mcp \
 
 ---
 
-## Example
+## Complete Examples
+
+**mark3labs SDK:**
+- [examples/mark3labs/simple/](../../examples/mark3labs/simple/) - Minimal HMAC setup
+- [examples/mark3labs/advanced/](../../examples/mark3labs/advanced/) - Full featured
+
+**Official SDK:**
+- [examples/official/simple/](../../examples/official/simple/) - Minimal HMAC setup
+- [examples/official/advanced/](../../examples/official/advanced/) - Full featured
 
-See [examples/simple/main.go](../../examples/simple/main.go) for a complete working example with HMAC provider.
+See [examples/README.md](../../examples/README.md) for setup instructions.
 
 ---
 
diff --git a/docs/providers/OKTA.md b/docs/providers/OKTA.md
index b47f417..f1a79de 100644
--- a/docs/providers/OKTA.md
+++ b/docs/providers/OKTA.md
@@ -1,5 +1,8 @@
 # Okta Provider Guide
 
+> **📢 v2.0.0:** This guide shows examples for both `mark3labs/mcp-go` and official `modelcontextprotocol/go-sdk`.
+> See [examples/README.md](../../examples/README.md) for complete Okta setup guide.
+
 ## Overview
 
 Okta provider uses OIDC/JWKS for JWT validation. Ideal for enterprise SSO, user management, and production deployments.
@@ -74,12 +77,28 @@ By default, Okta uses the org authorization server. For custom authorization ser
 
 **When:** Client handles OAuth (Claude Desktop, browser clients)
 
+**mark3labs SDK:**
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",
-    Issuer:   "https://yourcompany.okta.com",  // Your Okta domain
-    Audience: "api://your-mcp-server",         // Custom audience or Client ID
+    Issuer:   "https://yourcompany.okta.com",
+    Audience: "api://your-mcp-server",
 })
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+```
+
+**Official SDK:**
+```go
+import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider: "okta",
+    Issuer:   "https://yourcompany.okta.com",
+    Audience: "api://your-mcp-server",
+}, mcpServer)
+http.ListenAndServe(":8080", handler)
 ```
 
 Client configures OAuth directly with Okta. Server only validates tokens.
@@ -90,8 +109,11 @@ Client configures OAuth directly with Okta. Server only validates tokens.
 
 **When:** Client cannot do OAuth (simple CLI tools)
 
+**mark3labs SDK:**
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+import "github.com/tuannvm/oauth-mcp-proxy/mark3labs"
+
+_, oauthOption, _ := mark3labs.WithOAuth(mux, &oauth.Config{
     Provider:     "okta",
     Issuer:       "https://yourcompany.okta.com",
     Audience:     "api://your-mcp-server",
@@ -100,6 +122,23 @@ oauth.WithOAuth(mux, &oauth.Config{
     ServerURL:    "https://your-mcp-server.com",     // Your public URL
     RedirectURIs: "https://your-mcp-server.com/oauth/callback",
 })
+mcpServer := server.NewMCPServer("Server", "1.0.0", oauthOption)
+```
+
+**Official SDK:**
+```go
+import mcpoauth "github.com/tuannvm/oauth-mcp-proxy/mcp"
+
+_, handler, _ := mcpoauth.WithOAuth(mux, &oauth.Config{
+    Provider:     "okta",
+    Issuer:       "https://yourcompany.okta.com",
+    Audience:     "api://your-mcp-server",
+    ClientID:     "0oa...",                           // From Okta app
+    ClientSecret: "secret-from-okta",                 // From Okta app
+    ServerURL:    "https://your-mcp-server.com",     // Your public URL
+    RedirectURIs: "https://your-mcp-server.com/oauth/callback",
+}, mcpServer)
+http.ListenAndServe(":8080", handler)
 ```
 
 Server proxies OAuth flow. Client gets tokens from your server.
@@ -115,7 +154,8 @@ Okta tokens include `aud` (audience) claim. Configure it:
 Simplest approach:
 
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+// mark3labs or official SDK - same config
+mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",
     Issuer:   "https://yourcompany.okta.com",
     Audience: "0oa...",  // Same as ClientID
@@ -139,7 +179,8 @@ For custom audience (e.g., `api://my-server`):
 Then configure:
 
 ```go
-oauth.WithOAuth(mux, &oauth.Config{
+// mark3labs or official SDK - same config
+mark3labs.WithOAuth(mux, &oauth.Config{
     Provider: "okta",
     Issuer:   "https://yourcompany.okta.com",
     Audience: "api://my-server",  // Your custom audience
diff --git a/examples/README.md b/examples/README.md
index 608446d..d530c60 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -325,7 +325,7 @@ curl -H "Accept: application/json, text/event-stream" ...
 FROM golang:1.24 AS builder
 WORKDIR /app
 COPY . .
-RUN go build -o server ./examples/advanced
+RUN go build -o server ./examples/mark3labs/advanced
 
 FROM alpine:latest
 RUN apk --no-cache add ca-certificates

From 17e9c2c1edc5dd15f8d53be4a66a0011c0f394dc Mon Sep 17 00:00:00 2001
From: Tommy Nguyen <tuannvm@hotmail.com>
Date: Wed, 22 Oct 2025 18:36:48 -0700
Subject: [PATCH 9/9] refactor(advanced/main.go): Enhance TLS detection and
 Logging for server startup

Signed-off-by: Tommy Nguyen <tuannvm@hotmail.com>
---
 examples/official/advanced/main.go | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/examples/official/advanced/main.go b/examples/official/advanced/main.go
index 1dee456..4d89ee4 100644
--- a/examples/official/advanced/main.go
+++ b/examples/official/advanced/main.go
@@ -114,7 +114,9 @@ func main() {
 	}
 
 	// Feature 4: LogStartup - Displays OAuth endpoint information
-	useTLS := getEnv("HTTPS_CERT_FILE", "") != ""
+	certFile := getEnv("HTTPS_CERT_FILE", "")
+	keyFile := getEnv("HTTPS_KEY_FILE", "")
+	useTLS := certFile != "" && keyFile != ""
 	oauthServer.LogStartup(useTLS)
 
 	// Additional OAuth endpoints available via helper methods
@@ -130,15 +132,32 @@ func main() {
 	port := getEnv("MCP_PORT", "8080")
 	addr := ":" + port
 
-	log.Printf("\nStarting MCP server on %s", addr)
+	log.Printf("\nStarting MCP server (TLS=%t) on %s", useTLS, addr)
 	log.Println("\nMake sure to set your Okta environment variables:")
 	log.Println("  export OKTA_DOMAIN=dev-12345.okta.com")
 	log.Println("  export OKTA_AUDIENCE=api://my-mcp-server")
+	if useTLS {
+		log.Println("\nFor TLS, also set:")
+		log.Println("  export HTTPS_CERT_FILE=/path/to/cert.pem")
+		log.Println("  export HTTPS_KEY_FILE=/path/to/key.pem")
+	}
 	log.Println("\nTo test, obtain an access token from Okta and use:")
-	log.Printf("  curl -H 'Authorization: Bearer <okta-token>' http://localhost:%s", port)
+	protocol := "http"
+	if useTLS {
+		protocol = "https"
+	}
+	log.Printf("  curl -H 'Authorization: Bearer <okta-token>' %s://localhost:%s", protocol, port)
 
-	if err := http.ListenAndServe(addr, handler); err != nil {
-		log.Fatalf("Server failed: %v", err)
+	if useTLS {
+		log.Printf("\nStarting HTTPS server on %s", addr)
+		if err := http.ListenAndServeTLS(addr, certFile, keyFile, handler); err != nil {
+			log.Fatalf("HTTPS server failed: %v", err)
+		}
+	} else {
+		log.Printf("\nStarting HTTP server on %s", addr)
+		if err := http.ListenAndServe(addr, handler); err != nil {
+			log.Fatalf("HTTP server failed: %v", err)
+		}
 	}
 }