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

# Install via Homebrew (macOS)
brew install --cask claude

# Or download directly from Anthropic
# https://claude.ai/download

# Create a Claude workspace directory
mkdir -p ~/claude-workspace/{projects,templates,context}

# Add a context file for your coding standards
cat > ~/claude-workspace/context/coding-standards.md << 'EOF'
# My Coding Standards

- Use descriptive variable names
- Functions should do one thing well
- Write tests for all new features
- Handle errors explicitly
- Document complex logic
EOF

# Download TLA+ Tools from https://github.com/tlaplus/tlaplus/releases
# Or use Homebrew on macOS
brew install --cask tla-plus-toolbox

# For command-line usage (recommended for CI)
wget https://github.com/tlaplus/tlaplus/releases/download/v1.8.0/tla2tools.jar

code --install-extension alygin.vscode-tlaplus

--------------------------- MODULE SimpleCounter ---------------------------
VARIABLE counter

Init == counter = 0

Increment == counter' = counter + 1

Decrement == counter' = counter - 1

Next == Increment \/ Decrement

Spec == Init /\ [][Next]_counter

TypeInvariant == counter \in Int

=============================================================================

--------------------------- MODULE TaskManagement ---------------------------
EXTENDS Integers, Sequences, FiniteSets, TLC

CONSTANTS
    Users,          \* Set of users
    MaxTasks,       \* Maximum number of tasks
    MaxTime,        \* Maximum time value for simulation
    Titles,         \* Set of possible task titles
    Descriptions    \* Set of possible task descriptions

VARIABLES
    tasks,          \* Function from task ID to task record
    userTasks,      \* Function from user ID to set of task IDs
    nextTaskId,     \* Counter for generating unique task IDs
    currentUser,    \* Currently authenticated user
    clock,          \* Global clock for timestamps
    sessions        \* Active user sessions

\* Task states enumeration with valid transitions
TaskStates == {"pending", "in_progress", "completed", "cancelled", "blocked"}

\* Priority levels
Priorities == {"low", "medium", "high", "critical"}

\* Valid state transitions
ValidTransitions == {
    <<"pending", "in_progress">>,
    <<"pending", "cancelled">>,
    <<"pending", "blocked">>,
    <<"in_progress", "completed">>,
    <<"in_progress", "cancelled">>,
    <<"in_progress", "blocked">>,
    <<"in_progress", "pending">>,      \* Allow reverting to pending
    <<"blocked", "pending">>,
    <<"blocked", "in_progress">>,
    <<"blocked", "cancelled">>
}


TaskRecord == [
    id: Nat,
    title: STRING,
    description: STRING,
    status: TaskStates,
    priority: {"low", "medium", "high"},
    assignee: Users,
    createdAt: Nat,
    dueDate: Nat \cup {NULL}
]

\* Type invariants
TypeInvariant == 
    /\ tasks \in [Nat -> TaskRecord]
    /\ userTasks \in [Users -> SUBSET Nat]
    /\ nextTaskId \in Nat
    /\ currentUser \in Users \cup {NULL}

\* System initialization
Init ==
    /\ tasks = [i \in {} |-> CHOOSE x : FALSE]  \* Empty function
    /\ userTasks = [u \in Users |-> {}]
    /\ nextTaskId = 1
    /\ currentUser = "NULL"
    /\ clock = 0
    /\ sessions = [u \in Users |-> FALSE]

\* User authentication
Authenticate(user) ==
    /\ user \in Users
    /\ ~sessions[user]  \* User not already logged in
    /\ currentUser' = user
    /\ sessions' = [sessions EXCEPT ![user] = TRUE]
    /\ UNCHANGED <<tasks, userTasks, nextTaskId, clock>>

\* Create a new task
CreateTask(title, description, priority, dueDate) ==
    /\ currentUser # NULL
    /\ nextTaskId <= MaxTasks
    /\ LET newTask == [
           id |-> nextTaskId,
           title |-> title,
           description |-> description,
           status |-> "pending",
           priority |-> priority,
           assignee |-> currentUser,
           createdAt |-> nextTaskId, \* Simplified timestamp
           dueDate |-> dueDate
       ] IN
       /\ tasks' = tasks @@ (nextTaskId :> newTask)
       /\ userTasks' = [userTasks EXCEPT ![currentUser] = @ \cup {nextTaskId}]
       /\ nextTaskId' = nextTaskId + 1
       /\ UNCHANGED currentUser

\* Update task status
UpdateTaskStatus(taskId, newStatus) ==
    /\ currentUser # NULL
    /\ taskId \in DOMAIN tasks
    /\ taskId \in userTasks[currentUser]
    /\ newStatus \in TaskStates
    /\ tasks' = [tasks EXCEPT ![taskId].status = newStatus]
    /\ UNCHANGED <<userTasks, nextTaskId, currentUser>>

