Shahzad Bhatti Welcome to my ramblings and rants!

{
  "jsonrpc": "2.0",
  "id": 1,
  "method": "tools/call",
  "params": {
    "name": "hybrid_search",
    "arguments": {
      "query": "machine learning best practices",
      "limit": 5,
      "bm25_weight": 0.5,
      "vector_weight": 0.5
    }
  }
}

// ? BAD: Application-level filtering (can be bypassed)
func GetDocuments(tenantID string) ([]Document, error) {
    query := "SELECT * FROM documents WHERE tenant_id = ?"
    // What if someone forgets the WHERE clause?
    // What if there's a SQL injection?
    // What if a bug skips this check?
}

-- ? GOOD: Database-level Row-Level Security (impossible to bypass)
ALTER TABLE documents ENABLE ROW LEVEL SECURITY;

CREATE POLICY tenant_isolation ON documents
    USING (tenant_id = current_setting('app.current_tenant_id')::uuid);

// mcp-server/internal/auth/jwt.go
func (v *JWTValidator) ValidateToken(tokenString string) (*Claims, error) {
    token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
        if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
            return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
        }
        return v.publicKey, nil
    })

    if err != nil {
        return nil, fmt.Errorf("failed to parse token: %w", err)
    }

    claims, ok := token.Claims.(*Claims)
    if !ok || !token.Valid {
        return nil, fmt.Errorf("invalid token claims")
    }

    return claims, nil
}

Query: "SEC disclosure requirements GDPR data breach"

Vector-only results:
1. "Privacy Policy" (0.87 similarity)
2. "Data Protection Guide" (0.84 similarity)  
3. "General Security Practices" (0.81 similarity)
? Missed: Actual SEC regulation text

Hybrid results (0.5 BM25 + 0.5 Vector):
1. "SEC Rule 10b-5 Disclosure Requirements" (0.92 combined)
2. "GDPR Article 33 Breach Notification" (0.89 combined)
3. "Cross-Border Regulatory Compliance" (0.85 combined)
? Found: Exactly what we needed

// Hybrid search query using Reciprocal Rank Fusion
query := `
    WITH bm25_results AS (
        SELECT
            id,
            ts_rank_cd(
                to_tsvector('english', title || ' ' || content),
                plainto_tsquery('english', $1)
            ) AS bm25_score,
            ROW_NUMBER() OVER (ORDER BY ts_rank_cd(...) DESC) AS bm25_rank
        FROM documents
        WHERE to_tsvector('english', title || ' ' || content) @@ plainto_tsquery('english', $1)
    ),
    vector_results AS (
        SELECT
            id,
            1 - (embedding <=> $2) AS vector_score,
            ROW_NUMBER() OVER (ORDER BY embedding <=> $2) AS vector_rank
        FROM documents
        WHERE embedding IS NOT NULL
    ),
    combined AS (
        SELECT
            COALESCE(b.id, v.id) AS id,
            -- Reciprocal Rank Fusion score
            (
                COALESCE(1.0 / (60 + b.bm25_rank), 0) * $3 +
                COALESCE(1.0 / (60 + v.vector_rank), 0) * $4
            ) AS combined_score
        FROM bm25_results b
        FULL OUTER JOIN vector_results v ON b.id = v.id
    )
    SELECT * FROM combined
    ORDER BY combined_score DESC
    LIMIT $7
`

{
  "jsonrpc": "2.0",
  "id": 1,
  "method": "tools/call",
  "params": {
    "name": "hybrid_search",
    "arguments": {"query": "machine learning", "limit": 10}
  }
}

{
  "jsonrpc": "2.0",
  "id": 1,
  "result": {
    "content": [{
      "type": "text",
      "text": "[{\"doc_id\": \"123\", \"title\": \"ML Guide\", ...}]"
    }],
    "isError": false
  }
}

{
  "jsonrpc": "2.0",
  "id": 1,
  "error": {
    "code": -32602,
    "message": "Invalid params",
    "data": {"field": "query", "reason": "required"}
  }
}

// mcp-server/internal/tools/tool.go
type Tool interface {
    Definition() protocol.ToolDefinition
    Execute(ctx context.Context, args map[string]interface{}) (protocol.ToolCallResult, error)
}

// mcp-server/internal/tools/hybrid_search.go
type HybridSearchTool struct {
    db database.Store
}

func (t *HybridSearchTool) Execute(ctx context.Context, args map[string]interface{}) (protocol.ToolCallResult, error) {
    // 1. AUTHENTICATION: Extract tenant from JWT claims
    //    This happens at middleware level, but we verify here
    tenantID, ok := ctx.Value(auth.ContextKeyTenantID).(string)
    if !ok {
        return protocol.ToolCallResult{IsError: true}, fmt.Errorf("tenant ID not found in context")
    }

    // 2. PARAMETER PARSING: Extract and validate arguments
    query, _ := args["query"].(string)
    if query == "" {
        return protocol.ToolCallResult{IsError: true}, fmt.Errorf("query is required")
    }
    
    limit, _ := args["limit"].(float64)
    if limit <= 0 {
        limit = 10 // default
    }
    if limit > 50 {
        limit = 50 // max cap
    }
    
    bm25Weight, _ := args["bm25_weight"].(float64)
    vectorWeight, _ := args["vector_weight"].(float64)
    
    // 3. WEIGHT NORMALIZATION: Ensure weights sum to 1.0
    if bm25Weight == 0 && vectorWeight == 0 {
        bm25Weight = 0.5
        vectorWeight = 0.5
    }

    // 4. EMBEDDING GENERATION: Using Ollama for query embedding
    var embedding []float32
    if vectorWeight > 0 {
        embedding = generateEmbedding(query) // Calls Ollama API
    }

    // 5. DATABASE QUERY: Execute hybrid search with RLS
    params := database.HybridSearchParams{
        Query:        query,
        Embedding:    embedding,
        Limit:        int(limit),
        BM25Weight:   bm25Weight,
        VectorWeight: vectorWeight,
    }

    results, err := t.db.HybridSearch(ctx, tenantID, params)
    if err != nil {
        return protocol.ToolCallResult{IsError: true}, err
    }

    // 6. RESPONSE FORMATTING: Convert to JSON for client
    jsonData, _ := json.Marshal(results)
    return protocol.ToolCallResult{
        Content: []protocol.ContentBlock{{Type: "text", Text: string(jsonData)}},
        IsError: false,
    }, nil
}

