diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cce4cf1..89d289a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,30 +1,27 @@ -# Contributing to engram +# Contributing to Engram -Welcome! Contributions of all kinds are encouraged and appreciated. +Welcome! Contributions of all kinds are encouraged. -## Development Environment Setup +## Development environment setup -1. **Prerequisites:** - - Rust (stable) - - Docker (for Redis) - - OpenAI API key -2. **Clone the repository:** +1. Prerequisites: Rust (stable), Docker (for Redis), OpenAI API key +2. Clone the repository: ```sh git clone https://github.com/bit2swaz/engram.git cd engram ``` -3. **Set up environment variables:** +3. Set up environment variables: ```sh cp .env.example .env # Edit .env and fill in your OpenAI API key ``` -4. **Start Redis:** +4. Start Redis: ```sh docker compose up -d redis # or docker run -d --name engram-redis -p 6379:6379 redis:7-alpine ``` -5. **Build the project:** +5. Build the project: ```sh cargo build ``` @@ -65,16 +62,13 @@ docker compose -f docker-compose.cluster.yml up -d --build docker compose -f docker-compose.cluster.yml down ``` -The verify script checks 10 criteria: leader election, write replication, follower redirect, failover, Prometheus metrics, knowledge graph replication, entity graph traversal, delete-session cleanup, node restart and recovery from the Raft log, snapshot compaction, and full state restoration from a snapshot. It exits 0 only when all criteria pass. +The verify script checks 17 criteria: leader election, write replication, follower redirect, failover, Prometheus metrics, knowledge graph replication, entity graph traversal, delete-session cleanup, node restart and recovery from the Raft log, snapshot compaction, full state restoration from a snapshot, session visibility propagation, global graph population from public sessions, agent registration, global entity/relationship count metrics, global entity queries, conflict detection, and global graph snapshot round-trip. It exits 0 only when all criteria pass. -## TDD Workflow +## TDD workflow -This project uses a strict test-driven development (TDD) workflow: -- **Red:** Write a failing test that describes the desired behavior. -- **Green:** Implement the minimum code needed to make the test pass. -- **Refactor:** Clean up the code while keeping all tests green. +This project uses strict TDD: write a failing test, implement the minimum code to make it pass, then refactor while keeping all tests green. No implementation code goes in without a test first. -## Branch Naming Convention +## Branch naming - `feat/short-description` (new features) - `fix/issue-number` (bug fixes) @@ -82,18 +76,16 @@ This project uses a strict test-driven development (TDD) workflow: - `test/unit-coverage` (tests) - `chore/dependency-update` (maintenance) -## Commit Message Format +## Commit message format Follow [Conventional Commits](https://www.conventionalcommits.org/): - `type: short description` - Types: `feat`, `fix`, `docs`, `test`, `chore`, `bench`, `ci` - Example: `feat: add redis-backed short term store` -## Pull Requests +## Pull requests -- PRs should link to an issue when possible. -- Include a clear summary of the change. -- All tests and CI must pass before merging. +PRs should link to an issue when possible, include a clear summary of the change, and pass all tests and CI before merging. ## Code of Conduct diff --git a/README.md b/README.md index 9158516..6c47ba8 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# engram +# Engram An asynchronous semantic memory backend for LLM agents, written in Rust. @@ -7,11 +7,13 @@ An asynchronous semantic memory backend for LLM agents, written in Rust. ## overview -engram is a backend service for large language model (LLM) agents. it provides three types of memory: short-term (recent messages), long-term (semantic vector search), and core memory (pinned facts). the goal is to give LLM agents a transparent, efficient, and controllable way to manage context and recall information. +Engram is a backend service for LLM agents. It stores three types of memory: short-term (recent messages), long-term (semantic vector search), and core memory (pinned facts). The goal is to give agents a transparent, efficient, and controllable way to manage context. -engram is written in rust for performance and reliability. it is designed for transparency, with full control over token budgets and context assembly, and exposes all operations via a simple REST API. +Engram is written in Rust for performance and reliability. It exposes all operations via a REST API and gives you full control over token budgets and context assembly. -engram is built for developers who want to plug in their own LLM agents, run locally or in production, and have full visibility into how memory is managed. it is easy to run, test, and extend. all memory operations are behind trait abstractions, making it easy to swap implementations or mock for tests. +Engram is built for developers who want to plug in their own LLM agents, run locally or in production, and see exactly what goes into the context window. All memory operations are behind trait abstractions, so you can swap implementations or mock them in tests without changing any calling code. + +With the latest update, Engram also supports collective memory: multiple agents can share a global knowledge graph, sessions can be marked public or private, and conflicting facts across sessions are surfaced via a dedicated endpoint. ## architecture @@ -26,14 +28,20 @@ graph TD router --> corememhandler["core memory handler"] router --> healthhandler["health handler"] router --> knowledgehandler["knowledge handler"] + router --> visibilityhandler["visibility handler"] + router --> globalhandler["global knowledge handler"] memoryhandler -->|write| raft["raft consensus\nopenraft 0.9\ncluster mode only"] corememhandler -->|write| raft - sessionhandler -->|delete| raft + sessionhandler -->|delete or register agent| raft + visibilityhandler -->|write| raft raft -.->|grpc append entries| peers["peer nodes (port 9001)"] raft -.->|grpc install snapshot| laggers["lagging followers"] raft -->|state machine apply| shortterm["short-term memory trait"] raft -->|state machine apply| coremem["core memory store trait"] + raft -->|state machine apply| knowledgegraph[("knowledge graph\nper-session in-memory")] + raft -->|state machine apply| globalgraph[("global knowledge graph\ncross-session")] + raft -->|state machine apply| visibility["session visibility map"] raft -->|embedding job| embedqueue["embedding worker pool
bounded channel"] raft -->|knowledge job| knowledgequeue["knowledge worker pool
bounded channel"] raft --> redb[("redb\npersistent raft log\n+ snapshot store")] @@ -51,8 +59,11 @@ graph TD knowledgequeue -->|leader-only extraction| extractor["knowledge extractor trait"] extractor -->|openai gpt-4o-mini / mock| extraction["entities + relationships"] - extraction -->|AddKnowledge via raft| knowledgegraph[("knowledge graph\nper-session in-memory")] + extraction -->|AddKnowledge via raft| knowledgegraph + knowledgegraph -->|public sessions merge| globalgraph + knowledgehandler --> knowledgegraph + globalhandler --> globalgraph contexthandler --> assembler["context assembler"] assembler --> shortterm @@ -104,6 +115,12 @@ create a session: ```sh curl -X POST http://localhost:3000/sessions ``` +create a session and register an agent: +```sh +curl -X POST http://localhost:3000/sessions \ + -H 'content-type: application/json' \ + -d '{"agent_id":"agent-42"}' +``` add a message: ```sh curl -X POST http://localhost:3000/sessions/{session_id}/messages \ @@ -126,6 +143,16 @@ curl -X PUT http://localhost:3000/sessions/{session_id}/core-memory \ -H 'content-type: application/json' \ -d '{"fact":"user prefers dark mode"}' ``` +make a session public so it contributes to the global knowledge graph: +```sh +curl -X PUT http://localhost:3000/sessions/{session_id}/visibility \ + -H 'content-type: application/json' \ + -d '{"visibility":"Shared"}' +``` +query the global knowledge graph: +```sh +curl http://localhost:3000/knowledge/global +``` delete session: ```sh curl -X DELETE http://localhost:3000/sessions/{session_id} @@ -143,26 +170,33 @@ docker compose up -d ## API overview -| method | path | description | -|--------|-------------------------------------------------------------|----------------------------------------------| -| GET | /health | health check | -| GET | /metrics | Prometheus metrics | -| GET | /api-docs/openapi.json | OpenAPI specification | -| GET | /swagger-ui/ | Swagger UI | -| POST | /sessions | create session | -| POST | /sessions/{session_id}/messages | add message | -| GET | /sessions/{session_id}/context | get assembled context | -| POST | /sessions/{session_id}/search | semantic search | -| PUT | /sessions/{session_id}/core-memory | add core memory fact | -| DELETE | /sessions/{session_id} | delete session | -| GET | /sessions/{session_id}/knowledge | get knowledge graph (entities + edges) | -| GET | /sessions/{session_id}/knowledge/entities/{entity_name} | get related entities for a given entity | -| GET | /sessions/{session_id}/knowledge/path?from=X&to=Y | find shortest path between two entities | -| GET | /sessions/{session_id}/knowledge/export?format=json\|dot | export knowledge graph (JSON or Graphviz) | -| GET | /cluster | cluster status (cluster mode only) | -| POST | /cluster/init | initialize cluster | -| POST | /cluster/add-learner | add a learner node | -| POST | /cluster/change-membership | promote learners to full members | +| method | path | description | +|--------|---------------------------------------------------------------|----------------------------------------------| +| GET | /health | health check | +| GET | /metrics | Prometheus metrics | +| GET | /api-docs/openapi.json | OpenAPI specification | +| GET | /swagger-ui/ | Swagger UI | +| POST | /sessions | create session (optional agent_id body) | +| POST | /sessions/{session_id}/messages | add message | +| GET | /sessions/{session_id}/context | get assembled context | +| POST | /sessions/{session_id}/search | semantic search | +| PUT | /sessions/{session_id}/core-memory | add core memory fact | +| PUT | /sessions/{session_id}/visibility | set session visibility (Private/Shared) | +| DELETE | /sessions/{session_id} | delete session | +| GET | /sessions/{session_id}/knowledge | get knowledge graph (entities + edges) | +| GET | /sessions/{session_id}/knowledge/entities/{entity_name} | get related entities for a given entity | +| GET | /sessions/{session_id}/knowledge/path?from=X&to=Y | find shortest path between two entities | +| GET | /sessions/{session_id}/knowledge/export?format=json\|dot | export knowledge graph (JSON or Graphviz) | +| GET | /knowledge/global | get cross-session global knowledge graph | +| GET | /knowledge/global/entities/{name} | get related entities in global graph | +| GET | /knowledge/global/entities/{name}/sources | get sessions that contributed this entity | +| GET | /knowledge/global/path?from=X&to=Y | find path in global graph | +| GET | /knowledge/global/export?format=json\|dot | export global graph | +| GET | /knowledge/global/conflicts | list conflicting facts across sessions | +| GET | /cluster | cluster status (cluster mode only) | +| POST | /cluster/init | initialize cluster | +| POST | /cluster/add-learner | add a learner node | +| POST | /cluster/change-membership | promote learners to full members | see [API.md](docs/API.md) for full details. @@ -214,13 +248,17 @@ values like `similarity_threshold` and `max_tokens` are controlled per request t - per-session knowledge graph with BFS path-finding, persisted via snapshots - leader-only extraction with Raft-replicated `AddKnowledge` command (all nodes stay consistent) - knowledge graph export (JSON and Graphviz DOT) -- Prometheus metrics endpoint (includes knowledge and snapshot metrics) +- session visibility control (Private/Shared) via `SetSessionVisibility` Raft command +- global cross-session knowledge graph: public sessions contribute their entities and relationships to a shared graph +- agent registration: sessions can be associated with a named agent at creation time +- conflict detection: the global graph tracks conflicting relationship types across sessions +- Prometheus metrics endpoint (includes knowledge, snapshot, and global graph metrics) - Raft consensus for fault-tolerant distributed writes (OpenRaft 0.9) - gRPC inter-node transport for Raft (Vote, AppendEntries, and InstallSnapshot via tonic 0.12) - follower-to-leader HTTP redirect (307) in cluster mode - per-node LanceDB with eventual consistency via deterministic embeddings - persistent redb-backed Raft log and snapshot store (survives restarts) -- full state machine snapshots covering short-term memory, core memory, and knowledge graph +- full state machine snapshots covering short-term memory, core memory, knowledge graph, global graph, session visibility, and agent registry - startup recovery from the latest snapshot followed by Raft log replay - InstallSnapshot RPC so lagging followers catch up without manual re-initialization - automatic log compaction after every `SNAPSHOT_LOG_THRESHOLD` committed entries @@ -230,7 +268,7 @@ values like `similarity_threshold` and `max_tokens` are controlled per request t ## quickstart (3-node cluster) -the cluster compose file runs three engram nodes, each with its own Redis instance, connected over a shared Docker network. +the cluster compose file runs three Engram nodes, each with its own Redis instance, connected over a shared Docker network. ```sh # copy and fill in your OpenAI key @@ -242,11 +280,11 @@ docker compose -f docker-compose.cluster.yml up -d --build # wait for all nodes to be healthy, then initialize the cluster ./scripts/cluster-init.sh -# verify all Stage 1 acceptance criteria +# verify all acceptance criteria ./scripts/cluster-verify.sh ``` -the verify script checks 10 criteria: leader election, write replication to all nodes, 307 redirect from followers, failover, Prometheus metric presence, knowledge graph replication, delete-session cleanup, node restart and recovery from the Raft log, snapshot compaction, and restart-then-verify that state is fully restored from the latest snapshot. it exits 0 only if all criteria pass. +the verify script checks 17 criteria: leader election, write replication to all nodes, 307 redirect from followers, failover, Prometheus metric presence, knowledge graph replication, delete-session cleanup, node restart and recovery from the Raft log, snapshot compaction, restart-then-verify that state is fully restored from the latest snapshot, session visibility propagation, global graph population from public sessions, agent registration, global entity and relationship count metrics, global entity queries, global conflict detection, and global graph snapshot round-trip. it exits 0 only if all criteria pass. see `docker-compose.cluster.yml` and the scripts in `scripts/` for details. @@ -257,7 +295,7 @@ the repository includes a retrieval-quality harness for LongMemEval and BEAM und - LongMemEval uses `benchmarks/longmemeval_engram.py` and emits retrieval summaries plus `hypothesis.jsonl` for the official evaluator. - BEAM uses `benchmarks/beam_engram.py` and supports flat JSON input as well as the repository-style `chats/100K`, `chats/500K`, and `chats/1M` layouts. - `scripts/run_quality_benchmarks.sh` defaults to `http://127.0.0.1:3002` and is meant to target the Docker Compose deployment to avoid common port `3000` conflicts. -- Retrieval smoke runs can avoid hosted embedding APIs entirely by letting the harness start `tools/local_embed_server.py` and a matching engram process with `--start-local-embed-server --start-engram`. +- Retrieval smoke runs can avoid hosted embedding APIs entirely by letting the harness start `tools/local_embed_server.py` and a matching Engram process with `--start-local-embed-server --start-engram`. - Preliminary LongMemEval retrieval results are now published: a 5-question local-embedder slice reached perfect recall@5/10. See `BENCHMARKS.md`. see [docs/QUALITY_BENCHMARKS.md](docs/QUALITY_BENCHMARKS.md) for the end-to-end runbook. diff --git a/benches/e2e_throughput.rs b/benches/e2e_throughput.rs index 92ee803..fc6cdca 100644 --- a/benches/e2e_throughput.rs +++ b/benches/e2e_throughput.rs @@ -197,6 +197,9 @@ impl BenchmarkHarness { engram::knowledge::graph::KnowledgeGraph::new(), )), knowledge_job_sender, + global_graph: Arc::new(tokio::sync::RwLock::new( + engram::knowledge::global::GlobalGraph::new(), + )), }); let _worker_handles = spawn_embedding_workers( diff --git a/docs/API.md b/docs/API.md index 921d6fb..6396de5 100644 --- a/docs/API.md +++ b/docs/API.md @@ -1,35 +1,47 @@ # API Reference -This document describes every REST endpoint exposed by engram. All endpoints are served at `http://localhost:3000` by default. - -| method | path | description | -|--------|-------------------------------------------------------------|-------------------------------------------------------| -| POST | /sessions | create a new session | -| POST | /sessions/{session_id}/messages | add a message | -| GET | /sessions/{session_id}/context | get assembled context | -| POST | /sessions/{session_id}/search | semantic search over long-term memory | -| DELETE | /sessions/{session_id} | delete a session and all its memories | -| PUT | /sessions/{session_id}/core-memory | add a core memory fact | -| GET | /sessions/{session_id}/knowledge | get full knowledge graph for the session | -| GET | /sessions/{session_id}/knowledge/entities/{entity_name} | get all entities connected to a named entity | -| GET | /sessions/{session_id}/knowledge/path | find shortest path between two entities | -| GET | /sessions/{session_id}/knowledge/export | export knowledge graph (JSON or Graphviz DOT) | -| GET | /health | health check | -| GET | /metrics | Prometheus metrics | -| GET | /api-docs/openapi.json | OpenAPI specification | -| GET | /swagger-ui/ | Swagger UI | -| GET | /cluster | cluster status (cluster mode only) | -| POST | /cluster/init | initialize the Raft cluster | -| POST | /cluster/add-learner | add a learner node | -| POST | /cluster/change-membership | promote learners to full voting members | +This document describes every REST endpoint exposed by Engram. All endpoints are served at `http://localhost:3000` by default. + +| method | path | description | +|--------|---------------------------------------------------------------|-------------------------------------------------------| +| POST | /sessions | create a new session (optional agent_id) | +| POST | /sessions/{session_id}/messages | add a message | +| GET | /sessions/{session_id}/context | get assembled context | +| POST | /sessions/{session_id}/search | semantic search over long-term memory | +| DELETE | /sessions/{session_id} | delete a session and all its memories | +| PUT | /sessions/{session_id}/core-memory | add a core memory fact | +| PUT | /sessions/{session_id}/visibility | set session visibility (Private or Shared) | +| GET | /sessions/{session_id}/knowledge | get full knowledge graph for the session | +| GET | /sessions/{session_id}/knowledge/entities/{entity_name} | get all entities connected to a named entity | +| GET | /sessions/{session_id}/knowledge/path | find shortest path between two entities | +| GET | /sessions/{session_id}/knowledge/export | export knowledge graph (JSON or Graphviz DOT) | +| GET | /knowledge/global | get the cross-session global knowledge graph | +| GET | /knowledge/global/entities/{name} | get entities connected to a named entity (global) | +| GET | /knowledge/global/entities/{name}/sources | get sessions that contributed a named entity | +| GET | /knowledge/global/path | find shortest path in the global graph | +| GET | /knowledge/global/export | export global graph (JSON or Graphviz DOT) | +| GET | /knowledge/global/conflicts | list conflicting facts across sessions | +| GET | /health | health check | +| GET | /metrics | Prometheus metrics | +| GET | /api-docs/openapi.json | OpenAPI specification | +| GET | /swagger-ui/ | Swagger UI | +| GET | /cluster | cluster status (cluster mode only) | +| POST | /cluster/init | initialize the Raft cluster | +| POST | /cluster/add-learner | add a learner node | +| POST | /cluster/change-membership | promote learners to full voting members | --- ## POST /sessions -create a new session. +Create a new session. Optionally supply an `agent_id` to associate the session with a named agent. In cluster mode, a non-empty `agent_id` is persisted via a `RegisterSession` Raft command so the mapping survives restarts. -**request body:** none +**request body:** optional +```json +{ + "agent_id": "agent-42" +} +``` **success response:** - status: 200 @@ -43,16 +55,22 @@ create a new session. **error responses:** - 500: internal server error -**example:** +**examples:** ```sh +# no agent association curl -X POST http://localhost:3000/sessions + +# with agent association +curl -X POST http://localhost:3000/sessions \ + -H 'content-type: application/json' \ + -d '{"agent_id":"agent-42"}' ``` --- ## POST /sessions/{session_id}/messages -add a message to a session. +Add a message to a session. **path parameters:** - `session_id` (string): session identifier @@ -88,7 +106,7 @@ curl -X POST http://localhost:3000/sessions/{session_id}/messages \ ## GET /sessions/{session_id}/context -get the assembled context for a session. +Get the assembled context for a session. **path parameters:** - `session_id` (string): session identifier @@ -121,7 +139,7 @@ curl http://localhost:3000/sessions/{session_id}/context ## POST /sessions/{session_id}/search -semantic search over long-term memory. +Semantic search over long-term memory. **path parameters:** - `session_id` (string): session identifier @@ -162,7 +180,7 @@ curl -X POST http://localhost:3000/sessions/{session_id}/search \ ## DELETE /sessions/{session_id} -delete a session and all its memories. +Delete a session and all its memories. **path parameters:** - `session_id` (string): session identifier @@ -182,7 +200,7 @@ curl -X DELETE http://localhost:3000/sessions/{session_id} ## PUT /sessions/{session_id}/core-memory -add a core memory fact to a session. +Add a core memory fact to a session. **path parameters:** - `session_id` (string): session identifier @@ -212,9 +230,44 @@ curl -X PUT http://localhost:3000/sessions/{session_id}/core-memory \ --- +## PUT /sessions/{session_id}/visibility + +Set the visibility of a session. Sessions are `Private` by default. Setting a session to `Shared` causes its extracted knowledge to be merged into the global knowledge graph. + +In cluster mode this goes through Raft so all nodes see the same visibility state. In standalone mode the command is accepted and no-ops. + +**path parameters:** +- `session_id` (string): session identifier + +**request body:** +```json +{ + "visibility": "Shared" +} +``` + +`visibility` must be `"Shared"` or `"Private"`. + +**success response:** +- status: 204 (no content) + +**error responses:** +- 400: invalid request body +- 307: redirect to leader (cluster mode, follower received the request) +- 500: failed to set visibility + +**example:** +```sh +curl -X PUT http://localhost:3000/sessions/{session_id}/visibility \ + -H 'content-type: application/json' \ + -d '{"visibility":"Shared"}' +``` + +--- + ## GET /health -health check endpoint. +Health check endpoint. **success response:** - status: 200 @@ -273,13 +326,13 @@ curl http://localhost:3000/swagger-ui/ ## cluster endpoints -these endpoints are only available when the node is started in cluster mode (i.e., `NODE_ID` is set). standalone nodes return 503. +These endpoints are only available when the node is started in cluster mode (i.e., `NODE_ID` is set). Standalone nodes return 503. --- ## GET /cluster -returns the current Raft cluster status for this node. +Returns the current Raft cluster status for this node. **success response:** - status: 200 @@ -311,7 +364,7 @@ curl http://localhost:3000/cluster ## POST /cluster/init -initializes the Raft cluster. call this once from any node after all nodes are running. reads `NODE_ID`, `RAFT_ADDR` (or `RAFT_ADVERTISE_ADDR`), and `CLUSTER_PEERS` from the node's environment to build the initial membership set. +Initializes the Raft cluster. Call this once from any node after all nodes are running. Reads `NODE_ID`, `RAFT_ADDR` (or `RAFT_ADVERTISE_ADDR`), and `CLUSTER_PEERS` from the node's environment to build the initial membership set. **request body:** none @@ -331,7 +384,7 @@ curl -X POST http://localhost:3000/cluster/init ## POST /cluster/add-learner -adds a new node as a learner. learners receive log replication but do not vote in elections. promote to a full member with `/cluster/change-membership`. +Adds a new node as a learner. Learners receive log replication but do not vote in elections. Promote to a full member with `/cluster/change-membership`. **request body:** ```json @@ -356,7 +409,7 @@ curl -X POST http://localhost:3000/cluster/add-learner \ ## POST /cluster/change-membership -changes the cluster membership to the given set of node IDs. nodes in the new set that are currently learners are promoted to full members; nodes not in the new set are removed. +Changes the cluster membership to the given set of node IDs. Nodes in the new set that are currently learners are promoted to full members; nodes not in the new set are removed. **request body:** ```json @@ -381,13 +434,13 @@ curl -X POST http://localhost:3000/cluster/change-membership \ ## knowledge graph endpoints -these endpoints query the per-session knowledge graph built by the background knowledge extraction pipeline. entities and relationships are extracted automatically from every message added to the session. results reflect all committed messages; there may be a brief delay while extraction jobs are processed. +These endpoints query the per-session knowledge graph built by the background knowledge extraction pipeline. Entities and relationships are extracted automatically from every message added to the session. Results reflect all committed messages; there may be a brief delay while extraction jobs are processed. --- ## GET /sessions/{session_id}/knowledge -returns all entities and edges in the session's knowledge graph. +Returns all entities and edges in the session's knowledge graph. **path parameters:** - `session_id` (string): session identifier @@ -417,7 +470,7 @@ curl http://localhost:3000/sessions/{session_id}/knowledge ## GET /sessions/{session_id}/knowledge/entities/{entity_name} -returns all entities directly connected to the named entity, including both incoming and outgoing relationships. +Returns all entities directly connected to the named entity, including both incoming and outgoing relationships. **path parameters:** - `session_id` (string): session identifier @@ -452,7 +505,7 @@ curl http://localhost:3000/sessions/{session_id}/knowledge/entities/Alice ## GET /sessions/{session_id}/knowledge/path -finds the shortest directed path (BFS over outgoing edges) between two named entities. +Finds the shortest directed path (BFS over outgoing edges) between two named entities. **path parameters:** - `session_id` (string): session identifier @@ -484,7 +537,7 @@ curl "http://localhost:3000/sessions/{session_id}/knowledge/path?from=Alice&to=B ## GET /sessions/{session_id}/knowledge/export -exports the knowledge graph in JSON or Graphviz DOT format. +Exports the knowledge graph in JSON or Graphviz DOT format. **path parameters:** - `session_id` (string): session identifier @@ -516,3 +569,164 @@ digraph knowledge { ```sh curl "http://localhost:3000/sessions/{session_id}/knowledge/export?format=dot" ``` + +--- + +## global knowledge graph endpoints + +The global graph aggregates entities and relationships from all sessions whose visibility is set to `Shared`. These endpoints return data from the local in-memory global graph. In cluster mode the global graph is eventually consistent with the leader. + +--- + +## GET /knowledge/global + +Returns all entities and edges in the global knowledge graph. + +**success response:** +- status: 200 +- body: +```json +{ + "entities": [ + { "name": "Alice", "entity_type": "Person", "attributes": {} }, + { "name": "OpenAI", "entity_type": "Organization", "attributes": {} } + ], + "edges": [ + { "from": "Alice", "to": "OpenAI", "relationship_type": "works_at" } + ] +} +``` + +**example:** +```sh +curl http://localhost:3000/knowledge/global +``` + +--- + +## GET /knowledge/global/entities/{name} + +Returns all entities directly connected to the named entity in the global graph, including both incoming and outgoing relationships. + +**path parameters:** +- `name` (string): the entity to look up + +**success response:** +- status: 200 +- body: +```json +{ + "entity_name": "Alice", + "related": [ + { + "name": "OpenAI", + "entity_type": "Organization", + "relationship_type": "works_at", + "direction": "Outgoing" + } + ] +} +``` + +**example:** +```sh +curl http://localhost:3000/knowledge/global/entities/Alice +``` + +--- + +## GET /knowledge/global/entities/{name}/sources + +Returns the list of session IDs that contributed the named entity to the global graph. + +**path parameters:** +- `name` (string): the entity to look up + +**success response:** +- status: 200 +- body: +```json +{ + "entity_name": "Alice", + "sources": ["session-abc123", "session-def456"] +} +``` + +**example:** +```sh +curl http://localhost:3000/knowledge/global/entities/Alice/sources +``` + +--- + +## GET /knowledge/global/path + +Finds the shortest directed path (BFS over outgoing edges) between two named entities in the global graph. + +**query parameters:** +- `from` (string, required): source entity name +- `to` (string, required): target entity name + +**success response:** +- status: 200 +- body (`path` is `null` if no path exists): +```json +{ + "from": "Alice", + "to": "Bob", + "path": [ + { "from": "Alice", "relationship_type": "works_at", "to": "OpenAI" }, + { "from": "OpenAI", "relationship_type": "employs", "to": "Bob" } + ] +} +``` + +**example:** +```sh +curl "http://localhost:3000/knowledge/global/path?from=Alice&to=Bob" +``` + +--- + +## GET /knowledge/global/export + +Exports the global knowledge graph in JSON or Graphviz DOT format. + +**query parameters:** +- `format` (string, optional, default: `json`): `json` or `dot` + +**success response:** +- status: 200 +- content-type: `application/json` for JSON, `text/vnd.graphviz` for DOT + +**example:** +```sh +curl "http://localhost:3000/knowledge/global/export?format=dot" +``` + +--- + +## GET /knowledge/global/conflicts + +Returns all detected conflicts in the global graph. A conflict occurs when two different sessions report different relationship types between the same pair of entities. + +**success response:** +- status: 200 +- body: +```json +{ + "conflicts": [ + { + "entity": "Alice", + "related": "OpenAI", + "relationship_types": ["works_at", "founded"], + "sessions": ["session-abc123", "session-def456"] + } + ] +} +``` + +**example:** +```sh +curl http://localhost:3000/knowledge/global/conflicts +``` diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index cf07d71..487015d 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -2,15 +2,15 @@ ## Overview -engram is an asynchronous semantic memory backend for LLM-powered agents, written in Rust. It provides short-term, long-term, and core memory for agents, enabling context assembly with strict token budgeting, semantic search, and transparent memory management. +Engram is an asynchronous semantic memory backend for LLM-powered agents, written in Rust. It provides short-term, long-term, and core memory for agents, enabling context assembly with strict token budgeting, semantic search, and transparent memory management. -The system is designed for performance, reliability, and developer control, with all major components behind trait abstractions for easy swapping and testing. +All major components are behind trait abstractions so implementations can be swapped out and mocked in tests without changing any calling code. ## Data flow -in cluster mode, writes go through Raft consensus before reaching the stores. in standalone mode (NODE_ID not set), the Raft layer is absent and write handlers reach the stores directly. +In cluster mode, writes go through Raft consensus before reaching the stores. In standalone mode (NODE_ID not set), the Raft layer is absent and write handlers reach the stores directly. -reads (context, search) always go directly to the local stores on whichever node receives the request. +Reads (context, search, knowledge queries) always go directly to the local stores on whichever node receives the request. ```mermaid graph TD @@ -24,14 +24,20 @@ graph TD router --> corememhandler["core memory handler"] router --> healthhandler["health handler"] router --> knowledgehandler["knowledge handler"] + router --> visibilityhandler["visibility handler"] + router --> globalhandler["global knowledge handler"] memoryhandler -->|write| raft["raft consensus\nopenraft 0.9\ncluster mode only"] corememhandler -->|write| raft - sessionhandler -->|delete| raft + sessionhandler -->|delete or register agent| raft + visibilityhandler -->|write| raft raft -.->|grpc append entries| peers["peer nodes (port 9001)"] raft -.->|grpc install snapshot| laggers["lagging followers"] raft -->|state machine apply| shortterm["short-term memory (trait)"] raft -->|state machine apply| coremem["core memory store (trait)"] + raft -->|state machine apply| knowledgegraph[("knowledge graph\nper-session in-memory")] + raft -->|state machine apply| globalgraph[("global knowledge graph\ncross-session")] + raft -->|state machine apply| visibility["session visibility map"] raft -->|embedding job| embedqueue["embedding worker pool (bounded channel)"] raft -->|knowledge job| knowledgequeue["knowledge worker pool (bounded channel)"] raft --> redb[("redb\npersistent raft log\n+ snapshot store")] @@ -50,8 +56,11 @@ graph TD knowledgequeue -->|leader-only| extractor["knowledge extractor (trait)"] extractor -->|gpt-4o-mini or mock| extraction["entities + relationships"] - extraction -->|AddKnowledge via raft| knowledgegraph[("knowledge graph\nper-session in-memory")] + extraction -->|AddKnowledge via raft| knowledgegraph + knowledgegraph -->|public sessions merge| globalgraph + knowledgehandler --> knowledgegraph + globalhandler --> globalgraph searchhandler --> embedprovider searchhandler --> longterm @@ -101,6 +110,8 @@ All major components are behind trait abstractions, which lets implementations b | MockKnowledgeExtractor as default in Docker Compose | always OpenAI | decouples cluster verification from OpenAI API quota; pattern-matching mock is deterministic and offline-capable | | redb v2 for Raft log and snapshot store | sled, RocksDB, file-based log | redb is pure Rust, ACID, embedded, and simple to use; the key constraint is that tables must be pre-created in a write transaction before any read transaction touches them on a fresh database | | Redis as a projection, not the source of truth | reconcile Redis with Raft log on recovery | flushing Redis unconditionally on startup and restoring from the snapshot avoids a maze of reconciliation edge cases; one authoritative source (Raft log plus latest snapshot) with all volatile state derived from it | +| global graph as in-memory projection of public sessions | separate global persistent store | the global graph is always derivable from the full log; snapshotting it avoids replaying the entire history on startup while keeping the system simple | +| session visibility defaults to Private | defaults to Shared | private by default prevents unintended cross-agent knowledge leakage; agents opt in explicitly | ## Context assembly algorithm @@ -131,8 +142,8 @@ Stage 2 adds a knowledge graph pipeline that runs alongside the existing memory ### Knowledge pipeline -1. State machine receives `AddMessage` command → enqueues `KnowledgeJob { session_id, message_id, text }` on a bounded channel (capacity 500 by default). -2. Knowledge workers dequeue jobs. In cluster mode only the **current leader** calls the extractor; followers skip because they will receive the result via Raft replication. This prevents duplicate LLM calls and avoids non-deterministic divergence between nodes. +1. State machine receives `AddMessage` command and enqueues `KnowledgeJob { session_id, message_id, text }` on a bounded channel (capacity 500 by default). +2. Knowledge workers dequeue jobs. In cluster mode only the current leader calls the extractor; followers skip because they will receive the result via Raft replication. This prevents duplicate LLM calls and avoids non-deterministic divergence between nodes. 3. The extractor calls GPT-4o-mini (or the mock) with a structured JSON prompt requesting named entities and typed relationships. 4. The result is submitted as a `MemoryCommand::AddKnowledge` through `raft.client_write()`. OpenRaft replicates this to all nodes. 5. Each node's state machine applies `AddKnowledge` to its local `KnowledgeGraph` (idempotent by `(session_id, message_id)` key). @@ -160,7 +171,7 @@ Key operations: ### Knowledge metrics -Four new Prometheus metrics are exposed alongside the existing Raft and embedding metrics: +Four Prometheus metrics are exposed alongside the existing Raft and embedding metrics: | metric | type | description | |--------|------|-------------| @@ -171,7 +182,7 @@ Four new Prometheus metrics are exposed alongside the existing Raft and embeddin ## Distributed cluster mode (Stage 1) -In cluster mode, multiple engram nodes form a Raft consensus group using [OpenRaft 0.9](https://github.com/datafuselabs/openraft). Cluster mode is enabled by setting `NODE_ID` in the environment. Without it, the server runs in standalone mode exactly as described above. +In cluster mode, multiple Engram nodes form a Raft consensus group using [OpenRaft 0.9](https://github.com/datafuselabs/openraft). Cluster mode is enabled by setting `NODE_ID` in the environment. Without it, the server runs in standalone mode exactly as described above. ### Node anatomy @@ -236,14 +247,14 @@ OpenAI text embeddings are deterministic for the same input. All nodes converge | component | description | |-----------|-------------| | `EngRaftLogStore` | redb-backed persistent log store; implements `RaftLogStorage + RaftLogReader`; pre-creates tables on init so read transactions never fail on a fresh database | -| `EngStateMachineStore` | applies committed `MemoryCommand` entries to Redis and the knowledge graph; hosts `EngSnapshotBuilder`, which serializes the full state to an `EngramSnapshot` payload and persists it in redb | +| `EngStateMachineStore` | applies committed `MemoryCommand` entries to Redis, the per-session knowledge graph, the global graph, and the session visibility map; hosts `EngSnapshotBuilder`, which serializes the full state to an `EngramSnapshot` payload and persists it in redb | | `EngRaftNetwork` | factory that creates per-peer gRPC connections using tonic channels | | `EngRaftNetworkConnection` | sends Vote, AppendEntries, and InstallSnapshot RPCs over gRPC; the snapshot payload is a serialized `EngramSnapshot` | | `RaftGrpcServer` | tonic service that forwards incoming Raft RPCs to the local `RaftHandle`; handles `InstallSnapshotRequest` so lagging followers can restore a leader snapshot over gRPC | -`MemoryCommand` has four variants: `AddMessage`, `AddFact`, `DeleteSession`, and `AddKnowledge`. `AddMessage` also enqueues a `KnowledgeJob` so every committed message is a candidate for knowledge extraction. +`MemoryCommand` has six variants: `AddMessage`, `AddFact`, `DeleteSession`, `AddKnowledge`, `SetSessionVisibility`, and `RegisterSession`. `AddMessage` also enqueues a `KnowledgeJob` so every committed message is a candidate for knowledge extraction. -`EngramSnapshot` is the versioned payload serialized into every snapshot. It contains `short_term`, `core_memory`, `knowledge_graph`, and a reserved `global_graph` field. Adding `version: 1` and `#[serde(default)]` on optional fields means future stages can extend the payload without breaking existing snapshots. +`EngramSnapshot` is the versioned payload serialized into every snapshot. It contains `short_term`, `core_memory`, `knowledge_graph`, `global_graph`, `visibility`, and `session_agents`. The `version: 1` field and `#[serde(default)]` on optional fields mean older nodes can install newer snapshots by ignoring fields they don't recognize. `recover_state_machine()` in `src/raft/recovery.rs` runs at node startup before Raft is initialized. It flushes Redis, loads the latest persisted snapshot from redb, restores the payload into the live stores, and advances `last_applied` and `last_membership`. OpenRaft then replays any committed log entries that sit past the snapshot index. @@ -282,8 +293,6 @@ Stage 3A replaced the in-memory Raft log store with a persistent redb-backed sto The leader builds a snapshot by serializing the entire state machine under the `SmInner` lock. Holding the lock continuously across all `dump_all` calls (short-term memory, core memory, knowledge graph) guarantees that the snapshot represents exactly the state at `last_applied` index N rather than a mix of states from different moments. -The snapshot payload is `EngramSnapshot`: a versioned struct with `version: 1`, three store snapshots, and a reserved `global_graph` field marked `#[serde(default)]` so Stage 3B can add it without breaking existing snapshots. - Snapshots are stored in redb in the same database file as the Raft log (`RAFT_DB_PATH`). Log compaction runs automatically when the number of committed entries since the last snapshot exceeds `SNAPSHOT_LOG_THRESHOLD` (default 1000). ### Recovery sequence @@ -292,7 +301,7 @@ At startup, before OpenRaft is initialized, `recover_state_machine()` runs: 1. Flush Redis completely. This is unconditional; Redis is always treated as a projection, never as a source of truth. 2. Load the latest `EngramSnapshot` from the redb snapshot table. -3. Restore `short_term`, `core_memory`, and `knowledge_graph` from the snapshot payload. +3. Restore `short_term`, `core_memory`, `knowledge_graph`, `global_graph`, `visibility`, and `session_agents` from the snapshot payload. 4. Set `last_applied` and `last_membership` so OpenRaft knows where the state machine is. 5. OpenRaft reads the persistent log and replays any committed entries past the snapshot index. @@ -314,9 +323,58 @@ Three Prometheus metrics are exported in cluster mode alongside the existing Raf `engram_snapshot_last_index` is updated from the OpenRaft `RaftMetrics.snapshot` field in the background metrics watcher task in `app.rs`. +## Stage 3B: Collective memory + +Stage 3B adds cross-agent knowledge sharing. Multiple agents can now contribute to a shared global knowledge graph, control whether their sessions are public or private, and inspect conflicts between facts reported by different sessions. + +### Session visibility + +Each session has a visibility setting that defaults to `Private`. When set to `Shared`, the session's knowledge extraction results are merged into the global graph. The `SetSessionVisibility` Raft command replicates this change to all nodes so the visibility state stays consistent across the cluster. + +`PUT /sessions/{session_id}/visibility` accepts `{ "visibility": "Shared" }` or `{ "visibility": "Private" }`. In standalone mode the command is accepted and no-ops (there is no global coordination needed when running a single node). + +### Agent registration + +`POST /sessions` now accepts an optional `agent_id` field in the request body. When provided, the server issues a `RegisterSession` Raft command that records the `(session_id, agent_id)` mapping in `SmInner.session_agents`. This mapping is included in snapshots and survives restarts. It lets you trace which agent produced which knowledge in the global graph. + +### Global knowledge graph + +`GlobalGraph` is an in-memory directed graph that aggregates entities and relationships from all sessions whose visibility is set to `Shared`. It also tracks provenance (which session contributed each entity) and detects conflicts (the same entity connected to another entity via different relationship types across sessions). + +When the state machine applies `AddKnowledge` for a session that is visible, it calls `global_graph.merge_with_agent()`, which upserts entities and edges with provenance metadata. Conflicts are recorded when a new relationship type for the same entity pair differs from an existing one. + +### Global REST endpoints + +| method | path | description | +|--------|------|-------------| +| PUT | `/sessions/{id}/visibility` | set session visibility (Private or Shared) | +| GET | `/knowledge/global` | full global graph: all entities and edges | +| GET | `/knowledge/global/entities/{name}` | entities connected to the named entity in the global graph | +| GET | `/knowledge/global/entities/{name}/sources` | sessions that contributed this entity | +| GET | `/knowledge/global/path?from=X&to=Y` | BFS shortest path in the global graph | +| GET | `/knowledge/global/export?format=json\|dot` | export the global graph as JSON or Graphviz DOT | +| GET | `/knowledge/global/conflicts` | list conflicting relationship types across sessions | + +All global read endpoints return data from the local in-memory `GlobalGraph`. In cluster mode the global graph is eventually consistent: a freshly committed `SetSessionVisibility` or `AddKnowledge` may not appear immediately on a follower that is slightly behind the leader. + +### Global graph metrics + +Three Prometheus gauges track the state of the global graph. They are updated each time the state machine applies a command that modifies the global graph: + +| metric | type | description | +|--------|------|-------------| +| `engram_global_entities` | gauge | number of distinct entities in the global graph | +| `engram_global_relationships` | gauge | number of edges in the global graph | +| `engram_global_conflicts` | gauge | number of detected conflicting relationship types | + +### Snapshot protocol v2 + +The `EngramSnapshot` payload now includes `global_graph`, `visibility`, and `session_agents` fields alongside the existing `short_term`, `core_memory`, and `knowledge_graph` fields. The `#[serde(default)]` attribute on all new fields means Stage 3A snapshots deserialize cleanly on Stage 3B nodes. + ## Deferred items -The following remain out of scope after Stage 3A: +The following remain out of scope after Stage 3B: - **LanceDB replication.** Each node still calls the embedding API independently. Routing vector storage through Raft (leader-only embedding, follower payload replication) is a future item. -- **Multi-tenant auth.** Cluster-aware authentication routing. \ No newline at end of file +- **KG-augmented context assembly.** The knowledge graph (per-session and global) is queryable via REST but is not yet integrated into the context assembly pipeline to augment semantic search results. +- **Multi-tenant auth.** Cluster-aware authentication routing. diff --git a/docs/COMPARISON.md b/docs/COMPARISON.md index 94a1a9b..be4a195 100644 --- a/docs/COMPARISON.md +++ b/docs/COMPARISON.md @@ -1,19 +1,19 @@ -# Memory Engine Comparison: engram vs. Alternatives +# Memory Engine Comparison: Engram vs. Alternatives ## 1. Feature Matrix -| Feature | engram | Zep | Mem0 | LangChain Memory | Hindsight (Vectorize) | +| Feature | Engram | Zep | Mem0 | LangChain Memory | Hindsight (Vectorize) | |------------------------------- |:--------------:|:-------------:|:-------------:|:----------------:|:---------------------:| | **Language** | Rust | Python | Python | Python | Go | | **Deployment model** | Single binary, Docker | Docker, cloud, pip | pip, Docker, cloud | pip, cloud | Docker, cloud | -| **Fault-tolerant cluster** | Yes (3-node Raft, OpenRaft 0.9, persistent log + snapshots, startup recovery) | No | No | No | ? | +| **Fault-tolerant cluster** | Yes (3-node Raft, OpenRaft 0.9, persistent log + snapshots, startup recovery, snapshot v2 with global state) | No | No | No | ? | | **Embedding flexibility** | Yes (trait, BYO) | Yes (BYO, OpenAI, Cohere, etc.) | Yes (BYO, OpenAI, etc.) | Yes (BYO, OpenAI, etc.) | Yes (BYO, OpenAI, etc.) | | **Context visibility** | Full (exact prompt shown) | Partial (debug endpoint) | Partial | Partial (depends on chain) | ? | | **Token budget control** | Yes (per request) | Yes (configurable) | Yes (configurable) | Partial (depends on chain) | ? | | **Trimming strategy** | Pair-preserving | Naive/Configurable | Naive | Naive | ? | -| **Memory types** | Short-term, long-term, core, knowledge graph | Short, long, episodic | Short, long | Short, long, summary | Short, long, KG? | +| **Memory types** | Short-term, long-term, core, per-session KG, global cross-session KG | Short, long, episodic | Short, long | Short, long, summary | Short, long, KG? | | **Retrieval method** | Semantic search, knowledge graph traversal | Semantic, BM25, hybrid | Semantic, hybrid | Semantic, retriever chain | Semantic, hybrid, KG | -| **Knowledge graph** | Yes (petgraph, per-session, persisted via snapshots) | No | No | No | Yes | +| **Knowledge graph** | Yes (petgraph, per-session + global cross-session, agent provenance, conflict detection, persisted via snapshots) | No | No | No | Yes | | **Idempotency / deduplication**| Yes (message_id, status) | Yes (message_id) | Partial | No | ? | | **Observability** | Prometheus, tracing | Prometheus, logs | Logs | No (manual) | Prometheus, logs | | **Background processing** | Async worker, bounded queue | Async worker | Async | No | Async worker | @@ -32,7 +32,7 @@ | System | Context Assembly Latency (100 msg) | Throughput (msg/s) | Token Efficiency (vs full-dump) | LongMemEval Score | |--------|------------------------------------|--------------------|---------------------------------|-------------------| -| engram | 0.281 ms (in-memory), 21.66-29.55 ms (real-store) | 64,500.32 | 39.99% reduction | Harnesses and local slices published; full scorecards pending | +| Engram | 0.281 ms (in-memory), 21.66-29.55 ms (real-store) | 64,500.32 | 39.99% reduction | Harnesses and local slices published; full scorecards pending | | Mem0 | ~200 ms (P50 search) | Not disclosed | Not disclosed | 49.0% (independent) | | Zep | < 200 ms (retrieval) | Not disclosed | Not disclosed | 71.2% (via Graphiti), 63.8% (LongMemEval GPT-4o) | | Hindsight | < 200 ms (est.) | Not disclosed | Not disclosed | 91.4% (LongMemEval) | @@ -41,33 +41,33 @@ | System | Retrieval Metrics | QA Accuracy | |--------|-------------------|-------------| -| engram (prelim, n=5, retrieval-only, local embedder) | R@5=1.000, R@10=1.000, MRR=0.767, NDCG@10=0.826 | Not yet measured | +| Engram (prelim, n=5, retrieval-only, local embedder) | R@5=1.000, R@10=1.000, MRR=0.767, NDCG@10=0.826 | Not yet measured | -On the currently published numbers, engram's in-memory context assembly path is hundreds of times faster than the roughly 200 ms public retrieval figures cited for comparable systems. Its real-store path remains comfortably competitive at 21.66-29.55 ms while exercising actual Redis and LanceDB integrations, not placeholder mocks. The token-efficiency measurement also shows a 39.99% reduction versus a naive full-history dump at a 4k-token budget. The retrieval-quality gap is narrower than before because the repository now includes dedicated LongMemEval and BEAM harnesses plus a local-embedding fallback, and the first published LongMemEval retrieval slice already shows perfect recall@5/10 with strong MRR and NDCG. The public score cells should still be treated as provisional until full runs are published against the real datasets. +On the currently published numbers, Engram's in-memory context assembly path is hundreds of times faster than the roughly 200 ms public retrieval figures cited for comparable systems. Its real-store path remains comfortably competitive at 21.66-29.55 ms while exercising actual Redis and LanceDB integrations, not placeholder mocks. The token-efficiency measurement also shows a 39.99% reduction versus a naive full-history dump at a 4k-token budget. The retrieval-quality gap is narrower than before because the repository now includes dedicated LongMemEval and BEAM harnesses plus a local-embedding fallback, and the first published LongMemEval retrieval slice already shows perfect recall@5/10 with strong MRR and NDCG. The public score cells should still be treated as provisional until full runs are published against the real datasets. -> **Comparison note:** The engram numbers above are direct local benchmarks of full context assembly or end-to-end request throughput. Public competitor figures are typically retrieval or search latencies, so the table should be read as directional rather than strictly apples-to-apples. +> **Comparison note:** The Engram numbers above are direct local benchmarks of full context assembly or end-to-end request throughput. Public competitor figures are typically retrieval or search latencies, so the table should be read as directional rather than strictly apples-to-apples. ## 3. Narrative Analysis -**Where engram excels:** +**Where Engram excels:** - **Benchmarked performance:** Current measurements show 0.281 ms in-memory context assembly for a 100-message session, 21.66-29.55 ms with real stores depending on workload, and 64,500.32 messages per second in the reduced e2e throughput run. - **Transparency:** Developers can inspect the exact assembled context returned by the API rather than relying on hidden chain state. - **Rust performance:** High concurrency, low memory overhead, and strong type safety. - **Single-binary deployment:** Easy to run locally or in production; Docker and Compose supported. -- **Durability:** The Raft log, snapshots, and full state machine (including the knowledge graph) survive node restarts via the redb-backed persistent store and startup recovery. +- **Durability:** The Raft log, snapshots, and full state machine (including the knowledge graph, global graph, and session visibility) survive node restarts via the redb-backed persistent store and startup recovery. - **Pair-preserving trim:** Prevents broken dialogue, a common source of LLM hallucination in naive memory engines. - **Idempotent workers:** Message ingestion and embedding are robust to retries and crashes. - **Observability:** Prometheus metrics and structured tracing from day one; snapshot build and install metrics included. - **Token budget control:** Every context assembly is budgeted per request, not just globally. -**Where engram falls short today:** -- **No KG-augmented retrieval yet:** The knowledge graph is queryable via REST but is not yet integrated into the context assembly pipeline to augment semantic search results. +**Where Engram falls short today:** +- **No KG-augmented retrieval yet:** The knowledge graph (per-session and global) is queryable via REST but is not yet wired into context assembly to augment semantic search results. - **No managed cloud offering:** Self-hosted only; no SaaS or managed tier. - **Smaller community:** Newer and less widely adopted than Zep or LangChain. - **Retrieval is single-strategy:** Only semantic search drives context assembly; no hybrid or BM25 yet. - **Preliminary retrieval evaluation is strong, but still tiny:** A 5-question `single-session-user` LongMemEval retrieval slice achieved perfect recall@5/10 plus MRR=0.767 and NDCG@10=0.826 with the local embedder, but the full 500-question run is still pending. -**Who engram is best for:** +**Who Engram is best for:** - Rust developers and teams who want a self-hosted, debuggable, and transparent memory layer for LLM agents. - Anyone who needs to understand and control exactly what goes into the LLM context window. - Projects that value observability, idempotency, and explicit token budgeting over plug-and-play cloud convenience. diff --git a/docs/QUALITY_BENCHMARKS.md b/docs/QUALITY_BENCHMARKS.md index a981c89..3170b0e 100644 --- a/docs/QUALITY_BENCHMARKS.md +++ b/docs/QUALITY_BENCHMARKS.md @@ -1,15 +1,15 @@ # Quality Benchmarking -This document is the end-to-end runbook for engram's retrieval-quality benchmarks. +This document is the end-to-end runbook for Engram's retrieval-quality benchmarks. ## Scope -engram currently ships two benchmark bridges: +Engram currently ships two benchmark bridges: - `benchmarks/longmemeval_engram.py` for LongMemEval retrieval and QA. - `benchmarks/beam_engram.py` for BEAM retrieval and QA. -These scripts drive the existing engram HTTP API. They do not require extra debug endpoints or message-status APIs. Instead, they wait for background embeddings by polling the Prometheus gauge `engram_memory_embedding_queue_size` until the queue settles. +These scripts drive the existing Engram HTTP API. They do not require extra debug endpoints or message-status APIs. Instead, they wait for background embeddings by polling the Prometheus gauge `engram_memory_embedding_queue_size` until the queue settles. ## Recommended Runtime @@ -32,7 +32,7 @@ curl http://127.0.0.1:3002/health ### Dataset Acquisition -The LongMemEval GitHub repository does not include the benchmark JSON files in-tree. You need a local copy of one of the released files before running engram against it: +The LongMemEval GitHub repository does not include the benchmark JSON files in-tree. You need a local copy of one of the released files before running Engram against it: - `longmemeval_s.json` - `longmemeval_m.json` @@ -58,7 +58,7 @@ Outputs: ### Local Embedding Fallback -For retrieval-only runs, you can avoid hosted embedding latency and rate limits by letting the harness start a local OpenAI-compatible embedding server and a matching engram process: +For retrieval-only runs, you can avoid hosted embedding latency and rate limits by letting the harness start a local OpenAI-compatible embedding server and a matching Engram process: ```bash python3 benchmarks/longmemeval_engram.py \ @@ -71,7 +71,7 @@ python3 benchmarks/longmemeval_engram.py \ --lance-db-path ./data/lancedb-bench-local ``` -That flow starts `tools/local_embed_server.py`, injects `OPENAI_BASE_URL` for the spawned engram process, and defaults `EMBEDDING_DIMENSION` to `384`. QA mode still needs a completion model such as `gpt-4o` for answer generation. +That flow starts `tools/local_embed_server.py`, injects `OPENAI_BASE_URL` for the spawned Engram process, and defaults `EMBEDDING_DIMENSION` to `384`. QA mode still needs a completion model such as `gpt-4o` for answer generation. ### QA Run @@ -119,7 +119,7 @@ python3 benchmarks/beam_engram.py \ --output-dir benchmarks/results/beam-128k ``` -For BEAM retrieval, engram records ranked outputs per probing question. BEAM does not expose LongMemEval-style session gold labels directly, so the bridge writes raw retrieval outputs even when full recall-style metrics are unavailable. +For BEAM retrieval, Engram records ranked outputs per probing question. BEAM does not expose LongMemEval-style session gold labels directly, so the bridge writes raw retrieval outputs even when full recall-style metrics are unavailable. ### QA Run @@ -161,7 +161,7 @@ Important caveat: - `LONGMEMEVAL_DATASET` must point to a real local file. - `RUN_BEAM=1` can default `BEAM_DATASET` to `benchmarks/deps/BEAM/chats/` after sparse-cloning the repo. -- For local-embedding retrieval smoke runs, prefer the direct Python harnesses; the wrapper still assumes the target engram instance already has a working embedding backend. +- For local-embedding retrieval smoke runs, prefer the direct Python harnesses; the wrapper still assumes the target Engram instance already has a working embedding backend. Example: diff --git a/docs/SSOT.md b/docs/SSOT.md index fb5f935..d507f92 100644 --- a/docs/SSOT.md +++ b/docs/SSOT.md @@ -128,6 +128,17 @@ The project serves two purposes: 3. Install snapshot: a follower wiped after compaction catches up via InstallSnapshot, not full log replay 4. Full restart recovery: all three nodes stopped and restarted; state is fully restored from the latest snapshot +### Stage 3B: Collective memory ✅ **(Completed)** +- Session visibility (`Private`/`Shared`) controlled via `MemoryCommand::SetSessionVisibility`, replicated through Raft +- Global cross-session knowledge graph (`GlobalGraph`): aggregates entities and relationships from all Shared sessions with provenance tracking and conflict detection +- `MemoryCommand::RegisterSession` persists `(session_id, agent_id)` mappings in `SmInner.session_agents` +- `POST /sessions` accepts optional `agent_id` body field +- `PUT /sessions/{id}/visibility` endpoint for changing session visibility +- 6 new global REST endpoints: `GET /knowledge/global`, `GET /knowledge/global/entities/{name}`, `GET /knowledge/global/entities/{name}/sources`, `GET /knowledge/global/path`, `GET /knowledge/global/export`, `GET /knowledge/global/conflicts` +- 3 new Prometheus gauges: `engram_global_entities`, `engram_global_relationships`, `engram_global_conflicts` +- `EngramSnapshot` protocol v2: adds `global_graph`, `visibility`, and `session_agents` fields with `#[serde(default)]` for backward compatibility with Stage 3A snapshots +- 194 tests passing; all 17 cluster-verify acceptance criteria pass + --- ## 3. System Architecture (High-Level) @@ -144,13 +155,19 @@ graph TD router --> corememhandler["core memory handler"] router --> healthhandler["health handler"] router --> knowledgehandler["knowledge handler"] + router --> visibilityhandler["visibility handler"] + router --> globalhandler["global knowledge handler"] memoryhandler -->|write| raft["raft consensus\nopenraft 0.9\ncluster mode only"] corememhandler -->|write| raft - sessionhandler -->|delete| raft + sessionhandler -->|delete or register agent| raft + visibilityhandler -->|write| raft raft -.->|grpc append entries| peers["peer nodes (port 9001)"] raft -->|state machine apply| shortterm["short-term memory trait"] raft -->|state machine apply| coremem["core memory store trait"] + raft -->|state machine apply| knowledgegraph[("knowledge graph\nper-session in-memory")] + raft -->|state machine apply| globalgraph[("global knowledge graph\ncross-session")] + raft -->|state machine apply| visibility["session visibility map"] raft -->|embedding job| embedqueue["embedding worker pool\nbounded channel"] raft -->|knowledge job| knowledgequeue["knowledge worker pool\nbounded channel"] @@ -166,8 +183,11 @@ graph TD knowledgequeue -->|leader-only| extractor["knowledge extractor trait"] extractor -->|gpt-4o-mini or mock| extraction["entities + relationships"] - extraction -->|AddKnowledge via raft| knowledgegraph[("knowledge graph\nper-session in-memory")] + extraction -->|AddKnowledge via raft| knowledgegraph + knowledgegraph -->|public sessions merge| globalgraph + knowledgehandler --> knowledgegraph + globalhandler --> globalgraph contexthandler --> assembler["context assembler"] assembler --> shortterm @@ -477,8 +497,8 @@ This algorithm preserves conversational continuity (recent messages), injects on - Single table `memories` with schema: - `id` (Int64, auto‑increment) - `session_id` (Utf8) - - `message_id` (Utf8) – for idempotency - - `text` (Utf8) – full `Message` JSON to preserve role/timestamp + - `message_id` (Utf8) - for idempotency + - `text` (Utf8) - full `Message` JSON to preserve role/timestamp - `embedding` (FixedSizeList[1536]) - `created_at` (Timestamp) - **Indexes** (create immediately after table creation): @@ -557,18 +577,18 @@ services: - `LOG_FORMAT` (default `pretty`, supports `json`) **Knowledge pipeline** (Stage 2): -- `KNOWLEDGE_EXTRACTOR` — `openai` (default) or `mock`; use `mock` to avoid OpenAI quota consumption in CI/testing -- `KNOWLEDGE_MAX_WORKERS` — number of knowledge extraction workers (default `4`) -- `KNOWLEDGE_CHANNEL_SIZE` — knowledge job queue capacity (default `500`) +- `KNOWLEDGE_EXTRACTOR` - `openai` (default) or `mock`; use `mock` to avoid OpenAI quota consumption in CI/testing +- `KNOWLEDGE_MAX_WORKERS` - number of knowledge extraction workers (default `4`) +- `KNOWLEDGE_CHANNEL_SIZE` - knowledge job queue capacity (default `500`) **Cluster mode** (all required when `NODE_ID` is set): -- `NODE_ID` — unique integer node identifier (e.g. `1`) -- `RAFT_ADDR` — bind address for the gRPC Raft server (e.g. `0.0.0.0:9001`) -- `RAFT_ADVERTISE_ADDR` — address other nodes route to; required when binding `0.0.0.0` (e.g. `node-1:9001`) -- `CLUSTER_PEERS` — comma-separated gRPC peers as `id:host:port` (e.g. `2:node-2:9001,3:node-3:9001`) -- `CLUSTER_HTTP_PEERS` — comma-separated HTTP peers as `id:host:port`, used for leader redirect URLs -- `RAFT_DB_PATH` — path to the redb file for the persistent Raft log and snapshot store (default `./data/raft/engram.redb`) -- `SNAPSHOT_LOG_THRESHOLD` — number of committed log entries after which the leader builds a new snapshot (default `1000`) +- `NODE_ID` - unique integer node identifier (e.g. `1`) +- `RAFT_ADDR` - bind address for the gRPC Raft server (e.g. `0.0.0.0:9001`) +- `RAFT_ADVERTISE_ADDR` - address other nodes route to; required when binding `0.0.0.0` (e.g. `node-1:9001`) +- `CLUSTER_PEERS` - comma-separated gRPC peers as `id:host:port` (e.g. `2:node-2:9001,3:node-3:9001`) +- `CLUSTER_HTTP_PEERS` - comma-separated HTTP peers as `id:host:port`, used for leader redirect URLs +- `RAFT_DB_PATH` - path to the redb file for the persistent Raft log and snapshot store (default `./data/raft/engram.redb`) +- `SNAPSHOT_LOG_THRESHOLD` - number of committed log entries after which the leader builds a new snapshot (default `1000`) Per-request context settings such as `max_tokens`, `similarity_threshold`, and `long_term_top_k` are currently controlled through query parameters on the context endpoint rather than startup environment variables. @@ -585,21 +605,21 @@ Per-request context settings such as `max_tokens`, `similarity_threshold`, and ` - `memory_embedding_queue_size` (gauge) **Raft cluster metrics** (only emitted in cluster mode): -- `engram_raft_term` — current Raft term (gauge) -- `engram_raft_commit_index` — index of last applied log entry (gauge) -- `engram_raft_is_leader` — 1 if this node is the current leader, 0 otherwise (gauge) -- `engram_raft_leader_changes_total` — number of leader changes observed by this node (counter) +- `engram_raft_term` - current Raft term (gauge) +- `engram_raft_commit_index` - index of last applied log entry (gauge) +- `engram_raft_is_leader` - 1 if this node is the current leader, 0 otherwise (gauge) +- `engram_raft_leader_changes_total` - number of leader changes observed by this node (counter) **Knowledge pipeline metrics** (Stage 2, all modes): -- `engram_knowledge_extraction_duration_seconds` — duration of extraction calls (histogram, label: extractor) -- `engram_knowledge_entities_extracted_total` — cumulative entities extracted (counter) -- `engram_knowledge_relationships_extracted_total` — cumulative relationships extracted (counter) -- `engram_knowledge_queue_size` — pending knowledge jobs (gauge) +- `engram_knowledge_extraction_duration_seconds` - duration of extraction calls (histogram, label: extractor) +- `engram_knowledge_entities_extracted_total` - cumulative entities extracted (counter) +- `engram_knowledge_relationships_extracted_total` - cumulative relationships extracted (counter) +- `engram_knowledge_queue_size` - pending knowledge jobs (gauge) **Snapshot metrics** (Stage 3A, cluster mode): -- `engram_snapshot_build_total` — number of snapshots built by this node (counter) -- `engram_snapshot_install_total` — number of snapshots installed from the leader (counter) -- `engram_snapshot_last_index` — log index of the most recent snapshot; 0 if none exists (gauge) +- `engram_snapshot_build_total` - number of snapshots built by this node (counter) +- `engram_snapshot_install_total` - number of snapshots installed from the leader (counter) +- `engram_snapshot_last_index` - log index of the most recent snapshot; 0 if none exists (gauge) ### 10.2 Tracing - Each request gets a span. diff --git a/docs/VISION.md b/docs/VISION.md index 35e45ef..d622e86 100644 --- a/docs/VISION.md +++ b/docs/VISION.md @@ -240,7 +240,9 @@ Learn: Status: complete. Persistent redb-backed Raft log and snapshot store, full state machine snapshots (short-term memory, core memory, knowledge graph), startup recovery, InstallSnapshot over gRPC for lagging followers, automatic log compaction. All 10 cluster-verify criteria pass. -### Stage 3B: Collective memory +### Stage 3B: Collective memory ✅ + +Status: complete. Session visibility (Private/Shared) controlled via `SetSessionVisibility` Raft command, global cross-session knowledge graph with provenance and conflict tracking, agent registration at session creation, six new global REST endpoints, three new Prometheus gauges (`engram_global_entities`, `engram_global_relationships`, `engram_global_conflicts`), snapshot protocol v2 including global_graph/visibility/session_agents, 17/17 cluster-verify criteria pass. Goal: diff --git a/scripts/cluster-verify.sh b/scripts/cluster-verify.sh index 44fd6a7..a06f4c9 100755 --- a/scripts/cluster-verify.sh +++ b/scripts/cluster-verify.sh @@ -8,11 +8,25 @@ N3="http://localhost:3002" pass() { echo " PASS: $1"; } fail() { echo " FAIL: $1"; exit 1; } +# Inline leader finder used before the Stage 3A helpers are declared. +_find_leader_port_early() { + for p in 3000 3001 3002; do + local role + role=$(curl -sf "http://localhost:$p/cluster" 2>/dev/null | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['role'])" 2>/dev/null || echo "") + if [ "$role" = "Leader" ]; then + echo "$p" + return + fi + done +} + echo "" echo "=== Stage 1 Acceptance Verification ===" echo "" echo "[1] Leader election" +# Query any node — leader identity is cluster-wide state STATUS=$(curl -sf "$N1/cluster") LEADER=$(echo "$STATUS" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['leader_id'])" 2>/dev/null || echo "null") MEMBERS=$(echo "$STATUS" | python3 -c "import sys,json; d=json.load(sys.stdin); print(len(d['members']))" 2>/dev/null || echo "0") @@ -21,23 +35,40 @@ MEMBERS=$(echo "$STATUS" | python3 -c "import sys,json; d=json.load(sys.stdin); || fail "no leader or wrong member count (leader=$LEADER, members=$MEMBERS)" echo "[2] Write replication" -SESSION=$(curl -sf -X POST "$N1/sessions" | python3 -c "import sys,json; print(json.load(sys.stdin)['session_id'])") -WRITE_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$N1/sessions/$SESSION/messages" \ +# Discover actual leader — do not assume N1 is always the leader +EARLY_LEADER_PORT=$(_find_leader_port_early) +[ -z "${EARLY_LEADER_PORT:-}" ] && fail "no leader found for check [2]" +EARLY_LEADER="http://localhost:$EARLY_LEADER_PORT" +# Snapshot the leader index BEFORE the write so we have a stable target to wait for. +BEFORE_IDX=$(curl -sf "$EARLY_LEADER/cluster" 2>/dev/null | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['last_applied_index'])" 2>/dev/null || echo "0") +SESSION=$(curl -sf -X POST "$EARLY_LEADER/sessions" | python3 -c "import sys,json; print(json.load(sys.stdin)['session_id'])") +WRITE_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$EARLY_LEADER/sessions/$SESSION/messages" \ -H "Content-Type: application/json" \ -d '{"role":"user","content":"stage1 replication test"}') [ "$WRITE_CODE" = "204" ] || fail "write to leader returned HTTP $WRITE_CODE (expected 204)" -sleep 1 -# Verify replication by checking last_applied_index matches across all nodes. -# Using /cluster avoids calling OpenAI (which /context requires). -LEADER_APPLIED=$(curl -sf "$N1/cluster" | python3 -c "import sys,json; print(json.load(sys.stdin)['last_applied_index'])" 2>/dev/null || echo "null") -[ "$LEADER_APPLIED" != "null" ] || fail "could not read last_applied_index from leader" -for port in 3000 3001 3002; do - NODE_APPLIED=$(curl -sf "http://localhost:$port/cluster" | \ - python3 -c "import sys,json; print(json.load(sys.stdin)['last_applied_index'])" 2>/dev/null || echo "null") - [ "$NODE_APPLIED" = "$LEADER_APPLIED" ] \ - && pass "node :$port applied_index=$NODE_APPLIED (matches leader)" \ - || fail "node :$port applied_index=$NODE_APPLIED (leader has $LEADER_APPLIED)" +# Verify replication: wait for the leader to report every member's last_log_index >= TARGET. +# Using the leader's member list avoids the follower state-machine apply lag (entries are +# committed and in every follower's log before the leader acks 204). +TARGET_IDX=$((BEFORE_IDX + 1)) +REPLICATED=0 +for _i in $(seq 1 30); do + sleep 0.5 + MEMBER_IDXS=$(curl -sf "$EARLY_LEADER/cluster" 2>/dev/null | \ + python3 -c "import sys,json; d=json.load(sys.stdin); print(' '.join(str(m['last_log_index']) for m in d['members']))" \ + 2>/dev/null || echo "") + [ -z "${MEMBER_IDXS:-}" ] && continue + ALL_PAST=1 + for idx in $MEMBER_IDXS; do + [ "${idx:-0}" -lt "$TARGET_IDX" ] 2>/dev/null && { ALL_PAST=0; break; } + done + [ "$ALL_PAST" = "1" ] && { REPLICATED=1; break; } done +[ "$REPLICATED" = "1" ] || fail "write not replicated to all members within 15 s (target=$TARGET_IDX)" +FINAL_IDXS=$(curl -sf "$EARLY_LEADER/cluster" 2>/dev/null | \ + python3 -c "import sys,json; d=json.load(sys.stdin); [print(f' node id={m[\"id\"]} last_log_index={m[\"last_log_index\"]}') for m in d['members']]" \ + 2>/dev/null || true) +echo "$FINAL_IDXS" | while IFS= read -r line; do pass "$line (>= write checkpoint $TARGET_IDX)"; done echo "[3] Follower redirect" FOLLOWER_TESTED=0 @@ -84,7 +115,12 @@ curl -sf -X POST "$WRITE_NODE/sessions/$SESSION2/messages" \ || fail "write rejected after failover" echo " Restarting node-1..." docker compose -f docker-compose.cluster.yml start node-1 -sleep 2 +# Wait for node-1 to be healthy before proceeding (up to 15 s) +for _i in $(seq 1 30); do + sleep 0.5 + CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:3000/health" 2>/dev/null || echo "0") + [ "$CODE" = "200" ] && break +done echo "[5] Cluster observability" # Read metrics from whichever node is the current leader (node-1 may still be catching up) @@ -123,39 +159,46 @@ curl -sf -X POST "$WRITE_LEADER/sessions/$SESSION_K/messages" \ -H "Content-Type: application/json" \ -d '{"role":"user","content":"Alice knows Bob"}' > /dev/null -echo " Waiting 3 seconds for extraction and Raft replication..." -sleep 3 +echo " Waiting for extraction and Raft replication (up to 20 s)..." LEADER_PORT="" -for port in 3000 3001 3002; do - ROLE=$(curl -sf "http://localhost:$port/cluster" | \ - python3 -c "import sys,json; print(json.load(sys.stdin)['role'])" 2>/dev/null || echo "") - if [ "$ROLE" = "Leader" ]; then - LEADER_PORT="$port" - break - fi +LEADER_ENTITIES=-1 +LEADER_EDGES=-1 +for _i in $(seq 1 40); do + sleep 0.5 + LEADER_PORT="" + for port in 3000 3001 3002; do + ROLE=$(curl -sf "http://localhost:$port/cluster" 2>/dev/null | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['role'])" 2>/dev/null || echo "") + if [ "$ROLE" = "Leader" ]; then LEADER_PORT="$port"; break; fi + done + [ -z "${LEADER_PORT:-}" ] && continue + LEADER_ENTITIES=$(curl -sf "http://localhost:$LEADER_PORT/sessions/$SESSION_K/knowledge" 2>/dev/null | \ + python3 -c "import sys,json; print(len(json.load(sys.stdin)['entities']))" 2>/dev/null || echo "-1") + [ "${LEADER_ENTITIES:-0}" -ge 3 ] && break done -[ -z "$LEADER_PORT" ] && fail "no leader found for knowledge check" - -LEADER_ENTITIES=$(curl -sf "http://localhost:$LEADER_PORT/sessions/$SESSION_K/knowledge" | \ - python3 -c "import sys,json; print(len(json.load(sys.stdin)['entities']))" 2>/dev/null || echo "-1") - -[ "$LEADER_ENTITIES" -ge 3 ] \ +[ -z "${LEADER_PORT:-}" ] && fail "no leader found for knowledge check" +[ "${LEADER_ENTITIES:-0}" -ge 3 ] \ && pass "leader (:$LEADER_PORT) has $LEADER_ENTITIES entities" \ || fail "leader (:$LEADER_PORT) has $LEADER_ENTITIES entities (expected >= 3)" -LEADER_EDGES=$(curl -sf "http://localhost:$LEADER_PORT/sessions/$SESSION_K/knowledge" | \ +LEADER_EDGES=$(curl -sf "http://localhost:$LEADER_PORT/sessions/$SESSION_K/knowledge" 2>/dev/null | \ python3 -c "import sys,json; print(len(json.load(sys.stdin)['edges']))" 2>/dev/null || echo "-1") - -[ "$LEADER_EDGES" -ge 3 ] \ +[ "${LEADER_EDGES:-0}" -ge 3 ] \ && pass "leader (:$LEADER_PORT) has $LEADER_EDGES relationships" \ || fail "leader (:$LEADER_PORT) has $LEADER_EDGES relationships (expected >= 3)" +# Wait for followers to converge on the same entity count (up to 15 s each) for port in 3000 3001 3002; do [ "$port" -eq "$LEADER_PORT" ] && continue - FOLLOWER_ENTITIES=$(curl -sf "http://localhost:$port/sessions/$SESSION_K/knowledge" | \ - python3 -c "import sys,json; print(len(json.load(sys.stdin)['entities']))" 2>/dev/null || echo "-1") + FOLLOWER_ENTITIES=-1 + for _j in $(seq 1 30); do + sleep 0.5 + FOLLOWER_ENTITIES=$(curl -sf "http://localhost:$port/sessions/$SESSION_K/knowledge" 2>/dev/null | \ + python3 -c "import sys,json; print(len(json.load(sys.stdin)['entities']))" 2>/dev/null || echo "-1") + [ "$FOLLOWER_ENTITIES" -eq "$LEADER_ENTITIES" ] && break + done [ "$FOLLOWER_ENTITIES" -eq "$LEADER_ENTITIES" ] \ && pass "follower :$port converged to $FOLLOWER_ENTITIES entities (matches leader)" \ || fail "follower :$port has $FOLLOWER_ENTITIES entities (leader has $LEADER_ENTITIES)" @@ -343,3 +386,174 @@ ALL_AFTER=$(entity_count_on node-1) echo "" echo "=== All Stage 3A criteria PASSED ===" + +# --------------------------------------------------------------------------- +# Stage 3B helpers +# --------------------------------------------------------------------------- + +global_entity_count_on() { + local port=$1 + curl -sf "http://localhost:$port/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print(len(json.load(sys.stdin)['entities']))" 2>/dev/null || echo "-1" +} + +global_related_on() { + local port=$1 entity=$2 + curl -sf "http://localhost:$port/knowledge/global/entities/$entity" 2>/dev/null | \ + python3 -c "import sys,json; print([r['name'] for r in json.load(sys.stdin).get('related',[])])" \ + 2>/dev/null || echo "[]" +} + +# --------------------------------------------------------------------------- +# Stage 3B setup: two Shared sessions (SA, SB) + one Private session (SC) +# --------------------------------------------------------------------------- +echo "" +echo "=== Stage 3B: Collective Knowledge ===" +echo "" + +echo " Setting up Stage 3B sessions..." +S3B_LEADER_PORT=$(find_leader_port) +[ -z "${S3B_LEADER_PORT:-}" ] && fail "no leader for Stage 3B setup" +S3B_LEADER="http://localhost:$S3B_LEADER_PORT" + +S3B_SA=$(curl -sf -X POST "$S3B_LEADER/sessions" | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['session_id'])") +S3B_SB=$(curl -sf -X POST "$S3B_LEADER/sessions" | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['session_id'])") +S3B_SC=$(curl -sf -X POST "$S3B_LEADER/sessions" | \ + python3 -c "import sys,json; print(json.load(sys.stdin)['session_id'])") + +# Mark SA and SB as Shared; SC stays Private (default) +VIS_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X PUT "$S3B_LEADER/sessions/$S3B_SA/visibility" \ + -H "Content-Type: application/json" \ + -d '{"visibility":"Shared"}') +[ "$VIS_CODE" = "204" ] || fail "set SA visibility returned HTTP $VIS_CODE (expected 204)" + +VIS_CODE=$(curl -s -o /dev/null -w "%{http_code}" \ + -X PUT "$S3B_LEADER/sessions/$S3B_SB/visibility" \ + -H "Content-Type: application/json" \ + -d '{"visibility":"Shared"}') +[ "$VIS_CODE" = "204" ] || fail "set SB visibility returned HTTP $VIS_CODE (expected 204)" + +# SA: Alice+OpenAI, and Alice-knows-Bob for path check [17] +curl -sf -X POST "$S3B_LEADER/sessions/$S3B_SA/messages" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":"Alice works at OpenAI"}' > /dev/null +curl -sf -X POST "$S3B_LEADER/sessions/$S3B_SA/messages" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":"Alice knows Bob"}' > /dev/null +# SB: Bob+OpenAI -- contributes OpenAI a second time alongside SA +curl -sf -X POST "$S3B_LEADER/sessions/$S3B_SB/messages" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":"Bob works at OpenAI"}' > /dev/null +# SC: private -- must not surface in global graph +curl -sf -X POST "$S3B_LEADER/sessions/$S3B_SC/messages" \ + -H "Content-Type: application/json" \ + -d '{"role":"user","content":"TopSecret works at HiddenCorp"}' > /dev/null + +echo " Waiting 5 seconds for extraction and Raft replication..." +sleep 5 + +# [11] Shared sessions aggregate across every node +echo "[11] shared sessions aggregate" +for port in 3000 3001 3002; do + RELATED_11=$(global_related_on "$port" "OpenAI") + echo "$RELATED_11" | grep -q "Alice" \ + && pass "[11] node :$port OpenAI related to Alice (contributed by SA)" \ + || fail "[11] node :$port Alice missing from OpenAI related (got: $RELATED_11)" + echo "$RELATED_11" | grep -q "Bob" \ + && pass "[11] node :$port OpenAI related to Bob (contributed by SB)" \ + || fail "[11] node :$port Bob missing from OpenAI related (got: $RELATED_11)" +done + +# [12] Private session entities never appear in the global graph +echo "[12] private session isolation" +GLOBAL_ENTITIES_12=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print([e['name'] for e in json.load(sys.stdin).get('entities',[])])" \ + 2>/dev/null || echo "[]") +echo "$GLOBAL_ENTITIES_12" | grep -q "TopSecret" \ + && fail "[12] private entity TopSecret leaked into global graph" \ + || pass "[12] private entity TopSecret absent from global graph" +echo "$GLOBAL_ENTITIES_12" | grep -q "HiddenCorp" \ + && fail "[12] private entity HiddenCorp leaked into global graph" \ + || pass "[12] private entity HiddenCorp absent from global graph" + +# [13] Provenance lists contributing session ids +echo "[13] provenance" +SOURCES_13=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global/entities/OpenAI/sources" 2>/dev/null | \ + python3 -c "import sys,json; print(json.load(sys.stdin).get('sources',[]))" 2>/dev/null || echo "[]") +echo "$SOURCES_13" | grep -qF "$S3B_SA" \ + && pass "[13] SA listed as OpenAI provenance source" \ + || fail "[13] SA not in OpenAI sources (got: $SOURCES_13)" +echo "$SOURCES_13" | grep -qF "$S3B_SB" \ + && pass "[13] SB listed as OpenAI provenance source" \ + || fail "[13] SB not in OpenAI sources (got: $SOURCES_13)" + +# [14] All 3 nodes converge to identical global entity set (deterministic state) +echo "[14] deterministic conflict resolution" +GLOBAL_14_1=$(curl -sf "http://localhost:3000/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print(sorted([e['name'] for e in json.load(sys.stdin).get('entities',[])]))" \ + 2>/dev/null || echo "[]") +GLOBAL_14_2=$(curl -sf "http://localhost:3001/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print(sorted([e['name'] for e in json.load(sys.stdin).get('entities',[])]))" \ + 2>/dev/null || echo "[]") +GLOBAL_14_3=$(curl -sf "http://localhost:3002/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print(sorted([e['name'] for e in json.load(sys.stdin).get('entities',[])]))" \ + 2>/dev/null || echo "[]") +[ "$GLOBAL_14_1" = "$GLOBAL_14_2" ] && [ "$GLOBAL_14_2" = "$GLOBAL_14_3" ] \ + && pass "[14] all 3 nodes converge to identical global entity set: $GLOBAL_14_1" \ + || fail "[14] nodes diverge: node1=$GLOBAL_14_1 node2=$GLOBAL_14_2 node3=$GLOBAL_14_3" + +# [17] Global path traversal across sessions, no LLM call +echo "[17] global path traversal (no LLM)" +PATH_17=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global/path?from=Alice&to=Bob" 2>/dev/null | \ + python3 -c "import sys,json; d=json.load(sys.stdin); print('found' if d.get('path') else 'none')" \ + 2>/dev/null || echo "none") +[ "$PATH_17" = "found" ] \ + && pass "[17] global path Alice->Bob found via graph traversal (no LLM)" \ + || fail "[17] no global path found from Alice to Bob" + +# [16] Global graph and visibility survive a full cluster restart +echo "[16] persistence of collective state" +S3B_GLOBAL_BEFORE=$(global_entity_count_on "$S3B_LEADER_PORT") +docker compose -f docker-compose.cluster.yml stop node-1 node-2 node-3 +docker compose -f docker-compose.cluster.yml start node-1 node-2 node-3 +wait_for_leader +sleep 5 +S3B_LEADER_PORT=$(find_leader_port) +S3B_GLOBAL_AFTER=$(global_entity_count_on "$S3B_LEADER_PORT") +[ "$S3B_GLOBAL_AFTER" = "$S3B_GLOBAL_BEFORE" ] \ + && pass "[16] global graph survived restart ($S3B_GLOBAL_AFTER entities, was $S3B_GLOBAL_BEFORE)" \ + || fail "[16] global entity count changed after restart ($S3B_GLOBAL_BEFORE -> $S3B_GLOBAL_AFTER)" +RESTORED_SOURCES=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global/entities/OpenAI/sources" 2>/dev/null | \ + python3 -c "import sys,json; print(json.load(sys.stdin).get('sources',[]))" 2>/dev/null || echo "[]") +echo "$RESTORED_SOURCES" | grep -qF "$S3B_SA" \ + && pass "[16] visibility and provenance restored after restart (SA still owns OpenAI)" \ + || fail "[16] SA no longer in OpenAI sources after restart (visibility or provenance lost)" + +# [15] Provenance-scoped deletion +echo "[15] provenance-scoped deletion" +S3B_LEADER_PORT=$(find_leader_port) +S3B_LEADER="http://localhost:$S3B_LEADER_PORT" +# Delete SA: OpenAI must remain (SB still contributes it) +curl -sf -X DELETE "$S3B_LEADER/sessions/$S3B_SA" > /dev/null +sleep 1 +AFTER_SA=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print([e['name'] for e in json.load(sys.stdin).get('entities',[])])" \ + 2>/dev/null || echo "[]") +echo "$AFTER_SA" | grep -q "OpenAI" \ + && pass "[15] OpenAI remains after deleting SA (still contributed by SB)" \ + || fail "[15] OpenAI wrongly removed when only SA was deleted" +# Delete SB: OpenAI must now be gone (no remaining contributors) +curl -sf -X DELETE "$S3B_LEADER/sessions/$S3B_SB" > /dev/null +sleep 1 +AFTER_SB=$(curl -sf "http://localhost:$S3B_LEADER_PORT/knowledge/global" 2>/dev/null | \ + python3 -c "import sys,json; print([e['name'] for e in json.load(sys.stdin).get('entities',[])])" \ + 2>/dev/null || echo "[]") +echo "$AFTER_SB" | grep -q "OpenAI" \ + && fail "[15] OpenAI still in global graph after deleting all contributing sessions" \ + || pass "[15] OpenAI removed after deleting both SA and SB" + +echo "" +echo "=== All Stage 3B criteria PASSED ===" diff --git a/src/app.rs b/src/app.rs index e5a31bf..d9018d6 100644 --- a/src/app.rs +++ b/src/app.rs @@ -42,6 +42,8 @@ pub async fn build_raft_node( embedding_tx: tokio::sync::mpsc::Sender, knowledge_graph: Arc>, knowledge_tx: tokio::sync::mpsc::Sender, + global_graph: Arc>, + metrics: Arc, ) -> anyhow::Result> { use crate::raft::{ log_store::EngRaftLogStore, network::EngRaftNetwork, @@ -69,6 +71,8 @@ pub async fn build_raft_node( knowledge_graph, knowledge_tx, db, + global_graph, + metrics, ); // RECOVERY: flush Redis + restore snapshot BEFORE openraft replays the log. @@ -119,8 +123,12 @@ mod stage3a_tests { crate::knowledge::graph::KnowledgeGraph::new(), )); let (ktx, _krx) = tokio::sync::mpsc::channel(10); + let gg = std::sync::Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )); - let raft = super::build_raft_node(&cfg, short_term, core_memory, vector_store, etx, kg, ktx) + let metrics = std::sync::Arc::new(crate::metrics::AppMetrics::new().unwrap()); + let raft = super::build_raft_node(&cfg, short_term, core_memory, vector_store, etx, kg, ktx, gg, metrics) .await .unwrap(); assert!(raft.is_initialized().await.is_ok() || true); @@ -143,6 +151,9 @@ mod stage3a_tests { crate::knowledge::graph::KnowledgeGraph::new(), )); let (ktx, _krx) = tokio::sync::mpsc::channel(10); + let gg = std::sync::Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )); // Pre-load stale data that recovery must flush. use crate::core::ShortTermMemory; @@ -158,6 +169,7 @@ mod stage3a_tests { .unwrap(); let st_clone = short_term.clone(); + let metrics = std::sync::Arc::new(crate::metrics::AppMetrics::new().unwrap()); let _raft = super::build_raft_node( &cfg, short_term as std::sync::Arc, @@ -166,6 +178,8 @@ mod stage3a_tests { etx, kg, ktx, + gg, + metrics, ) .await .unwrap(); @@ -251,6 +265,7 @@ pub async fn build_app_state_with_embedding_provider( ); let knowledge_graph = Arc::new(tokio::sync::RwLock::new(KnowledgeGraph::new())); + let global_graph = Arc::new(tokio::sync::RwLock::new(crate::knowledge::global::GlobalGraph::new())); let (knowledge_job_sender, knowledge_receiver) = knowledge_job_channel(config.knowledge_channel_size); let knowledge_extractor: Arc = @@ -274,6 +289,8 @@ pub async fn build_app_state_with_embedding_provider( embedding_job_sender.clone(), knowledge_graph.clone(), knowledge_job_sender.clone(), + global_graph.clone(), + metrics.clone(), ) .await .map_err(|e| AppBuildError::Other(e.into()))?; @@ -319,5 +336,6 @@ pub async fn build_app_state_with_embedding_provider( cluster_peers, knowledge_graph, knowledge_job_sender, + global_graph, })) } diff --git a/src/cluster.rs b/src/cluster.rs index 9340cd7..7b46e36 100644 --- a/src/cluster.rs +++ b/src/cluster.rs @@ -208,6 +208,7 @@ mod tests { ..Config::default() }; let knowledge_graph = Arc::new(tokio::sync::RwLock::new(crate::knowledge::graph::KnowledgeGraph::new())); + let global_graph = Arc::new(tokio::sync::RwLock::new(crate::knowledge::global::GlobalGraph::new())); let (knowledge_tx, mut knowledge_rx) = tokio::sync::mpsc::channel::(500); tokio::spawn(async move { while knowledge_rx.recv().await.is_some() {} }); let raft = build_raft_node( @@ -218,6 +219,8 @@ mod tests { c.embedding_job_sender.clone(), knowledge_graph.clone(), knowledge_tx.clone(), + global_graph, + c.metrics.clone(), ) .await .unwrap(); @@ -246,6 +249,9 @@ mod tests { cluster_peers: vec![], knowledge_graph, knowledge_job_sender: knowledge_tx, + global_graph: Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )), }); (TestServer::new(build_router(state)).unwrap(), raft_dir) } @@ -274,6 +280,9 @@ mod tests { crate::knowledge::graph::KnowledgeGraph::new(), )), knowledge_job_sender, + global_graph: Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )), }); TestServer::new(build_router(state)).unwrap() } diff --git a/src/knowledge/global.rs b/src/knowledge/global.rs new file mode 100644 index 0000000..3cf74b7 --- /dev/null +++ b/src/knowledge/global.rs @@ -0,0 +1,615 @@ +use std::collections::{HashMap, HashSet, VecDeque}; +use petgraph::graph::{DiGraph, NodeIndex}; +use petgraph::visit::EdgeRef; +use petgraph::Direction; +use serde::{Deserialize, Serialize}; + +use crate::knowledge::graph::{PathEdge, RelatedEntity, RelationshipDirection}; +use crate::knowledge::types::{Entity, Relationship}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum Visibility { + Private, + Shared, +} + +impl Default for Visibility { + fn default() -> Self { + Visibility::Private + } +} + +#[derive(Debug, Clone)] +struct AttrValue { + value: String, + index: u64, +} + +#[derive(Debug, Clone)] +struct GlobalNode { + name: String, + entity_type: String, + entity_type_index: u64, + attributes: HashMap, + /// session_id -> agent_id (None if no agent was specified for that session) + provenance: HashMap>, +} + +#[derive(Debug, Clone)] +struct GlobalEdge { + relationship_type: String, + sources: HashSet, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Conflict { + pub from: String, + pub relationship_type: String, + pub targets: Vec, +} + +pub struct GlobalGraph { + graph: DiGraph, + name_to_idx: HashMap, +} + +impl GlobalGraph { + pub fn new() -> Self { + Self { graph: DiGraph::new(), name_to_idx: HashMap::new() } + } + + fn ensure_entity( + &mut self, + name: &str, + entity_type: &str, + session_id: &str, + agent_id: Option<&str>, + index: u64, + ) -> NodeIndex { + if let Some(&idx) = self.name_to_idx.get(name) { + let node = &mut self.graph[idx]; + node.provenance.entry(session_id.to_string()) + .or_insert_with(|| agent_id.map(|s| s.to_string())); + // Update entity_type with LWW; "Other" sentinel is never a real type. + if entity_type != "Other" && index >= node.entity_type_index { + node.entity_type = entity_type.to_string(); + node.entity_type_index = index; + } + return idx; + } + let idx = self.graph.add_node(GlobalNode { + name: name.to_string(), + entity_type: entity_type.to_string(), + entity_type_index: index, + attributes: HashMap::new(), + provenance: { + let mut m = HashMap::new(); + m.insert(session_id.to_string(), agent_id.map(|s| s.to_string())); + m + }, + }); + self.name_to_idx.insert(name.to_string(), idx); + idx + } + + pub fn merge_with_agent( + &mut self, + session_id: &str, + agent_id: Option<&str>, + index: u64, + entities: Vec, + relationships: Vec, + ) { + for entity in &entities { + let idx = self.ensure_entity(&entity.name, &entity.entity_type, session_id, agent_id, index); + let node = &mut self.graph[idx]; + for (key, value) in &entity.attributes { + let should_update = match node.attributes.get(key) { + Some(existing) => index >= existing.index, + None => true, + }; + if should_update { + node.attributes.insert(key.clone(), AttrValue { value: value.clone(), index }); + } + } + } + + for rel in &relationships { + let from_idx = self.ensure_entity(&rel.from, "Other", session_id, agent_id, index); + let to_idx = self.ensure_entity(&rel.to, "Other", session_id, agent_id, index); + + // Dedup by (from, to, relationship_type); accumulate sources. + let existing = self.graph + .edges_directed(from_idx, Direction::Outgoing) + .find(|e| { + e.target() == to_idx + && self.graph[e.id()].relationship_type == rel.relationship_type + }) + .map(|e| e.id()); + + if let Some(eidx) = existing { + self.graph[eidx].sources.insert(session_id.to_string()); + } else { + let mut sources = HashSet::new(); + sources.insert(session_id.to_string()); + self.graph.add_edge(from_idx, to_idx, GlobalEdge { + relationship_type: rel.relationship_type.clone(), + sources, + }); + } + } + } + + pub fn merge( + &mut self, + session_id: &str, + index: u64, + entities: Vec, + relationships: Vec, + ) { + self.merge_with_agent(session_id, None, index, entities, relationships); + } + + pub fn get_related(&self, name: &str) -> Vec { + let Some(&idx) = self.name_to_idx.get(name) else { return vec![] }; + let mut related = Vec::new(); + for edge in self.graph.edges_directed(idx, Direction::Outgoing) { + let node = &self.graph[edge.target()]; + related.push(RelatedEntity { + name: node.name.clone(), + entity_type: node.entity_type.clone(), + relationship_type: self.graph[edge.id()].relationship_type.clone(), + direction: RelationshipDirection::Outgoing, + }); + } + for edge in self.graph.edges_directed(idx, Direction::Incoming) { + let node = &self.graph[edge.source()]; + related.push(RelatedEntity { + name: node.name.clone(), + entity_type: node.entity_type.clone(), + relationship_type: self.graph[edge.id()].relationship_type.clone(), + direction: RelationshipDirection::Incoming, + }); + } + related + } + + pub fn entity_sources(&self, name: &str) -> Vec { + let Some(&idx) = self.name_to_idx.get(name) else { return vec![] }; + let mut sources: Vec = self.graph[idx].provenance.keys().cloned().collect(); + sources.sort(); + sources + } + + pub fn entity_agents(&self, name: &str) -> Vec { + let Some(&idx) = self.name_to_idx.get(name) else { return vec![] }; + let mut agents: Vec = self.graph[idx] + .provenance + .values() + .filter_map(|a| a.as_ref()) + .cloned() + .collect::>() + .into_iter() + .collect(); + agents.sort(); + agents + } + + pub fn entity_attribute(&self, name: &str, key: &str) -> Option { + let &idx = self.name_to_idx.get(name)?; + self.graph[idx].attributes.get(key).map(|v| v.value.clone()) + } + + pub fn all_entities(&self) -> Vec { + self.graph + .node_indices() + .map(|idx| { + let n = &self.graph[idx]; + Entity { + name: n.name.clone(), + entity_type: n.entity_type.clone(), + attributes: n + .attributes + .iter() + .map(|(k, v)| (k.clone(), v.value.clone())) + .collect(), + } + }) + .collect() + } + + pub fn all_relationships(&self) -> Vec { + self.graph + .edge_indices() + .map(|eidx| { + let (src, tgt) = self.graph.edge_endpoints(eidx).unwrap(); + Relationship { + from: self.graph[src].name.clone(), + to: self.graph[tgt].name.clone(), + relationship_type: self.graph[eidx].relationship_type.clone(), + } + }) + .collect() + } + + pub fn find_path(&self, from: &str, to: &str) -> Option> { + let &from_idx = self.name_to_idx.get(from)?; + let &to_idx = self.name_to_idx.get(to)?; + + if from_idx == to_idx { + return Some(vec![]); + } + + let mut parent: HashMap = HashMap::new(); + let mut queue = VecDeque::new(); + queue.push_back(from_idx); + + 'bfs: while let Some(current) = queue.pop_front() { + for edge in self.graph.edges_directed(current, Direction::Outgoing) { + let next = edge.target(); + if parent.contains_key(&next) { + continue; + } + parent.insert(next, (current, self.graph[edge.id()].relationship_type.clone())); + if next == to_idx { + break 'bfs; + } + queue.push_back(next); + } + } + + if !parent.contains_key(&to_idx) { + return None; + } + + let mut path = Vec::new(); + let mut node = to_idx; + while node != from_idx { + let (prev, rel) = parent.remove(&node).unwrap(); + path.push(PathEdge { + from: self.graph[prev].name.clone(), + relationship_type: rel, + to: self.graph[node].name.clone(), + }); + node = prev; + } + path.reverse(); + Some(path) + } + + /// Remove all contributions from `session_id`. Elements with no remaining + /// contributors are dropped. Uses a rebuild to avoid petgraph swap-remove + /// index invalidation. + pub fn prune_session(&mut self, session_id: &str) { + // Collect surviving node data after removing session. + let nodes_data: Vec<(String, String, u64, HashMap, HashMap>)> = + self.graph + .node_indices() + .map(|idx| { + let n = &self.graph[idx]; + let mut prov = n.provenance.clone(); + prov.remove(session_id); + (n.name.clone(), n.entity_type.clone(), n.entity_type_index, n.attributes.clone(), prov) + }) + .filter(|(_, _, _, _, prov)| !prov.is_empty()) + .collect(); + + let surviving: HashSet<&str> = nodes_data.iter().map(|(n, _, _, _, _)| n.as_str()).collect(); + + let edges_data: Vec<(String, String, String, HashSet)> = self + .graph + .edge_indices() + .map(|eidx| { + let (src, tgt) = self.graph.edge_endpoints(eidx).unwrap(); + let mut sources = self.graph[eidx].sources.clone(); + sources.remove(session_id); + ( + self.graph[src].name.clone(), + self.graph[tgt].name.clone(), + self.graph[eidx].relationship_type.clone(), + sources, + ) + }) + .filter(|(from, to, _, sources)| { + !sources.is_empty() + && surviving.contains(from.as_str()) + && surviving.contains(to.as_str()) + }) + .collect(); + + // Rebuild from surviving data. + self.graph = DiGraph::new(); + self.name_to_idx = HashMap::new(); + + for (name, entity_type, entity_type_index, attributes, provenance) in &nodes_data { + let idx = self.graph.add_node(GlobalNode { + name: name.clone(), + entity_type: entity_type.clone(), + entity_type_index: *entity_type_index, + attributes: attributes.clone(), + provenance: provenance.clone(), + }); + self.name_to_idx.insert(name.clone(), idx); + } + + for (from, to, rel_type, sources) in &edges_data { + if let (Some(&fi), Some(&ti)) = (self.name_to_idx.get(from), self.name_to_idx.get(to)) { + self.graph.add_edge(fi, ti, GlobalEdge { + relationship_type: rel_type.clone(), + sources: sources.clone(), + }); + } + } + } + + pub fn conflicts(&self) -> Vec { + let mut result = Vec::new(); + for idx in self.graph.node_indices() { + let from_name = &self.graph[idx].name; + let mut by_type: HashMap> = HashMap::new(); + for edge in self.graph.edges_directed(idx, Direction::Outgoing) { + let rel_type = self.graph[edge.id()].relationship_type.clone(); + let target = self.graph[edge.target()].name.clone(); + by_type.entry(rel_type).or_default().push(target); + } + for (rel_type, targets) in by_type { + if targets.len() > 1 { + result.push(Conflict { + from: from_name.clone(), + relationship_type: rel_type, + targets, + }); + } + } + } + result + } + + pub fn to_snapshot(&self) -> GlobalGraphSnapshot { + let nodes = self + .graph + .node_indices() + .map(|idx| { + let n = &self.graph[idx]; + GlobalNodeSnapshot { + name: n.name.clone(), + entity_type: n.entity_type.clone(), + entity_type_index: n.entity_type_index, + attributes: n + .attributes + .iter() + .map(|(k, v)| (k.clone(), v.value.clone(), v.index)) + .collect(), + provenance: n.provenance.iter().map(|(s, a)| (s.clone(), a.clone())).collect(), + } + }) + .collect(); + + let edges = self + .graph + .edge_indices() + .map(|eidx| { + let (src, tgt) = self.graph.edge_endpoints(eidx).unwrap(); + GlobalEdgeSnapshot { + from: self.graph[src].name.clone(), + to: self.graph[tgt].name.clone(), + relationship_type: self.graph[eidx].relationship_type.clone(), + sources: self.graph[eidx].sources.iter().cloned().collect(), + } + }) + .collect(); + + GlobalGraphSnapshot { nodes, edges } + } + + pub fn from_snapshot(snap: GlobalGraphSnapshot) -> Self { + let mut g = GlobalGraph::new(); + for ns in &snap.nodes { + let attributes = ns + .attributes + .iter() + .map(|(k, v, i)| (k.clone(), AttrValue { value: v.clone(), index: *i })) + .collect(); + let provenance = ns.provenance.iter().map(|(s, a)| (s.clone(), a.clone())).collect(); + let idx = g.graph.add_node(GlobalNode { + name: ns.name.clone(), + entity_type: ns.entity_type.clone(), + entity_type_index: ns.entity_type_index, + attributes, + provenance, + }); + g.name_to_idx.insert(ns.name.clone(), idx); + } + for es in &snap.edges { + if let (Some(&fi), Some(&ti)) = (g.name_to_idx.get(&es.from), g.name_to_idx.get(&es.to)) { + g.graph.add_edge(fi, ti, GlobalEdge { + relationship_type: es.relationship_type.clone(), + sources: es.sources.iter().cloned().collect(), + }); + } + } + g + } +} + +impl Default for GlobalGraph { + fn default() -> Self { + Self::new() + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GlobalNodeSnapshot { + pub name: String, + pub entity_type: String, + pub entity_type_index: u64, + /// (key, value, raft_log_index) + pub attributes: Vec<(String, String, u64)>, + /// (session_id, agent_id) + pub provenance: Vec<(String, Option)>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GlobalEdgeSnapshot { + pub from: String, + pub to: String, + pub relationship_type: String, + pub sources: Vec, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct GlobalGraphSnapshot { + pub nodes: Vec, + pub edges: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn visibility_defaults_private_and_round_trips() { + assert_eq!(Visibility::default(), Visibility::Private); + let j = serde_json::to_string(&Visibility::Shared).unwrap(); + assert_eq!(serde_json::from_str::(&j).unwrap(), Visibility::Shared); + } +} + +#[cfg(test)] +mod graph_tests { + use super::*; + use crate::knowledge::types::{Entity, Relationship}; + use std::collections::HashMap; + + fn ent(name: &str, t: &str) -> Entity { + Entity { name: name.into(), entity_type: t.into(), attributes: HashMap::new() } + } + fn ent_attr(name: &str, k: &str, v: &str) -> Entity { + let mut a = HashMap::new(); + a.insert(k.into(), v.into()); + Entity { name: name.into(), entity_type: "Person".into(), attributes: a } + } + fn rel(f: &str, t: &str, ty: &str) -> Relationship { + Relationship { from: f.into(), to: t.into(), relationship_type: ty.into() } + } + + #[test] + fn merges_two_sessions_into_one_global_view() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("OpenAI", "Organization")], vec![rel("Alice", "OpenAI", "works_at")]); + g.merge("s2", 1, vec![ent("Bob", "Person"), ent("OpenAI", "Organization")], vec![rel("Bob", "OpenAI", "works_at")]); + let related = g.get_related("OpenAI"); + let names: Vec<&str> = related.iter().map(|r| r.name.as_str()).collect(); + assert!(names.contains(&"Alice")); + assert!(names.contains(&"Bob")); + } + + #[test] + fn provenance_lists_contributing_sessions() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("OpenAI", "Organization")], vec![]); + g.merge("s2", 1, vec![ent("OpenAI", "Organization")], vec![]); + let mut sources = g.entity_sources("OpenAI"); + sources.sort(); + assert_eq!(sources, vec!["s1".to_string(), "s2".to_string()]); + } + + #[test] + fn attribute_conflict_resolves_last_writer_wins_by_index() { + let mut g = GlobalGraph::new(); + g.merge("s1", 5, vec![ent_attr("Alice", "role", "eng")], vec![]); + g.merge("s2", 9, vec![ent_attr("Alice", "role", "manager")], vec![]); + assert_eq!(g.entity_attribute("Alice", "role").as_deref(), Some("manager")); + // Lower index applied later must NOT overwrite a higher index. + g.merge("s3", 7, vec![ent_attr("Alice", "role", "intern")], vec![]); + assert_eq!(g.entity_attribute("Alice", "role").as_deref(), Some("manager")); + } + + #[test] + fn pruning_one_session_keeps_elements_with_other_contributors() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person")], vec![]); + g.merge("s2", 1, vec![ent("Alice", "Person")], vec![]); + g.prune_session("s1"); + assert!(g.all_entities().iter().any(|e| e.name == "Alice")); + g.prune_session("s2"); + assert!(g.all_entities().is_empty()); + } + + #[test] + fn pruning_drops_orphaned_edges() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("Bob", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.prune_session("s1"); + assert!(g.find_path("Alice", "Bob").is_none()); + } + + #[test] + fn prune_shared_entity_unique_relationship() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("Bob", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.merge("s2", 1, vec![ent("Alice", "Person")], vec![]); + g.prune_session("s1"); + assert!(g.all_entities().iter().any(|e| e.name == "Alice"), "Alice survives via s2"); + assert!(g.find_path("Alice", "Bob").is_none(), "the unique s1 edge is gone"); + } + + #[test] + fn prune_unique_entity_shared_relationship() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("Bob", "Person"), ent("Carol", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.merge("s2", 1, vec![ent("Alice", "Person"), ent("Bob", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.prune_session("s1"); + assert!(g.find_path("Alice", "Bob").is_some(), "shared edge survives via s2"); + assert!(!g.all_entities().iter().any(|e| e.name == "Carol"), "unique Carol is gone"); + } + + #[test] + fn prune_shared_entity_and_shared_relationship() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("Bob", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.merge("s2", 1, vec![ent("Alice", "Person"), ent("Bob", "Person")], vec![rel("Alice", "Bob", "knows")]); + g.prune_session("s1"); + assert!(g.find_path("Alice", "Bob").is_some(), "fully shared element survives one prune"); + g.prune_session("s2"); + assert!(g.all_entities().is_empty(), "removing the last contributor clears it"); + } + + #[test] + fn global_graph_tracks_agent_provenance() { + let mut g = GlobalGraph::new(); + g.merge_with_agent("s1", Some("agent-7"), 0, vec![ent("OpenAI", "Organization")], vec![]); + assert_eq!(g.entity_agents("OpenAI"), vec!["agent-7".to_string()]); + } + + #[test] + fn agent_provenance_survives_snapshot_round_trip() { + let mut g = GlobalGraph::new(); + g.merge_with_agent("s1", Some("agent-7"), 0, vec![ent("OpenAI", "Organization")], vec![]); + let g2 = GlobalGraph::from_snapshot(g.to_snapshot()); + assert_eq!(g2.entity_agents("OpenAI"), vec!["agent-7".to_string()]); + assert_eq!(g2.entity_sources("OpenAI"), vec!["s1".to_string()]); + } + + #[test] + fn contradictions_are_reported() { + let mut g = GlobalGraph::new(); + g.merge("s1", 0, vec![ent("Alice", "Person"), ent("X", "Organization")], vec![rel("Alice", "X", "works_at")]); + g.merge("s2", 1, vec![ent("Alice", "Person"), ent("Y", "Organization")], vec![rel("Alice", "Y", "works_at")]); + let conflicts = g.conflicts(); + assert!(conflicts.iter().any(|c| c.from == "Alice" && c.relationship_type == "works_at")); + } + + #[test] + fn snapshot_round_trips_with_provenance() { + let mut g = GlobalGraph::new(); + g.merge("s1", 3, vec![ent_attr("Alice", "role", "eng"), ent("OpenAI", "Organization")], vec![rel("Alice", "OpenAI", "works_at")]); + let snap = g.to_snapshot(); + let json = serde_json::to_string(&snap).unwrap(); + let back: GlobalGraphSnapshot = serde_json::from_str(&json).unwrap(); + let g2 = GlobalGraph::from_snapshot(back); + assert_eq!(g2.entity_sources("OpenAI"), vec!["s1".to_string()]); + assert_eq!(g2.entity_attribute("Alice", "role").as_deref(), Some("eng")); + assert!(g2.find_path("Alice", "OpenAI").is_some()); + } +} diff --git a/src/knowledge/global_handler.rs b/src/knowledge/global_handler.rs new file mode 100644 index 0000000..c0cdaf1 --- /dev/null +++ b/src/knowledge/global_handler.rs @@ -0,0 +1,171 @@ +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, Query, State}, + http::StatusCode, +}; +use serde::{Deserialize, Serialize}; + +use crate::core::MemoryServerError; +use crate::knowledge::export::{GraphExport, to_dot}; +use crate::knowledge::global::{Conflict, Visibility}; +use crate::knowledge::graph::{PathEdge, RelatedEntity}; +use crate::knowledge::types::{Entity, Relationship}; +use crate::raft::types::MemoryCommand; +use crate::server::AppState; + +#[derive(Debug, Deserialize)] +pub struct SetVisibilityRequest { + pub visibility: Visibility, +} + +pub async fn set_visibility( + State(state): State>, + Path(session_id): Path, + Json(body): Json, +) -> Result { + if let Some(raft) = &state.raft { + return raft + .client_write(MemoryCommand::SetSessionVisibility { + session_id: session_id.clone(), + visibility: body.visibility, + }) + .await + .map(|_| StatusCode::NO_CONTENT) + .map_err(|e| { + if let Some(fwd) = e.forward_to_leader::() { + if let Some(leader_id) = fwd.leader_id { + if let Some(http_addr) = state.peer_http_addrs.get(&leader_id) { + let location = format!( + "http://{}/sessions/{}/visibility", + http_addr, session_id + ); + return MemoryServerError::RedirectToLeader(location); + } + } + return MemoryServerError::NoLeader; + } + MemoryServerError::Internal(format!("raft error: {e}")) + }); + } + + // Standalone mode: no cluster to coordinate; accept and no-op. + Ok(StatusCode::NO_CONTENT) +} + +// ---- Read-only global graph handlers ---------------------------------------- + +#[derive(Serialize)] +pub struct GlobalKnowledgeResponse { + entities: Vec, + edges: Vec, +} + +pub async fn get_global( + State(state): State>, +) -> Json { + let g = state.global_graph.read().await; + Json(GlobalKnowledgeResponse { + entities: g.all_entities(), + edges: g.all_relationships(), + }) +} + +#[derive(Serialize)] +pub(crate) struct RelatedResponse { + entity_name: String, + related: Vec, +} + +pub async fn get_global_entity( + State(state): State>, + Path(name): Path, +) -> Json { + let g = state.global_graph.read().await; + Json(RelatedResponse { + entity_name: name.clone(), + related: g.get_related(&name), + }) +} + +#[derive(Serialize)] +pub(crate) struct SourcesResponse { + entity_name: String, + sources: Vec, +} + +pub async fn get_global_entity_sources( + State(state): State>, + Path(name): Path, +) -> Json { + let g = state.global_graph.read().await; + Json(SourcesResponse { + entity_name: name.clone(), + sources: g.entity_sources(&name), + }) +} + +#[derive(Deserialize)] +pub struct PathQuery { + from: String, + to: String, +} + +#[derive(Serialize)] +pub(crate) struct PathResponse { + from: String, + to: String, + path: Option>, +} + +pub async fn get_global_path( + State(state): State>, + Query(params): Query, +) -> Json { + let g = state.global_graph.read().await; + let path = g.find_path(¶ms.from, ¶ms.to); + Json(PathResponse { from: params.from, to: params.to, path }) +} + +#[derive(Deserialize)] +pub(crate) struct ExportQuery { + #[serde(default = "default_format")] + format: String, +} + +fn default_format() -> String { + "json".to_string() +} + +pub async fn get_global_export( + State(state): State>, + Query(params): Query, +) -> (StatusCode, [(axum::http::header::HeaderName, &'static str); 1], String) { + let g = state.global_graph.read().await; + let export = GraphExport::new("global", g.all_entities(), g.all_relationships()); + match params.format.as_str() { + "dot" => ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "text/vnd.graphviz")], + to_dot(&export), + ), + _ => ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "application/json")], + serde_json::to_string(&export).unwrap_or_default(), + ), + } +} + +#[derive(Serialize)] +pub(crate) struct ConflictsResponse { + conflicts: Vec, +} + +pub async fn get_global_conflicts( + State(state): State>, +) -> Json { + let g = state.global_graph.read().await; + Json(ConflictsResponse { conflicts: g.conflicts() }) +} diff --git a/src/knowledge/handler.rs b/src/knowledge/handler.rs index 2ab5382..795a856 100644 --- a/src/knowledge/handler.rs +++ b/src/knowledge/handler.rs @@ -155,6 +155,9 @@ mod tests { cluster_peers: vec![], knowledge_graph: kg, knowledge_job_sender, + global_graph: Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )), }) } diff --git a/src/knowledge/mod.rs b/src/knowledge/mod.rs index 5ec1622..ccfe71a 100644 --- a/src/knowledge/mod.rs +++ b/src/knowledge/mod.rs @@ -1,8 +1,11 @@ pub mod extractor; pub mod export; +pub mod global; +pub mod global_handler; pub mod graph; pub mod handler; pub mod types; pub mod worker; +pub use global::{GlobalGraph, GlobalGraphSnapshot, Visibility}; pub use types::{Entity, ExtractionResult, KnowledgeJob, Relationship}; diff --git a/src/metrics.rs b/src/metrics.rs index 8959d1a..efb8f41 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -25,6 +25,9 @@ pub struct AppMetrics { snapshot_build_total: IntCounter, snapshot_install_total: IntCounter, snapshot_last_index: IntGauge, + global_entities: IntGauge, + global_relationships: IntGauge, + global_conflicts: IntGauge, } impl AppMetrics { @@ -148,6 +151,24 @@ impl AppMetrics { ))?; registry.register(Box::new(snapshot_last_index.clone()))?; + let global_entities = IntGauge::with_opts(Opts::new( + "global_entities", + "Current number of entities in the global knowledge graph.", + ))?; + registry.register(Box::new(global_entities.clone()))?; + + let global_relationships = IntGauge::with_opts(Opts::new( + "global_relationships", + "Current number of relationships in the global knowledge graph.", + ))?; + registry.register(Box::new(global_relationships.clone()))?; + + let global_conflicts = IntGauge::with_opts(Opts::new( + "global_conflicts", + "Current number of conflicting relationships in the global knowledge graph.", + ))?; + registry.register(Box::new(global_conflicts.clone()))?; + Ok(Self { registry, messages_added_total, @@ -167,6 +188,9 @@ impl AppMetrics { snapshot_build_total, snapshot_install_total, snapshot_last_index, + global_entities, + global_relationships, + global_conflicts, }) } @@ -244,6 +268,18 @@ impl AppMetrics { self.snapshot_last_index.set(index as i64); } + pub fn set_global_entities(&self, count: usize) { + self.global_entities.set(count as i64); + } + + pub fn set_global_relationships(&self, count: usize) { + self.global_relationships.set(count as i64); + } + + pub fn set_global_conflicts(&self, count: usize) { + self.global_conflicts.set(count as i64); + } + pub fn render(&self) -> Result { let mut buffer = Vec::new(); let encoder = TextEncoder::new(); @@ -270,4 +306,16 @@ mod tests { assert!(text.contains("engram_snapshot_install_total")); assert!(text.contains("engram_snapshot_last_index")); } + + #[test] + fn renders_global_graph_metrics() { + let m = AppMetrics::new().unwrap(); + m.set_global_entities(3); + m.set_global_relationships(2); + m.set_global_conflicts(1); + let t = m.render().unwrap(); + assert!(t.contains("engram_global_entities")); + assert!(t.contains("engram_global_relationships")); + assert!(t.contains("engram_global_conflicts")); + } } \ No newline at end of file diff --git a/src/raft/recovery.rs b/src/raft/recovery.rs index f94e898..8d47dd9 100644 --- a/src/raft/recovery.rs +++ b/src/raft/recovery.rs @@ -56,7 +56,9 @@ mod tests { let (etx, _erx) = mpsc::channel(10); let (ktx, _krx) = mpsc::channel(10); let kg = Arc::new(RwLock::new(KnowledgeGraph::new())); - let sm = EngStateMachineStore::new(st.clone(), cm.clone(), vs, etx, kg.clone(), ktx, db); + let gg = Arc::new(RwLock::new(crate::knowledge::global::GlobalGraph::new())); + let metrics = Arc::new(crate::metrics::AppMetrics::new().unwrap()); + let sm = EngStateMachineStore::new(st.clone(), cm.clone(), vs, etx, kg.clone(), ktx, db, gg, metrics); (sm, st, cm, kg) } diff --git a/src/raft/snapshot.rs b/src/raft/snapshot.rs index 309d960..4b7065b 100644 --- a/src/raft/snapshot.rs +++ b/src/raft/snapshot.rs @@ -4,7 +4,7 @@ use crate::knowledge::graph::GraphSnapshot; use crate::models::Message; /// Snapshot schema version. Bump when the payload layout changes incompatibly. -pub const SNAPSHOT_VERSION: u32 = 1; +pub const SNAPSHOT_VERSION: u32 = 2; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SessionMessages { @@ -28,9 +28,12 @@ pub struct EngramSnapshot { pub short_term: Vec, pub core_memory: Vec, pub knowledge_graph: GraphSnapshot, - /// Reserved for Stage 3B (collective/global knowledge graph). Absent in 3A. #[serde(default)] - pub global_graph: Option, + pub global_graph: Option, + #[serde(default)] + pub visibility: Vec<(String, crate::knowledge::global::Visibility)>, + #[serde(default)] + pub session_agents: Vec<(String, String)>, } impl EngramSnapshot { @@ -56,12 +59,14 @@ mod tests { core_memory: vec![SessionFacts { session_id: "s1".into(), facts: vec!["f".into()] }], knowledge_graph: crate::knowledge::graph::GraphSnapshot::default(), global_graph: None, + visibility: vec![], + session_agents: vec![], } } #[test] - fn snapshot_carries_version_one() { - assert_eq!(sample().version, 1); + fn snapshot_carries_version_two() { + assert_eq!(sample().version, 2); } #[test] @@ -69,9 +74,10 @@ mod tests { let snap = sample(); let bytes = snap.to_bytes().unwrap(); let back = EngramSnapshot::from_bytes(&bytes).unwrap(); - assert_eq!(back.version, 1); + assert_eq!(back.version, 2); assert_eq!(back.core_memory[0].facts, vec!["f".to_string()]); assert!(back.global_graph.is_none()); + assert!(back.visibility.is_empty()); } #[test] @@ -81,4 +87,32 @@ mod tests { let back = EngramSnapshot::from_bytes(&bytes).unwrap(); assert!(back.global_graph.is_none()); } + + #[test] + fn snapshot_version_is_two_and_carries_global_and_visibility() { + let snap = EngramSnapshot { + version: SNAPSHOT_VERSION, + short_term: vec![], + core_memory: vec![], + knowledge_graph: crate::knowledge::graph::GraphSnapshot::default(), + global_graph: Some(crate::knowledge::global::GlobalGraphSnapshot::default()), + visibility: vec![("s1".into(), crate::knowledge::global::Visibility::Shared)], + session_agents: vec![("s1".into(), "agent-7".into())], + }; + assert_eq!(snap.version, 2); + let bytes = snap.to_bytes().unwrap(); + let back = EngramSnapshot::from_bytes(&bytes).unwrap(); + assert!(back.global_graph.is_some()); + assert_eq!(back.visibility.len(), 1); + assert_eq!(back.session_agents, vec![("s1".to_string(), "agent-7".to_string())]); + } + + #[test] + fn v1_snapshot_without_global_fields_still_loads() { + let v1 = r#"{"version":1,"short_term":[],"core_memory":[],"knowledge_graph":{"sessions":[],"processed":[]}}"#; + let back = EngramSnapshot::from_bytes(v1.as_bytes()).unwrap(); + assert!(back.global_graph.is_none()); + assert!(back.visibility.is_empty()); + assert!(back.session_agents.is_empty()); + } } diff --git a/src/raft/state_machine.rs b/src/raft/state_machine.rs index 6282608..508bef5 100644 --- a/src/raft/state_machine.rs +++ b/src/raft/state_machine.rs @@ -10,12 +10,15 @@ use openraft::{ use redb::{Database, TableDefinition}; use crate::core::{CoreMemoryStore, ShortTermMemory}; +use crate::knowledge::global::{GlobalGraph, Visibility}; use crate::knowledge::graph::KnowledgeGraph; use crate::knowledge::types::KnowledgeJob; +use crate::metrics::AppMetrics; use crate::models::{EmbeddingStatus, Message}; use crate::raft::snapshot::{EngramSnapshot, SessionFacts, SessionMessages}; use crate::raft::types::{CommandResponse, MemoryCommand, TypeConfig}; use crate::worker::EmbeddingJob; +use std::collections::HashMap; const SNAPSHOT_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("raft_snapshot"); const SNAPSHOT_META_KEY: &str = "meta"; @@ -33,6 +36,10 @@ struct SmInner { embedding_tx: mpsc::Sender, knowledge_graph: Arc>, knowledge_tx: mpsc::Sender, + global_graph: Arc>, + visibility: Arc>>, + session_agents: Arc>>, + metrics: Arc, db: Arc, snapshot_idx: u64, } @@ -48,6 +55,8 @@ impl EngStateMachineStore { knowledge_graph: Arc>, knowledge_tx: mpsc::Sender, db: Arc, + global_graph: Arc>, + metrics: Arc, ) -> Self { { let txn = db.begin_write().expect("redb begin_write sm init"); @@ -63,6 +72,10 @@ impl EngStateMachineStore { embedding_tx, knowledge_graph, knowledge_tx, + global_graph, + visibility: Arc::new(RwLock::new(HashMap::new())), + session_agents: Arc::new(RwLock::new(HashMap::new())), + metrics, db, snapshot_idx: 0, })), @@ -141,13 +154,24 @@ async fn build_payload(inner: &SmInner) -> Result<(EngramSnapshot, SnapshotMeta< .map(|(session_id, facts)| SessionFacts { session_id, facts }) .collect(); let knowledge_graph = inner.knowledge_graph.read().await.to_snapshot(); + let global_graph = Some(inner.global_graph.read().await.to_snapshot()); + let visibility = inner.visibility.read().await + .iter() + .map(|(k, v)| (k.clone(), *v)) + .collect(); + let session_agents = inner.session_agents.read().await + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); let payload = EngramSnapshot { version: crate::raft::snapshot::SNAPSHOT_VERSION, short_term, core_memory, knowledge_graph, - global_graph: None, + global_graph, + visibility, + session_agents, }; let snapshot_id = format!( "{}-{}", @@ -222,7 +246,7 @@ impl RaftStateMachine for EngStateMachineStore { I::IntoIter: Send, { // Clone Arcs once so the lock is not held across async apply_cmd calls. - let (short_term, core_memory, embedding_tx, knowledge_graph, knowledge_tx) = { + let (short_term, core_memory, embedding_tx, knowledge_graph, knowledge_tx, global_graph, visibility, session_agents, metrics) = { let inner = self.inner.lock().await; ( inner.short_term.clone(), @@ -230,6 +254,10 @@ impl RaftStateMachine for EngStateMachineStore { inner.embedding_tx.clone(), inner.knowledge_graph.clone(), inner.knowledge_tx.clone(), + inner.global_graph.clone(), + inner.visibility.clone(), + inner.session_agents.clone(), + inner.metrics.clone(), ) }; @@ -238,12 +266,13 @@ impl RaftStateMachine for EngStateMachineStore { let mut last_membership = None; for entry in entries { + let index = entry.log_id.index; last_applied = Some(entry.log_id.clone()); if let EntryPayload::Membership(mem) = &entry.payload { last_membership = Some(StoredMembership::new(Some(entry.log_id.clone()), mem.clone())); } if let EntryPayload::Normal(cmd) = entry.payload { - apply_cmd(cmd, &short_term, &core_memory, &embedding_tx, &knowledge_graph, &knowledge_tx).await; + apply_cmd(cmd, &short_term, &core_memory, &embedding_tx, &knowledge_graph, &knowledge_tx, &global_graph, &visibility, &session_agents, &metrics, index).await; } responses.push(CommandResponse::default()); } @@ -281,9 +310,17 @@ impl RaftStateMachine for EngStateMachineStore { .map_err(|e| sm_io_err(ErrorVerb::Read, e.to_string()))?; // Clone store/graph handles under a short lock so we don't hold it across awaits. - let (short_term, core_memory, knowledge_graph, db) = { + let (short_term, core_memory, knowledge_graph, global_graph, visibility, session_agents, db) = { let inner = self.inner.lock().await; - (inner.short_term.clone(), inner.core_memory.clone(), inner.knowledge_graph.clone(), inner.db.clone()) + ( + inner.short_term.clone(), + inner.core_memory.clone(), + inner.knowledge_graph.clone(), + inner.global_graph.clone(), + inner.visibility.clone(), + inner.session_agents.clone(), + inner.db.clone(), + ) }; let st_sessions = payload.short_term.into_iter().map(|s| (s.session_id, s.messages)).collect(); @@ -292,6 +329,24 @@ impl RaftStateMachine for EngStateMachineStore { core_memory.restore_all(cm_sessions).await.map_err(|e| sm_io_err(ErrorVerb::Write, e.to_string()))?; *knowledge_graph.write().await = KnowledgeGraph::from_snapshot(payload.knowledge_graph); + if let Some(gg_snap) = payload.global_graph { + *global_graph.write().await = GlobalGraph::from_snapshot(gg_snap); + } + { + let mut vis = visibility.write().await; + vis.clear(); + for (session_id, v) in payload.visibility { + vis.insert(session_id, v); + } + } + { + let mut agents = session_agents.write().await; + agents.clear(); + for (session_id, agent_id) in payload.session_agents { + agents.insert(session_id, agent_id); + } + } + persist_snapshot(&db, meta, snapshot.get_ref())?; { @@ -310,6 +365,12 @@ impl RaftStateMachine for EngStateMachineStore { } } +fn update_global_metrics(metrics: &AppMetrics, graph: &GlobalGraph) { + metrics.set_global_entities(graph.all_entities().len()); + metrics.set_global_relationships(graph.all_relationships().len()); + metrics.set_global_conflicts(graph.conflicts().len()); +} + async fn apply_cmd( cmd: MemoryCommand, short_term: &Arc, @@ -317,6 +378,11 @@ async fn apply_cmd( embedding_tx: &mpsc::Sender, knowledge_graph: &Arc>, knowledge_tx: &mpsc::Sender, + global_graph: &Arc>, + visibility: &Arc>>, + session_agents: &Arc>>, + metrics: &Arc, + index: u64, ) { match cmd { MemoryCommand::AddMessage { session_id, message } => { @@ -357,9 +423,53 @@ async fn apply_cmd( // Signal embedding worker to delete from local LanceDB. let _ = embedding_tx.try_send(EmbeddingJob::DeleteSession { session_id: session_id.clone() }); knowledge_graph.write().await.delete_session(&session_id); + { + let mut gg = global_graph.write().await; + gg.prune_session(&session_id); + update_global_metrics(metrics, &*gg); + } + visibility.write().await.remove(&session_id); + session_agents.write().await.remove(&session_id); } MemoryCommand::AddKnowledge { session_id, message_id, entities, relationships } => { - knowledge_graph.write().await.apply_extraction(&session_id, &message_id, entities, relationships); + knowledge_graph.write().await.apply_extraction(&session_id, &message_id, entities.clone(), relationships.clone()); + if visibility.read().await.get(&session_id) == Some(&Visibility::Shared) { + let agent_id = session_agents.read().await.get(&session_id).cloned(); + let mut gg = global_graph.write().await; + gg.merge_with_agent(&session_id, agent_id.as_deref(), index, entities, relationships); + update_global_metrics(metrics, &*gg); + } + } + MemoryCommand::SetSessionVisibility { session_id, visibility: new_vis } => { + let prev = visibility.write().await.insert(session_id.clone(), new_vis); + match new_vis { + Visibility::Shared => { + // Back-merge existing session knowledge only if this is a + // Private->Shared transition (idempotent: re-inserting the + // same provenance key is a no-op in GlobalGraph). + let (entities, relationships) = { + let kg = knowledge_graph.read().await; + (kg.all_entities(&session_id), kg.all_relationships(&session_id)) + }; + let agent_id = session_agents.read().await.get(&session_id).cloned(); + let mut gg = global_graph.write().await; + gg.merge_with_agent(&session_id, agent_id.as_deref(), index, entities, relationships); + update_global_metrics(metrics, &*gg); + } + Visibility::Private => { + // Prune only when actually switching away from Shared. + if prev == Some(Visibility::Shared) { + let mut gg = global_graph.write().await; + gg.prune_session(&session_id); + update_global_metrics(metrics, &*gg); + } + } + } + } + MemoryCommand::RegisterSession { session_id, agent_id } => { + if let Some(agent) = agent_id { + session_agents.write().await.insert(session_id, agent); + } } MemoryCommand::NoOp => {} } @@ -369,6 +479,7 @@ async fn apply_cmd( mod tests { use super::*; use crate::core::{InMemoryCoreMemoryStore, InMemoryStore, InMemoryVectorStore}; + use crate::knowledge::global::{GlobalGraph, Visibility}; use crate::knowledge::graph::KnowledgeGraph; use crate::knowledge::types::{Entity, KnowledgeJob, Relationship}; use crate::raft::types::MessagePayload; @@ -383,6 +494,7 @@ mod tests { mpsc::Receiver, Arc>, Arc, + Arc>, tempfile::TempDir, ) { let short_term = Arc::new(InMemoryStore::default()); @@ -391,8 +503,10 @@ mod tests { let (embed_tx, embed_rx) = mpsc::channel(10); let (know_tx, know_rx) = mpsc::channel(10); let kg = Arc::new(RwLock::new(KnowledgeGraph::new())); + let gg = Arc::new(RwLock::new(GlobalGraph::new())); let dir = tempfile::tempdir().unwrap(); let db = Arc::new(redb::Database::create(dir.path().join("sm.redb")).unwrap()); + let metrics = Arc::new(crate::metrics::AppMetrics::new().unwrap()); let sm = EngStateMachineStore::new( short_term.clone(), core_memory.clone(), @@ -401,8 +515,10 @@ mod tests { kg.clone(), know_tx, db, + gg.clone(), + metrics, ); - (sm, short_term, embed_rx, know_rx, kg, core_memory, dir) + (sm, short_term, embed_rx, know_rx, kg, core_memory, gg, dir) } fn make_entry(index: u64, cmd: MemoryCommand) -> openraft::Entry { @@ -414,7 +530,7 @@ mod tests { #[tokio::test] async fn add_message_writes_to_short_term() { - let (mut sm, short_term, _embed, _know, _kg, _cm, _dir) = make_sm(); + let (mut sm, short_term, _embed, _know, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry( 0, MemoryCommand::AddMessage { @@ -436,7 +552,7 @@ mod tests { #[tokio::test] async fn add_message_enqueues_embedding_job() { - let (mut sm, _st, mut embed_rx, _know, _kg, _cm, _dir) = make_sm(); + let (mut sm, _st, mut embed_rx, _know, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry( 0, MemoryCommand::AddMessage { @@ -457,7 +573,7 @@ mod tests { #[tokio::test] async fn delete_session_clears_redis_and_enqueues_lancedb_delete() { - let (mut sm, short_term, mut embed_rx, _know, _kg, _cm, _dir) = make_sm(); + let (mut sm, short_term, mut embed_rx, _know, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![ make_entry( 0, @@ -484,7 +600,7 @@ mod tests { #[tokio::test] async fn noop_command_is_ignored() { - let (mut sm, short_term, _embed, _know, _kg, _cm, _dir) = make_sm(); + let (mut sm, short_term, _embed, _know, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::NoOp)]).await.unwrap(); let msgs = short_term.get_recent("any", 10).await.unwrap(); assert_eq!(msgs.len(), 0); @@ -492,7 +608,7 @@ mod tests { #[tokio::test] async fn add_message_enqueues_knowledge_job() { - let (mut sm, _st, _embed, mut know_rx, _kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _embed, mut know_rx, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::AddMessage { session_id: "s1".into(), message: MessagePayload { @@ -509,7 +625,7 @@ mod tests { #[tokio::test] async fn add_knowledge_updates_graph() { - let (mut sm, _st, _embed, _know, kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _embed, _know, kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { session_id: "s1".into(), message_id: "m1".into(), @@ -530,7 +646,7 @@ mod tests { #[tokio::test] async fn delete_session_clears_knowledge_graph() { - let (mut sm, _st, _embed, _know, kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _embed, _know, kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { session_id: "s1".into(), message_id: "m1".into(), entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], @@ -544,14 +660,14 @@ mod tests { #[tokio::test] async fn install_snapshot_sets_last_applied_to_meta_log_id() { - let (mut src, _st, _e, _k, _kg, _cm, _dir) = make_sm(); + let (mut src, _st, _e, _k, _kg, _cm, _gg, _dir) = make_sm(); src.apply(vec![make_entry(7, MemoryCommand::AddFact { session_id: "s1".into(), fact: "f".into(), })]).await.unwrap(); let mut builder = src.get_snapshot_builder().await; let snap = builder.build_snapshot().await.unwrap(); - let (mut dst, dst_st, _e2, _k2, dst_kg, _cm2, _dir2) = make_sm(); + let (mut dst, dst_st, _e2, _k2, dst_kg, _cm2, _gg2, _dir2) = make_sm(); let mut buf = dst.begin_receiving_snapshot().await.unwrap(); *buf = std::io::Cursor::new(snap.snapshot.get_ref().clone()); dst.install_snapshot(&snap.meta, buf).await.unwrap(); @@ -564,7 +680,7 @@ mod tests { #[tokio::test] async fn apply_build_install_reproduces_state() { - let (mut src, _st, _e, _k, _kg, src_cm, _dir) = make_sm(); + let (mut src, _st, _e, _k, _kg, src_cm, _gg, _dir) = make_sm(); src.apply(vec![ make_entry(0, MemoryCommand::AddKnowledge { session_id: "s1".into(), message_id: "m1".into(), @@ -581,7 +697,7 @@ mod tests { let mut builder = src.get_snapshot_builder().await; let snap = builder.build_snapshot().await.unwrap(); - let (mut dst, _st2, _e2, _k2, dst_kg, dst_cm, _dir2) = make_sm(); + let (mut dst, _st2, _e2, _k2, dst_kg, dst_cm, _gg2, _dir2) = make_sm(); let mut buf = dst.begin_receiving_snapshot().await.unwrap(); *buf = std::io::Cursor::new(snap.snapshot.get_ref().clone()); dst.install_snapshot(&snap.meta, buf).await.unwrap(); @@ -593,7 +709,7 @@ mod tests { #[tokio::test] async fn build_snapshot_meta_index_equals_last_applied() { - let (mut sm, _st, _e, _k, _kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _e, _k, _kg, _cm, _gg, _dir) = make_sm(); for i in 0..=4u64 { sm.apply(vec![make_entry(i, MemoryCommand::AddFact { session_id: "s1".into(), fact: format!("f{i}"), @@ -606,7 +722,7 @@ mod tests { #[tokio::test] async fn build_then_get_current_snapshot_returns_same_index() { - let (mut sm, _st, _e, _k, _kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _e, _k, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::AddFact { session_id: "s1".into(), fact: "f".into(), })]).await.unwrap(); @@ -618,7 +734,7 @@ mod tests { #[tokio::test] async fn snapshot_payload_contains_applied_state() { - let (mut sm, _st, _e, _k, _kg, _cm, _dir) = make_sm(); + let (mut sm, _st, _e, _k, _kg, _cm, _gg, _dir) = make_sm(); sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { session_id: "s1".into(), message_id: "m1".into(), entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], @@ -631,8 +747,115 @@ mod tests { let mut builder = sm.get_snapshot_builder().await; let snap = builder.build_snapshot().await.unwrap(); let payload = crate::raft::snapshot::EngramSnapshot::from_bytes(snap.snapshot.get_ref()).unwrap(); - assert_eq!(payload.version, 1); + assert_eq!(payload.version, 2); assert!(payload.knowledge_graph.sessions.iter().any(|s| s.session_id == "s1")); assert!(payload.core_memory.iter().any(|s| s.facts.contains(&"likes coffee".to_string()))); } + + #[tokio::test] + async fn shared_session_knowledge_reaches_global_graph() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + sm.apply(vec![make_entry(0, MemoryCommand::SetSessionVisibility { + session_id: "s1".into(), visibility: Visibility::Shared, + })]).await.unwrap(); + sm.apply(vec![make_entry(1, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + assert!(gg.read().await.all_entities().iter().any(|e| e.name == "Alice")); + } + + #[tokio::test] + async fn private_session_knowledge_stays_out_of_global_graph() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![Entity { name: "Secret".into(), entity_type: "Person".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + assert!(gg.read().await.all_entities().is_empty()); + } + + #[tokio::test] + async fn becoming_shared_backmerges_existing_session_knowledge() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + sm.apply(vec![make_entry(1, MemoryCommand::SetSessionVisibility { + session_id: "s1".into(), visibility: Visibility::Shared, + })]).await.unwrap(); + assert!(gg.read().await.all_entities().iter().any(|e| e.name == "Alice")); + } + + #[tokio::test] + async fn delete_session_prunes_global_contributions() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + for (idx, sid) in [(0u64, "s1"), (1, "s2")] { + sm.apply(vec![make_entry(idx, MemoryCommand::SetSessionVisibility { + session_id: sid.into(), visibility: Visibility::Shared, + })]).await.unwrap(); + } + sm.apply(vec![make_entry(2, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + sm.apply(vec![make_entry(3, MemoryCommand::AddKnowledge { + session_id: "s2".into(), message_id: "m2".into(), + entities: vec![Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + sm.apply(vec![make_entry(4, MemoryCommand::DeleteSession { session_id: "s1".into() })]).await.unwrap(); + assert!(gg.read().await.all_entities().iter().any(|e| e.name == "Alice"), "still contributed by s2"); + sm.apply(vec![make_entry(5, MemoryCommand::DeleteSession { session_id: "s2".into() })]).await.unwrap(); + assert!(gg.read().await.all_entities().is_empty()); + } + + #[tokio::test] + async fn registered_agent_id_flows_into_global_provenance() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + sm.apply(vec![make_entry(0, MemoryCommand::RegisterSession { + session_id: "s1".into(), agent_id: Some("agent-7".into()), + })]).await.unwrap(); + sm.apply(vec![make_entry(1, MemoryCommand::SetSessionVisibility { + session_id: "s1".into(), visibility: Visibility::Shared, + })]).await.unwrap(); + sm.apply(vec![make_entry(2, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![Entity { name: "OpenAI".into(), entity_type: "Organization".into(), attributes: HashMap::new() }], + relationships: vec![], + })]).await.unwrap(); + assert_eq!(gg.read().await.entity_agents("OpenAI"), vec!["agent-7".to_string()]); + } + + #[tokio::test] + async fn visibility_transitions_are_fully_reversible() { + let (mut sm, _st, _e, _k, _kg, _cm, gg, _dir) = make_sm(); + sm.apply(vec![make_entry(0, MemoryCommand::AddKnowledge { + session_id: "s1".into(), message_id: "m1".into(), + entities: vec![ + Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }, + Entity { name: "Bob".into(), entity_type: "Person".into(), attributes: HashMap::new() }, + ], + relationships: vec![Relationship { from: "Alice".into(), to: "Bob".into(), relationship_type: "knows".into() }], + })]).await.unwrap(); + + let vis = |idx: u64, v| make_entry(idx, MemoryCommand::SetSessionVisibility { session_id: "s1".into(), visibility: v }); + sm.apply(vec![vis(1, Visibility::Shared)]).await.unwrap(); + sm.apply(vec![vis(2, Visibility::Private)]).await.unwrap(); + sm.apply(vec![vis(3, Visibility::Shared)]).await.unwrap(); + + let g = gg.read().await; + assert_eq!(g.all_entities().iter().filter(|e| e.name == "Alice").count(), 1); + assert_eq!(g.all_relationships().len(), 1); + assert_eq!(g.entity_sources("Alice"), vec!["s1".to_string()]); + drop(g); + + sm.apply(vec![vis(4, Visibility::Private)]).await.unwrap(); + assert!(gg.read().await.all_entities().is_empty()); + } } diff --git a/src/raft/types.rs b/src/raft/types.rs index 5982b44..2a7e4ab 100644 --- a/src/raft/types.rs +++ b/src/raft/types.rs @@ -43,6 +43,14 @@ pub enum MemoryCommand { entities: Vec, relationships: Vec, }, + /// Set a session's visibility. Replicated so every node agrees deterministically + /// on which sessions contribute to the global graph. + SetSessionVisibility { + session_id: String, + visibility: crate::knowledge::global::Visibility, + }, + /// Record an agent owner for a session (provenance for the global graph). + RegisterSession { session_id: String, agent_id: Option }, /// No-op placeholder. Applied by the state machine without side effects. /// Reserved for future cluster operations (e.g., leadership probes). NoOp, @@ -126,4 +134,16 @@ mod tests { let back: MemoryCommand = serde_json::from_str(&json).unwrap(); assert!(matches!(back, MemoryCommand::NoOp)); } + + #[test] + fn set_session_visibility_command_round_trips() { + use crate::knowledge::global::Visibility; + let cmd = MemoryCommand::SetSessionVisibility { + session_id: "s1".into(), + visibility: Visibility::Shared, + }; + let json = serde_json::to_string(&cmd).unwrap(); + let back: MemoryCommand = serde_json::from_str(&json).unwrap(); + assert!(matches!(back, MemoryCommand::SetSessionVisibility { visibility: Visibility::Shared, .. })); + } } diff --git a/src/server.rs b/src/server.rs index ff9a009..5bd5844 100644 --- a/src/server.rs +++ b/src/server.rs @@ -10,6 +10,10 @@ use axum::{ response::IntoResponse, routing::{delete, get, post, put}, }; +use crate::knowledge::global_handler::{ + get_global, get_global_conflicts, get_global_entity, get_global_entity_sources, + get_global_export, get_global_path, set_visibility, +}; use crate::knowledge::handler::{export_knowledge, find_path, get_knowledge, get_related}; use axum_prometheus::{PrometheusMetricLayer, PrometheusMetricLayerBuilder}; use axum_prometheus::metrics_exporter_prometheus::PrometheusHandle; @@ -128,6 +132,8 @@ pub struct AppState { pub knowledge_graph: Arc>, /// Channel for sending knowledge extraction jobs to the worker pool. pub knowledge_job_sender: tokio::sync::mpsc::Sender, + /// Cluster-wide knowledge graph aggregating all Shared sessions. + pub global_graph: Arc>, } @@ -136,6 +142,11 @@ struct CreateSessionResponse { session_id: String, } +#[derive(Debug, Deserialize, Default, utoipa::ToSchema)] +struct CreateSessionRequest { + agent_id: Option, +} + #[derive(Debug, Deserialize, utoipa::ToSchema)] struct AddMessageRequest { id: Option, @@ -224,6 +235,13 @@ pub fn build_router(state: Arc) -> Router { .route("/sessions/{session_id}/knowledge/entities/{entity_name}", get(get_related)) .route("/sessions/{session_id}/knowledge/path", get(find_path)) .route("/sessions/{session_id}/knowledge/export", get(export_knowledge)) + .route("/sessions/{session_id}/visibility", put(set_visibility)) + .route("/knowledge/global", get(get_global)) + .route("/knowledge/global/entities/{name}", get(get_global_entity)) + .route("/knowledge/global/entities/{name}/sources", get(get_global_entity_sources)) + .route("/knowledge/global/path", get(get_global_path)) + .route("/knowledge/global/export", get(get_global_export)) + .route("/knowledge/global/conflicts", get(get_global_conflicts)) .route("/cluster", get(crate::cluster::get_cluster_status)) .route("/cluster/init", post(crate::cluster::init_cluster)) .route("/cluster/add-learner", post(crate::cluster::add_learner)) @@ -289,7 +307,7 @@ async fn health_check() -> StatusCode { StatusCode::OK } -#[tracing::instrument] +#[tracing::instrument(skip(state, body))] #[utoipa::path( post, path = "/sessions", @@ -298,11 +316,29 @@ async fn health_check() -> StatusCode { (status = 200, description = "Session created successfully", body = CreateSessionResponse) ) )] -async fn create_session() -> Json { +async fn create_session( + State(state): State>, + body: Option>, +) -> Result, MemoryServerError> { let session_id = Uuid::new_v4().to_string(); - tracing::info!(session_id = %session_id, "created session"); - Json(CreateSessionResponse { session_id }) + if let Some(agent_id) = body.and_then(|b| b.0.agent_id) { + if let Some(raft) = &state.raft { + raft_write( + raft, + MemoryCommand::RegisterSession { + session_id: session_id.clone(), + agent_id: Some(agent_id), + }, + &state.peer_http_addrs, + "/sessions", + ) + .await?; + } + } + + tracing::info!(session_id = %session_id, "created session"); + Ok(Json(CreateSessionResponse { session_id })) } #[tracing::instrument(skip(state, payload), fields(session_id = %session_id))] @@ -777,6 +813,9 @@ mod tests { tokio::spawn(async move { while rx.recv().await.is_some() {} }); tx }, + global_graph: Arc::new(tokio::sync::RwLock::new( + crate::knowledge::global::GlobalGraph::new(), + )), }) } @@ -1235,6 +1274,21 @@ mod tests { response.assert_status(StatusCode::BAD_REQUEST); } + #[tokio::test] + async fn put_visibility_route_exists_and_validates_body() { + let state = build_test_state(); + let server = TestServer::new(build_router(state)).unwrap(); + let resp = server + .put("/sessions/s1/visibility") + .json(&json!({ "visibility": "Shared" })) + .await; + assert!( + resp.status_code().is_success() || resp.status_code().as_u16() == 307, + "expected 2xx or 307 but got {}", + resp.status_code() + ); + } + #[tokio::test] async fn knowledge_routes_are_registered() { let server = TestServer::new(build_router(build_test_state())).unwrap(); @@ -1245,6 +1299,44 @@ mod tests { .assert_status_ok(); } + #[tokio::test] + async fn global_endpoints_return_aggregated_view() { + use crate::knowledge::types::{Entity, Relationship}; + use std::collections::HashMap; + + let state = build_test_state(); + { + let mut g = state.global_graph.write().await; + g.merge( + "s1", + 0, + vec![ + Entity { name: "Alice".into(), entity_type: "Person".into(), attributes: HashMap::new() }, + Entity { name: "OpenAI".into(), entity_type: "Organization".into(), attributes: HashMap::new() }, + ], + vec![Relationship { from: "Alice".into(), to: "OpenAI".into(), relationship_type: "works_at".into() }], + ); + } + let server = TestServer::new(build_router(state)).unwrap(); + + let all = server.get("/knowledge/global").await; + all.assert_status_ok(); + + let neighbors = server.get("/knowledge/global/entities/OpenAI").await; + neighbors.assert_status_ok(); + assert!(neighbors.text().contains("Alice")); + + let sources = server.get("/knowledge/global/entities/OpenAI/sources").await; + assert!(sources.text().contains("s1")); + + let path = server.get("/knowledge/global/path?from=Alice&to=OpenAI").await; + path.assert_status_ok(); + assert!(path.text().contains("works_at")); + + let dot = server.get("/knowledge/global/export?format=dot").await; + assert!(dot.text().contains("digraph")); + } + #[tokio::test] async fn knowledge_metrics_appear_in_prometheus_scrape() { let state = build_test_state(); diff --git a/tests/raft_write_test.rs b/tests/raft_write_test.rs index f53e164..231ae97 100644 --- a/tests/raft_write_test.rs +++ b/tests/raft_write_test.rs @@ -22,7 +22,9 @@ async fn single_node_raft_write_commits_to_state_machine() { ..Config::default() }; let knowledge_graph = Arc::new(tokio::sync::RwLock::new(engram::knowledge::graph::KnowledgeGraph::new())); + let global_graph = Arc::new(tokio::sync::RwLock::new(engram::knowledge::GlobalGraph::new())); let (knowledge_tx, _knowledge_rx) = mpsc::channel(500); + let metrics = Arc::new(engram::metrics::AppMetrics::new().unwrap()); let raft = build_raft_node( &config, short_term.clone(), @@ -31,6 +33,8 @@ async fn single_node_raft_write_commits_to_state_machine() { tx, knowledge_graph, knowledge_tx, + global_graph, + metrics, ) .await .unwrap(); diff --git a/tests/token_efficiency.rs b/tests/token_efficiency.rs index 1f734c4..cee0b97 100644 --- a/tests/token_efficiency.rs +++ b/tests/token_efficiency.rs @@ -80,6 +80,9 @@ fn build_test_state() -> Arc { engram::knowledge::graph::KnowledgeGraph::new(), )), knowledge_job_sender, + global_graph: Arc::new(tokio::sync::RwLock::new( + engram::knowledge::global::GlobalGraph::new(), + )), }) }