oikos/cmd/nomos/agent.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"time"

	"github.com/google/uuid"
	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"github.com/openai/openai-go/shared"
)

const maxIterations = 15

type agent struct {
	client   *mcpClient
	system   string
	provider *openai.Client
	model    string
	store    *store
	agentID  uuid.UUID
	reqOpts  []option.RequestOption
}

func newAgent(ctx context.Context, mcpClient *mcpClient, st *store, agentSlug string) (*agent, error) {
	system := loadSoul()
	apiKey := os.Getenv("OPENROUTER_API_KEY")
	model := os.Getenv("NOMOS_MODEL")
	if model == "" {
		model = "deepseek/deepseek-v4-flash"
	}

	provider := openai.NewClient(
		option.WithBaseURL("https://openrouter.ai/api/v1"),
		option.WithAPIKey(apiKey),
	)

	agentID := st.resolveAgentID(ctx, agentSlug)
	if agentID == uuid.Nil {
		slog.Warn("nomos: agent entity not found; tool-call activity will not be logged", "slug", agentSlug)
	}

	// OpenRouter provider routing. data_collection=deny pins to zero-data-
	// retention providers (privacy: conversations + tool results transit
	// OpenRouter); require_parameters ensures the routed provider actually
	// supports tool calling. NOMOS_PROVIDER_SORT (price|throughput|latency)
	// and Exacto tool-accuracy routing are opt-in — the latter via a model
	// suffix in NOMOS_MODEL (e.g. "deepseek/deepseek-v4-flash:exacto"), so an
	// unsupported value never silently breaks the confirmed routing below.
	providerRouting := map[string]any{
		"data_collection":    "deny",
		"require_parameters": true,
	}
	if sort := os.Getenv("NOMOS_PROVIDER_SORT"); sort != "" {
		providerRouting["sort"] = sort
	}
	reqOpts := []option.RequestOption{option.WithJSONSet("provider", providerRouting)}

	return &agent{
		client:   mcpClient,
		system:   system,
		provider: &provider,
		model:    model,
		store:    st,
		agentID:  agentID,
		reqOpts:  reqOpts,
	}, nil
}

func loadSoul() string {
	paths := []string{"/app/nomos/SOUL.md", "nomos/SOUL.md"}
	for _, p := range paths {
		if data, err := os.ReadFile(p); err == nil {
			return string(data)
		}
	}
	return `You are Nomos, the steward of the oikos — the AI agent for the hubris homelab.
You have access to MCP tools to query topology, health, knowledge, and request
gated mutations through request_execution. Be concise. Prefer tools over guessing.`
}

type toolDef struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

type agentEvent struct {
	Type       string `json:"type"`
	Data       any    `json:"data,omitempty"`
	SessionID  string `json:"session_id,omitempty"`
	Iteration  int    `json:"iteration,omitempty"`
}