// ? This crashes with NULL embeddings
var embedding pgvector.Vector

err = tx.QueryRow(ctx, query, docID).Scan(
    &doc.ID,
    &doc.TenantID,
    &doc.Title,
    &doc.Content,
    &doc.Metadata,
    &embedding, // CRASH: can't scan <nil> into pgvector.Vector
    &doc.CreatedAt,
    &doc.UpdatedAt,
)

can't scan into dest[5]: unsupported data type: <nil>

// ? Use pointer types for nullable fields
var embedding *pgvector.Vector // Pointer allows NULL

err = tx.QueryRow(ctx, query, docID).Scan(
    &doc.ID,
    &doc.TenantID,
    &doc.Title,
    &doc.Content,
    &doc.Metadata,
    &embedding, // Can be NULL now
    &doc.CreatedAt,
    &doc.UpdatedAt,
)

// Handle NULL embeddings gracefully
if embedding != nil && embedding.Slice() != nil {
    doc.Embedding = embedding.Slice()
} else {
    doc.Embedding = nil // Explicitly set to nil
}

return doc, nil

-- Hybrid search handles NULL embeddings gracefully
WITH bm25_results AS (
    SELECT id, ts_rank(to_tsvector('english', content), query) AS bm25_score
    FROM documents
    WHERE to_tsvector('english', content) @@ query
),
vector_results AS (
    SELECT id, 1 - (embedding <=> $1) AS vector_score
    FROM documents
    WHERE embedding IS NOT NULL  -- ? Skip NULL embeddings
)
SELECT
    d.*,
    COALESCE(b.bm25_score, 0) AS bm25_score,
    COALESCE(v.vector_score, 0) AS vector_score,
    ($2 * COALESCE(b.bm25_score, 0) + $3 * COALESCE(v.vector_score, 0)) AS combined_score
FROM documents d
LEFT JOIN bm25_results b ON d.id = b.id
LEFT JOIN vector_results v ON d.id = v.id
WHERE COALESCE(b.bm25_score, 0) > 0 OR COALESCE(v.vector_score, 0) > 0
ORDER BY combined_score DESC
LIMIT $4;

// mcp-server/internal/database/postgres.go
func (db *DB) SetTenantContext(ctx context.Context, tx pgx.Tx, tenantID string) error {
    // Note: SET commands don't support parameter binding
    // TenantID is validated as UUID by JWT validator, so this is safe
    query := fmt.Sprintf("SET LOCAL app.current_tenant_id = '%s'", tenantID)
    _, err := tx.Exec(ctx, query)
    return err
}

// Integration test: Verify tenant isolation
func TestTenantIsolation(t *testing.T) {
    // Create documents for two tenants
    tenant1Doc := createDocument(t, db, "tenant-1", "Secret Data A")
    tenant2Doc := createDocument(t, db, "tenant-2", "Secret Data B")
    
    // Query as tenant-1
    ctx1 := contextWithTenant(ctx, "tenant-1")
    results1, _ := db.ListDocuments(ctx1, "tenant-1", ListParams{Limit: 100})
    
    // Query as tenant-2
    ctx2 := contextWithTenant(ctx, "tenant-2")
    results2, _ := db.ListDocuments(ctx2, "tenant-2", ListParams{Limit: 100})
    
    // Assertions
    assert.Contains(t, results1, tenant1Doc)
    assert.NotContains(t, results1, tenant2Doc) // ? Cannot see other tenant
    
    assert.Contains(t, results2, tenant2Doc)
    assert.NotContains(t, results2, tenant1Doc) // ? Cannot see other tenant
}

// a2a-server/internal/handlers/tasks.go
func (h *TaskHandler) StreamEvents(w http.ResponseWriter, r *http.Request) {
    taskID := chi.URLParam(r, "taskId")

    // Set SSE headers
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.Header().Set("Access-Control-Allow-Origin", "*")

    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming not supported", http.StatusInternalServerError)
        return
    }

    // Stream task events
    for {
        event := h.taskManager.GetNextEvent(taskID)
        if event == nil {
            break // Task complete
        }

        // Format as SSE event
        data, _ := json.Marshal(event)
        fmt.Fprintf(w, "event: task_update\n")
        fmt.Fprintf(w, "data: %s\n\n", data)
        flusher.Flush()

        if event.Status == "completed" || event.Status == "failed" {
            break
        }
    }
}

# streamlit-ui/pages/3_?_A2A_Tasks.py
def stream_task_events(task_id: str):
    url = f"{A2A_BASE_URL}/tasks/{task_id}/events"

    with requests.get(url, stream=True) as response:
        for line in response.iter_lines():
            if line.startswith(b'data:'):
                data = json.loads(line[5:])
                st.write(f"Update: {data['message']}")
                yield data

