llmgw/internal/router/memory.go

package router

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strconv"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// agentHeader names the agent whose notes a request reads and writes.
//
// A namespace inside a project, not a credential: with auth on, the key decides
// the project and this only separates agents sharing it; with auth off it is
// just a name.
const agentHeader = "x-llmgw-agent"

const (
	defaultAgent   = "default"
	maxAgentID     = 100
	maxMemoryBody  = 64 << 10
	mcpProtocolVer = "2025-03-26"
)

// errMemoryOff is what a memory call meets when nothing is configured to keep
// notes. Refused rather than pretended: an agent that believes it saved
// something and finds nothing later has been lied to.
var errMemoryOff = errors.New(
	"this gateway keeps no memory notes: enable auth with a project that keeps memory, " +
		"or set threads.store and threads.retain_hours")

// memoryScope says whose notes this request is about.
func (r *Router) memoryScope(req *http.Request) (project, agent string, err error) {
	var caller string
	if p, ok := auth.FromContext(req.Context()); ok {
		caller = p.ProjectID
	}
	project, err = r.threadProject(caller)
	if err != nil {
		if errors.Is(err, errThreadsOff) {
			err = errMemoryOff
		}
		return "", "", err
	}
	agent = req.Header.Get(agentHeader)
	if agent == "" {
		agent = defaultAgent
	}
	if len(agent) > maxAgentID {
		return "", "", errors.New(agentHeader + " is too long")
	}
	return project, agent, nil
}

type noteJSON struct {
	Key       string   `json:"key"`
	Text      string   `json:"text"`
	Tags      []string `json:"tags"`
	CreatedAt string   `json:"created_at"`
	UpdatedAt string   `json:"updated_at"`
}

func toNoteJSON(n store.Note) noteJSON {
	tags := n.Tags
	if tags == nil {
		tags = []string{}
	}
	return noteJSON{
		Key: n.Key, Text: n.Text, Tags: tags,
		CreatedAt: n.CreatedAt.UTC().Format(time.RFC3339),
		UpdatedAt: n.UpdatedAt.UTC().Format(time.RFC3339),
	}
}

// memoryStatus maps a store error to what a client should be told.
func memoryStatus(err error) int {
	switch {
	case errors.Is(err, store.ErrNoNote):
		return http.StatusNotFound
	case errors.Is(err, errThreadsNeedAuth):
		return http.StatusUnauthorized
	case errors.Is(err, errMemoryOff), errors.Is(err, store.ErrMemoryNotEnabled),
		errors.Is(err, store.ErrMemoryFull):
		return http.StatusBadRequest
	default:
		return 0
	}
}

func writeMemoryError(w http.ResponseWriter, err error) {
	if code := memoryStatus(err); code != 0 {
		msg := err.Error()
		if code == http.StatusNotFound {
			// One answer for missing, another agent's and another project's, so
			// a caller cannot map which keys exist where they cannot read.
			msg = "no such note"
		}
		writeError(w, code, "invalid_request_error", msg)
		return
	}
	writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
}

// handlePutNote writes a note: PUT /v1/memory/{key}.
func (r *Router) handlePutNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	if r.noStore(callerSubject(req)) {
		writeErrorCode(w, http.StatusForbidden, "invalid_request_error", "no_store", errNoStoreMsg)
		return
	}
	var body struct {
		Text string   `json:"text"`
		Tags []string `json:"tags"`
	}
	raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
	if err != nil || json.Unmarshal(raw, &body) != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			`body must be JSON like {"text": "...", "tags": ["..."]}`)
		return
	}
	note := store.Note{Key: req.PathValue("key"), Text: body.Text, Tags: body.Tags}
	if err := r.threads.PutNote(req.Context(), project, agent, note); err != nil {
		writeMemoryError(w, err)
		return
	}
	saved, err := r.threads.GetNote(req.Context(), project, agent, note.Key)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, toNoteJSON(saved))
}

// handleGetNote reads one note: GET /v1/memory/{key}.
func (r *Router) handleGetNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	n, err := r.threads.GetNote(req.Context(), project, agent, req.PathValue("key"))
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, toNoteJSON(n))
}

// handleSearchNotes lists or searches: GET /v1/memory?q=&tag=&limit=.
func (r *Router) handleSearchNotes(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	q := req.URL.Query()
	limit, _ := strconv.Atoi(q.Get("limit"))
	notes, err := r.threads.SearchNotes(req.Context(), project, agent, q.Get("q"), q.Get("tag"), limit)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	out := make([]noteJSON, 0, len(notes))
	for _, n := range notes {
		out = append(out, toNoteJSON(n))
	}
	writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": out})
}

