oikos/cmd/nomos/main.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintln(os.Stderr, "usage: nomos serve")
		os.Exit(1)
	}

	mcpURL := os.Getenv("NOMOS_MCP_URL")
	if mcpURL == "" {
		mcpURL = "http://localhost:8090/mcp"
	}

	agentSlug := os.Getenv("NOMOS_AGENT_SLUG")
	if agentSlug == "" {
		agentSlug = "agent:nomos"
	}

	switch os.Args[1] {
	case "serve":
		ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
		defer cancel()

		client, err := newMCPClient(mcpURL)
		if err != nil {
			slog.Error("nomos: mcp connect", "url", mcpURL, "error", err)
			os.Exit(1)
		}

		mux := http.NewServeMux()
		mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(200)
			w.Write([]byte("ok"))
		})
		mux.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
			handleQuery(w, r, client, agentSlug, mcpURL)
		})

		addr := os.Getenv("NOMOS_LISTEN")
		if addr == "" {
			addr = ":8092"
		}

		srv := &http.Server{Addr: addr, Handler: mux}
		go func() {
			slog.Info("nomos: gateway listening", "addr", addr, "mcp", mcpURL)
			if err := srv.ListenAndServe(); err != http.ErrServerClosed {
				slog.Error("nomos: serve", "error", err)
			}
		}()

		<-ctx.Done()
		slog.Info("nomos: shutting down")
		srv.Shutdown(context.Background())
		client.close()

	default:
		fmt.Fprintf(os.Stderr, "unknown command: %s\n", os.Args[1])
		os.Exit(1)
	}
}