# orchestration/workflows/rag_workflow.py
class RAGWorkflow:
    def __init__(self, mcp_url: str):
        self.mcp_client = MCPClient(mcp_url)
        self.workflow = self.build_workflow()

    def build_workflow(self) -> StateGraph:
        workflow = StateGraph(RAGState)

        # Define workflow steps
        workflow.add_node("search", self.search_documents)
        workflow.add_node("rank", self.rank_results)
        workflow.add_node("generate", self.generate_answer)
        workflow.add_node("verify", self.verify_sources)

        # Define edges (workflow graph)
        workflow.add_edge(START, "search")
        workflow.add_edge("search", "rank")
        workflow.add_edge("rank", "generate")
        workflow.add_edge("generate", "verify")
        workflow.add_edge("verify", END)

        return workflow.compile()

    def search_documents(self, state: RAGState) -> RAGState:
        """Search for relevant documents using MCP hybrid search"""
        # This is where MCP and A2A integrate!
        results = self.mcp_client.hybrid_search(
            query=state["query"],
            limit=10,
            bm25_weight=0.5,
            vector_weight=0.5
        )

        state["documents"] = results
        state["progress"] = f"Found {len(results)} documents"
        
        # Emit progress event via A2A
        emit_progress_event(state["task_id"], "search_complete", state["progress"])
        
        return state

    def rank_results(self, state: RAGState) -> RAGState:
        """Rank results by combined score"""
        docs = sorted(
            state["documents"],
            key=lambda x: x["score"],
            reverse=True
        )[:5]

        state["ranked_docs"] = docs
        state["progress"] = "Ranked top 5 documents"
        
        emit_progress_event(state["task_id"], "ranking_complete", state["progress"])
        
        return state

    def generate_answer(self, state: RAGState) -> RAGState:
        """Generate answer using retrieved context"""
        context = "\n\n".join([
            f"Document: {doc['title']}\n{doc['content']}"
            for doc in state["ranked_docs"]
        ])

        prompt = f"""Based on the following documents, answer the question.

Context:
{context}

Question: {state['query']}

Answer:"""

        # Call Ollama for local inference
        response = ollama.generate(
            model="llama3.2",
            prompt=prompt
        )

        state["answer"] = response["response"]
        state["progress"] = "Generated final answer"
        
        emit_progress_event(state["task_id"], "generation_complete", state["progress"])
        
        return state
        
    def verify_sources(self, state: RAGState) -> RAGState:
        """Verify sources are accurately cited"""
        # Check each cited document exists in ranked_docs
        cited_docs = extract_citations(state["answer"])
        verified = all(doc in state["ranked_docs"] for doc in cited_docs)
        
        state["verified"] = verified
        state["progress"] = "Verified sources" if verified else "Source verification failed"
        
        emit_progress_event(state["task_id"], "verification_complete", state["progress"])
        
        return state

# streamlit-ui/pages/1_?_Authentication.py
def generate_jwt_token(tenant_id: str, user_id: str, ttl: int = 3600) -> str:
    """Generate RS256 JWT token with proper claims"""
    now = datetime.now(timezone.utc)

    payload = {
        "tenant_id": tenant_id,
        "user_id": user_id,
        "iat": now,              # Issued at
        "exp": now + timedelta(seconds=ttl),  # Expiration
        "nbf": now,              # Not before
        "jti": str(uuid.uuid4()), # JWT ID (for revocation)
        "iss": "mcp-demo-ui",    # Issuer
        "aud": "mcp-server"      # Audience
    }

    # Sign with RSA private key
    with open("/app/certs/private_key.pem", "rb") as f:
        private_key = serialization.load_pem_private_key(
            f.read(),
            password=None
        )

    token = jwt.encode(payload, private_key, algorithm="RS256")
    return token

// mcp-server/internal/middleware/auth.go
func AuthMiddleware(validator *auth.JWTValidator) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            // 1. Extract token from Authorization header
            authHeader := r.Header.Get("Authorization")
            if authHeader == "" {
                http.Error(w, "missing authorization header", http.StatusUnauthorized)
                return
            }

            tokenString := strings.TrimPrefix(authHeader, "Bearer ")
            
            // 2. Validate token signature and claims
            claims, err := validator.ValidateToken(tokenString)
            if err != nil {
                log.Printf("Token validation failed: %v", err)
                http.Error(w, "invalid token", http.StatusUnauthorized)
                return
            }

            // 3. Check token expiration
            if claims.ExpiresAt.Before(time.Now()) {
                http.Error(w, "token expired", http.StatusUnauthorized)
                return
            }

            // 4. Check token not used before nbf
            if claims.NotBefore.After(time.Now()) {
                http.Error(w, "token not yet valid", http.StatusUnauthorized)
                return
            }

            // 5. Verify audience (prevent token reuse across services)
            if claims.Audience != "mcp-server" {
                http.Error(w, "invalid token audience", http.StatusUnauthorized)
                return
            }

            // 6. Add claims to context for downstream handlers
            ctx := context.WithValue(r.Context(), auth.ContextKeyTenantID, claims.TenantID)
            ctx = context.WithValue(ctx, auth.ContextKeyUserID, claims.UserID)
            ctx = context.WithValue(ctx, auth.ContextKeyJTI, claims.JTI)

            next.ServeHTTP(w, r.WithContext(ctx))
        })
    }
}

# streamlit-ui/pages/4_?_Cost_Tracking.py
class CostTracker:
    def __init__(self):
        self.costs = []
        self.pricing = {
            # Local models (Ollama)
            "llama3.2": 0.0001,      # $0.0001 per 1K tokens
            "mistral": 0.0001,
            
            # OpenAI models
            "gpt-4": 0.03,           # $0.03 per 1K tokens
            "gpt-3.5-turbo": 0.002,  # $0.002 per 1K tokens
            
            # Anthropic models
            "claude-3": 0.015,       # $0.015 per 1K tokens
            "claude-3-haiku": 0.0025,
        }

    def track_request(self, user_id: str, model: str, 
                     input_tokens: int, output_tokens: int,
                     metadata: dict = None):
        """Track a single request with detailed token breakdown"""
        
        # Calculate costs
        input_cost = (input_tokens / 1000) * self.pricing.get(model, 0)
        output_cost = (output_tokens / 1000) * self.pricing.get(model, 0)
        total_cost = input_cost + output_cost

        # Store record
        self.costs.append({
            "timestamp": datetime.now(),
            "user_id": user_id,
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "input_cost": input_cost,
            "output_cost": output_cost,
            "total_cost": total_cost,
            "metadata": metadata or {}
        })
        
        return total_cost

    def check_budget(self, user_id: str, budget: float) -> tuple[bool, float]:
        """Check if user is within budget"""
        user_costs = [
            c["total_cost"] for c in self.costs
            if c["user_id"] == user_id
        ]

        total_spent = sum(user_costs)
        remaining = budget - total_spent
        
        return remaining > 0, remaining

    def get_usage_by_model(self, user_id: str) -> dict:
        """Get cost breakdown by model"""
        model_costs = {}
        
        for cost in self.costs:
            if cost["user_id"] == user_id:
                model = cost["model"]
                if model not in model_costs:
                    model_costs[model] = {
                        "requests": 0,
                        "total_tokens": 0,
                        "total_cost": 0.0
                    }
                
                model_costs[model]["requests"] += 1
                model_costs[model]["total_tokens"] += cost["input_tokens"] + cost["output_tokens"]
                model_costs[model]["total_cost"] += cost["total_cost"]
        
        return model_costs

