oikos/cmd/nomos/main.go

package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"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"
	}

	databaseURL := os.Getenv("DATABASE_URL")
	if databaseURL == "" {
		databaseURL = os.Getenv("OIKOS_DATABASE_URL")
	}

	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)
		}

		st, err := newStore(ctx, databaseURL)
		if err != nil {
			slog.Error("nomos: db connect", "error", err)
			os.Exit(1)
		}
		if st != nil {
			defer st.close()
		}

		nAgent, err := newAgent(ctx, client, st, agentSlug)
		if err != nil {
			slog.Error("nomos: agent init", "error", err)
			os.Exit(1)
		}

		// Event-driven auto-continuation: feed finished async executions back
		// into the agent so an approved plan runs to completion (and recovers
		// from failures) without the operator ticking it forward each step.
		go nAgent.runContinuationWorker(ctx)

		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)
		})
		mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
			handleChat(w, r, nAgent, st)
		})
		mux.HandleFunc("/sessions", func(w http.ResponseWriter, r *http.Request) {
			handleSessionsList(w, r, st)
		})
		mux.HandleFunc("/sessions/", func(w http.ResponseWriter, r *http.Request) {
			handleSessionDetail(w, r, st)
		})

		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, "db", databaseURL != "")
			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)
	}
}

func sseEvent(w http.ResponseWriter, flusher http.Flusher, event agentEvent) {
	data, _ := json.Marshal(event)
	fmt.Fprintf(w, "data: %s\n\n", data)
	flusher.Flush()
}

func handleChat(w http.ResponseWriter, r *http.Request, a *agent, st *store) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", 405)
		return
	}

	var req struct {
		SessionID string `json:"session_id"`
		Message   string `json:"message"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "bad request: "+err.Error(), 400)
		return
	}
	if req.Message == "" {
		http.Error(w, "message is required", 400)
		return
	}

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

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(200)

	ctx := r.Context()
	sessionID := req.SessionID

	if sessionID == "" {
		title := truncate(req.Message, 80)
		sess, err := st.createSession(ctx, title)
		if err != nil {
			slog.Error("nomos: create session", "error", err)
			sessionID = "ephemeral"
		} else {
			sessionID = sess.ID
		}
	} else {
		st.touchSession(ctx, sessionID)
	}

	slog.Info("nomos: chat", "session", sessionID, "message", truncate(req.Message, 100))

	userMsg, _ := json.Marshal(map[string]any{"role": "user", "text": req.Message})
	st.saveMessage(ctx, sessionID, "user", userMsg)

	sseEvent(w, flusher, agentEvent{Type: "session", Data: sessionID, SessionID: sessionID})

	toolCalls := []map[string]any{}
	var finalText string

	a.chat(ctx, sessionID, req.Message, func(ev agentEvent) {
		if ev.Type == "tool_use" || ev.Type == "tool_result" {
			if m, ok := ev.Data.(map[string]any); ok {
				m["type"] = ev.Type
				toolCalls = append(toolCalls, m)
			}
		}
		if ev.Type == "text" {
			finalText, _ = ev.Data.(string)
		}
		sseEvent(w, flusher, ev)
	})

	assistantMsg, _ := json.Marshal(map[string]any{
		"role":       "assistant",
		"text":       finalText,
		"tool_calls": toolCalls,
	})
	st.saveMessage(ctx, sessionID, "assistant", assistantMsg)

	// Generate a meaningful title from the assistant's first answer
	// instead of reusing the raw user message for every session.
	if finalText != "" && sessionID != "ephemeral" {
		title := truncate(finalText, 80)
		if title != "" {
			st.updateSessionTitle(ctx, sessionID, title)
		}
	}
}

func handleSessionsList(w http.ResponseWriter, r *http.Request, st *store) {
	if st == nil {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{"sessions": []any{}})
		return
	}

	if r.Method == http.MethodOptions {
		return
	}

	sessions, err := st.listSessions(r.Context())
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{"sessions": sessions})
}

func handleSessionDetail(w http.ResponseWriter, r *http.Request, st *store) {
	if st == nil {
		http.Error(w, "not found", 404)
		return
	}

	id := strings.TrimPrefix(r.URL.Path, "/sessions/")
	if id == "" {
		http.Error(w, "session id required", 400)
		return
	}

	switch r.Method {
	case http.MethodDelete:
		if err := st.deleteSession(r.Context(), id); err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		w.WriteHeader(204)

	case http.MethodGet:
		messages, err := st.getMessages(r.Context(), id)
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{"session_id": id, "messages": messages})

	default:
		http.Error(w, "method not allowed", 405)
	}
}

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()

	if req.Tool != "" {
		result, err := client.callTool(req.Tool, req.Args)
		duration := time.Since(start).Milliseconds()
		if err != nil {
			slog.Error("nomos: query failed", "tool", req.Tool, "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,
		})
		return
	}

	if req.Query != "" {
		if strings.Contains(strings.ToLower(req.Query), "what can you do") ||
			strings.Contains(strings.ToLower(req.Query), "help") {

			tools, err := client.listTools()
			duration := time.Since(start).Milliseconds()
			if err != nil {
				w.Header().Set("Content-Type", "application/json")
				json.NewEncoder(w).Encode(map[string]any{
					"error":      err.Error(),
					"elapsed_ms": duration,
				})
				return
			}
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]any{
				"message":    "natural language queries belong to /chat. Use structured /query with 'tool' for direct MCP calls.",
				"tools":      tools,
				"elapsed_ms": duration,
			})
			return
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{
			"message":    "natural language queries belong to /chat. Use structured /query with 'tool' for direct MCP calls.",
			"elapsed_ms": time.Since(start).Milliseconds(),
		})
		return
	}

	http.Error(w, "either 'tool' or 'query' required", 400)
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}

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

type mcpClient struct {
	baseURL   string
	sessionID string
	http      *http.Client
	nextID    int
	mu        sync.Mutex // MCP is one stateful session; serialize concurrent calls
}

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

	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

	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"`
}