\* Delete a task
DeleteTask(taskId) ==
    /\ currentUser # NULL
    /\ taskId \in DOMAIN tasks
    /\ taskId \in userTasks[currentUser]
    /\ tasks' = [id \in (DOMAIN tasks \ {taskId}) |-> tasks[id]]
    /\ userTasks' = [userTasks EXCEPT ![currentUser] = @ \ {taskId}]
    /\ UNCHANGED <<nextTaskId, currentUser>>

\* Safety properties
NoOrphanTasks ==
    \A taskId \in DOMAIN tasks :
        \E user \in Users : taskId \in GetUserTasks(user)

TaskOwnership ==
    \A taskId \in DOMAIN tasks :
        tasks[taskId].assignee \in Users /\
        taskId \in GetUserTasks(tasks[taskId].assignee)

ValidTaskIds ==
    \A taskId \in DOMAIN tasks : 
        /\ taskId < nextTaskId
        /\ taskId >= 1

NoDuplicateTaskIds ==
    \A t1, t2 \in DOMAIN tasks :
        t1 = t2 \/ tasks[t1].id # tasks[t2].id

ValidStateTransitionsInvariant ==
    \A taskId \in DOMAIN tasks :
        tasks[taskId].status \in TaskStates

ConsistentTimestamps ==
    \A taskId \in DOMAIN tasks :
        /\ tasks[taskId].createdAt <= tasks[taskId].updatedAt
        /\ tasks[taskId].updatedAt <= clock

NoCyclicDependencies ==
    LET
        \* Transitive closure of dependencies
        RECURSIVE TransitiveDeps(_)
        TransitiveDeps(taskId) ==
            IF ~TaskExists(taskId) THEN {}
            ELSE LET directDeps == tasks[taskId].dependencies IN
                 directDeps \cup 
                 UNION {TransitiveDeps(dep) : dep \in directDeps}
    IN
    \A taskId \in DOMAIN tasks :
        taskId \notin TransitiveDeps(taskId)

AuthenticationRequired ==
    \* All task operations require authentication
    \A taskId \in DOMAIN tasks :
        tasks[taskId].createdBy \in Users

SafetyInvariant ==
    /\ NoOrphanTasks
    /\ TaskOwnership
    /\ ValidTaskIds
    /\ NoDuplicateTaskIds
    /\ ValidStateTransitionsInvariant
    /\ ConsistentTimestamps
    /\ NoCyclicDependencies
    /\ AuthenticationRequired

\* Next state relation
Next ==
    \/ AdvanceTime
    \/ \E user \in Users : Authenticate(user)
    \/ Logout
    \/ \E t \in Titles, d \in Descriptions, p \in Priorities, 
         u \in Users, dd \in 0..MaxTime \cup {"NULL"},
         tags \in SUBSET {"bug", "feature", "enhancement", "documentation"},
         deps \in SUBSET DOMAIN tasks :
       CreateTask(t, d, p, u, dd, tags, deps)
    \/ \E taskId \in DOMAIN tasks, newStatus \in TaskStates :
       UpdateTaskStatus(taskId, newStatus)
    \/ \E taskId \in DOMAIN tasks, newPriority \in Priorities :
       UpdateTaskPriority(taskId, newPriority)
    \/ \E taskId \in DOMAIN tasks, newAssignee \in Users :
       ReassignTask(taskId, newAssignee)
    \/ \E taskId \in DOMAIN tasks, t \in Titles, 
         d \in Descriptions, dd \in 0..MaxTime \cup {"NULL"} :
       UpdateTaskDetails(taskId, t, d, dd)
    \/ \E taskId \in DOMAIN tasks : DeleteTask(taskId)
    \/ CheckDependencies
    \/ \E taskIds \in SUBSET DOMAIN tasks, newStatus \in TaskStates :
       taskIds # {} /\ BulkUpdateStatus(taskIds, newStatus)

\* Properties to check
THEOREM TypeCorrectness == Spec => []TypeInvariant
THEOREM SafetyHolds == Spec => []SafetyInvariant
THEOREM LivenessHolds == Spec => (EventualCompletion /\ FairProgress)
THEOREM NoDeadlock == Spec => []<>Next
THEOREM Termination == Spec => <>(\A taskId \in DOMAIN tasks : 
                                    tasks[taskId].status \in {"completed", "cancelled"})
=============================================================================

\* Model configuration for TaskManagementImproved module
SPECIFICATION Spec

\* Constants definition
CONSTANTS
    Users = {alice, bob, charlie}
    MaxTasks = 5
    MaxTime = 20
    Titles = {task1, task2, task3, task4, task5}
    Descriptions = {desc1, desc2, desc3}

\* Model values for special constants
CONSTANT
    NULL = NULL
    EMPTY_STRING = EMPTY_STRING

\* Initial state constraint
CONSTRAINT
    /\ nextTaskId <= MaxTasks + 1
    /\ clock <= MaxTime
    /\ Cardinality(DOMAIN tasks) <= MaxTasks

\* State space reduction (optional, for faster checking)
ACTION_CONSTRAINT
    \* Limit number of active sessions
    /\ Cardinality({u \in Users : sessions[u] = TRUE}) <= 2
    \* Prevent creating too many tasks at once
    /\ nextTaskId <= MaxTasks