# Budget enforcement by user tier
BUDGET_TIERS = {
    "free": {
        "monthly_budget": 0.50,      # $0.50/month
        "rate_limit": 10,            # 10 req/min
        "models": ["llama3.2"]       # Local only
    },
    "pro": {
        "monthly_budget": 25.00,     # $25/month
        "rate_limit": 100,           # 100 req/min
        "models": ["llama3.2", "gpt-3.5-turbo", "claude-3-haiku"]
    },
    "enterprise": {
        "monthly_budget": 500.00,    # $500/month
        "rate_limit": 1000,          # 1000 req/min
        "models": ["*"]              # All models
    }
}

# orchestration/workflows/rag_workflow.py
try:
    from langfuse.decorators import observe, langfuse_context
    LANGFUSE_AVAILABLE = True
except ImportError:
    LANGFUSE_AVAILABLE = False
    # Create no-op decorator for local dev
    def observe(*args, **kwargs):
        def decorator(func):
            return func
        return decorator if not args else decorator(args[0])

@observe(name="rag_workflow")
def run_rag_workflow(query: str, user_id: str, tenant_id: str) -> str:
    """Run RAG workflow with observability"""
    workflow = RAGWorkflow(mcp_url="http://mcp-server:8080")

    result = workflow.run({
        "query": query,
        "user_id": user_id,
        "tenant_id": tenant_id
    })

    if LANGFUSE_AVAILABLE:
        # Add metadata for debugging
        langfuse_context.update_current_trace(
            metadata={
                "documents_found": len(result["documents"]),
                "top_score": result["ranked_docs"][0]["score"],
                "model": "llama3.2",
                "tenant_id": tenant_id,
                "user_id": user_id
            },
            tags=["rag", "production", tenant_id]
        )

    return result["answer"]

// mcp-server/internal/middleware/ratelimit.go
import "golang.org/x/time/rate"

func RateLimitMiddleware(limiter *rate.Limiter) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            if !limiter.Allow() {
                http.Error(w, "rate limit exceeded", http.StatusTooManyRequests)
                return
            }
            next.ServeHTTP(w, r)
        })
    }
}

// Usage: 100 requests per second per tenant
limiter := rate.NewLimiter(100, 200) // 100 req/sec burst of 200

// mcp-server/internal/middleware/ratelimit_redis.go
type RedisRateLimiter struct {
    client *redis.Client
    limit  int
    window time.Duration
}

func (r *RedisRateLimiter) Allow(ctx context.Context, tenantID string) (bool, error) {
    key := fmt.Sprintf("ratelimit:tenant:%s", tenantID)
    
    // Increment counter
    count, err := r.client.Incr(ctx, key).Result()
    if err != nil {
        return false, err
    }
    
    // Set expiration on first request
    if count == 1 {
        r.client.Expire(ctx, key, r.window)
    }
    
    // Check limit
    return count <= int64(r.limit), nil
}

// mcp-server/internal/database/postgres_integration_test.go
func TestGetDocument_WithNullEmbedding(t *testing.T) {
    db := setupTestDB(t)
    defer db.Close()

    ctx := context.Background()

    // Insert document WITHOUT embedding (common in real world)
    testDoc := &Document{
        TenantID:  testTenantID,
        Title:     "Test Document Without Embedding",
        Content:   "This document has no embedding vector",
        Metadata:  map[string]interface{}{"test": true},
        Embedding: nil, // Explicitly no embedding
    }

    err := db.InsertDocument(ctx, testTenantID, testDoc)
    require.NoError(t, err)

    // Retrieve - should NOT fail with NULL scan error
    retrieved, err := db.GetDocument(ctx, testTenantID, testDoc.ID)
    require.NoError(t, err)
    assert.NotNil(t, retrieved)
    assert.Nil(t, retrieved.Embedding) // Embedding is NULL
    assert.Equal(t, testDoc.Title, retrieved.Title)
    assert.Equal(t, testDoc.Content, retrieved.Content)

    // Cleanup
    db.DeleteDocument(ctx, testTenantID, testDoc.ID)
}

func TestHybridSearch_HandlesNullEmbeddings(t *testing.T) {
    db := setupTestDB(t)
    defer db.Close()

    ctx := context.Background()

    // Insert documents with and without embeddings
    docWithEmbedding := createDocumentWithEmbedding(t, db, testTenantID, "AI Guide")
    docWithoutEmbedding := createDocumentWithoutEmbedding(t, db, testTenantID, "ML Tutorial")

    // Create query embedding
    queryEmbedding := make([]float32, 1536)
    for i := range queryEmbedding {
        queryEmbedding[i] = 0.1
    }

    params := HybridSearchParams{
        Query:        "artificial intelligence machine learning",
        Embedding:    queryEmbedding,
        Limit:        10,
        BM25Weight:   0.5,
        VectorWeight: 0.5,
    }

    // Should work even with NULL embeddings
    results, err := db.HybridSearch(ctx, testTenantID, params)
    require.NoError(t, err)
    assert.NotNil(t, results)
    assert.Greater(t, len(results), 0)

    // Documents without embeddings get vector_score = 0
    for _, result := range results {
        if result.Document.Embedding == nil {
            assert.Equal(t, 0.0, result.VectorScore)
            assert.Greater(t, result.BM25Score, 0.0) // But BM25 should work
        }
    }
}

