# Production-Grade MCP & A2A Implementation in Go

A comprehensive, tutorial-focused implementation demonstrating production-quality integration of **Model Context Protocol (MCP)** and **Agent-to-Agent (A2A)** protocols with complete observability, security, and cost control features.

Built with **Go servers** and an interactive **Streamlit UI** for hands-on exploration of all features.

## 🎯 Overview

This repository showcases two production-ready use cases:

1. **Multi-Tenant RAG Pipeline (MCP)**: Secure document search with hybrid search (BM25 + Vector), JWT authentication, rate limiting, and tenant isolation
2. **Cost-Controlled Research Assistant (A2A)**: Budget-aware research agent with task management, real-time streaming, and multi-tier cost tracking

## ✨ Key Features

### 🔐 Security & Multi-Tenancy
- **JWT Authentication**: RS256 tokens with tenant and user claims
- **Multi-tenant Isolation**: Row-Level Security (RLS) in PostgreSQL
- **Rate Limiting**: Redis-backed per-tenant request throttling
- **Scope-based Authorization**: Fine-grained access control

### 🔍 Search & Retrieval
- **Hybrid Search**: BM25 (keyword) + Vector (semantic) with Reciprocal Rank Fusion
- **pgvector**: Efficient similarity search with HNSW indexing
- **Document Management**: Full CRUD operations with tenant isolation
- **Pagination**: Efficient cursor-based pagination for large result sets

### 💰 Cost Control & Budgeting
- **Token Tracking**: Accurate per-request token counting for GPT-4, GPT-3.5, Claude
- **Budget Enforcement**: Pre-flight checks prevent exceeding limits
- **Multi-tier Plans**: Basic ($10), Pro ($50), Enterprise ($200) monthly budgets
- **Cost Attribution**: Per-user and per-task cost tracking

### 📊 Observability & Monitoring
- **Distributed Tracing**: OpenTelemetry + Jaeger for end-to-end request visibility
- **Metrics**: Prometheus-compatible metrics for all operations
- **Health Checks**: Readiness and liveness probes for all services
- **Structured Logging**: JSON logs with trace context propagation

### 🚀 Real-time Streaming
- **Server-Sent Events (SSE)**: Real-time task updates
- **Task Lifecycle**: Pending → Running → Completed/Failed/Cancelled
- **Event Broadcasting**: Pub/sub pattern for task state changes

### 🧪 Test Coverage
- **MCP Server**: 95% average coverage, 125+ unit tests
- **A2A Server**: 92.6% average coverage, 75+ unit tests
- **Integration Tests**: Redis, database, and middleware testing
- **Total**: 200+ passing tests

## 🏗️ Architecture

```
┌─────────────────────────────────────────────────────────┐
│                    Streamlit UI                         │
│  - Authentication & JWT token management                │
│  - MCP RAG testing (hybrid search, documents)           │
│  - A2A task creation & monitoring                       │
│  - Cost tracking & budget visualization                 │
│  - Metrics dashboards (Prometheus)                      │
│  - Distributed tracing (Jaeger)                         │
└─────────────────┬───────────────────────────────────────┘
                  │
      ┌───────────┴───────────┐
      │                       │
┌─────▼────────┐   ┌──────────▼──────────┐
│ MCP Server   │   │   A2A Server        │
│ Port 8080    │   │   Port 8081         │
│              │   │                     │
│ - JSON-RPC   │   │ - REST API          │
│ - JWT Auth   │   │ - Agent Cards       │
│ - Hybrid     │   │ - Task Management   │
│   Search     │   │ - SSE Streaming     │
│ - Rate Limit │   │ - Cost Tracking     │
└─────┬────────┘   └─────────┬───────────┘
      │                      │
  ┌───┴───┬────────┬─────────┴─────┬──────────┐
  │       │        │               │          │
┌─▼──┐ ┌─▼───┐ ┌──▼──────┐  ┌─────▼────┐ ┌──▼──────┐
│PG  │ │Redis│ │ Jaeger  │  │Prometheus│ │ Grafana │
│+pgv│ │     │ │ Tracing │  │ Metrics  │ │Dashboard│
└────┘ └─────┘ └─────────┘  └──────────┘ └─────────┘
```

## 📁 Project Structure