\* Invariants to check
INVARIANT TypeInvariant
INVARIANT SafetyInvariant
INVARIANT NoOrphanTasks
INVARIANT TaskOwnership
INVARIANT ValidTaskIds
INVARIANT NoDuplicateTaskIds
INVARIANT ValidStateTransitionsInvariant
INVARIANT ConsistentTimestamps
INVARIANT NoCyclicDependencies
INVARIANT AuthenticationRequired

\* Properties to check
PROPERTY EventualCompletion
PROPERTY FairProgress
PROPERTY EventualUnblocking
PROPERTY EventualAuthentication
PROPERTY NoStarvation

\* Check for deadlocks
CHECK_DEADLOCK TRUE

\* View for debugging (optional)
VIEW <<nextTaskId, Cardinality(DOMAIN tasks), clock>>

\* Alias for better state visualization
ALIAS TaskSummary == [
    totalTasks |-> Cardinality(DOMAIN tasks),
    pendingTasks |-> Cardinality({t \in DOMAIN tasks : tasks[t].status = "pending"}),
    inProgressTasks |-> Cardinality({t \in DOMAIN tasks : tasks[t].status = "in_progress"}),
    completedTasks |-> Cardinality({t \in DOMAIN tasks : tasks[t].status = "completed"}),
    blockedTasks |-> Cardinality({t \in DOMAIN tasks : tasks[t].status = "blocked"}),
    activeSessions |-> {u \in Users : sessions[u] = TRUE},
    currentUser |-> currentUser,
    time |-> clock
]

# Run the model checker
java -jar tla2tools.jar -config TaskManagement.cfg TaskManagement.tla

Using this TLA+ specification for a task management system, implement a Go REST API 
that satisfies all the specified behaviors and properties.

[Include the full TLA+ specification here]

Requirements:
1. Implement all actions defined in the TLA+ spec
2. Ensure all safety properties hold at runtime
3. Use Go's type system to enforce the data structures
4. Include comprehensive tests that cover the execution traces
5. Add proper error handling for violation of preconditions

Focus on behavioral correctness over performance optimization.

// Package domain contains the core business entities matching the TLA+ specification
package domain

import (
	"fmt"
	"time"
)

// TaskID represents a unique task identifier (maps to TLA+ task ID)
type TaskID int

// UserID represents a user identifier (maps to TLA+ Users set)
type UserID string

// TaskStatus represents the state of a task (maps to TLA+ TaskStates)
type TaskStatus string

const (
	StatusPending    TaskStatus = "pending"
	StatusInProgress TaskStatus = "in_progress"
	StatusCompleted  TaskStatus = "completed"
	StatusCancelled  TaskStatus = "cancelled"
	StatusBlocked    TaskStatus = "blocked"
)

// Priority represents task priority levels (maps to TLA+ Priorities)
type Priority string

const (
	PriorityLow      Priority = "low"
	PriorityMedium   Priority = "medium"
	PriorityHigh     Priority = "high"
	PriorityCritical Priority = "critical"
)

// Tag represents task categories (maps to TLA+ tags subset)
type Tag string

const (
	TagBug           Tag = "bug"
	TagFeature       Tag = "feature"
	TagEnhancement   Tag = "enhancement"
	TagDocumentation Tag = "documentation"
)

// Task represents a task entity (maps to TLA+ task record)
type Task struct {
	ID           TaskID            `json:"id"`
	Title        string            `json:"title"`
	Description  string            `json:"description"`
	Status       TaskStatus        `json:"status"`
	Priority     Priority          `json:"priority"`
	Assignee     UserID            `json:"assignee"`
	CreatedBy    UserID            `json:"created_by"`
	CreatedAt    time.Time         `json:"created_at"`
	UpdatedAt    time.Time         `json:"updated_at"`
	DueDate      *time.Time        `json:"due_date,omitempty"`
	Tags         []Tag             `json:"tags"`
	Dependencies map[TaskID]bool   `json:"dependencies"`
}

// ValidTransition represents a valid state transition (maps to TLA+ ValidTransitions)
type ValidTransition struct {
	From TaskStatus
	To   TaskStatus
}

// ValidTransitions defines all allowed state transitions
var ValidTransitions = map[ValidTransition]bool{
	{StatusPending, StatusInProgress}:    true,
	{StatusPending, StatusCancelled}:     true,
	{StatusPending, StatusBlocked}:       true,
	{StatusInProgress, StatusCompleted}:  true,
	{StatusInProgress, StatusCancelled}:  true,
	{StatusInProgress, StatusBlocked}:    true,
	{StatusInProgress, StatusPending}:    true, // Allow reverting
	{StatusBlocked, StatusPending}:       true,
	{StatusBlocked, StatusInProgress}:    true,
	{StatusBlocked, StatusCancelled}:     true,
}

// IsValidTransition checks if a state transition is valid (maps to TLA+ IsValidTransition)
func IsValidTransition(from, to TaskStatus) bool {
	return ValidTransitions[ValidTransition{From: from, To: to}]
}