func (a *agent) chat(ctx context.Context, sessionID, message string, emit func(agentEvent)) {
	correlationID := uuid.New().String()

	tools, err := a.buildTools()
	if err != nil {
		emit(agentEvent{Type: "error", Data: fmt.Sprintf("build tools: %v", err), SessionID: sessionID})
		return
	}

	// Rebuild conversation context from persisted history so sessions are
	// multi-turn. The current user turn is saved by the HTTP handler before
	// this runs, so it is already included in the history for real sessions.
	// Prior tool_use/tool_result pairs are replayed as a tool-calling
	// assistant message followed by matching tool-role results, so the agent
	// starts each turn already knowing what it already checked instead of
	// re-querying the same tools from scratch. Ephemeral sessions (no store)
	// fall back to the single incoming message.
	system := a.system
	if snapshot := a.fleetSnapshot(); snapshot != "" {
		system += "\n\n" + snapshot
	}
	messages := []openai.ChatCompletionMessageParamUnion{openai.SystemMessage(system)}
	history, _ := a.store.getMessages(ctx, sessionID)
	for _, m := range history {
		text := extractText(m.Content)
		switch m.Role {
		case "user":
			messages = append(messages, openai.UserMessage(text))
		case "assistant":
			if calls := extractToolCalls(m.Content); len(calls) > 0 {
				messages = append(messages, assistantToolCallMessage(calls))
				for _, c := range calls {
					messages = append(messages, openai.ToolMessage(c.resultText(), c.id))
				}
			}
			if text != "" {
				messages = append(messages, openai.AssistantMessage(text))
			}
		}
	}
	if len(history) == 0 {
		messages = append(messages, openai.UserMessage(message))
	}

	for i := 0; i < maxIterations; i++ {
		params := openai.ChatCompletionNewParams{
			Model:    openai.ChatModel(a.model),
			Messages: messages,
			Tools:    tools,
		}

		// Stream the completion, emitting token deltas as they arrive. The
		// accumulator reassembles the full message (content + tool calls) for
		// the loop's control flow.
		stream := a.provider.Chat.Completions.NewStreaming(ctx, params, a.reqOpts...)
		acc := openai.ChatCompletionAccumulator{}
		for stream.Next() {
			chunk := stream.Current()
			acc.AddChunk(chunk)
			if len(chunk.Choices) > 0 {
				if delta := chunk.Choices[0].Delta.Content; delta != "" {
					emit(agentEvent{Type: "text_delta", Data: delta, SessionID: sessionID, Iteration: i + 1})
				}
			}
		}
		if err := stream.Err(); err != nil {
			emit(agentEvent{Type: "error", Data: fmt.Sprintf("llm: %v", err), SessionID: sessionID})
			return
		}
		if len(acc.Choices) == 0 {
			emit(agentEvent{Type: "error", Data: "no choices in response", SessionID: sessionID})
			return
		}

		msg := acc.Choices[0].Message

		if len(msg.ToolCalls) == 0 {
			emit(agentEvent{Type: "text", Data: msg.Content, SessionID: sessionID})
			emit(agentEvent{Type: "done", Data: map[string]any{
				"session_id":     sessionID,
				"usage":          acc.Usage,
				"correlation_id": correlationID,
				"iterations":     i + 1,
			}, SessionID: sessionID})
			return
		}

		slog.Info("nomos: tool calls", "count", len(msg.ToolCalls), "iter", i+1, "correlation", correlationID)

		messages = append(messages, msg.ToParam())

		for _, tc := range msg.ToolCalls {
			var args map[string]any
			if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil {
				args = map[string]any{}
			}

			emit(agentEvent{
				Type:      "tool_use",
				Data:      map[string]any{"name": tc.Function.Name, "args": args, "id": tc.ID},
				SessionID: sessionID,
				Iteration: i + 1,
			})

			start := time.Now()
			result, callErr := a.client.callTool(tc.Function.Name, args)
			elapsed := int(time.Since(start).Milliseconds())

			inputJSON, _ := json.Marshal(args)
			inputStr := string(inputJSON)

			if callErr != nil {
				a.store.logActivity(ctx, a.agentID, sessionID, tc.Function.Name, inputStr, callErr.Error(), elapsed, false, correlationID)

				emit(agentEvent{
					Type:      "tool_result",
					Data:      map[string]any{"name": tc.Function.Name, "error": callErr.Error(), "id": tc.ID},
					SessionID: sessionID,
					Iteration: i + 1,
				})
				messages = append(messages, openai.ToolMessage(callErr.Error(), tc.ID))
				slog.Error("nomos: tool error", "tool", tc.Function.Name, "error", callErr, "ms", elapsed)
				continue
			}

			resultJSON, _ := json.Marshal(result)
			a.store.logActivity(ctx, a.agentID, sessionID, tc.Function.Name, inputStr, string(resultJSON), elapsed, true, correlationID)

			emit(agentEvent{
				Type:      "tool_result",
				Data:      map[string]any{"name": tc.Function.Name, "result": result, "id": tc.ID},
				SessionID: sessionID,
				Iteration: i + 1,
			})
			messages = append(messages, openai.ToolMessage(string(resultJSON), tc.ID))
			slog.Info("nomos: tool success", "tool", tc.Function.Name, "ms", elapsed)
		}
	}

	emit(agentEvent{Type: "text", Data: "Agent loop: max iterations reached without final answer.", SessionID: sessionID})
	emit(agentEvent{Type: "done", Data: map[string]any{
		"session_id":     sessionID,
		"correlation_id": correlationID,
		"iterations":     maxIterations,
	}, SessionID: sessionID})
}

// extractText pulls the "text" field from a persisted message's JSONB content.
func extractText(content json.RawMessage) string {
	var m struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(content, &m); err != nil {
		return ""
	}
	return m.Text
}

// persistedCall is one merged tool_use+tool_result pair from a persisted
// assistant message's tool_calls array. The store keeps them as two entries
// sharing the same id (mirroring the SSE event pair); replay needs one
// entry per id to build a valid tool-calling assistant message.
type persistedCall struct {
	id     string
	name   string
	args   json.RawMessage
	result json.RawMessage
	errMsg string
}

func (c persistedCall) resultText() string {
	if c.errMsg != "" {
		return c.errMsg
	}
	if len(c.result) > 0 {
		return string(c.result)
	}
	return "null"
}