// errStaleSession signals that the MCP server rejected our session id (e.g.
// after an api/MCP restart), so the client should re-initialize and retry.
var errStaleSession = fmt.Errorf("mcp session stale")

// doRequest serializes MCP calls and transparently re-initializes the session
// if the server has forgotten it (common after an api redeploy), retrying the
// original call once. Without this, an api restart permanently breaks nomos
// until it is itself restarted.
func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPCResponse, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	resp, err := c.send(method, params)
	if err != nil && method != "initialize" && errors.Is(err, errStaleSession) {
		slog.Warn("nomos: mcp session stale, reconnecting")
		if rerr := c.reconnectLocked(); rerr != nil {
			return nil, fmt.Errorf("mcp reconnect: %w (original: %v)", rerr, err)
		}
		return c.send(method, params)
	}
	return resp, err
}

// reconnectLocked re-initializes the MCP session. The caller must hold c.mu.
func (c *mcpClient) reconnectLocked() error {
	c.sessionID = ""
	resp, err := c.send("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 err
	}
	if resp.sessionID == "" {
		return fmt.Errorf("no session ID on re-initialize")
	}
	c.sessionID = resp.sessionID
	_, _ = c.send("notifications/initialized", map[string]any{})
	slog.Info("nomos: mcp reconnected", "session", c.sessionID[:16]+"...")
	return nil
}

// send performs one MCP round-trip. It does not lock; callers hold c.mu.
func (c *mcpClient) send(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()

	// A rejected/unknown session comes back as 4xx (commonly 400/404).
	if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusBadRequest {
		return nil, errStaleSession
	}

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

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	gotData := false
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			gotData = true
			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))
	}

	// Empty body with no result and a session set: the server likely dropped
	// our session. Notifications legitimately return no data, so exempt them.
	if !gotData && result.Result == nil && method != "notifications/initialized" {
		return nil, errStaleSession
	}

	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
	}

	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" {
			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() {
}