// handleDeleteNote forgets one: DELETE /v1/memory/{key}.
func (r *Router) handleDeleteNote(w http.ResponseWriter, req *http.Request) {
	project, agent, err := r.memoryScope(req)
	if err != nil {
		writeMemoryError(w, err)
		return
	}
	if err := r.threads.DeleteNote(req.Context(), project, agent, req.PathValue("key")); err != nil {
		writeMemoryError(w, err)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// --- MCP ---
//
// The same notes, offered as tools an agent can be given. Only the small part
// of the protocol a tool server needs: initialize, tools/list and tools/call,
// answered as plain JSON over POST. It never streams and never calls the
// client, so it needs no session.

type rpcRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type rpcResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  any             `json:"result,omitempty"`
	Error   *rpcError       `json:"error,omitempty"`
}

var mcpTools = []map[string]any{
	{
		"name":        "remember",
		"description": "Save a note that will still be there in later conversations. Writing to an existing key replaces it.",
		"inputSchema": map[string]any{
			"type":     "object",
			"required": []string{"key", "text"},
			"properties": map[string]any{
				"key":  map[string]any{"type": "string", "description": "Short unique name for the note"},
				"text": map[string]any{"type": "string", "description": "What to remember"},
				"tags": map[string]any{"type": "array", "items": map[string]any{"type": "string"}, "description": "Optional lowercase labels"},
			},
		},
	},
	{
		"name":        "recall",
		"description": "Find saved notes by words, by tag, or by exact key. With nothing given, returns the most recent notes.",
		"inputSchema": map[string]any{
			"type": "object",
			"properties": map[string]any{
				"query": map[string]any{"type": "string", "description": "Words to search for"},
				"tag":   map[string]any{"type": "string", "description": "Only notes with this tag"},
				"key":   map[string]any{"type": "string", "description": "Exact key to read"},
				"limit": map[string]any{"type": "integer", "description": "Most notes to return (up to 100)"},
			},
		},
	},
	{
		"name":        "forget",
		"description": "Delete a saved note.",
		"inputSchema": map[string]any{
			"type":       "object",
			"required":   []string{"key"},
			"properties": map[string]any{"key": map[string]any{"type": "string"}},
		},
	},
}

// handleMCP serves POST /mcp.
func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
	// Pictures to make a mesh from or to look at come in the call, as
	// base64.
	raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, mcpBody))
	if mbe := (*http.MaxBytesError)(nil); errors.As(err, &mbe) {
		writeJSON(w, http.StatusRequestEntityTooLarge, rpcResponse{
			JSONRPC: "2.0", ID: json.RawMessage("null"),
			Error: &rpcError{Code: -32600, Message: fmt.Sprintf("the call is larger than %d MB: send pictures smaller, "+
				"or fewer at once", mcpBody>>20)},
		})
		return
	}
	var rpc rpcRequest
	if err != nil || json.Unmarshal(raw, &rpc) != nil || rpc.Method == "" {
		writeJSON(w, http.StatusBadRequest, rpcResponse{
			JSONRPC: "2.0", ID: json.RawMessage("null"),
			Error: &rpcError{Code: -32600, Message: "expected one JSON-RPC request object"},
		})
		return
	}
	// A notification has no id and gets no answer.
	if len(rpc.ID) == 0 {
		w.WriteHeader(http.StatusAccepted)
		return
	}

	resp := rpcResponse{JSONRPC: "2.0", ID: rpc.ID}
	switch rpc.Method {
	case "initialize":
		result := map[string]any{
			"protocolVersion": mcpVersionFor(rpc.Params),
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": "llmgw", "version": "1"},
		}
		if r.mediaModels().any() {
			result["instructions"] = mcpMediaInstructions
		}
		resp.Result = result
	case "ping":
		resp.Result = map[string]any{}
	case "tools/list":
		tools := append(slices.Clone(mcpTools), r.mediaTools()...)
		for _, t := range r.mcpgw.list(req) {
			tools = append(tools, t)
		}
		resp.Result = map[string]any{"tools": tools}
	case "tools/call":
		var call struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
			Meta      struct {
				ProgressToken json.RawMessage `json:"progressToken"`
			} `json:"_meta"`
		}
		_ = json.Unmarshal(rpc.Params, &call)
		if isCompanyTool(call.Name) {
			if acceptsStream(req) {
				// The server's progress passed on while it works.
				w.Header().Set("Content-Type", "text/event-stream")
				w.Header().Set("Cache-Control", "no-cache")
				w.WriteHeader(http.StatusOK)
				sw := &streamWriter{w: w, token: call.Meta.ProgressToken}
				stop := make(chan struct{})
				go sw.keepAlive(stop, 15*time.Second)
				resp.Result = r.mcpgw.call(req, call.Name, call.Arguments, func(p json.RawMessage) {
					var pr struct {
						Progress, Total float64
						Message         string
					}
					if json.Unmarshal(p, &pr) == nil && pr.Total > 0 {
						sw.progress(100*pr.Progress/pr.Total, pr.Message)
					}
				})
				close(stop)
				sw.event(resp)
				return
			}
			resp.Result = r.mcpgw.call(req, call.Name, call.Arguments, nil)
			break
		}
		if !isMediaTool(call.Name) {
			resp.Result = r.callMemoryTool(req, rpc.Params)
			break
		}
		if acceptsStream(req) {
			// Answered as a stream: how far along it is while it runs, which
			// a mesh's minutes want, then the answer.
			w.Header().Set("Content-Type", "text/event-stream")
			w.Header().Set("Cache-Control", "no-cache")
			w.WriteHeader(http.StatusOK)
			sw := &streamWriter{w: w, token: call.Meta.ProgressToken}
			stop := make(chan struct{})
			go sw.keepAlive(stop, 15*time.Second)
			resp.Result = r.callMediaTool(mediaCall{req: req, progress: sw.progress, streamed: true}, call.Name, call.Arguments)
			close(stop)
			sw.event(resp)
			return
		}
		resp.Result = r.callMediaTool(mediaCall{req: req, progress: func(float64, string) {}}, call.Name, call.Arguments)
	default:
		resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
	}
	writeJSON(w, http.StatusOK, resp)
}