// CanDelete checks if a task can be deleted (only completed or cancelled)
func (t *Task) CanDelete() bool {
	return t.Status == StatusCompleted || t.Status == StatusCancelled
}

// IsBlocked checks if task should be blocked based on dependencies
func (t *Task) IsBlocked(allTasks map[TaskID]*Task) bool {
	if len(t.Dependencies) == 0 {
		return false
	}
	
	for depID := range t.Dependencies {
		if dep, exists := allTasks[depID]; exists {
			if dep.Status != StatusCompleted {
				return true
			}
		}
	}
	return false
}

// ShouldUnblock checks if a blocked task can be unblocked
func (t *Task) ShouldUnblock(allTasks map[TaskID]*Task) bool {
	if t.Status != StatusBlocked {
		return false
	}
	
	for depID := range t.Dependencies {
		if dep, exists := allTasks[depID]; exists {
			if dep.Status != StatusCompleted {
				return false
			}
		}
	}
	return true
}

// Validate performs domain validation on the task
func (t *Task) Validate() error {
	if t.Title == "" {
		return fmt.Errorf("task title cannot be empty")
	}
	if t.Description == "" {
		return fmt.Errorf("task description cannot be empty")
	}
	if !isValidStatus(t.Status) {
		return fmt.Errorf("invalid task status: %s", t.Status)
	}
	if !isValidPriority(t.Priority) {
		return fmt.Errorf("invalid task priority: %s", t.Priority)
	}
	if t.Assignee == "" {
		return fmt.Errorf("task must have an assignee")
	}
	if t.CreatedBy == "" {
		return fmt.Errorf("task must have a creator")
	}
	if t.CreatedAt.After(t.UpdatedAt) {
		return fmt.Errorf("created time cannot be after updated time")
	}
	for _, tag := range t.Tags {
		if !isValidTag(tag) {
			return fmt.Errorf("invalid tag: %s", tag)
		}
	}
	return nil
}

func isValidStatus(status TaskStatus) bool {
	switch status {
	case StatusPending, StatusInProgress, StatusCompleted, StatusCancelled, StatusBlocked:
		return true
	default:
		return false
	}
}

func isValidPriority(priority Priority) bool {
	switch priority {
	case PriorityLow, PriorityMedium, PriorityHigh, PriorityCritical:
		return true
	default:
		return false
	}
}

func isValidTag(tag Tag) bool {
	switch tag {
	case TagBug, TagFeature, TagEnhancement, TagDocumentation:
		return true
	default:
		return false
	}
}

// Package usecase implements the TLA+ actions as use cases
package usecase

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"time"
	
	"github.com/bhatti/sample-task-management/internal/domain"
	"github.com/bhatti/sample-task-management/internal/repository"
)

// TaskUseCase implements task-related TLA+ actions
type TaskUseCase struct {
	uow              repository.UnitOfWork
	invariantChecker InvariantChecker
}

// InvariantChecker interface for runtime invariant validation
type InvariantChecker interface {
	CheckAllInvariants(state *domain.SystemState) error
	CheckTaskInvariants(task *domain.Task, state *domain.SystemState) error
	CheckTransitionInvariant(from, to domain.TaskStatus) error
}

// NewTaskUseCase creates a new task use case
func NewTaskUseCase(uow repository.UnitOfWork, checker InvariantChecker) *TaskUseCase {
	return &TaskUseCase{
		uow:              uow,
		invariantChecker: checker,
	}
}

// Authenticate implements TLA+ Authenticate action
func (uc *TaskUseCase) Authenticate(userID domain.UserID) (*domain.Session, error) {
	// Preconditions from TLA+:
	// - user \in Users
	// - ~sessions[user]
	
	user, err := uc.uow.Users().GetUser(userID)
	if err != nil {
		return nil, fmt.Errorf("user not found: %w", err)
	}
	
	// Check if user already has an active session
	existingSession, _ := uc.uow.Sessions().GetSessionByUser(userID)
	if existingSession != nil && existingSession.IsValid() {
		return nil, fmt.Errorf("user %s already has an active session", userID)
	}
	
	// Create new session
	token := generateToken()
	session := &domain.Session{
		UserID:    user.ID,
		Token:     token,
		Active:    true,
		CreatedAt: time.Now(),
		ExpiresAt: time.Now().Add(24 * time.Hour),
	}
	
	// Update state
	if err := uc.uow.Sessions().CreateSession(session); err != nil {
		return nil, fmt.Errorf("failed to create session: %w", err)
	}
	
	if err := uc.uow.SystemState().SetCurrentUser(&userID); err != nil {
		return nil, fmt.Errorf("failed to set current user: %w", err)
	}
	
	// Check invariants
	state, _ := uc.uow.SystemState().GetSystemState()
	if err := uc.invariantChecker.CheckAllInvariants(state); err != nil {
		uc.uow.Rollback()
		return nil, fmt.Errorf("invariant violation: %w", err)
	}
	
	return session, nil
}


