From 2eb79f3400b160388991bf318fbd9f3878029764 Mon Sep 17 00:00:00 2001 From: SirmaXX Date: Thu, 24 Sep 2026 04:04:03 +0300 Subject: [PATCH] agenta memory eklendi --- CHANGELOG.md | 18 + docs/install.md | 28 ++ docs/quickstart.md | 375 ++++++++++++++++++ internal/cli/keys.go | 1 + internal/cli/memory.go | 62 +++ internal/config/config.go | 24 ++ internal/router/handlers.go | 11 + internal/router/memory.go | 361 +++++++++++++++++ internal/router/router.go | 40 +- internal/router/threads.go | 62 ++- internal/store/memory.go | 284 +++++++++++++ internal/store/memory_test.go | 208 ++++++++++ .../store/migrations/sqlite/0007_memory.sql | 45 +++ internal/store/threads.go | 22 + 14 files changed, 1523 insertions(+), 18 deletions(-) create mode 100644 docs/quickstart.md create mode 100644 internal/cli/memory.go create mode 100644 internal/router/memory.go create mode 100644 internal/store/memory.go create mode 100644 internal/store/memory_test.go create mode 100644 internal/store/migrations/sqlite/0007_memory.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index e0f35f4..7896b68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,24 @@ APIs change between minors without a deprecation period. ## Unreleased +### Memory notes for agents, and threads without API keys + +Notes an agent writes for itself and reads back in later conversations, offered +as REST (`/v1/memory`) and as MCP tools at `/mcp` (`remember`, `recall`, +`forget`). A note is a key, text and tags, searched by words or tag, private to +a project and to the agent id in `x-llmgw-agent`. llmgw does not inject the +tools into chat requests: an agent that is not given the server never sees +memory. + +Off for every project until someone runs `llmgw keys memory --retain-hours N`, +and the retention is required, as it is for threads. + +On a gateway with no authentication, a `threads:` block (`store`, +`retain_hours`) now keeps conversations and notes in one shared pool. Before, +`x-llmgw-thread` was silently ignored there; now a gateway configured to keep +nothing refuses it with an explanation instead. The shared pool is for one +trusted machine: anyone who can reach the port and knows an id can read it. + ### Conversations that can move between models A thread a client names with `x-llmgw-thread`, carrying only its new turn. diff --git a/docs/install.md b/docs/install.md index 200abd3..e3de8fe 100644 --- a/docs/install.md +++ b/docs/install.md @@ -644,6 +644,34 @@ next hourly prune: llmgw keys threads --project default --off ``` +## Memory notes for agents + +Also off until a project asks, for the same reason. Where conversations are what +was said, notes are what an agent chose to keep, and they persist across +conversations: + +``` +llmgw keys memory --project default --retain-hours 720 +``` + +The retention is required and is measured from the last write to a note. An agent +then saves, searches and deletes notes through `/v1/memory` or, as tools, through +the MCP server at `/mcp` (`remember`, `recall`, `forget`). Both need the API key +of the project, and notes are private to the project and to the agent id sent in +`x-llmgw-agent`. A note in another project or agent answers the same as one that +does not exist. llmgw does not add these tools to requests itself. + +Turning it off stops new storage, and existing notes go at the next hourly +prune: + +``` +llmgw keys memory --project default --off +``` + +Without authentication, `threads.store` and `threads.retain_hours` in +`server.yaml` keep both conversations and notes in one shared pool. That is for +one trusted machine only; see [quickstart.md](quickstart.md). + ## Reclaiming disk `llmgw ctl nodes` has a DISK column: free space against the size it is free diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 0000000..5440b55 --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,375 @@ +# Quickstart: one machine, one model, ten minutes + +Written for someone who wants to see llmgw answer a request on a single +workstation before reading anything else. No systemd, no `sudo`. For a real +cluster, see [install.md](install.md). + +What you end up with: an OpenAI-compatible endpoint on `localhost:8080` in front +of a real model served by llama.cpp. + +## Two ways to run + +Steps 1 to 4 use the simple one, **router-only**: you start the engines yourself +and tell the router where they are with `--endpoint`. No certificates, no node +agent. + +Step 5 uses the **control plane and node agent**, the same architecture as a real +cluster on one machine: the agent starts and supervises the engine, and the +control plane places the model. + +## 1. Build + +``` +git clone https://gitea.mynodes.xyz/ilker/llmgw && cd llmgw +make build +make test +./bin/llmgw --version +``` + +`make build` writes `./bin/llmgw`. Go 1.26 or later is the only requirement. + +## 2. Check the path with a fake engine + +Before involving a GPU, prove the router works. `sim-engine` is an +OpenAI-compatible fake that streams placeholder words. + +``` +./bin/llmgw sim-engine --listen :8001 --model llama & +./bin/llmgw server --endpoint e1=http://127.0.0.1:8001=llama & + +curl -N localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' \ + -H 'x-session-id: demo' \ + -d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}' +``` + +Chunks of `alpha beta gamma ...` ending in `data: [DONE]` mean the router is +fine. Stop both before continuing (see [Stopping](#stopping)). + +## 3. Serve a real model + +Two processes, two ports. Clients only ever talk to the router; the router +knows where the engine is because you tell it with `--endpoint`: + +``` +curl :8080 -> llmgw router :8080 -> llama serve :8082 -> model + (client) (the gateway) (the engine) +``` + +Requests go to 8080. Port 8082 is the engine's own, and you only touch it to +check that the engine is up. + +llmgw does not install engines or weights. You need: + +- a `llama-server` build from llama.cpp (CUDA build if you want the GPU), and +- a `.gguf` file. + +Check the card first with `nvidia-smi`. An 8 GB card holds a 1B model at bf16 +(2.4 GB) or a 9B model at `q4_k_m` (5.4 GB). + +Start the engine on its own port. The model name given to `--alias` is the name +clients will use: + +``` +llama serve \ + -m /path/to/Llama-3.2-1B-Instruct-BF16.gguf \ + --host 0.0.0.0 --port 8082 -ngl 99 --alias llama +``` + +`-ngl 99` offloads every layer to the GPU. Confirm it is up and knows the model: + +``` +curl -s localhost:8082/v1/models +``` + +Then point the router at it. The format is `id=URL=model`: + +``` +./bin/llmgw server --endpoint e1=http://127.0.0.1:8082=llama +``` + +## 4. Ask it something + +Streaming: + +``` +curl -N localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' \ + -H 'x-session-id: demo' \ + -d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"stream":true}' +``` + +One reply at once: + +``` +curl -s localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' \ + -d '{"model":"llama","messages":[{"role":"user","content":"hello"}],"max_tokens":100}' +``` + +Also useful: `curl -s localhost:8080/v1/models` and +`curl -s localhost:8080/health`. `llmgw ctl` does not work in this mode: it +talks to the control plane's admin API, which router-only mode does not have, +so `ctl status` answers `404`. + +On the machine this was written on, a 1B bf16 model on an RTX 4070 laptop GPU +answered at roughly 85 tokens per second. + +## Several replicas of one model + +Give the router more than one `--endpoint` with the same model name and it +spreads requests across them. `x-session-id` keeps a conversation on the same +replica so its cache is reused: + +``` +./bin/llmgw server \ + --endpoint e1=http://127.0.0.1:8082=llama \ + --endpoint e2=http://127.0.0.1:8083=llama +``` + +## Stopping + +``` +pkill -f "llmgw server" +pkill -f "llmgw sim-engine" +``` + +Stop the engine the way you started it (Ctrl-C in its terminal). + +## What goes wrong + +**`listening on :8080: bind: address already in use`.** An earlier router is +still running. `ss -ltnp | grep 8080` names the process; stop it first. This is +the usual result of moving from step 2 to step 3 without stopping the fake +engine's router. + +**`/v1/models` on the engine returns an empty list.** The engine was started +without `-m`, so it is a router with nothing to load. Restart it with a model. + +**`error: nothing to serve`.** `llmgw server` with no `--endpoint` and no model +carrying `profiles` has nothing to do. In router-only mode, always pass at least +one `--endpoint`. + +**Requests fail with a model-not-found error.** The name in the request must +match the name after the last `=` in `--endpoint`, and it must be a name the +engine actually serves (its `--alias`). + +**Answers are poor.** A 1B model is a plumbing test, not a chat assistant. Swap +in a larger `.gguf`; nothing else changes. + +## 5. Let llmgw run the engine (control plane and node agent) + +Steps 2 and 3 start the engine by hand. Here a **node agent** starts and +supervises llama.cpp for you, and a control plane decides where the model runs. +Still one machine, still no `sudo`. Stop everything from steps 2 and 3 first: +the GPU has to be free and port 8080 unused. + +``` +curl :8080 -> llmgw server (control plane + router) <- llmgw agent (this node) + | + starts llama-server :20000 +``` + +The agent dials the control plane over mutual TLS, so this needs a small private +certificate authority. Everything goes in a scratch directory, `devcluster/`, +which is in `.gitignore`. + +### Certificates + +``` +mkdir -p devcluster/pki +./bin/llmgw ctl cert init-ca --dir devcluster/pki +./bin/llmgw ctl cert issue-server llmgw-ctl --host 127.0.0.1 --host localhost \ + --dir devcluster/pki --out devcluster +./bin/llmgw ctl cert issue-client llmgw-node01 --dir devcluster/pki --out devcluster +``` + +The name given to `issue-client` becomes the certificate's common name and must +equal `node_id` in the agent config. The control plane refuses an agent whose +claimed id differs from its certificate. + +### The engine wrapper + +The agent starts the engine as ` --model ... --port ...`. A `llama` binary +that wants `llama serve ...` refuses every flag, so put a wrapper in front of it: + +``` +cat > devcluster/llama-server <<'EOF' +#!/bin/sh +exec /path/to/llama serve "$@" +EOF +chmod +x devcluster/llama-server + +./bin/llmgw ctl engine-check llamacpp --bin devcluster/llama-server +``` + +`engine-check` should end with `every flag llmgw builds is accepted`. If your +build already ships a `llama-server` binary, point `bin` at that and skip the +wrapper. + +### Control plane config: `devcluster/server.yaml` + +```yaml +listen: "127.0.0.1:8080" +admin_listen: "127.0.0.1:9090" +agent_listen: "127.0.0.1:7443" + +agent_tls: + ca: devcluster/pki/ca.pem + cert: devcluster/server.pem + key: devcluster/server-key.pem + +models: + - name: llama + source: /path/to/Llama-3.2-1B-Instruct-BF16.gguf + routing_profile: chat + context_length: 8192 + scaling: { min: 1, max: 1 } + profiles: + - name: gpu + engine: llamacpp + gpu_class: 4070-8g + weight: 1 + priority: 100 + startup_timeout: 5m +``` + +Three things in there are easy to get wrong: + +- **The model must have at least one profile.** That is what switches the control + plane on. With `models: []` the server starts as a plain router and exits with + `nothing to serve`. +- **`gpu_class` decides CPU or GPU.** Without it the replica is placed in the CPU + pool and runs on the CPU. Get the class name from `./bin/llmgw agent inventory` + (the `CLASS` column). +- **There is no `auth:` block.** Leave it out for a local trial; with auth on, + requests need an API key. + +### Agent config: `devcluster/agent.yaml` + +```yaml +node_id: llmgw-node01 +server: "https://127.0.0.1:7443" +mgmt_addr: "127.0.0.1" + +tls: + ca: devcluster/pki/ca.pem + cert: devcluster/llmgw-node01.pem + key: devcluster/llmgw-node01-key.pem + +engines: + llamacpp: + bin: /absolute/path/to/devcluster/llama-server + +supervisor: + work_dir: devcluster/agent-work +``` + +### Run it + +Two terminals: + +``` +./bin/llmgw server --config devcluster/server.yaml +./bin/llmgw agent --config devcluster/agent.yaml +``` + +The model is not placed at once. The control plane waits about 35 seconds for the +fleet to settle after it starts, then places the replica, and the agent then +loads the weights. Watch it with: + +``` +./bin/llmgw ctl nodes +./bin/llmgw ctl models +./bin/llmgw ctl replicas +./bin/llmgw ctl explain-placement llama +``` + +`STATUS ready` in `ctl models` means it is serving. `nvidia-smi` should now show +the model in GPU memory. The same curl from step 4 works against port 8080. + +`ctl` works here, unlike in router-only mode, because there is a control plane +for it to talk to. + +Stop it with Ctrl-C in each terminal. The agent leaves the engine running when +the agent itself stops, so if you restart the agent and the GPU stays full, +`pgrep -af "llama serve"` will show the engine still up. + +## 6. Memory: conversations and notes + +Off unless you ask. llmgw otherwise stores nothing anyone typed. On a single +machine with no API keys, one block in the server config turns on two things: + +```yaml +threads: + store: devcluster/data/threads.db + retain_hours: 24 +``` + +Restart the server after adding it. `retain_hours` is required: everything kept +expires that long after it was last written. + +Everything kept this way sits in one shared pool with no owner, and anyone who +can reach the port and knows an id can read it. That is fine on one machine and +wrong on a shared network. With auth on, threads and notes are per project +instead and the block is ignored; see [install.md](install.md). + +### Conversations (threads) + +Send the same `x-llmgw-thread` id and only the new turn. llmgw puts the stored +history in front of it: + +``` +curl -s localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' -H 'x-llmgw-thread: demo' \ + -d '{"model":"llama","messages":[{"role":"user","content":"My name is Deniz."}]}' + +curl -s localhost:8080/v1/chat/completions \ + -H 'content-type: application/json' -H 'x-llmgw-thread: demo' \ + -d '{"model":"llama","messages":[{"role":"user","content":"What is my name?"}]}' + +curl -X DELETE localhost:8080/v1/threads/demo # forget it +``` + +Without the block, a request carrying `x-llmgw-thread` is refused with +`this gateway keeps no conversations` rather than answered without its history. + +### Notes an agent keeps for itself + +Notes outlive any one conversation. Each is a key, some text and optional tags, +kept per agent; name the agent with `x-llmgw-agent` (default `default`), so two +agents using the same key do not read each other's notes. + +``` +curl -X PUT localhost:8080/v1/memory/favorite-color \ + -H 'content-type: application/json' \ + -d '{"text":"The user likes dark green.","tags":["preference"]}' + +curl 'localhost:8080/v1/memory?q=green' # search by words +curl 'localhost:8080/v1/memory?tag=preference' # or by tag +curl localhost:8080/v1/memory/favorite-color +curl -X DELETE localhost:8080/v1/memory/favorite-color +``` + +The same notes are offered to an agent as tools over MCP at +`http://localhost:8080/mcp`: `remember`, `recall` and `forget`. Add that URL to +your agent's MCP server list (HTTP transport; how you set request headers such +as `x-llmgw-agent` depends on the client). llmgw does not put these tools into +chat requests on its own, so an agent that is not given the server never sees +memory. + +Limits: a key is up to 200 bytes, a note's text up to 16 KiB, at most 16 tags, +1000 notes per agent. Tags are lowercase letters, digits, `-` and `_`. + +## Where this stops + +Router-only mode (steps 2 to 4) has no authentication, no placement and no model +management: you start and restart engines yourself. Step 5 adds placement, model +management and a node agent, and can add authentication as `install.md` +describes. For a real cluster, with several nodes, systemd and a persistent +certificate authority, follow [install.md](install.md). + +One caveat found while writing this page: `install.md` recommends starting with +an empty `models: []` and adding models later with `ctl model add`. The control +plane only turns on when the config declares a model with `profiles`, so that +first start does not produce a control plane. diff --git a/internal/cli/keys.go b/internal/cli/keys.go index fc60352..595b0a2 100644 --- a/internal/cli/keys.go +++ b/internal/cli/keys.go @@ -36,6 +36,7 @@ func newTenancyCmd() *cobra.Command { newKeyRevokeCmd(open), newLimitsCmd(open), newThreadsCmd(open), + newMemoryCmd(open), ) return cmd } diff --git a/internal/cli/memory.go b/internal/cli/memory.go new file mode 100644 index 0000000..a11bab7 --- /dev/null +++ b/internal/cli/memory.go @@ -0,0 +1,62 @@ +package cli + +import ( + "fmt" + + "github.com/spf13/cobra" +) + +// newMemoryCmd turns note storage on for a project, or off. +// +// A command rather than a configuration key for the same reason threads are: +// keeping what an agent decided to remember is a decision about data, not about +// deployment, and it should look like an act rather than a default nobody read. +func newMemoryCmd(open opener) *cobra.Command { + var ( + project string + hours int + off bool + ) + + cmd := &cobra.Command{ + Use: "memory", + Short: "Keep notes an agent writes for itself, across conversations", + Long: "With this on, an agent can save and search notes through /v1/memory or the\n" + + "MCP tools at /mcp (remember, recall, forget), and find them again in later\n" + + "conversations.\n\n" + + "Off for every project until this is run. The retention is required: notes\n" + + "expire that many hours after they were last written, so a store nobody\n" + + "chose a number for does not keep everything forever.", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + if !off && hours <= 0 { + return fmt.Errorf("--retain-hours is required: say how long notes are kept, or pass --off") + } + s, err := open(cmd.Context()) + if err != nil { + return err + } + defer s.Close() + + keep := hours + if off { + keep = 0 + } + if err := s.SetMemoryRetention(cmd.Context(), project, keep); err != nil { + return err + } + if keep == 0 { + fmt.Fprintf(cmd.OutOrStdout(), + "%s keeps no notes. Existing ones are removed at the next prune.\n", project) + return nil + } + fmt.Fprintf(cmd.OutOrStdout(), + "%s keeps notes for %d hours after they were last written.\n", project, keep) + return nil + }, + } + cmd.Flags().StringVar(&project, "project", "default", "project id") + cmd.Flags().IntVar(&hours, "retain-hours", 0, "how long a note is kept after it was last written") + cmd.Flags().BoolVar(&off, "off", false, "stop keeping notes for this project") + return cmd +} diff --git a/internal/config/config.go b/internal/config/config.go index ab9a83f..6071484 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -52,6 +52,11 @@ type Server struct { // Auth turns on API keys on the client-facing port. Auth Auth `koanf:"auth"` + // Threads keeps conversations and agent memory without API keys, for a + // gateway on one trusted machine. With auth on they live in the auth store, + // per project, and this block is ignored. + Threads Threads `koanf:"threads"` + // Landing configures the page served at "/". Landing Landing `koanf:"landing"` @@ -84,6 +89,22 @@ type Landing struct { Notice string `koanf:"notice"` } +// Threads configures conversation and memory storage for a gateway without +// authentication. +// +// Everything stored is in one pool with no owner, readable by anyone who can +// reach the client port and knows an id. That is acceptable on one machine and +// wrong on a shared network, which is why it takes a setting rather than being +// the default. +type Threads struct { + // Store is the SQLite file. Empty means nothing is kept. + Store string `koanf:"store"` + // RetainHours is how long idle conversations and notes are kept. Required + // with Store: a store with no expiry keeps everything anyone typed, by + // omission. + RetainHours int `koanf:"retain_hours"` +} + // Auth configures client authentication (plan ยง7.9). // // Off by default, because a gateway on a trusted network with no keys issued @@ -391,6 +412,9 @@ func (s Server) Validate() error { if s.Auth.Enabled && s.Auth.Store == "" { return fmt.Errorf("auth.store is required when auth is enabled") } + if s.Threads.Store != "" && s.Threads.RetainHours <= 0 { + return fmt.Errorf("threads.retain_hours must be positive when threads.store is set") + } seen := make(map[string]struct{}, len(s.Endpoints)) for i, e := range s.Endpoints { diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 1209c7e..4974d74 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -99,6 +99,14 @@ func (r *Router) Handler() http.Handler { // client endpoint, and scoped to the caller's project by the store. mux.HandleFunc("DELETE /v1/threads/{id}", r.handleDeleteThread) + // Notes an agent keeps for itself, over REST and as MCP tools. Off unless a + // project or the threads block turned storage on. + mux.HandleFunc("PUT /v1/memory/{key}", r.handlePutNote) + mux.HandleFunc("GET /v1/memory/{key}", r.handleGetNote) + mux.HandleFunc("GET /v1/memory", r.handleSearchNotes) + mux.HandleFunc("DELETE /v1/memory/{key}", r.handleDeleteNote) + mux.HandleFunc("POST /mcp", r.handleMCP) + if r.store != nil { // The one-time link that delivers a key to the person it was issued // for. Served without a key, because the person opening it is the one @@ -175,6 +183,9 @@ type parsedRequest struct { // exchange can be stored after a reply without re-reading a body that has // since had the history spliced into it. ThreadTurns []store.Message + // ThreadProject is the project the thread is stored under, resolved once by + // loadThread so saving it afterwards cannot disagree. + ThreadProject string // reservation is carried from the limit check to the point where the // request is settled, so the record counts against the same budget window diff --git a/internal/router/memory.go b/internal/router/memory.go new file mode 100644 index 0000000..6cd8802 --- /dev/null +++ b/internal/router/memory.go @@ -0,0 +1,361 @@ +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) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 7b8bf06..193119d 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -76,7 +76,19 @@ type Router struct { accounts *account.Service store *storepkg.SQLite acct *accounting.Accountant - recorder *accounting.Recorder + + // threads holds conversations and memory notes. It is store when auth is on, + // its own file when only the threads block is configured, and nil when + // nothing is kept. Separate from store because a non-nil store also switches + // on the accounts, web pages and usage recording, none of which belong to a + // gateway that only wants to remember things. + threads *storepkg.SQLite + // anonymous means threads has no project to scope to: everything shares the + // one pool named by anonymousProject. + anonymous bool + // ownsThreads is true when threads is not store and so must be closed here. + ownsThreads bool + recorder *accounting.Recorder // prefix is nil unless approximate prefix routing is turned on. prefix *scheduler.PrefixScorer @@ -220,6 +232,17 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k, }), r.recorder, log) + r.threads = r.store + } else if cfg.Threads.Store != "" { + st, err := storepkg.OpenSQLite(context.Background(), cfg.Threads.Store) + if err != nil { + return nil, fmt.Errorf("threads store: %w", err) + } + if err := st.EnsureAnonymousProject(context.Background(), anonymousProject, cfg.Threads.RetainHours); err != nil { + _ = st.Close() + return nil, fmt.Errorf("threads store: %w", err) + } + r.threads, r.anonymous, r.ownsThreads = st, true, true } if cfg.Router.Admission.Enabled { @@ -333,7 +356,7 @@ func (r *Router) Start(ctx context.Context) { r.log.Warn("could not restore budgets; they start from zero this period", "err", err) } } - if r.store != nil { + if r.threads != nil { go r.pruneThreads(ctx) } } @@ -349,7 +372,7 @@ func (r *Router) Start(ctx context.Context) { // an hour is fine granularity for a limit expressed in hours. func (r *Router) pruneThreads(ctx context.Context) { prune := func() { - n, err := r.store.PruneThreads(ctx, time.Now()) + n, err := r.threads.PruneThreads(ctx, time.Now()) if err != nil { r.log.Warn("could not prune expired conversations", "err", err) return @@ -357,6 +380,14 @@ func (r *Router) pruneThreads(ctx context.Context) { if n > 0 { r.log.Info("removed conversations past their retention", "threads", n) } + m, err := r.threads.PruneMemory(ctx, time.Now()) + if err != nil { + r.log.Warn("could not prune expired memory notes", "err", err) + return + } + if m > 0 { + r.log.Info("removed memory notes past their retention", "notes", m) + } } prune() @@ -385,6 +416,9 @@ func (r *Router) Stop() { if r.store != nil { _ = r.store.Close() } + if r.ownsThreads { + _ = r.threads.Close() + } } // Registry exposes the endpoint registry, for tests and admin endpoints. diff --git a/internal/router/threads.go b/internal/router/threads.go index f5bcd37..0bedc33 100644 --- a/internal/router/threads.go +++ b/internal/router/threads.go @@ -212,15 +212,14 @@ func (p *parsedRequest) captureBytes() int { // never told โ€” which reads as the model being wrong rather than as a setting // being off. func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error) { - if req.ThreadID == "" || r.store == nil { + if req.ThreadID == "" { return 0, nil } - if req.ProjectID == "" { - // No identity, so no project to scope the thread to. Storing - // conversations for an unauthenticated gateway would put every - // caller's history in one pool, readable by anyone who guesses an id. - return 0, errThreadsNeedAuth + project, err := r.threadProject(req.ProjectID) + if err != nil { + return 0, err } + req.ThreadProject = project // Asked before anything else, because a project that keeps no // conversations must be told so rather than served. Without this the @@ -228,7 +227,7 @@ func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error // the model has no idea what they are referring to โ€” which reads as the // model being broken. The e2e test that caught this asserted the refusal // and got a cheerful 200. - retention, err := r.store.ThreadRetention(ctx, req.ProjectID) + retention, err := r.threads.ThreadRetention(ctx, project) if err != nil { return 0, err } @@ -240,7 +239,7 @@ func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error // can be stored afterwards without re-reading a body that has changed. req.ThreadTurns = ownMessages(req.UpstreamBody) - history, err := r.store.Thread(ctx, req.ProjectID, req.ThreadID) + history, err := r.threads.Thread(ctx, project, req.ThreadID) if err != nil && !errors.Is(err, store.ErrNoThread) { return 0, err } @@ -279,23 +278,52 @@ func (r *Router) loadThread(ctx context.Context, req *parsedRequest) (int, error // The failure is logged rather than swallowed, because a thread silently not // growing is the kind of fault someone discovers days later. func (r *Router) saveThread(ctx context.Context, req *parsedRequest, reply string, model string) { - if req.ThreadID == "" || r.store == nil || reply == "" { + if req.ThreadID == "" || r.threads == nil || req.ThreadProject == "" || reply == "" { return } turns := append([]store.Message(nil), req.ThreadTurns...) turns = append(turns, store.Message{Role: "assistant", Content: reply, Model: model}) - if err := r.store.AppendThread(ctx, req.ProjectID, req.ThreadID, turns); err != nil { + if err := r.threads.AppendThread(ctx, req.ThreadProject, req.ThreadID, turns); err != nil { r.log.Warn("could not store a conversation turn", - "thread", req.ThreadID, "project", req.ProjectID, "err", err) + "thread", req.ThreadID, "project", req.ThreadProject, "err", err) } } +// anonymousProject owns everything kept by a gateway with no authentication. +const anonymousProject = "_anonymous" + // errThreadsNeedAuth is refused rather than ignored, for the same reason a // disabled project is. var errThreadsNeedAuth = errors.New( "a conversation needs an API key: without one there is no project to keep it in") +// errThreadsOff is what a thread header meets when nothing is configured to +// keep threads. Refused rather than ignored: ignoring it sends the model a +// conversation with no history, and it answers confidently about something it +// was never told. +var errThreadsOff = errors.New( + "this gateway keeps no conversations: enable auth with a project that keeps " + + "threads, or set threads.store and threads.retain_hours") + +// threadProject says which project a request's threads and notes belong to. +// +// The caller's own project when authenticated. The shared anonymous one only +// when the operator configured it, because with no identity every caller's +// history lands in one pool readable by anyone who knows an id. +func (r *Router) threadProject(callerProject string) (string, error) { + switch { + case r.threads == nil: + return "", errThreadsOff + case callerProject != "": + return callerProject, nil + case r.anonymous: + return anonymousProject, nil + default: + return "", errThreadsNeedAuth + } +} + // handleDeleteThread forgets a conversation. // // A client endpoint rather than an admin one: the person who said the thing is @@ -303,12 +331,16 @@ var errThreadsNeedAuth = errors.New( // would mean the operator reads it first. func (r *Router) handleDeleteThread(w http.ResponseWriter, req *http.Request) { id := req.PathValue("id") - if r.store == nil { + if r.threads == nil { writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread") return } - p, ok := auth.FromContext(req.Context()) - if !ok || p.ProjectID == "" { + var caller string + if p, ok := auth.FromContext(req.Context()); ok { + caller = p.ProjectID + } + project, err := r.threadProject(caller) + if err != nil { writeError(w, http.StatusUnauthorized, "invalid_request_error", "deleting a conversation needs the key that owns it") return @@ -318,7 +350,7 @@ func (r *Router) handleDeleteThread(w http.ResponseWriter, req *http.Request) { // another one is "no such thread" โ€” the same answer as one that never // existed, because two answers would let a caller map which ids are real // in projects they cannot read. - if err := r.store.DeleteThread(req.Context(), p.ProjectID, id); err != nil { + if err := r.threads.DeleteThread(req.Context(), project, id); err != nil { if errors.Is(err, store.ErrNoThread) { writeError(w, http.StatusNotFound, "invalid_request_error", "no such thread") return diff --git a/internal/store/memory.go b/internal/store/memory.go new file mode 100644 index 0000000..cc00bc7 --- /dev/null +++ b/internal/store/memory.go @@ -0,0 +1,284 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "fmt" + "regexp" + "strings" + "time" +) + +// ErrMemoryNotEnabled means the project has not asked llmgw to keep notes. +// +// Returned rather than ignoring the write: an agent that believes it saved +// something and finds nothing later has been lied to. +var ErrMemoryNotEnabled = errors.New("store: this project does not keep memory notes") + +// ErrNoNote means the key does not exist, or belongs to another project or agent. +var ErrNoNote = errors.New("store: no such note") + +// ErrMemoryFull means the agent already holds as many notes as one may. +var ErrMemoryFull = errors.New("store: this agent has reached its limit of notes") + +// Limits on what one note and one agent may hold. Notes are memory held on +// behalf of a caller, so each is bounded rather than trusted. +const ( + MaxNoteKey = 200 + MaxNoteText = 16 << 10 + MaxNoteTags = 16 + MaxNotesPerAgent = 1000 + MaxNoteResults = 100 +) + +var tagPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,63}$`) + +// Note is one thing an agent chose to remember. +type Note struct { + Key string + Text string + Tags []string + CreatedAt time.Time + UpdatedAt time.Time +} + +// NormalizeTag lowercases and validates a tag, so "Work" and "work" are one tag +// and nothing outside a small alphabet can reach the stored form. +func NormalizeTag(tag string) (string, error) { + t := strings.ToLower(strings.TrimSpace(tag)) + if !tagPattern.MatchString(t) { + return "", fmt.Errorf("invalid tag %q: use lowercase letters, digits, - and _ (up to 64)", tag) + } + return t, nil +} + +func encodeTags(tags []string) (string, error) { + seen := map[string]bool{} + var out []string + for _, tag := range tags { + t, err := NormalizeTag(tag) + if err != nil { + return "", err + } + if !seen[t] { + seen[t] = true + out = append(out, t) + } + } + if len(out) > MaxNoteTags { + return "", fmt.Errorf("too many tags: at most %d", MaxNoteTags) + } + if len(out) == 0 { + return "", nil + } + return "," + strings.Join(out, ",") + ",", nil +} + +func decodeTags(s string) []string { + s = strings.Trim(s, ",") + if s == "" { + return nil + } + return strings.Split(s, ",") +} + +// MemoryRetention reports how long a project keeps notes. Zero means none. +func (s *SQLite) MemoryRetention(ctx context.Context, projectID string) (time.Duration, error) { + var hours sql.NullInt64 + err := s.db.QueryRowContext(ctx, + `SELECT memory_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours) + if errors.Is(err, sql.ErrNoRows) { + return 0, nil + } + if err != nil { + return 0, err + } + if !hours.Valid || hours.Int64 <= 0 { + return 0, nil + } + return time.Duration(hours.Int64) * time.Hour, nil +} + +// SetMemoryRetention turns note storage on for a project, or off. +func (s *SQLite) SetMemoryRetention(ctx context.Context, projectID string, hours int) error { + var v any + if hours > 0 { + v = hours + } + res, err := s.db.ExecContext(ctx, + `UPDATE projects SET memory_retention_hours = ? WHERE id = ?`, v, projectID) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return fmt.Errorf("store: no project %q", projectID) + } + return nil +} + +// PutNote writes a note, replacing any note the agent already holds under the +// same key. +func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note) error { + if n.Key == "" || len(n.Key) > MaxNoteKey { + return fmt.Errorf("key must be 1 to %d bytes", MaxNoteKey) + } + if n.Text == "" || len(n.Text) > MaxNoteText { + return fmt.Errorf("text must be 1 to %d bytes", MaxNoteText) + } + tags, err := encodeTags(n.Tags) + if err != nil { + return err + } + retention, err := s.MemoryRetention(ctx, projectID) + if err != nil { + return err + } + if retention == 0 { + return ErrMemoryNotEnabled + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + + now := time.Now().UnixMilli() + res, err := tx.ExecContext(ctx, + `UPDATE memory_notes SET text = ?, tags = ?, updated_at = ? + WHERE project_id = ? AND agent_id = ? AND key = ?`, + n.Text, tags, now, projectID, agentID, n.Key) + if err != nil { + return err + } + if updated, _ := res.RowsAffected(); updated == 0 { + var count int + if err := tx.QueryRowContext(ctx, + `SELECT COUNT(*) FROM memory_notes WHERE project_id = ? AND agent_id = ?`, + projectID, agentID).Scan(&count); err != nil { + return err + } + if count >= MaxNotesPerAgent { + return ErrMemoryFull + } + if _, err := tx.ExecContext(ctx, + `INSERT INTO memory_notes (project_id, agent_id, key, text, tags, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + projectID, agentID, n.Key, n.Text, tags, now, now); err != nil { + return err + } + } + return tx.Commit() +} + +func scanNote(rows interface{ Scan(...any) error }) (Note, error) { + var ( + n Note + tags string + created int64 + updatedRaw int64 + ) + if err := rows.Scan(&n.Key, &n.Text, &tags, &created, &updatedRaw); err != nil { + return Note{}, err + } + n.Tags = decodeTags(tags) + n.CreatedAt = time.UnixMilli(created) + n.UpdatedAt = time.UnixMilli(updatedRaw) + return n, nil +} + +const noteColumns = `key, text, tags, created_at, updated_at` + +// GetNote reads one note. Scoped to the project and agent in the query itself. +func (s *SQLite) GetNote(ctx context.Context, projectID, agentID, key string) (Note, error) { + n, err := scanNote(s.db.QueryRowContext(ctx, + `SELECT `+noteColumns+` FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`, + projectID, agentID, key)) + if errors.Is(err, sql.ErrNoRows) { + return Note{}, ErrNoNote + } + return n, err +} + +// ftsQuery turns free text into a query that cannot be a syntax error: every +// word becomes a quoted term, and they must all match. FTS5 treats characters +// such as - : * and " as operators, and an agent's search words are not written +// for that syntax. +func ftsQuery(q string) string { + words := strings.Fields(q) + for i, w := range words { + words[i] = `"` + strings.ReplaceAll(w, `"`, `""`) + `"` + } + return strings.Join(words, " ") +} + +// SearchNotes finds an agent's notes by words and/or tag, newest first. With +// neither it lists the most recent. +func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag string, limit int) ([]Note, error) { + if limit <= 0 || limit > MaxNoteResults { + limit = MaxNoteResults + } + q := `SELECT ` + noteColumns + ` FROM memory_notes WHERE project_id = ? AND agent_id = ?` + args := []any{projectID, agentID} + + if match := ftsQuery(query); match != "" { + q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)` + args = append(args, match) + } + if tag != "" { + t, err := NormalizeTag(tag) + if err != nil { + return nil, err + } + q += ` AND tags LIKE ?` + args = append(args, "%,"+t+",%") + } + q += ` ORDER BY updated_at DESC, id DESC LIMIT ?` + args = append(args, limit) + + return queryAll(ctx, s.db, q, args, func(rows *sql.Rows) (Note, error) { return scanNote(rows) }) +} + +// DeleteNote forgets one note. +func (s *SQLite) DeleteNote(ctx context.Context, projectID, agentID, key string) error { + res, err := s.db.ExecContext(ctx, + `DELETE FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`, + projectID, agentID, key) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrNoNote + } + return nil +} + +// PruneMemory deletes notes not refreshed within their project's retention. +func (s *SQLite) PruneMemory(ctx context.Context, now time.Time) (int, error) { + type window struct { + project string + hours int64 + } + windows, err := queryAll(ctx, s.db, + `SELECT id, memory_retention_hours FROM projects WHERE memory_retention_hours > 0`, + nil, func(rows *sql.Rows) (window, error) { + var w window + return w, rows.Scan(&w.project, &w.hours) + }) + if err != nil { + return 0, err + } + var removed int + for _, w := range windows { + cutoff := now.Add(-time.Duration(w.hours) * time.Hour).UnixMilli() + res, err := s.db.ExecContext(ctx, + `DELETE FROM memory_notes WHERE project_id = ? AND updated_at < ?`, w.project, cutoff) + if err != nil { + return removed, err + } + n, _ := res.RowsAffected() + removed += int(n) + } + return removed, nil +} diff --git a/internal/store/memory_test.go b/internal/store/memory_test.go new file mode 100644 index 0000000..fd52065 --- /dev/null +++ b/internal/store/memory_test.go @@ -0,0 +1,208 @@ +package store + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func memoryProject(t *testing.T, s *SQLite) string { + t.Helper() + proj := tree(t, s) + if err := s.SetMemoryRetention(context.Background(), proj, 24); err != nil { + t.Fatal(err) + } + return proj +} + +func TestAProjectKeepsNoNotesUnlessItAsks(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := tree(t, s) + + err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "secret"}) + if !errors.Is(err, ErrMemoryNotEnabled) { + t.Fatalf("writing into a project that never asked returned %v", err) + } + if _, err := s.GetNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) { + t.Errorf("a refused write left a note behind: %v", err) + } +} + +func TestANoteIsWrittenReadBackAndReplaced(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + + if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Ankara", Tags: []string{"Geo", "geo", "facts"}}); err != nil { + t.Fatal(err) + } + got, err := s.GetNote(ctx, proj, "a", "city") + if err != nil { + t.Fatal(err) + } + if got.Text != "Ankara" || len(got.Tags) != 2 || got.Tags[0] != "geo" || got.Tags[1] != "facts" { + t.Fatalf("note came back as %+v: tags should be lowercased and deduplicated", got) + } + + if err := s.PutNote(ctx, proj, "a", Note{Key: "city", Text: "Istanbul"}); err != nil { + t.Fatal(err) + } + got, _ = s.GetNote(ctx, proj, "a", "city") + if got.Text != "Istanbul" || len(got.Tags) != 0 { + t.Errorf("rewriting a key should replace text and tags: %+v", got) + } + notes, _ := s.SearchNotes(ctx, proj, "a", "", "", 0) + if len(notes) != 1 { + t.Errorf("rewriting a key created %d notes, want 1", len(notes)) + } +} + +func TestNotesAreFoundByWordsAndTagsAndSpecialCharactersAreHarmless(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + put := func(key, text string, tags ...string) { + t.Helper() + if err := s.PutNote(ctx, proj, "a", Note{Key: key, Text: text, Tags: tags}); err != nil { + t.Fatal(err) + } + } + put("deploy", "the staging cluster is on port 8443", "ops") + put("food", "she prefers green tea", "personal") + put("old", "deploy notes from last year", "ops", "archive") + + byWord, err := s.SearchNotes(ctx, proj, "a", "staging cluster", "", 0) + if err != nil || len(byWord) != 1 || byWord[0].Key != "deploy" { + t.Fatalf("search by words: %+v (err %v)", byWord, err) + } + byTag, _ := s.SearchNotes(ctx, proj, "a", "", "ops", 0) + if len(byTag) != 2 { + t.Errorf("tag ops matched %d notes, want 2", len(byTag)) + } + both, _ := s.SearchNotes(ctx, proj, "a", "deploy", "archive", 0) + if len(both) != 1 || both[0].Key != "old" { + t.Errorf("words and tag together: %+v", both) + } + // A tag is matched whole: "op" must not match "ops". + if part, _ := s.SearchNotes(ctx, proj, "a", "", "op", 0); len(part) != 0 { + t.Errorf("a partial tag matched %d notes", len(part)) + } + // FTS operators in an agent's words must not become a syntax error. + for _, q := range []string{`"`, `a -b`, `x:y`, `foo*`, `AND OR NOT`, `(`} { + if _, err := s.SearchNotes(ctx, proj, "a", q, "", 0); err != nil { + t.Errorf("query %q failed: %v", q, err) + } + } +} + +func TestOneAgentOrProjectCannotReadAnothersNotes(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + if err := s.PutOrg(ctx, Org{ID: "o2", Name: "Other"}); err != nil { + t.Fatal(err) + } + if err := s.PutTeam(ctx, Team{ID: "t2", OrgID: "o2", Name: "T"}); err != nil { + t.Fatal(err) + } + if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "t2", Name: "P"}); err != nil { + t.Fatal(err) + } + if err := s.SetMemoryRetention(ctx, "proj2", 24); err != nil { + t.Fatal(err) + } + + if err := s.PutNote(ctx, proj, "agent1", Note{Key: "k", Text: "private"}); err != nil { + t.Fatal(err) + } + if _, err := s.GetNote(ctx, proj, "agent2", "k"); !errors.Is(err, ErrNoNote) { + t.Errorf("another agent in the same project read the note: %v", err) + } + if _, err := s.GetNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) { + t.Errorf("another project read the note: %v", err) + } + if got, _ := s.SearchNotes(ctx, "proj2", "agent1", "private", "", 0); len(got) != 0 { + t.Errorf("another project found the note by search: %+v", got) + } + if err := s.DeleteNote(ctx, "proj2", "agent1", "k"); !errors.Is(err, ErrNoNote) { + t.Errorf("another project deleted the note: %v", err) + } +} + +func TestNotesAreBoundedAndBadInputIsRefused(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + + bad := []Note{ + {Key: "", Text: "x"}, + {Key: "k", Text: ""}, + {Key: strings.Repeat("k", MaxNoteKey+1), Text: "x"}, + {Key: "k", Text: strings.Repeat("x", MaxNoteText+1)}, + {Key: "k", Text: "x", Tags: []string{"has space"}}, + } + for i, n := range bad { + if err := s.PutNote(ctx, proj, "a", n); err == nil { + t.Errorf("bad note %d was accepted", i) + } + } + + for i := 0; i < MaxNotesPerAgent; i++ { + if err := s.PutNote(ctx, proj, "a", Note{Key: fmt.Sprintf("note-%d", i), Text: "x"}); err != nil { + t.Fatalf("note %d: %v", i, err) + } + } + if err := s.PutNote(ctx, proj, "a", Note{Key: "one-too-many", Text: "x"}); !errors.Is(err, ErrMemoryFull) { + t.Errorf("exceeding the per-agent limit returned %v", err) + } + // Replacing an existing note is not growth and must still work when full. + if err := s.PutNote(ctx, proj, "a", Note{Key: "note-0", Text: "y"}); err != nil { + t.Errorf("replacing a note in a full agent failed: %v", err) + } + // Another agent has its own allowance. + if err := s.PutNote(ctx, proj, "b", Note{Key: "note-0", Text: "x"}); err != nil { + t.Errorf("a second agent was blocked by the first one's notes: %v", err) + } +} + +func TestAForgottenNoteIsGoneFromSearchToo(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + + if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "needle in a haystack"}); err != nil { + t.Fatal(err) + } + if err := s.DeleteNote(ctx, proj, "a", "k"); err != nil { + t.Fatal(err) + } + if got, _ := s.SearchNotes(ctx, proj, "a", "needle", "", 0); len(got) != 0 { + t.Errorf("a deleted note is still searchable: %+v", got) + } + if err := s.DeleteNote(ctx, proj, "a", "k"); !errors.Is(err, ErrNoNote) { + t.Errorf("deleting twice returned %v", err) + } +} + +func TestNotesNotRefreshedWithinRetentionArePruned(t *testing.T) { + ctx := context.Background() + s := openDB(t) + proj := memoryProject(t, s) + + if err := s.PutNote(ctx, proj, "a", Note{Key: "k", Text: "x"}); err != nil { + t.Fatal(err) + } + if n, err := s.PruneMemory(ctx, time.Now()); err != nil || n != 0 { + t.Fatalf("a fresh note was pruned: %d, %v", n, err) + } + if n, err := s.PruneMemory(ctx, time.Now().Add(25*time.Hour)); err != nil || n != 1 { + t.Fatalf("an expired note survived: %d, %v", n, err) + } + if got, _ := s.SearchNotes(ctx, proj, "a", "", "", 0); len(got) != 0 { + t.Errorf("pruned note still listed: %+v", got) + } +} diff --git a/internal/store/migrations/sqlite/0007_memory.sql b/internal/store/migrations/sqlite/0007_memory.sql new file mode 100644 index 0000000..2754cb8 --- /dev/null +++ b/internal/store/migrations/sqlite/0007_memory.sql @@ -0,0 +1,45 @@ +-- Notes an agent writes for itself and reads back later, across threads. +-- +-- Off unless a project turns it on, for the same reason threads are: llmgw +-- stores no message content unless somebody decided it should. NULL means this +-- project keeps no notes. The number is hours, and is required to turn it on, +-- because a store with no expiry is a transcript kept forever by omission. +ALTER TABLE projects ADD COLUMN memory_retention_hours INTEGER; + +CREATE TABLE memory_notes ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + -- Scoped to a project and checked in SQL on every read, like threads. + project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + -- A namespace inside the project, so two agents sharing a key do not read + -- each other's notes. + agent_id TEXT NOT NULL, + key TEXT NOT NULL, + text TEXT NOT NULL, + -- Normalised to ",a,b," so a tag matches whole, never as part of another. + tags TEXT NOT NULL DEFAULT '', + created_at INTEGER NOT NULL, + -- Touched on every write, so retention measures how long a note has gone + -- unrefreshed rather than how old it is. + updated_at INTEGER NOT NULL, + UNIQUE (project_id, agent_id, key) +); +CREATE INDEX memory_notes_by_updated ON memory_notes(updated_at); + +-- Full-text search over the fields an agent would search by. An external +-- content table, so the text is stored once. +CREATE VIRTUAL TABLE memory_notes_fts USING fts5( + key, text, tags, content = 'memory_notes', content_rowid = 'id' +); + +CREATE TRIGGER memory_notes_ai AFTER INSERT ON memory_notes BEGIN + INSERT INTO memory_notes_fts(rowid, key, text, tags) VALUES (new.id, new.key, new.text, new.tags); +END; +CREATE TRIGGER memory_notes_ad AFTER DELETE ON memory_notes BEGIN + INSERT INTO memory_notes_fts(memory_notes_fts, rowid, key, text, tags) + VALUES ('delete', old.id, old.key, old.text, old.tags); +END; +CREATE TRIGGER memory_notes_au AFTER UPDATE ON memory_notes BEGIN + INSERT INTO memory_notes_fts(memory_notes_fts, rowid, key, text, tags) + VALUES ('delete', old.id, old.key, old.text, old.tags); + INSERT INTO memory_notes_fts(rowid, key, text, tags) VALUES (new.id, new.key, new.text, new.tags); +END; diff --git a/internal/store/threads.go b/internal/store/threads.go index 944aa30..7c2d042 100644 --- a/internal/store/threads.go +++ b/internal/store/threads.go @@ -33,6 +33,28 @@ type Message struct { Model string } +// EnsureAnonymousProject creates the single project that owns everything kept +// by a gateway with no authentication, and sets how long it keeps it. +// +// A real project rather than a special case in every query: threads and notes +// are scoped by project id and enforce it in SQL, and an empty id would have to +// be an exception in each of those places. +func (s *SQLite) EnsureAnonymousProject(ctx context.Context, id string, retainHours int) error { + if err := s.PutOrg(ctx, Org{ID: id, Name: id}); err != nil { + return err + } + if err := s.PutTeam(ctx, Team{ID: id, OrgID: id, Name: id}); err != nil { + return err + } + if err := s.PutProject(ctx, Project{ID: id, TeamID: id, Name: id}); err != nil { + return err + } + if err := s.SetThreadRetention(ctx, id, retainHours); err != nil { + return err + } + return s.SetMemoryRetention(ctx, id, retainHours) +} + // ThreadRetention reports how long a project keeps conversations. // // Zero means it keeps none, which is the default and was the only behaviour -- 2.47.3