```
.
├── mcp-server/                    # Go MCP server (95% coverage)
│   ├── cmd/server/main.go         # Entry point
│   ├── internal/
│   │   ├── auth/                  # JWT validation (93.1% coverage)
│   │   ├── database/              # PostgreSQL + pgvector
│   │   ├── protocol/              # JSON-RPC 2.0 (100% coverage)
│   │   ├── tools/                 # MCP tools (97.4% coverage)
│   │   ├── middleware/            # Logging, rate limiting (92.9%)
│   │   └── server/                # HTTP server (94.8% coverage)
│   ├── Dockerfile                 # Multi-stage build
│   └── go.mod
│
├── a2a-server/                    # Go A2A server (92.6% coverage)
│   ├── cmd/server/main.go         # Entry point with 3 capabilities
│   ├── internal/
│   │   ├── protocol/              # A2A types (100% coverage)
│   │   ├── agentcard/             # Agent Card store (100% coverage)
│   │   ├── tasks/                 # Task lifecycle (98.3% coverage)
│   │   ├── cost/                  # Cost tracking (91.5% coverage)
│   │   └── server/                # HTTP + SSE server (81.8%)
│   ├── Dockerfile
│   └── go.mod
│
├── streamlit-ui/                  # Interactive testing UI
│   ├── app.py                     # Main dashboard
│   ├── pages/
│   │   ├── 1_🔐_Authentication.py # JWT generation & testing
│   │   ├── 2_📄_MCP_RAG.py        # Hybrid search, documents
│   │   ├── 3_🤖_A2A_Tasks.py      # Task creation & monitoring
│   │   ├── 4_💰_Cost_Tracking.py  # Budget analytics
│   │   ├── 5_📊_Metrics.py        # Prometheus dashboards
│   │   └── 6_🔍_Tracing.py        # Jaeger integration guide
│   ├── utils/
│   │   ├── mcp_client.py          # JSON-RPC 2.0 client
│   │   ├── a2a_client.py          # REST + SSE client
│   │   └── auth.py                # JWT token generation
│   ├── Dockerfile
│   └── requirements.txt
│
├── docker compose.yml             # Complete local stack
├── README.md                      # This file
├── DESIGN.md                      # Detailed architecture 
└── go.work                        # Go workspace
```

## 🚀 Quick Start

### Prerequisites

- **Docker** & **Docker Compose** (required)
- **Go 1.23+** (for local development)
- **Python 3.11+** (for Streamlit UI development)

### Start Everything with Docker Compose

```bash
# Clone the repository
git clone <repo-url>
cd mcp-a2a-go

# Start all services (PostgreSQL, Redis, Jaeger, Prometheus, MCP, A2A, Streamlit)
docker compose up --build

# Wait for services to be healthy (check logs)
# When you see "Streamlit UI available at http://localhost:8501"
```

**Access the services:**

- **Streamlit UI**: http://localhost:8501 (Interactive testing dashboard)
- **MCP Server**: http://localhost:8080 (JSON-RPC endpoint: /mcp)
- **A2A Server**: http://localhost:8081 (REST API)
- **Jaeger**: http://localhost:16686 (Distributed tracing)
- **Prometheus**: http://localhost:9090 (Metrics)
- **Grafana**: http://localhost:3000 (Dashboards - admin/admin)

### Using the Streamlit UI

The Streamlit UI provides a complete interactive environment for testing all features:

#### 1. **🔐 Authentication Page**
- Generate JWT tokens for 3 demo tenants (acme-corp, globex, initech)
- Configure scopes (read, write, admin)
- Test token validation
- View decoded token claims

#### 2. **📄 MCP RAG Page**
- Initialize MCP session
- List available tools
- Test hybrid search with adjustable BM25/Vector weights
- List documents with pagination
- Retrieve specific documents
- Verify multi-tenant isolation

#### 3. **🤖 A2A Tasks Page**
- View agent card with capabilities
- Create tasks for 3 capabilities:
  - `search_papers`: Search academic papers
  - `analyze_code`: Analyze code repositories
  - `summarize_research`: Summarize research topics
- Select budget tier (Basic $10, Pro $50, Enterprise $200)
- Monitor real-time task progress via SSE
- View task lifecycle and status

#### 4. **💰 Cost Tracking Page**
- View budget overview by tier
- Monitor cost by model (GPT-4, GPT-3.5, Claude)
- Analyze usage timeline
- Track token consumption
- Compare model pricing

#### 5. **📊 Metrics Page**
- View system health metrics
- Request rate and error analysis
- Response time distribution
- Rate limiting status
- Database and cache metrics

#### 6. **🔍 Tracing Page**
- Learn about distributed tracing
- View sample traces
- Understand span hierarchy
- Performance debugging guide
- Link to Jaeger UI

### Manual Testing (Without UI)

#### Test MCP Server