func TestTenantIsolation_CannotAccessOtherTenant(t *testing.T) {
    db := setupTestDB(t)
    defer db.Close()

    tenant1ID := "tenant-1-" + uuid.New().String()
    tenant2ID := "tenant-2-" + uuid.New().String()

    // Create documents for both tenants
    doc1 := createDocument(t, db, tenant1ID, "Tenant 1 Secret Data")
    doc2 := createDocument(t, db, tenant2ID, "Tenant 2 Secret Data")

    // Query as tenant-1
    ctx1 := context.Background()
    results1, err := db.ListDocuments(ctx1, tenant1ID, ListParams{Limit: 100})
    require.NoError(t, err)

    // Query as tenant-2
    ctx2 := context.Background()
    results2, err := db.ListDocuments(ctx2, tenant2ID, ListParams{Limit: 100})
    require.NoError(t, err)

    // Verify isolation
    assert.Contains(t, results1, doc1)
    assert.NotContains(t, results1, doc2) // ? Cannot see other tenant

    assert.Contains(t, results2, doc2)
    assert.NotContains(t, results2, doc1) // ? Cannot see other tenant
}

func TestConcurrentRetrievals_NoRaceConditions(t *testing.T) {
    db := setupTestDB(t)
    defer db.Close()

    // Create test documents
    docs := make([]*Document, 50)
    for i := 0; i < 50; i++ {
        docs[i] = createDocument(t, db, testTenantID, fmt.Sprintf("Document %d", i))
    }

    // Concurrent retrievals
    var wg sync.WaitGroup
    errors := make(chan error, 500)

    for worker := 0; worker < 10; worker++ {
        wg.Add(1)
        go func() {
            defer wg.Done()

            for i := 0; i < 50; i++ {
                doc := docs[i]
                retrieved, err := db.GetDocument(context.Background(), testTenantID, doc.ID)
                if err != nil {
                    errors <- err
                    return
                }
                if retrieved.ID != doc.ID {
                    errors <- fmt.Errorf("document mismatch: got %s, want %s", retrieved.ID, doc.ID)
                    return
                }
            }
        }()
    }

    wg.Wait()
    close(errors)

    // Check for errors
    for err := range errors {
        t.Error(err)
    }
}

# Unit tests (fast, no dependencies)
cd mcp-server
go test -v ./...

# Integration tests (requires PostgreSQL)
./scripts/run-integration-tests.sh

? Running MCP Server Integration Tests
========================================
? PostgreSQL is ready

? Running integration tests...

=== RUN   TestGetDocument_WithNullEmbedding
--- PASS: TestGetDocument_WithNullEmbedding (0.05s)
=== RUN   TestGetDocument_WithEmbedding
--- PASS: TestGetDocument_WithEmbedding (0.04s)
=== RUN   TestHybridSearch_HandlesNullEmbeddings
--- PASS: TestHybridSearch_HandlesNullEmbeddings (0.12s)
=== RUN   TestTenantIsolation
--- PASS: TestTenantIsolation (0.08s)
=== RUN   TestConcurrentRetrievals
--- PASS: TestConcurrentRetrievals (2.34s)

PASS
coverage: 95.3% of statements
ok  	github.com/bhatti/mcp-a2a-go/mcp-server/internal/database	3.456s

? Integration tests completed!

# Client code
results = mcp_client.hybrid_search(
    query="data breach notification requirements GDPR Article 33",
    limit=10,
    bm25_weight=0.6,  # Favor exact keyword matches for legal terms
    vector_weight=0.4  # But include semantic similarity
)

for result in results:
    print(f"Document: {result['title']}")
    print(f"BM25 Score: {result['bm25_score']:.2f}")
    print(f"Vector Score: {result['vector_score']:.2f}")
    print(f"Combined: {result['score']:.2f}")
    print(f"Tenant: {result['tenant_id']}")
    print("---")

Document: GDPR Compliance Framework - Article 33 Analysis
BM25 Score: 0.89  (matched "GDPR", "Article 33", "notification")
Vector Score: 0.76  (understood "data breach requirements")
Combined: 0.84
Tenant: client-acme-legal

Document: Data Breach Response Procedures
BM25 Score: 0.45  (matched "data breach", "notification")
Vector Score: 0.91  (strong semantic match)
Combined: 0.65
Tenant: client-acme-legal

Document: SEC Disclosure Requirements
BM25 Score: 0.78  (matched "requirements", "notification")
Vector Score: 0.52  (weak semantic match)
Combined: 0.67
Tenant: client-acme-legal