// mcpVersionFor is the protocol version answered: the client's, of those
// this server speaks, else its own.
func mcpVersionFor(params json.RawMessage) string {
	var p struct {
		ProtocolVersion string `json:"protocolVersion"`
	}
	_ = json.Unmarshal(params, &p)
	if slices.Contains([]string{"2025-03-26", "2025-06-18"}, p.ProtocolVersion) {
		return p.ProtocolVersion
	}
	return mcpProtocolVer
}

func toolText(text string, isError bool) map[string]any {
	return map[string]any{
		"content": []map[string]any{{"type": "text", "text": text}},
		"isError": isError,
	}
}

func toolJSON(v any) map[string]any {
	b, _ := json.MarshalIndent(v, "", "  ")
	return toolText(string(b), false)
}

// callMemoryTool runs one tool. Failures are tool results with isError set, not
// protocol errors, because that is how a model gets to see what went wrong.
func (r *Router) callMemoryTool(req *http.Request, params json.RawMessage) map[string]any {
	var call struct {
		Name      string `json:"name"`
		Arguments struct {
			Key   string   `json:"key"`
			Text  string   `json:"text"`
			Tags  []string `json:"tags"`
			Query string   `json:"query"`
			Tag   string   `json:"tag"`
			Limit int      `json:"limit"`
		} `json:"arguments"`
	}
	if err := json.Unmarshal(params, &call); err != nil {
		return toolText("invalid tool call", true)
	}
	project, agent, err := r.memoryScope(req)
	if err != nil {
		return toolText(err.Error(), true)
	}
	a, ctx := call.Arguments, req.Context()

	switch call.Name {
	case "remember":
		if r.noStore(callerSubject(req)) {
			return toolText(errNoStoreMsg, true)
		}
		if err := r.threads.PutNote(ctx, project, agent, store.Note{Key: a.Key, Text: a.Text, Tags: a.Tags}); err != nil {
			return toolText(err.Error(), true)
		}
		return toolText("saved "+a.Key, false)
	case "recall":
		if a.Key != "" {
			n, err := r.threads.GetNote(ctx, project, agent, a.Key)
			if errors.Is(err, store.ErrNoNote) {
				return toolText("no note with key "+a.Key, false)
			}
			if err != nil {
				return toolText(err.Error(), true)
			}
			return toolJSON([]noteJSON{toNoteJSON(n)})
		}
		notes, err := r.threads.SearchNotes(ctx, project, agent, a.Query, a.Tag, a.Limit)
		if err != nil {
			return toolText(err.Error(), true)
		}
		out := make([]noteJSON, 0, len(notes))
		for _, n := range notes {
			out = append(out, toNoteJSON(n))
		}
		return toolJSON(out)
	case "forget":
		if err := r.threads.DeleteNote(ctx, project, agent, a.Key); err != nil {
			if errors.Is(err, store.ErrNoNote) {
				return toolText("no note with key "+a.Key, false)
			}
			return toolText(err.Error(), true)
		}
		return toolText("forgot "+a.Key, false)
	default:
		return toolText("unknown tool "+call.Name, true)
	}
}