// CreateTask implements TLA+ CreateTask action
func (uc *TaskUseCase) CreateTask(
	title, description string,
	priority domain.Priority,
	assignee domain.UserID,
	dueDate *time.Time,
	tags []domain.Tag,
	dependencies []domain.TaskID,
) (*domain.Task, error) {
	// Preconditions from TLA+:
	// - currentUser # NULL
	// - currentUser \in Users
	// - nextTaskId <= MaxTasks
	// - deps \subseteq DOMAIN tasks
	// - \A dep \in deps : tasks[dep].status # "cancelled"
	
	currentUser, err := uc.uow.SystemState().GetCurrentUser()
	if err != nil || currentUser == nil {
		return nil, fmt.Errorf("authentication required")
	}
	
	// Check max tasks limit
	nextID, err := uc.uow.SystemState().GetNextTaskID()
	if err != nil {
		return nil, fmt.Errorf("failed to get next task ID: %w", err)
	}
	
	if nextID > domain.MaxTasks {
		return nil, fmt.Errorf("maximum number of tasks (%d) reached", domain.MaxTasks)
	}
	
	// Validate dependencies
	allTasks, err := uc.uow.Tasks().GetAllTasks()
	if err != nil {
		return nil, fmt.Errorf("failed to get tasks: %w", err)
	}
	
	depMap := make(map[domain.TaskID]bool)
	for _, depID := range dependencies {
		depTask, exists := allTasks[depID]
		if !exists {
			return nil, fmt.Errorf("dependency task %d does not exist", depID)
		}
		if depTask.Status == domain.StatusCancelled {
			return nil, fmt.Errorf("cannot depend on cancelled task %d", depID)
		}
		depMap[depID] = true
	}
	
	// Check for cyclic dependencies
	if err := uc.checkCyclicDependencies(nextID, depMap, allTasks); err != nil {
		return nil, err
	}
	
	// Determine initial status based on dependencies
	status := domain.StatusPending
	if len(dependencies) > 0 {
		// Check if all dependencies are completed
		allCompleted := true
		for depID := range depMap {
			if allTasks[depID].Status != domain.StatusCompleted {
				allCompleted = false
				break
			}
		}
		if !allCompleted {
			status = domain.StatusBlocked
		}
	}
	
	// Create task
	task := &domain.Task{
		ID:           nextID,
		Title:        title,
		Description:  description,
		Status:       status,
		Priority:     priority,
		Assignee:     assignee,
		CreatedBy:    *currentUser,
		CreatedAt:    time.Now(),
		UpdatedAt:    time.Now(),
		DueDate:      dueDate,
		Tags:         tags,
		Dependencies: depMap,
	}
	
	// Validate task
	if err := task.Validate(); err != nil {
		return nil, fmt.Errorf("task validation failed: %w", err)
	}
	
	// Save task
	if err := uc.uow.Tasks().CreateTask(task); err != nil {
		return nil, fmt.Errorf("failed to create task: %w", err)
	}
	
	// Increment next task ID
	if _, err := uc.uow.SystemState().IncrementNextTaskID(); err != nil {
		return nil, fmt.Errorf("failed to increment task ID: %w", err)
	}
	
	// Check invariants
	state, _ := uc.uow.SystemState().GetSystemState()
	if err := uc.invariantChecker.CheckAllInvariants(state); err != nil {
		uc.uow.Rollback()
		return nil, fmt.Errorf("invariant violation after task creation: %w", err)
	}
	
	return task, nil
}

// UpdateTaskStatus implements TLA+ UpdateTaskStatus action
func (uc *TaskUseCase) UpdateTaskStatus(taskID domain.TaskID, newStatus domain.TaskStatus) error {
	// Preconditions from TLA+:
	// - currentUser # NULL
	// - TaskExists(taskId)
	// - taskId \in GetUserTasks(currentUser)
	// - IsValidTransition(tasks[taskId].status, newStatus)
	// - newStatus = "in_progress" => all dependencies completed
	
	currentUser, err := uc.uow.SystemState().GetCurrentUser()
	if err != nil || currentUser == nil {
		return fmt.Errorf("authentication required")
	}
	
	task, err := uc.uow.Tasks().GetTask(taskID)
	if err != nil {
		return fmt.Errorf("task not found: %w", err)
	}
	
	// Check user owns the task
	userTasks, err := uc.uow.SystemState().GetUserTasks(*currentUser)
	if err != nil {
		return fmt.Errorf("failed to get user tasks: %w", err)
	}
	
	hasTask := false
	for _, id := range userTasks {
		if id == taskID {
			hasTask = true
			break
		}
	}
	
	if !hasTask {
		return fmt.Errorf("user does not have access to task %d", taskID)
	}
	
	// Check valid transition
	if !domain.IsValidTransition(task.Status, newStatus) {
		return fmt.Errorf("invalid transition from %s to %s", task.Status, newStatus)
	}
	
	// Check dependencies if moving to in_progress
	if newStatus == domain.StatusInProgress {
		allTasks, _ := uc.uow.Tasks().GetAllTasks()
		for depID := range task.Dependencies {
			if depTask, exists := allTasks[depID]; exists {
				if depTask.Status != domain.StatusCompleted {
					return fmt.Errorf("cannot start task: dependency %d is not completed", depID)
				}
			}
		}
	}
	
	// Update status
	task.Status = newStatus
	task.UpdatedAt = time.Now()
	
	if err := uc.uow.Tasks().UpdateTask(task); err != nil {
		return fmt.Errorf("failed to update task: %w", err)
	}
	
	// Check invariants
	state, _ := uc.uow.SystemState().GetSystemState()
	if err := uc.invariantChecker.CheckAllInvariants(state); err != nil {
		uc.uow.Rollback()
		return fmt.Errorf("invariant violation: %w", err)
	}
	
	return nil
}
...

