oikos/internal/mcp/server.go

// Package mcp implements the Oikos MCP interface (plan R3-10).
// Uses the official MCP Go SDK with Streamable HTTP transport.
package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"github.com/dtoro/oikos/internal/db"
	"github.com/google/jsonschema-go/jsonschema"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// prop is one input-schema property (name → type + description).
type prop struct {
	name, typ, desc string
}

// objSchema builds an "object" JSON Schema from a list of properties. The
// MCP SDK requires every tool to declare an object input schema so tools
// are self-describing to the agent; a nil schema panics at registration.
func objSchema(props ...prop) *jsonschema.Schema {
	s := &jsonschema.Schema{Type: "object", Properties: map[string]*jsonschema.Schema{}}
	for _, p := range props {
		s.Properties[p.name] = &jsonschema.Schema{Type: p.typ, Description: p.desc}
	}
	return s
}

// NewHandler creates an http.Handler that serves the Oikos MCP server.
func NewHandler(pool *db.Pool, token string) http.Handler {
	s := newServer(pool)
	handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
		if token != "" {
			if r.Header.Get("Authorization") != "Bearer "+token {
				return nil
			}
		}
		return s
	}, nil)
	return handler
}

func newServer(pool *db.Pool) *mcp.Server {
	s := mcp.NewServer(&mcp.Implementation{Name: "oikos", Version: "dev"}, &mcp.ServerOptions{
		Logger: slog.Default(),
	})

	// All tools use the untyped handler (s.AddTool) for simplicity.
	// Arguments are accessed via req.Parameters.Arguments.(map[string]any).

	s.AddTool(&mcp.Tool{Name: "get_entity", Description: "Get an entity by slug or UUID",
		InputSchema: objSchema(prop{"slug_or_id", "string", "Entity slug (e.g. host:hubris) or UUID"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		idOrSlug, _ := args["slug_or_id"].(string)
		return queryEntity(ctx, pool, idOrSlug), nil
	})

	s.AddTool(&mcp.Tool{Name: "list_entities", Description: "List entities filtered by type, state, or search",
		InputSchema: objSchema(
			prop{"type", "string", "Filter by entity type"},
			prop{"state", "string", "Filter by lifecycle state"},
			prop{"q", "string", "Substring match on slug or name"},
			prop{"limit", "integer", "Max rows (default 50)"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		limit := int(getFloat(args, "limit", 50))
		return queryRows(ctx, pool, `
			SELECT e.slug, e.type, e.name, e.state, e.version, e.created_at, e.updated_at
			FROM entities e
			WHERE ($1::text IS NULL OR e.type = $1)
			  AND ($2::text IS NULL OR e.state = $2)
			  AND ($3::text IS NULL OR e.slug ILIKE '%'||$3||'%' OR e.name ILIKE '%'||$3||'%')
			ORDER BY e.slug LIMIT $4`,
			nStr(args["type"]), nStr(args["state"]), nStr(args["q"]), limit), nil
	})

	s.AddTool(&mcp.Tool{Name: "get_relations", Description: "Get relationships for an entity",
		InputSchema: objSchema(prop{"entity_id", "string", "Entity slug"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		slug, _ := args["entity_id"].(string)
		return queryRows(ctx, pool, `
			SELECT r.type, src.slug AS source, tgt.slug AS target
			FROM relationships r
			JOIN entities src ON src.id = r.source_id
			JOIN entities tgt ON tgt.id = r.target_id
			WHERE (src.slug = $1 OR tgt.slug = $1) AND r.valid_to IS NULL
			ORDER BY r.type`, slug), nil
	})

	s.AddTool(&mcp.Tool{Name: "get_blast_radius", Description: "Find entities affected if this entity goes down",
		InputSchema: objSchema(
			prop{"entity_id", "string", "Entity slug"},
			prop{"depth", "integer", "Traversal depth (default 3)"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		slug, _ := args["entity_id"].(string)
		depth := int(getFloat(args, "depth", 3))
		return queryRows(ctx, pool,
			"SELECT slug, CAST(depth AS int) FROM blast_radius((SELECT id FROM entities WHERE slug = $1), $2)",
			slug, depth), nil
	})

	s.AddTool(&mcp.Tool{Name: "get_health_summary", Description: "Current fleet health summary",
		InputSchema: objSchema(),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return queryRows(ctx, pool, `
			SELECT e.slug, e.type, st.health, st.last_check_at
			FROM entity_status st JOIN entities e ON e.id = st.entity_id
			ORDER BY e.slug`), nil
	})

	s.AddTool(&mcp.Tool{Name: "get_audit_trail", Description: "Query the audit log",
		InputSchema: objSchema(prop{"entity_id", "string", "Filter by affected entity UUID"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		return queryRows(ctx, pool, `
			SELECT id, ts, actor_type, action, entity_id::text, method, path, correlation_id
			FROM audit_log
			WHERE ($1::text IS NULL OR entity_id::text = $1)
			ORDER BY ts DESC LIMIT 50`, nStr(args["entity_id"])), nil
	})

	s.AddTool(&mcp.Tool{Name: "search_knowledge", Description: "Full-text search across documentation",
		InputSchema: objSchema(prop{"query", "string", "Search terms"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		q, _ := args["query"].(string)
		return queryRows(ctx, pool, `
			SELECT id, title, LEFT(content, 500) AS preview
			FROM knowledge_entities
			WHERE title ILIKE '%'||$1||'%' OR content ILIKE '%'||$1||'%'
			ORDER BY title LIMIT 20`, q), nil
	})

	s.AddTool(&mcp.Tool{Name: "query_metrics", Description: "Query time-series metrics",
		InputSchema: objSchema(prop{"hours", "integer", "Look-back window in hours (default 24)"}),
	}, func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args := argsMap(req)
		hours := int(getFloat(args, "hours", 24))
		return queryRows(ctx, pool, `
			SELECT time_bucket('1 hour', ts) AS bucket,
			       entity_id::text, metric,
			       ROUND(avg(value)::numeric, 2) AS avg,
			       ROUND(min(value)::numeric, 2) AS min,
			       ROUND(max(value)::numeric, 2) AS max
			FROM metric_samples
			WHERE ts > now() - make_interval(hours => $1)
			GROUP BY bucket, entity_id, metric
			ORDER BY bucket DESC LIMIT 100`, hours), nil
	})

	return s
}

// ─── Helpers ──────────────────────────────────────────────────────────

func argsMap(req *mcp.CallToolRequest) map[string]any {
	if req == nil || len(req.Params.Arguments) == 0 {
		return nil
	}
	var m map[string]any
	json.Unmarshal(req.Params.Arguments, &m)
	return m
}

func getFloat(m map[string]any, key string, def float64) float64 {
	if m == nil {
		return def
	}
	switch v := m[key].(type) {
	case float64:
		return v
	case int:
		return float64(v)
	}
	return def
}

func nStr(v any) any {
	if v == nil {
		return nil
	}
	s, _ := v.(string)
	if s == "" {
		return nil
	}
	return s
}

func textResult(s string) *mcp.CallToolResult {
	return &mcp.CallToolResult{
		Content: []mcp.Content{&mcp.TextContent{Text: s}},
	}
}

func queryEntity(ctx context.Context, pool *db.Pool, idOrSlug string) *mcp.CallToolResult {
	var id uuid.UUID
	if u, err := uuid.Parse(idOrSlug); err == nil {
		id = u
	} else {
		pool.QueryRow(ctx, "SELECT id FROM entities WHERE slug = $1", idOrSlug).Scan(&id)
	}
	if id == uuid.Nil {
		return textResult(fmt.Sprintf("entity not found: %s", idOrSlug))
	}
	return queryRows(ctx, pool, `
		SELECT slug, type, name, state, attributes,
		       maintenance_until::text, version, created_at, updated_at
		FROM entities WHERE id = $1`, id)
}

func queryRows(ctx context.Context, pool *db.Pool, query string, args ...any) *mcp.CallToolResult {
	rows, err := pool.Query(ctx, query, args...)
	if err != nil {
		return textResult(fmt.Sprintf("error: %v", err))
	}
	defer rows.Close()

	cols := rows.FieldDescriptions()
	var items []map[string]any
	for rows.Next() {
		vals, err := rows.Values()
		if err != nil {
			continue
		}
		m := make(map[string]any)
		for i, col := range cols {
			m[string(col.Name)] = fmt.Sprintf("%v", vals[i])
		}
		items = append(items, m)
	}
	if err := rows.Err(); err != nil {
		return textResult(fmt.Sprintf("error: %v", err))
	}

	data, _ := json.MarshalIndent(items, "", "  ")
	return textResult(string(data))
}

var _ = pgx.ErrNoRows
var _ = time.Now