# 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.

**Actual Outcome**: ✅ Completed. File created with 68 lines. 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.

**Actual Outcome**: ✅ Completed. File created with 46 lines including WithUser() function. 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.

**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 ✅

**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.

**Actual Outcome**: ✅ Completed. Added ValidateTokenCached() and WithUser() to core. All tests pass.

***

## 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.

**Actual Outcome**: ✅ Completed. Created mark3labs/ directory with oauth.go and middleware.go files.

***

### 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.

**Actual Outcome**: ✅ Completed. Created mark3labs/oauth.go with 45 lines. 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.

**Actual Outcome**: ✅ Completed. Created mark3labs/middleware.go with 38 lines using ValidateTokenCached(). Package builds successfully.

***

### Checkpoint 2.4: Update examples to use mark3labs package ✅

**Task**: Update example code to import from SDK-specific packages and create examples for both SDKs.

**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
for example in $(find examples -name main.go); do
  echo "Building $example..." && go build "$example"
done
```

**Expected Outcome**: All 4 examples build and run successfully.

**Actual Outcome**: ✅ Completed. Created 4 examples (2 per SDK). All examples build successfully with Okta configuration.

***

## 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.

**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 ✅

**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.

**Actual Outcome**: ✅ Completed. Created mcp/ directory with oauth.go file.

***

### 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.

**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 ⏭️

**Task**: Create example demonstrating official SDK integration.

**Status**: Skipped for initial release. Can be added later.

**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.

**Status**: Core tests passing. Adapter-specific tests pending.

### 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.

**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`

**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.

**Status**: Pending - to be added in follow-up PR.

**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.

**Actual Outcome**: ✅ Completed. All core tests pass. Build successful across all packages (core, mark3labs, mcp, provider, examples).

***

## 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.

**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 v1.0.0 release.

### Checkpoint 6.1: Update version and changelog ⬜

**Task**: Prepare release artifacts.

**Files to Update/Create**:

* `CHANGELOG.md` (document v1.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

# Test all examples build
for example in $(find examples -name main.go); do
  echo "Building $example..." && go build "$example"
done
```

**Expected Outcome**: Everything works.

***

### Checkpoint 6.3: Create release PR ⬜

**Task**: Create pull request for v1.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         | ✅ 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**                              | **🟢 Implementation Complete** | **82%**    |

***

## 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**: 2025-10-22 **Implementation Completion Date**: 2025-10-22

**Current Status**: ✅ Core implementation complete (Phases 0-3 + core testing). Documentation updates pending (Phase 5).


---

# Agent Instructions: Querying This Documentation

If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.

Perform an HTTP GET request on the current page URL with the `ask` query parameter:

```
GET https://docs.tuannvm.com/oauth-mcp-proxy/docs/generic-implementation.md?ask=<question>
```

The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.

Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.