Generate Go tests that verify the implementation satisfies the TLA+ specification.
Create test cases that:

1. Test all TLA+ actions with valid preconditions
2. Test safety property violations
3. Test edge cases from the TLA+ model boundary conditions
4. Use property-based testing where appropriate

Include tests that would catch the execution traces TLA+ model checker explores.

// Package property implements property-based tests for TLA+ invariants
package property

import (
	"math/rand"
	"testing"
	"time"

	"github.com/bhatti/sample-task-management/internal/domain"
	"github.com/bhatti/sample-task-management/internal/infrastructure/memory"
	"github.com/bhatti/sample-task-management/internal/usecase"
	"github.com/bhatti/sample-task-management/pkg/invariants"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestInvariantsHoldAfterOperations verifies invariants hold after each operation
func TestInvariantsHoldAfterOperations(t *testing.T) {
	repo := memory.NewMemoryRepository()
	uow := memory.NewMemoryUnitOfWork(repo)
	checker := invariants.NewInvariantChecker()
	uc := usecase.NewTaskUseCase(uow, checker)

	// Setup initial users
	users := []domain.UserID{"alice", "bob", "charlie"}
	for _, userID := range users {
		user := &domain.User{
			ID:       userID,
			Name:     string(userID),
			Email:    string(userID) + "@example.com",
			JoinedAt: time.Now(),
		}
		require.NoError(t, repo.CreateUser(user))
	}

	// Property: Invariants hold after authentication
	t.Run("InvariantsAfterAuthentication", func(t *testing.T) {
		for _, userID := range users {
			session, err := uc.Authenticate(userID)
			assert.NoError(t, err)
			assert.NotNil(t, session)

			state, _ := repo.GetSystemState()
			assert.NoError(t, checker.CheckAllInvariants(state))

			// Cleanup
			_ = uc.Logout(userID)
		}
	})

	// Property: Invariants hold after task creation
	t.Run("InvariantsAfterTaskCreation", func(t *testing.T) {
		uc.Authenticate("alice")

		for i := 0; i < 10; i++ {
			task, err := uc.CreateTask(
				"Task "+string(rune(i)),
				"Description",
				randomPriority(),
				randomUser(users),
				randomDueDate(),
				randomTags(),
				[]domain.TaskID{}, // No dependencies initially
			)

			assert.NoError(t, err)
			assert.NotNil(t, task)

			state, _ := repo.GetSystemState()
			assert.NoError(t, checker.CheckAllInvariants(state))
		}
	})

	// Property: Invariants hold after status transitions
	t.Run("InvariantsAfterStatusTransitions", func(t *testing.T) {
		uc.Authenticate("alice")

		// Create a task
		task, _ := uc.CreateTask(
			"Test Task",
			"Description",
			domain.PriorityMedium,
			"alice",
			nil,
			[]domain.Tag{domain.TagFeature},
			[]domain.TaskID{},
		)

		// Valid transitions
		validTransitions := []domain.TaskStatus{
			domain.StatusInProgress,
			domain.StatusCompleted,
		}

		for _, status := range validTransitions {
			err := uc.UpdateTaskStatus(task.ID, status)
			if err == nil {
				state, _ := repo.GetSystemState()
				assert.NoError(t, checker.CheckAllInvariants(state))
			}
		}
	})

	// Property: No cyclic dependencies can be created
	t.Run("NoCyclicDependencies", func(t *testing.T) {
		uc.Authenticate("alice")

		// Create tasks with potential cycles
		task1, _ := uc.CreateTask("Task1", "Desc", domain.PriorityLow, "alice", nil, nil, []domain.TaskID{})
		task2, _ := uc.CreateTask("Task2", "Desc", domain.PriorityLow, "alice", nil, nil, []domain.TaskID{task1.ID})
		task3, _ := uc.CreateTask("Task3", "Desc", domain.PriorityLow, "alice", nil, nil, []domain.TaskID{task2.ID})

		// Attempting to create a cycle should fail
		_, err := uc.CreateTask("Task4", "Desc", domain.PriorityLow, "alice", nil, nil,
			[]domain.TaskID{task3.ID, task1.ID}) // This would create a cycle
		assert.NoError(t, err)

		// Even if it doesn't fail explicitly, invariants should catch it
		state, _ := repo.GetSystemState()
		assert.NoError(t, checker.CheckAllInvariants(state))
	})
}

// TestTransitionInvariants tests state transition validity
func TestTransitionInvariants(t *testing.T) {
	checker := invariants.NewInvariantChecker()

	// Test all valid transitions
	validTransitions := []struct {
		from domain.TaskStatus
		to   domain.TaskStatus
	}{
		{domain.StatusPending, domain.StatusInProgress},
		{domain.StatusPending, domain.StatusCancelled},
		{domain.StatusInProgress, domain.StatusCompleted},
		{domain.StatusInProgress, domain.StatusCancelled},
		{domain.StatusBlocked, domain.StatusPending},
		{domain.StatusBlocked, domain.StatusCancelled},
	}

	for _, trans := range validTransitions {
		t.Run(string(trans.from)+"_to_"+string(trans.to), func(t *testing.T) {
			err := checker.CheckTransitionInvariant(trans.from, trans.to)
			assert.NoError(t, err)
		})
	}

	// Test invalid transitions
	invalidTransitions := []struct {
		from domain.TaskStatus
		to   domain.TaskStatus
	}{
		{domain.StatusCompleted, domain.StatusPending},
		{domain.StatusCompleted, domain.StatusInProgress},
		{domain.StatusCancelled, domain.StatusInProgress},
		{domain.StatusPending, domain.StatusCompleted}, // Must go through in_progress
	}

	for _, trans := range invalidTransitions {
		t.Run("Invalid_"+string(trans.from)+"_to_"+string(trans.to), func(t *testing.T) {
			err := checker.CheckTransitionInvariant(trans.from, trans.to)
			assert.Error(t, err)
		})
	}
}

// TestPropertyTaskOwnership verifies task ownership invariants
func TestPropertyTaskOwnership(t *testing.T) {
	repo := memory.NewMemoryRepository()
	uow := memory.NewMemoryUnitOfWork(repo)
	checker := invariants.NewInvariantChecker()
	uc := usecase.NewTaskUseCase(uow, checker)

	// Setup users
	users := []domain.UserID{"alice", "bob"}
	for _, userID := range users {
		user := &domain.User{
			ID:       userID,
			Name:     string(userID),
			Email:    string(userID) + "@example.com",
			JoinedAt: time.Now(),
		}
		repo.CreateUser(user)
	}

	// Property: Task reassignment maintains ownership invariants
	t.Run("ReassignmentMaintainsOwnership", func(t *testing.T) {
		uc.Authenticate("alice")

		// Create task assigned to Alice
		task, err := uc.CreateTask(
			"Test Task",
			"Description",
			domain.PriorityHigh,
			"alice",
			nil,
			[]domain.Tag{domain.TagBug},
			[]domain.TaskID{},
		)
		require.NoError(t, err)

		// Check initial ownership
		state, _ := repo.GetSystemState()
		assert.NoError(t, checker.CheckAllInvariants(state))

		aliceTasks := state.GetUserTasks("alice")
		assert.Contains(t, aliceTasks, task.ID)

		// Reassign to Bob
		err = uc.ReassignTask(task.ID, "bob")
		require.NoError(t, err)

		// Check ownership after reassignment
		state, _ = repo.GetSystemState()
		assert.NoError(t, checker.CheckAllInvariants(state))

		aliceTasks = state.GetUserTasks("alice")
		bobTasks := state.GetUserTasks("bob")
		assert.NotContains(t, aliceTasks, task.ID)
		assert.Contains(t, bobTasks, task.ID)
	})
}

// TestPropertyConcurrentOperations tests invariants under concurrent operations
func TestPropertyConcurrentOperations(t *testing.T) {
	repo := memory.NewMemoryRepository()
	uow := memory.NewMemoryUnitOfWork(repo)
	checker := invariants.NewInvariantChecker()

	// Setup users
	users := []domain.UserID{"user1", "user2", "user3"}
	for _, userID := range users {
		user := &domain.User{
			ID:       userID,
			Name:     string(userID),
			Email:    string(userID) + "@example.com",
			JoinedAt: time.Now(),
		}
		repo.CreateUser(user)
	}

	// Run concurrent operations
	done := make(chan bool, len(users))

	for _, userID := range users {
		go func(uid domain.UserID) {
			uc := usecase.NewTaskUseCase(uow, checker)

			// Authenticate
			uc.Authenticate(uid)

			// Create multiple tasks
			for i := 0; i < 5; i++ {
				uc.CreateTask(
					"Task",
					"Description",
					randomPriority(),
					uid,
					nil,
					randomTags(),
					[]domain.TaskID{},
				)

				// Random delay
				time.Sleep(time.Duration(rand.Intn(10)) * time.Millisecond)
			}

			done <- true
		}(userID)
	}

	// Wait for all goroutines
	for i := 0; i < len(users); i++ {
		<-done
	}

	// Check invariants after concurrent operations
	state, _ := repo.GetSystemState()
	assert.NoError(t, checker.CheckAllInvariants(state))
}

// Helper functions

func randomPriority() domain.Priority {
	priorities := []domain.Priority{
		domain.PriorityLow,
		domain.PriorityMedium,
		domain.PriorityHigh,
		domain.PriorityCritical,
	}
	return priorities[rand.Intn(len(priorities))]
}

func randomUser(users []domain.UserID) domain.UserID {
	return users[rand.Intn(len(users))]
}

func randomDueDate() *time.Time {
	if rand.Float32() < 0.5 {
		return nil
	}
	due := time.Now().Add(time.Duration(rand.Intn(30)) * 24 * time.Hour)
	return &due
}

func randomTags() []domain.Tag {
	allTags := []domain.Tag{
		domain.TagBug,
		domain.TagFeature,
		domain.TagEnhancement,
		domain.TagDocumentation,
	}

	numTags := rand.Intn(len(allTags) + 1)
	if numTags == 0 {
		return nil
	}

	tags := make([]domain.Tag, 0, numTags)
	used := make(map[domain.Tag]bool)

	for len(tags) < numTags {
		tag := allTags[rand.Intn(len(allTags))]
		if !used[tag] {
			tags = append(tags, tag)
			used[tag] = true
		}
	}

	return tags
}

\* Concurrent task updates with conflict resolution
ConcurrentUpdateTask(taskId, newStatus, version) ==
    /\ currentUser # NULL
    /\ taskId \in DOMAIN tasks
    /\ taskId \in userTasks[currentUser]
    /\ tasks[taskId].version = version  \* Optimistic concurrency control
    /\ tasks' = [tasks EXCEPT ![taskId] = [
                     @ EXCEPT 
                     !.status = newStatus,
                     !.version = @ + 1,
                     !.lastModified = currentUser
                 ]]
    /\ UNCHANGED <<userTasks, nextTaskId, currentUser>>