// extractToolCalls parses and merges a persisted message's tool_calls array,
// preserving first-seen order across ids.
func extractToolCalls(content json.RawMessage) []persistedCall {
	var m struct {
		ToolCalls []struct {
			ID     string          `json:"id"`
			Type   string          `json:"type"`
			Name   string          `json:"name"`
			Args   json.RawMessage `json:"args"`
			Result json.RawMessage `json:"result"`
			Error  string          `json:"error"`
		} `json:"tool_calls"`
	}
	if err := json.Unmarshal(content, &m); err != nil || len(m.ToolCalls) == 0 {
		return nil
	}

	byID := make(map[string]*persistedCall, len(m.ToolCalls))
	var order []string
	for _, tc := range m.ToolCalls {
		if tc.ID == "" {
			continue
		}
		pc, ok := byID[tc.ID]
		if !ok {
			pc = &persistedCall{id: tc.ID}
			byID[tc.ID] = pc
			order = append(order, tc.ID)
		}
		if tc.Name != "" {
			pc.name = tc.Name
		}
		if len(tc.Args) > 0 && string(tc.Args) != "null" {
			pc.args = tc.Args
		}
		if tc.Type == "tool_result" {
			pc.errMsg = tc.Error
			pc.result = tc.Result
		}
	}

	calls := make([]persistedCall, 0, len(order))
	for _, id := range order {
		calls = append(calls, *byID[id])
	}
	return calls
}

// assistantToolCallMessage builds the tool-calling assistant message that
// must precede the tool-role results being replayed.
func assistantToolCallMessage(calls []persistedCall) openai.ChatCompletionMessageParamUnion {
	toolCalls := make([]openai.ChatCompletionMessageToolCallParam, 0, len(calls))
	for _, c := range calls {
		args := string(c.args)
		if args == "" {
			args = "{}"
		}
		toolCalls = append(toolCalls, openai.ChatCompletionMessageToolCallParam{
			ID: c.id,
			Function: openai.ChatCompletionMessageToolCallFunctionParam{
				Name:      c.name,
				Arguments: args,
			},
		})
	}
	return openai.ChatCompletionMessageParamUnion{
		OfAssistant: &openai.ChatCompletionAssistantMessageParam{ToolCalls: toolCalls},
	}
}

// fleetSnapshot returns a compact, current-as-of-now fleet health line for
// the system prompt so the agent starts each turn already oriented instead
// of spending its first iteration rediscovering topology it already has
// tools to query. Best-effort: an empty string on any failure just means no
// snapshot, not an error for the turn.
func (a *agent) fleetSnapshot() string {
	result, err := a.client.callTool("get_health_summary", map[string]any{})
	if err != nil {
		return ""
	}
	rows, ok := result.([]any)
	if !ok {
		return ""
	}

	counts := map[string]int{}
	var attention []string
	for _, r := range rows {
		row, ok := r.(map[string]any)
		if !ok {
			continue
		}
		health, _ := row["health"].(string)
		counts[health]++
		if health != "healthy" && health != "" {
			if slug, ok := row["slug"].(string); ok && len(attention) < 10 {
				attention = append(attention, fmt.Sprintf("%s(%s)", slug, health))
			}
		}
	}
	if len(counts) == 0 {
		return ""
	}

	summary := fmt.Sprintf("Current fleet snapshot (as of now): healthy=%d degraded=%d down=%d stale=%d unknown=%d.",
		counts["healthy"], counts["degraded"], counts["down"], counts["stale"], counts["unknown"])
	if len(attention) > 0 {
		summary += " Needs attention: " + fmt.Sprintf("%v", attention) + "."
	}
	return summary
}

func (a *agent) buildTools() ([]openai.ChatCompletionToolParam, error) {
	defs, err := a.client.listToolsFull()
	if err != nil {
		return nil, err
	}

	var tools []openai.ChatCompletionToolParam
	for _, d := range defs {
		params := shared.FunctionParameters(d.InputSchema)
		if params == nil {
			params = shared.FunctionParameters{"type": "object", "properties": map[string]any{}}
		}

		tools = append(tools, openai.ChatCompletionToolParam{
			Type: "function",
			Function: shared.FunctionDefinitionParam{
				Name:        d.Name,
				Description: openai.String(d.Description),
				Parameters:  params,
			},
		})
	}
	return tools, nil
}

func (c *mcpClient) listToolsFull() ([]toolDef, 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"`
			InputSchema map[string]any `json:"inputSchema"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(resp.Result, &tr); err != nil {
		return nil, err
	}
	out := make([]toolDef, len(tr.Tools))
	for i, t := range tr.Tools {
		out[i] = toolDef{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: t.InputSchema,
		}
	}
	return out, nil
}