362 lines
11 KiB
Go
362 lines
11 KiB
Go
package router
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"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
|
|
}
|
|
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) {
|
|
raw, err := io.ReadAll(http.MaxBytesReader(w, req.Body, maxMemoryBody))
|
|
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":
|
|
resp.Result = map[string]any{
|
|
"protocolVersion": mcpProtocolVer,
|
|
"capabilities": map[string]any{"tools": map[string]any{}},
|
|
"serverInfo": map[string]any{"name": "llmgw-memory", "version": "1"},
|
|
}
|
|
case "ping":
|
|
resp.Result = map[string]any{}
|
|
case "tools/list":
|
|
resp.Result = map[string]any{"tools": mcpTools}
|
|
case "tools/call":
|
|
resp.Result = r.callMemoryTool(req, rpc.Params)
|
|
default:
|
|
resp.Error = &rpcError{Code: -32601, Message: "method not found: " + rpc.Method}
|
|
}
|
|
writeJSON(w, http.StatusOK, resp)
|
|
}
|
|
|
|
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 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)
|
|
}
|
|
}
|