Implement optimistic concurrency control for the task updates based on this 
TLA+ specification. Include version tracking and conflict detection.

\* Business rule: High priority tasks cannot be cancelled directly
ValidStatusTransition(currentStatus, newStatus, priority) ==
    \/ newStatus = currentStatus
    \/ /\ currentStatus = "pending" 
       /\ newStatus \in {"in_progress", "cancelled"}
    \/ /\ currentStatus = "in_progress"
       /\ newStatus \in {"completed", "pending"}
    \/ /\ currentStatus = "in_progress"
       /\ newStatus = "cancelled"
       /\ priority # "high"  \* High priority tasks cannot be cancelled

\* Start with basic CRUD
Init, CreateTask, UpdateTask, DeleteTask

\* Add complexity incrementally  
Authentication, Authorization, Concurrency, Business Rules

\* This property might fail, revealing a design issue
ConsistencyProperty == 
    \A user \in Users:
        \A taskId \in userTasks[user]:
            /\ taskId \in DOMAIN tasks
            /\ tasks[taskId].assignee = user
            /\ tasks[taskId].status # "deleted"  \* Soft delete consideration

# TLA+ finds this counterexample:
# Step 1: User1 creates Task1
# Step 2: User1 deletes Task1  
# Step 3: User2 creates Task2 (gets same ID due to reuse)
# Step 4: User1 tries to update Task1 -> Security violation!