```bash
# Generate a JWT token (you'll need to create a test token)
# See streamlit-ui/utils/auth.py for token generation

# 1. Initialize MCP session
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer YOUR_JWT_TOKEN" \
  -d '{
    "jsonrpc": "2.0",
    "id": "1",
    "method": "initialize",
    "params": {
      "protocolVersion": "2024-11-05",
      "clientInfo": {"name": "test-client", "version": "1.0.0"}
    }
  }'

# 2. List tools
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer YOUR_JWT_TOKEN" \
  -d '{
    "jsonrpc": "2.0",
    "id": "2",
    "method": "tools/list"
  }'

# 3. Hybrid search
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer YOUR_JWT_TOKEN" \
  -d '{
    "jsonrpc": "2.0",
    "id": "3",
    "method": "tools/call",
    "params": {
      "name": "hybrid_search",
      "arguments": {
        "query": "machine learning",
        "limit": 10,
        "bm25_weight": 0.5,
        "vector_weight": 0.5
      }
    }
  }'
```

#### Test A2A Server

```bash
# 1. Get agent card
curl http://localhost:8081/agent

# 2. Create a task
curl -X POST http://localhost:8081/tasks \
  -H "Content-Type: application/json" \
  -d '{
    "user_id": "demo-user-pro",
    "agent_id": "research-assistant",
    "capability": "search_papers",
    "input": {
      "query": "transformer architecture",
      "limit": 5
    }
  }'

# 3. Get task status
curl http://localhost:8081/tasks/{task_id}

# 4. Stream task events (SSE)
curl -N http://localhost:8081/tasks/{task_id}/events
```

## 🧪 Running Tests

### All Tests

```bash
# From project root
./scripts/run-tests.sh
```

### MCP Server Tests

```bash
cd mcp-server

# All tests with coverage
go test ./... -coverprofile=coverage.out

# View coverage report
go tool cover -html=coverage.out

# Specific package
go test ./internal/auth/...
go test ./internal/tools/...
go test ./internal/middleware/...
```

### A2A Server Tests

```bash
cd a2a-server

# All tests with coverage
go test ./... -coverprofile=coverage.out

# View coverage report
go tool cover -html=coverage.out

# Specific package
go test ./internal/protocol/...
go test ./internal/tasks/...
go test ./internal/cost/...
```

### Test Coverage Summary

| Package | Coverage | Tests |
|---------|----------|-------|
| mcp-server/internal/protocol | 100% | 25+ |
| mcp-server/internal/auth | 93.1% | 20+ |
| mcp-server/internal/tools | 97.4% | 30+ |
| mcp-server/internal/middleware | 92.9% | 25+ |
| mcp-server/internal/server | 94.8% | 25+ |
| **MCP Server Average** | **95%** | **125+** |
| a2a-server/internal/protocol | 100% | 12 |
| a2a-server/internal/agentcard | 100% | 10 |
| a2a-server/internal/tasks | 98.3% | 13 |
| a2a-server/internal/cost | 91.5% | 13 |
| a2a-server/internal/server | 81.8% | 27 |
| **A2A Server Average** | **92.6%** | **75+** |
| **Total** | **94%** | **200+** |

## 📚 Use Case Walkthroughs

### Use Case 1: Multi-Tenant RAG Pipeline

**Scenario**: You have multiple teams (acme-corp, globex, initech) that need isolated document search.

**Steps**:
1. Go to **Authentication** page in Streamlit
2. Generate token for `acme-corp` tenant with `read` and `write` scopes
3. Go to **MCP RAG** page
4. Initialize session with the token
5. Test hybrid search with query: "security policy"
6. Adjust BM25/Vector weights to see different results
7. List documents to see tenant-isolated data
8. Generate token for different tenant and verify isolation

**Key Learnings**:
- JWT claims enforce tenant isolation at database level
- Hybrid search combines keyword and semantic matching
- Rate limiting prevents abuse
- All requests are traced in Jaeger

### Use Case 3: Cost-Controlled Research Assistant

**Scenario**: Research team needs AI assistance with budget constraints.

**Steps**:
1. Go to **A2A Tasks** page in Streamlit
2. View agent card showing 3 capabilities
3. Select budget tier: Basic ($10/month), Pro ($50/month), or Enterprise ($200/month)
4. Create task: "Search papers on transformer architecture"
5. Watch real-time SSE events as task executes
6. Go to **Cost Tracking** page to see budget usage
7. Try creating tasks until budget is exceeded
8. See budget enforcement in action

**Key Learnings**:
- Pre-flight budget checks prevent overspending
- Different models have different costs (GPT-4 vs GPT-3.5)
- SSE provides real-time task updates
- Cost attribution tracks spending by user and task

## 🔧 Local Development (Without Docker)