// handleQuery maps structured queries to MCP tool calls.
func handleQuery(w http.ResponseWriter, r *http.Request, client *mcpClient, agentSlug, mcpURL string) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", 405)
		return
	}

	var req struct {
		Query string `json:"query"`
		Tool  string `json:"tool"`
		Args  map[string]any `json:"args"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "bad request: "+err.Error(), 400)
		return
	}

	start := time.Now()
	var result any
	var err error

	// Direct tool call (structured)
	if req.Tool != "" {
		result, err = client.callTool(req.Tool, req.Args)
	} else {
		// Natural-language-ish query routing
		q := strings.ToLower(req.Query)
		result, err = routeQuery(client, q, agentSlug)
	}

	duration := time.Since(start).Milliseconds()

	if err != nil {
		slog.Error("nomos: query failed", "query", req.Query, "error", err)
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{
			"error":       err.Error(),
			"elapsed_ms":  duration,
			"agent_slug":  agentSlug,
		})
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{
		"result":      result,
		"elapsed_ms":  duration,
		"agent_slug":  agentSlug,
		"mcp_url":     mcpURL,
	})
}

// routeQuery maps natural-language-style queries to MCP tool calls.
func routeQuery(client *mcpClient, query, agentSlug string) (any, error) {
	switch {
	case strings.Contains(query, "depends on") || strings.Contains(query, "depend on"):
		entity := extractEntity(query)
		if entity == "" {
			return nil, fmt.Errorf("no entity found in query: %s", query)
		}
		return client.callTool("get_blast_radius", map[string]any{
			"entity_id": entity,
		})

	case strings.Contains(query, "what is") || strings.Contains(query, "describe"):
		entity := extractEntity(query)
		if entity == "" {
			entity = query
		}
		return client.callTool("get_entity", map[string]any{
			"slug_or_id": entity,
		})

	case strings.Contains(query, "health") || strings.Contains(query, "status"):
		return client.callTool("get_health_summary", map[string]any{})

	case strings.Contains(query, "restart") || strings.Contains(query, "reload"):
		entity := extractEntity(query)
		if entity == "" {
			return nil, fmt.Errorf("no entity found in query: %s", query)
		}
		return client.callTool("request_execution", map[string]any{
			"target": entity,
			"action": "restart",
		})

	case strings.Contains(query, "what can you do") || strings.Contains(query, "help"):
		return client.callTool("tools/list", nil)

	default:
		return client.callTool("get_health_summary", map[string]any{})
	}
}

// extractEntity guesses an entity slug from a query.
func extractEntity(query string) string {
	for _, slug := range []string{"authentik", "caddy", "vaultwarden", "gitea", "immich"} {
		if strings.Contains(query, slug) {
			return "service:" + slug
		}
	}
	if strings.Contains(query, "mac-mini") {
		return "host:mac-mini"
	}
	if strings.Contains(query, "hubris") {
		return "host:hubris"
	}
	return ""
}

// ─── MCP Streamable HTTP client ────────────────────────────────────────

type mcpClient struct {
	baseURL    string
	sessionID  string
	http       *http.Client
	nextID     int
}

func newMCPClient(baseURL string) (*mcpClient, error) {
	c := &mcpClient{
		baseURL: baseURL,
		http:    &http.Client{Timeout: 30 * time.Second},
	}

	// Initialize session
	resp, err := c.doRequest("initialize", map[string]any{
		"protocolVersion": "2024-11-05",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "nomos", "version": "2.0"},
	})
	if err != nil {
		return nil, fmt.Errorf("initialize: %w", err)
	}
	if resp.sessionID == "" {
		return nil, fmt.Errorf("no session ID in initialize response")
	}
	c.sessionID = resp.sessionID

	// Send initialized notification
	c.doRequest("notifications/initialized", map[string]any{})

	slog.Info("nomos: mcp connected", "session", c.sessionID[:16]+"...")
	return c, nil
}

type mcpJSONRPCResponse struct {
	sessionID string
	Result    json.RawMessage `json:"result"`
	Error     json.RawMessage `json:"error"`
}

func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPCResponse, error) {
	c.nextID++
	body, _ := json.Marshal(map[string]any{
		"jsonrpc": "2.0",
		"method":  method,
		"params":  params,
		"id":      c.nextID,
	})

	req, err := http.NewRequest(http.MethodPost, c.baseURL, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	if c.sessionID != "" {
		req.Header.Set("Mcp-Session-Id", c.sessionID)
	}

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	result := &mcpJSONRPCResponse{}
	result.sessionID = resp.Header.Get("Mcp-Session-Id")

	// Parse SSE stream: "event: message\ndata: <json>\n\n"
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			data := line[6:]
			if err := json.Unmarshal([]byte(data), result); err != nil {
				return nil, fmt.Errorf("parse response: %w", err)
			}
		}
	}

	if result.Error != nil {
		return nil, fmt.Errorf("rpc error: %s", string(result.Error))
	}

	if result.sessionID != "" {
		c.sessionID = result.sessionID
	}

	return result, nil
}

func (c *mcpClient) callTool(name string, args map[string]any) (any, error) {
	resp, err := c.doRequest("tools/call", map[string]any{
		"name":      name,
		"arguments": args,
	})
	if err != nil {
		return nil, err
	}

	// Parse MCP content: { "content": [{ "type": "text", "text": "..." }] }
	var toolResult struct {
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
	}
	if err := json.Unmarshal(resp.Result, &toolResult); err != nil {
		return string(resp.Result), nil
	}

	var texts []string
	for _, c := range toolResult.Content {
		if c.Type == "text" {
			// Try to parse as JSON for structured display
			var parsed any
			if json.Unmarshal([]byte(c.Text), &parsed) == nil {
				return parsed, nil
			}
			texts = append(texts, c.Text)
		}
	}
	if len(texts) == 1 {
		return texts[0], nil
	}
	return texts, nil
}

func (c *mcpClient) listTools() ([]string, error) {
	resp, err := c.doRequest("tools/list", map[string]any{})
	if err != nil {
		return nil, err
	}
	var tr struct {
		Tools []struct {
			Name        string `json:"name"`
			Description string `json:"description"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(resp.Result, &tr); err != nil {
		return nil, err
	}
	var names []string
	for _, t := range tr.Tools {
		names = append(names, t.Name)
	}
	return names, nil
}

func (c *mcpClient) close() {
	// MCP sessions are ephemeral; no explicit close needed
}