# orchestration/workflows/research_workflow.py
class ResearchWorkflow:
    def build_workflow(self):
        workflow = StateGraph(ResearchState)

        # Define research steps
        workflow.add_node("search_company", self.search_company_docs)
        workflow.add_node("search_competitors", self.search_competitors)
        workflow.add_node("search_financials", self.search_financial_data)
        workflow.add_node("analyze_trends", self.analyze_market_trends)
        workflow.add_node("verify_facts", self.verify_with_sources)
        workflow.add_node("generate_report", self.generate_final_report)

        # Define workflow graph
        workflow.add_edge(START, "search_company")
        workflow.add_edge("search_company", "search_competitors")
        workflow.add_edge("search_competitors", "search_financials")
        workflow.add_edge("search_financials", "analyze_trends")
        workflow.add_edge("analyze_trends", "verify_facts")
        workflow.add_edge("verify_facts", "generate_report")
        workflow.add_edge("generate_report", END)

        return workflow.compile()
    
    def search_company_docs(self, state: ResearchState) -> ResearchState:
        """Step 1: Search company documents via MCP"""
        company = state["company_name"]
        
        # Call MCP hybrid search
        results = self.mcp_client.hybrid_search(
            query=f"{company} business operations revenue products",
            limit=20,
            bm25_weight=0.5,
            vector_weight=0.5
        )
        
        state["company_docs"] = results
        state["progress"] = f"Found {len(results)} company documents"
        
        # Emit progress via A2A SSE
        emit_progress("search_company_complete", state["progress"])
        
        return state
    
    def search_competitors(self, state: ResearchState) -> ResearchState:
        """Step 2: Identify and search competitors"""
        company = state["company_name"]
        
        # Extract competitors from company docs
        competitors = self.extract_competitors(state["company_docs"])
        
        # Search each competitor
        competitor_data = {}
        for competitor in competitors:
            results = self.mcp_client.hybrid_search(
                query=f"{competitor} market share products revenue",
                limit=10
            )
            competitor_data[competitor] = results
        
        state["competitors"] = competitor_data
        state["progress"] = f"Analyzed {len(competitors)} competitors"
        
        emit_progress("search_competitors_complete", state["progress"])
        
        return state
    
    def search_financial_data(self, state: ResearchState) -> ResearchState:
        """Step 3: Extract financial metrics"""
        company = state["company_name"]
        
        # Search for financial documents
        results = self.mcp_client.hybrid_search(
            query=f"{company} revenue earnings profit margin cash flow",
            limit=15,
            bm25_weight=0.7,  # Favor exact financial terms
            vector_weight=0.3
        )
        
        # Extract key metrics
        metrics = self.extract_financial_metrics(results)
        
        state["financials"] = metrics
        state["progress"] = f"Extracted {len(metrics)} financial metrics"
        
        emit_progress("search_financials_complete", state["progress"])
        
        return state
    
    def verify_facts(self, state: ResearchState) -> ResearchState:
        """Step 5: Verify all facts with sources"""
        # Check each claim has supporting document
        claims = self.extract_claims(state["report_draft"])
        
        verified_claims = []
        for claim in claims:
            sources = self.find_supporting_docs(claim, state)
            if sources:
                verified_claims.append({
                    "claim": claim,
                    "sources": sources,
                    "verified": True
                })
        
        state["verified_claims"] = verified_claims
        state["progress"] = f"Verified {len(verified_claims)} claims"
        
        emit_progress("verification_complete", state["progress"])
        
        return state

# Before each request
tier = get_user_tier(user_id)
budget = BUDGET_TIERS[tier]["monthly_budget"]
allowed, remaining = cost_tracker.check_budget(user_id, budget)

if not allowed:
    raise BudgetExceededError(
        f"Monthly budget of ${budget} exceeded. "
        f"Upgrade to {next_tier} for higher limits."
    )

# Track the request
response = llm.generate(prompt)
cost = cost_tracker.track_request(
    user_id=user_id,
    model="llama3.2",
    input_tokens=len(prompt.split()),
    output_tokens=len(response.split())
)

# Alert when approaching limit
if remaining < 5.0:  # $5 remaining
    send_alert(user_id, f"Budget alert: ${remaining:.2f} remaining")

# streamlit-ui/pages/4_?_Cost_Tracking.py
def enforce_budget_limits():
    """Check budget before task creation"""
    
    user_tier = st.session_state.get("user_tier", "free")
    budget = BUDGET_TIERS[user_tier]["monthly_budget"]
    
    # Calculate current spend
    spent = cost_tracker.get_total_cost(user_id)
    remaining = budget - spent
    
    # Display budget status
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.metric("Budget", f"${budget:.2f}")
    
    with col2:
        st.metric("Spent", f"${spent:.2f}", 
                 delta=f"-${spent:.2f}", delta_color="inverse")
    
    with col3:
        progress = (spent / budget) * 100
        st.metric("Remaining", f"${remaining:.2f}")
        st.progress(progress / 100)
    
    # Block if exceeded
    if remaining <= 0:
        st.error("? Monthly budget exceeded. Upgrade to continue.")
        st.button("Upgrade to Pro ($25/month)", on_click=upgrade_tier)
        return False
    
    # Warn if close
    if remaining < 5.0:
        st.warning(f"?? Budget alert: Only ${remaining:.2f} remaining this month")
    
    return True

# docker-compose.yml
version: '3.8'

services:
  postgres:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_DB: mcp_db
      POSTGRES_USER: mcp_user
      POSTGRES_PASSWORD: ${DB_PASSWORD:-mcp_secure_pass}
    volumes:
      - ./scripts/init-db.sql:/docker-entrypoint-initdb.d/init.sql
      - postgres_data:/var/lib/postgresql/data
    ports:
      - "5432:5432"
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mcp_user -d mcp_db"]
      interval: 5s
      timeout: 5s
      retries: 5

  ollama:
    image: ollama/ollama:latest
    volumes:
      - ollama_data:/root/.ollama
    ports:
      - "11434:11434"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:11434/api/tags"]
      interval: 10s
      timeout: 5s
      retries: 3

  mcp-server:
    build:
      context: ./mcp-server
      dockerfile: Dockerfile
    environment:
      DB_HOST: postgres
      DB_PORT: 5432
      DB_USER: mcp_user
      DB_PASSWORD: ${DB_PASSWORD:-mcp_secure_pass}
      DB_NAME: mcp_db
      JWT_PUBLIC_KEY_PATH: /app/certs/public_key.pem
      OLLAMA_URL: http://ollama:11434
      LOG_LEVEL: ${LOG_LEVEL:-info}
    ports:
      - "8080:8080"
    depends_on:
      postgres:
        condition: service_healthy
      ollama:
        condition: service_healthy
    volumes:
      - ./certs:/app/certs:ro
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
      interval: 10s
      timeout: 5s
      retries: 3

  a2a-server:
    build:
      context: ./a2a-server
      dockerfile: Dockerfile
    environment:
      MCP_SERVER_URL: http://mcp-server:8080
      OLLAMA_URL: http://ollama:11434
      LOG_LEVEL: ${LOG_LEVEL:-info}
    ports:
      - "8082:8082"
    depends_on:
      - mcp-server
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8082/health"]
      interval: 10s
      timeout: 5s
      retries: 3

  streamlit-ui:
    build:
      context: ./streamlit-ui
      dockerfile: Dockerfile
    environment:
      MCP_SERVER_URL: http://mcp-server:8080
      A2A_SERVER_URL: http://a2a-server:8082
    ports:
      - "8501:8501"
    volumes:
      - ./certs:/app/certs:ro
    depends_on:
      - mcp-server
      - a2a-server

volumes:
  postgres_data:
  ollama_data:

# Start all services
docker compose up -d

# Check status
docker compose ps

# Expected output:
# NAME              STATUS        PORTS
# postgres          Up (healthy)  0.0.0.0:5432->5432/tcp
# ollama            Up (healthy)  0.0.0.0:11434->11434/tcp
# mcp-server        Up (healthy)  0.0.0.0:8080->8080/tcp
# a2a-server        Up (healthy)  0.0.0.0:8082->8082/tcp
# streamlit-ui      Up            0.0.0.0:8501->8501/tcp

# View logs
docker compose logs -f mcp-server
docker compose logs -f a2a-server

# Run health checks
curl http://localhost:8080/health  # MCP server
curl http://localhost:8082/health  # A2A server

# Pull Ollama model
docker compose exec ollama ollama pull llama3.2

# Initialize database with sample data
docker compose exec postgres psql -U mcp_user -d mcp_db -f /docker-entrypoint-initdb.d/init.sql

# .env.production
DB_PASSWORD=$(openssl rand -base64 32)
JWT_PRIVATE_KEY_PATH=/secrets/jwt_private_key.pem
JWT_PUBLIC_KEY_PATH=/secrets/jwt_public_key.pem
LANGFUSE_SECRET_KEY=${LANGFUSE_SECRET_KEY}
LANGFUSE_PUBLIC_KEY=${LANGFUSE_PUBLIC_KEY}
OLLAMA_URL=http://ollama:11434
LOG_LEVEL=info
SENTRY_DSN=${SENTRY_DSN}

# Install migrate
curl -L https://github.com/golang-migrate/migrate/releases/download/v4.16.2/migrate.linux-amd64.tar.gz | tar xvz
mv migrate /usr/local/bin/

# Create migration
migrate create -ext sql -dir db/migrations -seq add_embeddings_index

# Apply migrations
migrate -path db/migrations \
        -database "postgresql://user:pass@localhost:5432/db?sslmode=disable" \
        up

# Rollback if needed
migrate -path db/migrations \
        -database "postgresql://user:pass@localhost:5432/db?sslmode=disable" \
        down 1

# k8s/mcp-server-deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: mcp-server
  namespace: mcp-a2a
spec:
  replicas: 3  # High availability
  selector:
    matchLabels:
      app: mcp-server
  template:
    metadata:
      labels:
        app: mcp-server
    spec:
      containers:
      - name: mcp-server
        image: ghcr.io/bhatti/mcp-server:latest
        ports:
        - containerPort: 8080
        env:
        - name: DB_HOST
          value: postgres-service
        - name: DB_PASSWORD
          valueFrom:
            secretKeyRef:
              name: db-credentials
              key: password
        - name: JWT_PUBLIC_KEY_PATH
          value: /certs/public_key.pem
        volumeMounts:
        - name: certs
          mountPath: /certs
          readOnly: true
        resources:
          requests:
            memory: "128Mi"
            cpu: "100m"
          limits:
            memory: "512Mi"
            cpu: "500m"
        livenessProbe:
          httpGet:
            path: /health
            port: 8080
          initialDelaySeconds: 10
          periodSeconds: 5
        readinessProbe:
          httpGet:
            path: /health
            port: 8080
          initialDelaySeconds: 5
          periodSeconds: 3
      volumes:
      - name: certs
        secret:
          secretName: jwt-certs

# Create namespace
kubectl create namespace mcp-a2a

# Apply secrets
kubectl create secret generic db-credentials \
  --from-literal=password=$(openssl rand -base64 32) \
  -n mcp-a2a

kubectl create secret generic jwt-certs \
  --from-file=public_key.pem=./certs/public_key.pem \
  --from-file=private_key.pem=./certs/private_key.pem \
  -n mcp-a2a

# Apply manifests
kubectl apply -f k8s/namespace.yaml
kubectl apply -f k8s/postgres.yaml
kubectl apply -f k8s/mcp-server.yaml
kubectl apply -f k8s/a2a-server.yaml
kubectl apply -f k8s/streamlit-ui.yaml

# Check pods
kubectl get pods -n mcp-a2a

# View logs
kubectl logs -f deployment/mcp-server -n mcp-a2a

# Scale up
kubectl scale deployment mcp-server --replicas=5 -n mcp-a2a

// mcp-server/internal/metrics/prometheus.go
var (
    requestDuration = prometheus.NewHistogramVec(
        prometheus.HistogramOpts{
            Name: "mcp_request_duration_seconds",
            Help: "MCP request duration in seconds",
            Buckets: prometheus.DefBuckets,
        },
        []string{"method", "status"},
    )

    activeRequests = prometheus.NewGauge(
        prometheus.GaugeOpts{
            Name: "mcp_active_requests",
            Help: "Number of active MCP requests",
        },
    )
    
    hybridSearchQueries = prometheus.NewCounterVec(
        prometheus.CounterOpts{
            Name: "mcp_hybrid_search_queries_total",
            Help: "Total number of hybrid search queries",
        },
        []string{"tenant_id"},
    )
    
    budgetExceeded = prometheus.NewCounterVec(
        prometheus.CounterOpts{
            Name: "mcp_budget_exceeded_total",
            Help: "Number of requests blocked due to budget limits",
        },
        []string{"user_id", "tier"},
    )
)

func init() {
    prometheus.MustRegister(requestDuration)
    prometheus.MustRegister(activeRequests)
    prometheus.MustRegister(hybridSearchQueries)
    prometheus.MustRegister(budgetExceeded)
}

# prometheus/alerts.yml
groups:
- name: mcp_alerts
  rules:
  - alert: HighErrorRate
    expr: rate(mcp_request_duration_seconds_count{status="error"}[5m]) > 0.1
    for: 5m
    annotations:
      summary: "High error rate on MCP server"
      description: "Error rate is {{ $value }} errors/sec"
  
  - alert: BudgetExceededRate
    expr: rate(mcp_budget_exceeded_total[1h]) > 100
    annotations:
      summary: "High budget exceeded rate"
      description: "{{ $value }} users hitting budget limits per hour"
  
  - alert: DatabaseLatency
    expr: mcp_request_duration_seconds{method="hybrid_search"} > 1.0
    for: 2m
    annotations:
      summary: "Slow hybrid search queries"
      description: "Hybrid search taking {{ $value }}s (should be <1s)"