### Prerequisites
- PostgreSQL 16 with pgvector extension
- Redis 7+
- Go 1.23+

### Setup PostgreSQL

```bash
# Install pgvector extension
psql -U postgres -c "CREATE EXTENSION IF NOT EXISTS vector;"

# Create database
psql -U postgres -c "CREATE DATABASE mcp_dev;"

# Set environment variables
export DB_HOST=localhost
export DB_PORT=5432
export DB_USER=postgres
export DB_PASSWORD=postgres
export DB_NAME=mcp_dev
export DB_SSLMODE=disable
```

### Setup Redis

```bash
# Start Redis
redis-server

# Set environment variable
export REDIS_ADDR=localhost:6379
```

### Run MCP Server

```bash
cd mcp-server

# Install dependencies
go mod download

# Run server
go run cmd/server/main.go

# Server starts on http://localhost:8080
```

### Run A2A Server

```bash
cd a2a-server

# Install dependencies
go mod download

# Run server
go run cmd/server/main.go

# Server starts on http://localhost:8081
```

### Run Streamlit UI

```bash
cd streamlit-ui

# Install dependencies
pip install -r requirements.txt

# Set environment variables
export MCP_SERVER_URL=http://localhost:8080
export A2A_SERVER_URL=http://localhost:8081
export JAEGER_URL=http://localhost:16686
export PROMETHEUS_URL=http://localhost:9090

# Run Streamlit
streamlit run app.py

# UI starts on http://localhost:8501
```

## 🔐 Security Features

### Authentication & Authorization

- **JWT Tokens**: RS256 algorithm with public/private key pairs
- **Token Claims**: `tenant_id`, `user_id`, `scopes`, `exp`, `iat`, `nbf`
- **Scope Validation**: Endpoints require specific scopes (read, write, admin)
- **Token Expiry**: Configurable expiration with automatic validation

### Multi-Tenancy

- **Row-Level Security (RLS)**: PostgreSQL policies enforce tenant isolation
- **Context Propagation**: Tenant ID from JWT flows through all operations
- **Isolated Rate Limits**: Each tenant has separate rate limit counters
- **Data Isolation**: Queries automatically filtered by tenant_id

### Rate Limiting

- **Algorithm**: Token bucket with Redis backend
- **Configuration**: Per-tenant and per-endpoint limits
- **Response**: HTTP 429 with `Retry-After` header
- **Monitoring**: Prometheus metrics for rate limit hits

## 📊 Observability

This implementation provides **production-grade observability** with dual instrumentation:
- **OpenTelemetry**: Service-to-service tracing, infrastructure metrics
- **Langfuse**: LLM-specific observability (prompts, tokens, costs)

### Distributed Tracing (OpenTelemetry + Jaeger)

**Full stack instrumentation:**
- ✅ **Go Servers**: HTTP middleware, tool execution, database queries
- ✅ **Python Workflows**: LangGraph nodes, MCP calls, LLM invocations
- ✅ **End-to-End Traces**: Streamlit UI → Python → Go → Database → LLM

**Features:**
- **Trace Propagation**: W3C Trace Context via HTTP headers
- **Span Attributes**: tenant_id, user_id (traces only), tool names, query params
- **Sampling**: Configurable (100% by default for dev)
- **Auto-instrumentation**: HTTP clients (requests, httpx) in Python

**Example trace:**
```
[Streamlit] user.query (1200ms)
  ├─ [Python] rag_workflow.execute (1150ms)
  │  ├─ [Python] mcp.hybrid_search (300ms)
  │  │  └─ [Go MCP] mcp.request → mcp.tool.call → mcp.db.hybrid_search (280ms)
  │  ├─ [Python] llm.generate (800ms) + [Langfuse] tracks tokens/cost
  │  └─ [Python] format.response (50ms)
```

**View Traces**: http://localhost:16686

### Metrics (OpenTelemetry + Prometheus)

**MCP Server Metrics** (`/metrics`):
- `mcp.request.count`, `mcp.request.duration` - HTTP request metrics
- `mcp.tool.execution.duration` - Tool execution time by tool name
- `mcp.db.query.duration` - Database query performance
- `mcp.search.results` - Search result count distribution

**A2A Server Metrics** (`/metrics`):
- `a2a.task.count`, `a2a.task.duration` - Task lifecycle metrics
- `a2a.cost.total`, `a2a.tokens.total` - Cost tracking by model
- `a2a.budget.remaining` - Budget utilization by tier
- `a2a.sse.connections` - Active SSE connections