Generate API documentation from this TLA+ specification that includes:
1. Endpoint descriptions derived from TLA+ actions
2. Request/response schemas from TLA+ data structures  
3. Error conditions from TLA+ preconditions
4. Behavioral guarantees from TLA+ properties

\* Wrong - this looks like pseudocode
CreateTask == 
    if currentUser != null then
        task = new Task()

\* Right - mathematical specification  
CreateTask == 
    /\ currentUser # NULL
    /\ tasks' = tasks @@ (nextTaskId :> newTask)

# ~/.claude/tla-prompts/implementation.md
Implement [language] code that satisfies this TLA+ specification:

[SPEC]

Requirements:
- All TLA+ actions become functions/methods
- All preconditions become runtime checks  
- All data structures match TLA+ types
- Include comprehensive tests covering specification traces

--------------------------- MODULE [ModuleName] ---------------------------
EXTENDS Integers, Sequences, FiniteSets

CONSTANTS [Constants]

VARIABLES [StateVariables]

[TypeDefinitions]

Init == [InitialConditions]

[Actions]

Next == [ActionDisjunction]

Spec == Init /\ [][Next]_[StateVariables]

[SafetyProperties]

[LivenessProperties]

=============================================================================

Explain how you would implement real-time task updates using WebSockets. 
What are the trade-offs between Socket.io and native WebSockets?
What state management challenges should I consider?

1. First: "Create a basic task model with CRUD operations"
2. Then: "Add validation and error handling"
3. Then: "Add authentication and authorization"
4. Finally: "Add real-time updates and notifications"

Review the task API implementation for:
- Security vulnerabilities
- Performance issues
- Code style consistency
- Missing error cases

Be critical and suggest improvements.