#!/bin/bash
# scripts/backup-database.sh

BACKUP_DIR="/backups"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DB_NAME="mcp_db"
DB_USER="mcp_user"

# Create backup directory
mkdir -p ${BACKUP_DIR}

# Dump database
docker compose exec -T postgres pg_dump -U ${DB_USER} ${DB_NAME} | \
    gzip > ${BACKUP_DIR}/${DB_NAME}_${TIMESTAMP}.sql.gz

# Upload to S3 (optional)
aws s3 cp ${BACKUP_DIR}/${DB_NAME}_${TIMESTAMP}.sql.gz \
    s3://my-backups/mcp-db/

# Keep last 7 days locally
find ${BACKUP_DIR} -name "${DB_NAME}_*.sql.gz" -mtime +7 -delete

echo "Backup completed: ${DB_NAME}_${TIMESTAMP}.sql.gz"

Benchmark: Hybrid Search (10 results, 1536-dim embeddings)
- Requests/sec: 5,247
- P50 latency: 12ms
- P95 latency: 45ms
- P99 latency: 89ms
- Memory: 52MB baseline, 89MB under load
- CPU: 23% average (4 cores)

Benchmark: Vector search (cosine similarity)
- Documents: 100,000
- Embedding dimensions: 1536
- Index: HNSW (m=16, ef_construction=64)
- Query time: <100ms (P95)
- Throughput: 150 queries/sec (single connection)
- Concurrent queries: 100+ simultaneous

# Using hey (HTTP load generator)
hey -n 10000 -c 100 -m POST \
    -H "Authorization: Bearer $TOKEN" \
    -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"hybrid_search","arguments":{"query":"machine learning","limit":10}}}' \
    http://localhost:8080/mcp

Summary:
  Total:        19.8421 secs
  Slowest:      0.2847 secs
  Fastest:      0.0089 secs
  Average:      0.1974 secs
  Requests/sec: 503.98
  
  Status code distribution:
    [200]	10000 responses

Latency distribution:
  10% in 0.0234 secs
  25% in 0.0456 secs
  50% in 0.1842 secs
  75% in 0.3123 secs
  90% in 0.4234 secs
  95% in 0.4867 secs
  99% in 0.5634 secs

// ? Structured logging
log.Info().
    Str("tenant_id", tenantID).
    Str("user_id", userID).
    Int("results_count", len(results)).
    Float64("duration_ms", duration.Milliseconds()).
    Msg("hybrid search completed")

// ? Per-tenant rate limiting
type RedisRateLimiter struct {
    client *redis.Client
}

func (r *RedisRateLimiter) Allow(ctx context.Context, tenantID string, limit int) (bool, error) {
    key := fmt.Sprintf("ratelimit:tenant:%s", tenantID)
    
    pipe := r.client.Pipeline()
    incr := pipe.Incr(ctx, key)
    pipe.Expire(ctx, key, time.Minute)
    _, err := pipe.Exec(ctx)
    if err != nil {
        return false, err
    }
    
    count, err := incr.Result()
    if err != nil {
        return false, err
    }
    
    return count <= int64(limit), nil
}

# embeddings-service/app.py (what I should have built)
from fastapi import FastAPI
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = SentenceTransformer('all-MiniLM-L6-v2')

@app.post("/embed")
async def embed(texts: list[str]):
    embeddings = model.encode(texts, batch_size=32)
    return {"embeddings": embeddings.tolist()}

// ? Circuit breaker pattern
type CircuitBreaker struct {
    maxFailures int
    timeout     time.Duration
    failures    int
    lastFail    time.Time
    state       string  // "closed", "open", "half-open"
}

func (cb *CircuitBreaker) Call(fn func() error) error {
    if cb.state == "open" {
        if time.Since(cb.lastFail) > cb.timeout {
            cb.state = "half-open"
        } else {
            return fmt.Errorf("circuit breaker open")
        }
    }
    
    err := fn()
    if err != nil {
        cb.failures++
        cb.lastFail = time.Now()
        
        if cb.failures >= cb.maxFailures {
            cb.state = "open"
        }
        return err
    }
    
    cb.failures = 0
    cb.state = "closed"
    return nil
}

# ? Typical MCP tutorial
def handle_request(request):
    tool = request["params"]["name"]
    args = request["params"]["arguments"]
    return execute_tool(tool, args)  # No auth, no validation, no limits

// ? Production reality
func (h *MCPHandler) handleToolsCall(ctx context.Context, req *protocol.Request) {
    // 1. Authenticate (JWT validation)
    // 2. Authorize (check permissions)
    // 3. Rate limit (per-tenant quotas)
    // 4. Validate input (prevent injection)
    // 5. Inject tenant context (RLS)
    // 6. Trace request (observability)
    // 7. Track cost (budget enforcement)
    // 8. Circuit breaker (fail fast)
    // 9. Retry logic (handle transients)
    // 10. Audit log (compliance)
    
    return h.toolRegistry.Execute(ctx, toolReq.Name, toolReq.Arguments)
}

embedding = ollama.embed(text)  # Ollama down ? hangs forever ? system freezes

ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
embedding, err := ollama.Embed(ctx, text)
if err != nil {
    return db.BM25Search(ctx, query)  // Degrade gracefully, skip embeddings
}

for task in tasks:
    result = external_api.call(task)  # Fails 1000 times, wastes time/money

if circuitBreaker.IsOpen() {
    return cachedResult  // Fail fast, don't waste resources
}

Tenant A: 10,000 req/sec ? Database crashes ? ALL tenants down

if !rateLimiter.Allow(tenantID) {
    return ErrRateLimitExceeded  // Other tenants unaffected
}

git clone https://github.com/bhatti/mcp-a2a-go
cd mcp-a2a-go
docker compose up -d
./scripts/run-integration-tests.sh
open http://localhost:8501