**Configuration:**
```bash
# Enable/disable observability
OTEL_ENABLE_TRACING=true
OTEL_ENABLE_METRICS=true

# OTLP endpoint (Jaeger)
OTEL_EXPORTER_OTLP_ENDPOINT=jaeger:4318

# Sampling rate (0.0 to 1.0)
OTEL_TRACES_SAMPLER_ARG=1.0  # 100% sampling

# Environment
ENVIRONMENT=development  # or production
```

**View Metrics**: http://localhost:9090

### LLM Observability (Langfuse)

**Complementary to OpenTelemetry:**
- **Prompt Tracking**: Full prompt/response history
- **Token Usage**: Per-call token counts and costs
- **Model Performance**: Latency by model type
- **Quality Metrics**: User feedback and ratings

**Integration**: Python workflows use both `@observe` (Langfuse) and OTel spans

### Testing Observability

**Quick automated test:**
```bash
./scripts/test-observability.sh
```

**Manual verification:**
```bash
# 1. Check metrics endpoints
curl http://localhost:8080/metrics  # MCP server
curl http://localhost:8081/metrics  # A2A server

# 2. Query Prometheus
curl -G http://localhost:9090/api/v1/query \
  --data-urlencode 'query=mcp_request_count'

# 3. Make test requests to generate traces
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":1,"method":"tools/list"}'

# 4. View traces in Jaeger UI
open http://localhost:16686
# Select service: mcp-server, a2a-server, or rag-workflow
# Click "Find Traces" to see distributed traces
```

**Common Prometheus queries:**
```promql
# Request rate by service
sum by (service) (rate(mcp_request_count[5m]))

# Error rate
rate(mcp_request_count{status="error"}[5m])

# Tool execution time (95th percentile)
histogram_quantile(0.95, rate(mcp_tool_execution_duration_bucket[5m]))

# Total cost by model
sum by (model) (a2a_cost_total)
```

**Troubleshooting:**
- **No metrics?** Check Prometheus targets: http://localhost:9090/targets (all should be "UP")
- **No traces?** Verify `OTEL_ENABLE_TRACING=true` in docker-compose.yml
- **Context not propagating?** Check HTTP headers include `traceparent`

## 💡 Configuration

### Environment Variables

#### MCP Server

```bash
# Database
DB_HOST=postgres
DB_PORT=5432
DB_USER=postgres
DB_PASSWORD=postgres
DB_NAME=mcp_dev
DB_SSLMODE=disable

# Redis
REDIS_ADDR=redis:6379

# Server
MCP_PORT=8080
MCP_LOG_LEVEL=info

# JWT
JWT_PUBLIC_KEY_PATH=/path/to/public.pem
JWT_ISSUER=mcp-server-demo
JWT_AUDIENCE=mcp-server

# Rate Limiting
RATE_LIMIT_REQUESTS=100
RATE_LIMIT_WINDOW=60s

# Observability
OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger:14268/api/traces
```

#### A2A Server

```bash
# Redis
REDIS_ADDR=redis:6379

# Server
A2A_PORT=8081
A2A_LOG_LEVEL=info

# Cost Limits (monthly budgets in USD)
BUDGET_BASIC=10.0
BUDGET_PRO=50.0
BUDGET_ENTERPRISE=200.0

# Observability
OTEL_EXPORTER_JAEGER_ENDPOINT=http://jaeger:14268/api/traces
```

## 🐛 Troubleshooting

### MCP Server won't start

```bash
# Check database connection
psql -h localhost -U postgres -d mcp_dev -c "SELECT 1;"

# Check pgvector extension
psql -h localhost -U postgres -d mcp_dev -c "SELECT * FROM pg_extension WHERE extname='vector';"

# Check Redis connection
redis-cli ping
```

### JWT tokens rejected

- Verify token hasn't expired (check `exp` claim)
- Ensure `issuer` matches server configuration
- Ensure `audience` matches server configuration
- Check token signature (must be RS256)

### Rate limiting too strict

- Adjust `RATE_LIMIT_REQUESTS` and `RATE_LIMIT_WINDOW`
- Check Redis for stuck counters: `redis-cli KEYS "rate_limit:*"`

### SSE streaming not working

- Verify task ID is correct
- Check CORS headers if accessing from browser
- Ensure connection timeout is long enough
- Check server logs for subscription errors

## 📖 Further Reading

- **[Design Document](DESIGN.md)**: Detailed architecture and implementation guide
- **MCP Specification**: https://spec.modelcontextprotocol.io/
- **OpenTelemetry Go**: https://opentelemetry.io/docs/languages/go/
- **pgvector**: https://github.com/pgvector/pgvector

## 📝 License

MIT License - see LICENSE file for details