diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ddc088..9e42dee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -107,6 +107,23 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). claimed a "background cleanup worker" that did not exist; it now describes what runs. +- **Optional API-key authentication** (`amp_server.auth` + + `AMP_API_KEYS_FILE`). The spec carries agent identity in `X-AMP-Agent-ID` and + defines no credential (RFC §6.1), so the header was an assertion: anyone who + could reach the port could claim any agent id, and every access rule was decided + from that claim. With a key store configured, the id must now be proven with a + matching key in `X-AMP-API-Key`; the file stores `sha256:` digests rather than + keys, and a wrong key, a missing key and an unknown agent id all answer the same + `401 UNAUTHENTICATED` so the endpoints cannot enumerate agent ids. Off by + default, so the spec's binding keeps working unchanged for anyone who has not + opted in; a key store that cannot be read stops the server instead of silently + falling back to trusting the header. This also closes `POST /memories`'s + `identity.created_by` fallback whenever keys are configured - without that, the + fallback would have been an authentication bypass. Both SDKs take the key as a + constructor argument, and identity resolution now happens in one FastAPI + dependency rather than in each handler, so a route cannot resolve an agent + without proving it. + ### Changed - **Every endpoint returns one error shape.** `PATCH /memories/{id}` answered a conflict with `{"detail": ...}` while `DELETE` answered with diff --git a/docs/api-reference.md b/docs/api-reference.md index 039467a..c1406b2 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -4,7 +4,7 @@ **Protocol version:** `0.1.0` **Content-Type:** `application/json` -All endpoints accept and return JSON. Memory-cell access control is expressed via `access_policy` on each cell; the only endpoint with its own auth is `POST /lifecycle/run`, gated on an admin token. +All endpoints accept and return JSON. Memory-cell access control is expressed via `access_policy` on each cell. Identity travels in the `X-AMP-Agent-ID` header; see [Authentication](#authentication) for how that is proven when the server is run with API keys, and `POST /lifecycle/run` for the one endpoint with a separate admin token. --- @@ -25,6 +25,55 @@ All endpoints accept and return JSON. Memory-cell access control is expressed vi --- +## Authentication + +Three separate things, which are easy to confuse: + +**Agent identity.** `X-AMP-Agent-ID` names the agent making the request. The spec +defines identity here and defines no credential +([RFC-AMP-001 §6.1](https://github.com/glatinone/agent-memory-protocol/blob/master/spec/rfcs/RFC-AMP-001.md)), +so by default the header is taken at its word - every access rule downstream +(`readable_by`, `writable_by`, the owner check) is decided from it. A server run +this way is exactly as the spec describes, and is only safe on a network you +control. + +**API keys (optional).** Set `AMP_API_KEYS_FILE` to a JSON file mapping agent id +to a key digest, and the header must then be proven: + +```json +{ + "agent_assistant": "sha256:2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae" +} +``` + +```bash +python -m amp_server.auth hash 'the-agent-key' # prints the value to paste +export AMP_API_KEYS_FILE=/etc/amp/api-keys.json +``` + +The file holds digests, never keys, so a leaked file does not hand over working +credentials. With keys configured, a request must send the matching key in +`X-AMP-API-Key`; a missing key, a wrong key and a key for an agent that does not +exist all answer the same `401 UNAUTHENTICATED`, so the endpoints cannot be used +to enumerate agent ids. Two consequences worth stating plainly: + +- `POST /memories` no longer falls back to `identity.created_by` in the body. + That fallback is the documented default without keys, and with keys configured + it would let any caller create a cell as any agent. +- A store that cannot be read stops the server from starting. Falling back to + trusting the header would silently remove the protection an operator asked for. + +`GET /health` and `GET /spec` stay open: clients and probes read them before they +have a key. The SDKs accept the key as a constructor argument +(`AMPClient(url, agent_id, api_key=...)`, `new AMPClient(url, agentId, apiKey)`). + +**Admin token.** `POST /lifecycle/run` is gated on `AMP_ADMIN_TOKEN`, separately +from the above, because it mutates lifecycle state for every cell in storage and +erases data outright when `AMP_PURGE_RETENTION` is on. With no token configured +the endpoint answers `403 ADMIN_DISABLED` rather than being open. + +--- + ## Machine-readable contract The table above is prose. The contract itself is @@ -86,6 +135,7 @@ curl http://localhost:8765/amp/v1/spec "capabilities": { "mcp_compatible": false, "storage_backends": ["chroma"], + "api_keys_required": false, "embedding": {"provider": "chroma-default", "dimensions": 384}, "max_cell_size_bytes": 65536, "retention_days": 30, @@ -108,6 +158,7 @@ suite checks it against the server's own numbers rather than a fixed value: | Capability | What it commits the server to | |---|---| | `mcp_compatible` | Whether **this HTTP server** speaks MCP directly. It is `false`: the MCP integration ships as a separate stdio process (`amp-mcp`, see `examples/mcp-claude-desktop/`), not as an endpoint on this API. | +| `api_keys_required` | Whether this server requires `X-AMP-API-Key` (`AMP_API_KEYS_FILE` is set). Reported here so a client learns it needs a key before a call fails with `401`. | | `storage_backends` | The adapter actually wired in (`chroma` or `postgres`), selected with `AMP_STORAGE_BACKEND`; see [getting started](getting-started.md#6-choosing-a-storage-backend). | | `embedding` | Which provider turns text into vectors, and the width of the vectors it produces (`null` when the service decides per request). Configured with `AMP_EMBEDDING_PROVIDER`; see [getting started](getting-started.md#5-choosing-an-embedding-provider). | | `max_cell_size_bytes` | The largest serialized cell the server will accept. Enforced on create and on update; a larger cell is refused with `413 CELL_TOO_LARGE` before anything is written, and the number here is the number the check uses. | diff --git a/docs/getting-started.md b/docs/getting-started.md index 905892b..1e15887 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -227,3 +227,37 @@ order. `GET /spec` reports which one is in use under `storage_backends`. The embedding-provider rules in step 5 apply to either backend: switching provider invalidates the vectors already stored. + +--- + +## 7. Turning on API keys (optional) + +By default the server trusts the `X-AMP-Agent-ID` header, which is what the spec's +binding describes: the header names the agent, and every access rule is decided +from it. Anyone who can reach the port can therefore claim any agent id. + +To require proof, point the server at a key store: + +```bash +python -m amp_server.auth hash 'the-agent-key' # prints the value to paste +``` + +```json title="api-keys.json" +{ + "agent_assistant": "sha256:2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae" +} +``` + +```bash +export AMP_API_KEYS_FILE=/etc/amp/api-keys.json +``` + +Clients then send the key alongside the identity header: + +```python +client = AMPClient("http://localhost:8765", "agent_assistant", api_key="the-agent-key") +``` + +The file stores digests rather than keys, and a store that cannot be read stops +the server rather than falling back to trusting the header. Full detail, including +the two rules this changes, is in the [API reference](api-reference.md#authentication). diff --git a/sdk/README.md b/sdk/README.md index ed3ac41..e2ae240 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -26,6 +26,10 @@ Store and retrieve memories with the synchronous client in under 5 lines: from amp_client import AMPClient client = AMPClient("http://localhost:8765", agent_id="agent_assistant") + +# Only when the server runs with AMP_API_KEYS_FILE; without it the agent id is +# accepted on its own, which is the binding the spec describes. +client = AMPClient("http://localhost:8765", agent_id="agent_assistant", api_key="...") client.remember(content="User prefers email correspondence.", owner_id="user_123") memories = client.recall(query="communication preferences", owner_id="user_123") print(memories[0]["content"]["text"]) diff --git a/sdk/amp_client/async_client.py b/sdk/amp_client/async_client.py index 82070b2..c7a64c6 100644 --- a/sdk/amp_client/async_client.py +++ b/sdk/amp_client/async_client.py @@ -11,13 +11,26 @@ class AsyncAMPClient: """Asynchronous AMP Client using httpx.""" - def __init__(self, server_url: str, agent_id: str): + def __init__(self, server_url: str, agent_id: str, api_key: str | None = None): + """Initialize the async AMP client. + + `api_key` is the key belonging to `agent_id`, sent as `X-AMP-API-Key`. It + is only needed when the server was started with `AMP_API_KEYS_FILE`. + """ self.server_url = server_url.rstrip("/") if not self.server_url.endswith("/amp/v1"): self.server_url += "/amp/v1" self.agent_id = agent_id + self.api_key = api_key self._client: httpx.AsyncClient | None = None + def identity_headers(self) -> dict[str, str]: + """See AMPClient.identity_headers: one place, so no call path is missed.""" + headers = {"X-AMP-Agent-ID": self.agent_id} + if self.api_key: + headers["X-AMP-API-Key"] = self.api_key + return headers + async def __aenter__(self) -> AsyncAMPClient: self._client = httpx.AsyncClient() return self @@ -81,7 +94,7 @@ async def remember( resp = await client.post( f"{self.server_url}/memories", json=body, - headers={"X-AMP-Agent-ID": self.agent_id}, + headers=self.identity_headers(), ) except httpx.HTTPError as exc: raise AMPError(f"HTTP request failed: {exc}") from exc @@ -108,7 +121,7 @@ async def recall( resp = await client.post( f"{self.server_url}/memories/search", json=body, - headers={"X-AMP-Agent-ID": self.agent_id}, + headers=self.identity_headers(), ) except httpx.HTTPError as exc: raise AMPError(f"HTTP request failed: {exc}") from exc @@ -122,7 +135,7 @@ async def forget(self, memory_id: str) -> bool: try: resp = await client.delete( f"{self.server_url}/memories/{memory_id}", - headers={"X-AMP-Agent-ID": self.agent_id}, + headers=self.identity_headers(), ) except httpx.HTTPError as exc: raise AMPError(f"HTTP request failed: {exc}") from exc @@ -149,7 +162,7 @@ async def list_memories( resp = await client.get( f"{self.server_url}/memories", params=params, - headers={"X-AMP-Agent-ID": self.agent_id}, + headers=self.identity_headers(), ) except httpx.HTTPError as exc: raise AMPError(f"HTTP request failed: {exc}") from exc diff --git a/sdk/amp_client/client.py b/sdk/amp_client/client.py index b0f2e15..f5c0b9c 100644 --- a/sdk/amp_client/client.py +++ b/sdk/amp_client/client.py @@ -10,12 +10,18 @@ class AMPClient: """Synchronous client for the Agent Memory Protocol (AMP) server.""" - def __init__(self, server_url: str, agent_id: str) -> None: + def __init__( + self, server_url: str, agent_id: str, api_key: str | None = None + ) -> None: """Initialize the AMP client. Args: server_url: The base URL of the AMP server. agent_id: The ID of the agent using the client. + api_key: The key belonging to `agent_id`, sent as `X-AMP-API-Key`. + Only needed when the server was started with + `AMP_API_KEYS_FILE`; without that, the server accepts the agent + id on its own, which is the binding the spec describes. """ # Normalize server_url (strip trailing slash and append /amp/v1 if not present) normalized_url = server_url.rstrip("/") @@ -24,16 +30,27 @@ def __init__(self, server_url: str, agent_id: str) -> None: self.server_url = normalized_url self.agent_id = agent_id + self.api_key = api_key self.session = requests.Session() + def identity_headers(self) -> dict[str, str]: + """The headers that identify this client, built in one place. + + One source, for the same reason the server resolves identity in one + dependency: a credential added to one call path and forgotten on another + fails as a confusing 401 rather than as a code error. + """ + headers = {"X-AMP-Agent-ID": self.agent_id} + if self.api_key: + headers["X-AMP-API-Key"] = self.api_key + return headers + def _request(self, method: str, path: str, **kwargs: Any) -> requests.Response: """Internal helper to execute HTTP requests with error handling.""" url = f"{self.server_url}{path}" - # Ensure standard headers are present - headers = kwargs.pop("headers", {}) - if "X-AMP-Agent-ID" not in headers: - headers["X-AMP-Agent-ID"] = self.agent_id + # Identity first, so a caller-supplied header can still override it. + headers = {**self.identity_headers(), **kwargs.pop("headers", {})} try: response = self.session.request(method, url, headers=headers, **kwargs) diff --git a/sdk/node/README.md b/sdk/node/README.md index 17f8549..5ae4250 100644 --- a/sdk/node/README.md +++ b/sdk/node/README.md @@ -36,10 +36,13 @@ All methods are async. `serverUrl` is normalized the same way as the Python client: a bare host, a host with a trailing slash, or an explicit `/amp/v1` all resolve to the same endpoint prefix. -### `new AMPClient(serverUrl, agentId)` +### `new AMPClient(serverUrl, agentId, apiKey?)` - `serverUrl` - base URL of the AMP server. - `agentId` - this agent's identifier, sent as the `X-AMP-Agent-ID` header. +- `apiKey` - optional; this agent's key, sent as `X-AMP-API-Key`. Only needed when + the server is run with `AMP_API_KEYS_FILE`, in which case a request without it is + answered `401 UNAUTHENTICATED`. ### `remember(content, ownerId, options?)` diff --git a/sdk/node/src/client.js b/sdk/node/src/client.js index 2ba7137..d2db131 100644 --- a/sdk/node/src/client.js +++ b/sdk/node/src/client.js @@ -37,8 +37,11 @@ export class AMPClient { /** * @param {string} serverUrl Base URL of the AMP server. * @param {string} agentId Identifier of this agent, sent as `X-AMP-Agent-ID`. + * @param {string} [apiKey] Key belonging to `agentId`, sent as + * `X-AMP-API-Key`. Only needed when the server is run with + * `AMP_API_KEYS_FILE`; otherwise the agent id alone is accepted. */ - constructor(serverUrl, agentId) { + constructor(serverUrl, agentId, apiKey) { if (!serverUrl) throw new AMPError("serverUrl is required"); if (!agentId) throw new AMPError("agentId is required"); @@ -50,6 +53,20 @@ export class AMPClient { ? normalized : `${normalized}/amp/v1`; this.agentId = agentId; + this.apiKey = apiKey; + } + + /** + * The headers that identify this client, built in one place so no call path + * can be missing the credential. + * + * @returns {Record} + * @private + */ + _identityHeaders() { + const headers = { "X-AMP-Agent-ID": this.agentId }; + if (this.apiKey) headers["X-AMP-API-Key"] = this.apiKey; + return headers; } /** @@ -66,7 +83,7 @@ export class AMPClient { async _request(method, path, { body, headers = {} } = {}) { const url = new URL(`${this.serverUrl}${path}`); - const finalHeaders = { "X-AMP-Agent-ID": this.agentId, ...headers }; + const finalHeaders = { ...this._identityHeaders(), ...headers }; let payload; if (body !== undefined) { finalHeaders["Content-Type"] = "application/json"; diff --git a/sdk/node/test/client.test.js b/sdk/node/test/client.test.js index 01bcd97..df2525f 100644 --- a/sdk/node/test/client.test.js +++ b/sdk/node/test/client.test.js @@ -23,6 +23,26 @@ const OWNER = `node-sdk-test-user-${Date.now()}`; // being collected, before any hook runs, and would always skip. let serverUp = false; +describe("API keys", () => { + test("omits the key header unless one was given", () => { + const client = new AMPClient("http://localhost:8000", "agent-1"); + assert.deepEqual(client._identityHeaders(), { "X-AMP-Agent-ID": "agent-1" }); + }); + + test("carries the key header when one was given", () => { + const client = new AMPClient("http://localhost:8000", "agent-1", "agent-one-key"); + assert.deepEqual(client._identityHeaders(), { + "X-AMP-Agent-ID": "agent-1", + "X-AMP-API-Key": "agent-one-key", + }); + }); + + test("lets a caller-supplied header win", () => { + const client = new AMPClient("http://localhost:8000", "agent-1"); + assert.deepEqual(client._identityHeaders(), { "X-AMP-Agent-ID": "agent-1" }); + }); +}); + describe("URL normalization", () => { test("appends /amp/v1 when missing", () => { const client = new AMPClient("http://localhost:8000", "a"); diff --git a/sdk/python/tests/test_async_client.py b/sdk/python/tests/test_async_client.py index ae43644..49aceb5 100644 --- a/sdk/python/tests/test_async_client.py +++ b/sdk/python/tests/test_async_client.py @@ -121,3 +121,22 @@ async def test_async_health_ok(mock_get): async with AsyncAMPClient("http://localhost:8000", "test_agent") as client: assert await client.health() is True + + +# --------------------------------------------------------------------------- +# API keys (X-AMP-API-Key) +# --------------------------------------------------------------------------- + + +def test_async_identity_headers_omit_the_key_unless_one_was_given(): + client = AsyncAMPClient("http://localhost:8000", "agent-1") + assert client.identity_headers() == {"X-AMP-Agent-ID": "agent-1"} + + +def test_async_identity_headers_carry_the_key_when_one_was_given(): + """Every async call site uses this builder, so one test covers them all.""" + client = AsyncAMPClient("http://localhost:8000", "agent-1", api_key="agent-one-key") + assert client.identity_headers() == { + "X-AMP-Agent-ID": "agent-1", + "X-AMP-API-Key": "agent-one-key", + } diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py index d3826b6..e89d2bd 100644 --- a/sdk/python/tests/test_client.py +++ b/sdk/python/tests/test_client.py @@ -186,3 +186,40 @@ def test_health_error(mock_request): client = AMPClient("http://localhost:8000", "test_agent") assert client.health() is False + + +# --------------------------------------------------------------------------- +# API keys (X-AMP-API-Key) +# --------------------------------------------------------------------------- + + +def test_identity_headers_omit_the_key_unless_one_was_given(): + """No key means no key header: the server trusts the agent id by default.""" + assert AMPClient("http://localhost:8000", "agent-1").identity_headers() == { + "X-AMP-Agent-ID": "agent-1" + } + assert AMPClient("http://localhost:8000", "agent-1").api_key is None + + +def test_identity_headers_carry_the_key_when_one_was_given(): + client = AMPClient("http://localhost:8000", "agent-1", api_key="agent-one-key") + assert client.identity_headers() == { + "X-AMP-Agent-ID": "agent-1", + "X-AMP-API-Key": "agent-one-key", + } + + +@patch("requests.Session.request") +def test_the_key_is_sent_on_requests(mock_request): + """The header has to reach the wire, not just the constructor.""" + mock_response = MagicMock() + mock_response.status_code = 201 + mock_response.json.return_value = {"id": "mem_123"} + mock_request.return_value = mock_response + + client = AMPClient("http://localhost:8000", "agent-1", api_key="agent-one-key") + client.remember(content="a fact", owner_id="user_abc") + + sent = mock_request.call_args.kwargs["headers"] + assert sent["X-AMP-API-Key"] == "agent-one-key" + assert sent["X-AMP-Agent-ID"] == "agent-1" diff --git a/server/amp_server/auth.py b/server/amp_server/auth.py new file mode 100644 index 0000000..897ff5f --- /dev/null +++ b/server/amp_server/auth.py @@ -0,0 +1,149 @@ +"""Optional API-key authentication for the HTTP binding. + +`spec/rfcs/RFC-AMP-001.md` §6.1 carries agent identity in the `X-AMP-Agent-ID` +header and defines no credential, so today a client can claim to be any agent it +likes: the header is an assertion, not proof. Every access rule downstream - +`readable_by`, `writable_by`, the owner check - is decided from that assertion. + +This module adds proof without changing the binding. The agent id keeps +travelling in the header exactly where the spec puts it; when a key store is +configured, the caller must additionally present a key that belongs to that id. + +It is off unless `AMP_API_KEYS_FILE` is set. A reference server must stay usable +with no configuration, and the spec's binding must keep working as written - but +a deployment that has configured keys and then finds them unreadable must not +quietly fall back to trusting the header, so a broken store stops the server. + +Two deliberate choices: + +- **Hashes, not keys.** The file stores `sha256:` digests, so a leaked file does + not hand over working credentials. `python -m amp_server.auth hash ` + prints the line to paste. +- **One 401 for every failure.** A wrong key and a key for an agent that does not + exist answer identically, so the endpoint cannot be used to enumerate agent + ids. The same reasoning as the uniform 403 on deleted cells in spec §8.4. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import logging +import os +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + +logger = logging.getLogger(__name__) + +SHA256_PREFIX = "sha256:" +ENV_VAR = "AMP_API_KEYS_FILE" + + +def digest(api_key: str) -> str: + """The stored form of a key. Never the key itself.""" + return SHA256_PREFIX + hashlib.sha256(api_key.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class ApiKeyStore: + """Which key belongs to which agent id. + + One key per agent id rather than several: the store answers "is this key + this agent's?", and a list of keys per agent would only widen what a leaked + file is worth. + """ + + digests: Mapping[str, str] + + def verify(self, agent_id: str, presented_key: str) -> bool: + """True when `presented_key` is the key registered for `agent_id`. + + Compared with `hmac.compare_digest` so the check does not reveal how much + of a wrong key was right. An unknown agent id is compared against a + throwaway digest rather than returning early, so a wrong-key attempt + costs the same whether or not the agent exists. + """ + expected = self.digests.get(agent_id, digest("no-such-agent")) + return hmac.compare_digest(expected, digest(presented_key)) + + def __len__(self) -> int: + return len(self.digests) + + +def load_store(path: str | Path) -> ApiKeyStore: + """Read a key store, refusing anything that is not a clean mapping. + + Raises ValueError naming the file and the offending entry. A store that is + half-readable is worse than none: the entries that parsed would be enforced + and the others silently would not. + """ + store_path = Path(path) + try: + raw = json.loads(store_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise ValueError( + f"{ENV_VAR} points at {store_path}, which does not exist" + ) from exc + except json.JSONDecodeError as exc: + raise ValueError(f"{store_path} is not valid JSON: {exc}") from exc + + if not isinstance(raw, dict): + raise ValueError( + f"{store_path} must be a JSON object mapping agent id to key digest" + ) + + digests: dict[str, str] = {} + for agent_id, value in raw.items(): + if not isinstance(agent_id, str) or not agent_id: + raise ValueError(f"{store_path} has an empty or non-string agent id") + if not isinstance(value, str) or not value.startswith(SHA256_PREFIX): + raise ValueError( + f"{store_path}: entry {agent_id!r} must be a " + f"'{SHA256_PREFIX}' digest, not a raw key" + ) + if len(value) != len(SHA256_PREFIX) + 64: + raise ValueError(f"{store_path}: entry {agent_id!r} is not a sha256 digest") + digests[agent_id] = value + + if not digests: + raise ValueError( + f"{store_path} contains no agents, which would authenticate nobody; " + f"unset {ENV_VAR} to run without keys" + ) + return ApiKeyStore(digests=digests) + + +def store_from_env() -> ApiKeyStore | None: + """The configured store, or None when this deployment has no keys. + + Called from the lifespan, so a malformed file stops startup instead of + surfacing as a 401 on a request an operator expected to work. + """ + path = os.environ.get(ENV_VAR) + if not path: + return None + store = load_store(path) + logger.info("API key authentication enabled for %d agent(s)", len(store)) + return store + + +def _main(argv: list[str]) -> int: + """`python -m amp_server.auth hash `: the value to put in the store. + + Prints the digest alone, because the agent id is the key in the store file + and only the operator knows which id the key belongs to. + """ + if len(argv) != 2 or argv[0] != "hash": + print("usage: python -m amp_server.auth hash ") + print("writes the sha256 value to use in the file named by AMP_API_KEYS_FILE") + return 2 + print(digest(argv[1])) + return 0 + + +if __name__ == "__main__": # pragma: no cover - operator convenience + import sys + + raise SystemExit(_main(sys.argv[1:])) diff --git a/server/amp_server/errors.py b/server/amp_server/errors.py index 0fbde31..8ccbee5 100644 --- a/server/amp_server/errors.py +++ b/server/amp_server/errors.py @@ -57,6 +57,20 @@ def access_denied() -> AMPError: return AMPError(403, "ACCESS_DENIED", "Access denied") +def unauthenticated() -> AMPError: + """No key, or a key that does not belong to the claimed agent id. + + One response for both, so the endpoint cannot be used to find out which agent + ids exist. Raised only when a key store is configured; without one the header + is trusted, as the spec's binding describes. + """ + return AMPError( + 401, + "UNAUTHENTICATED", + "X-AMP-API-Key is missing or not valid for this agent id", + ) + + def admin_disabled() -> AMPError: """Manual lifecycle runs are opt-in: no token configured means no access.""" return AMPError( diff --git a/server/amp_server/main.py b/server/amp_server/main.py index 1499bb9..f839546 100644 --- a/server/amp_server/main.py +++ b/server/amp_server/main.py @@ -9,10 +9,11 @@ from datetime import UTC, datetime from typing import Any -from fastapi import APIRouter, FastAPI, Header, Request +from fastapi import APIRouter, Depends, FastAPI, Header, Request from fastapi.responses import JSONResponse, Response from amp_server.access_control import check_read_access, check_write_access +from amp_server.auth import ApiKeyStore, store_from_env from amp_server.embeddings import provider_from_env from amp_server.errors import ( AMPError, @@ -20,6 +21,7 @@ admin_disabled, invalid_transition, missing_agent_id, + unauthenticated, ) from amp_server.lifecycle import LifecycleEngine from amp_server.limits import MAX_CELL_SIZE_BYTES @@ -68,6 +70,10 @@ def configure_logging() -> None: _storage: StorageAdapter | None = None _lifecycle: LifecycleEngine | None = None _lifecycle_settings = settings_from_env() +# None means this deployment runs without API keys, and the identity header is +# trusted as the spec's binding describes. Set from the lifespan, so a key store +# that cannot be read stops the server rather than turning into a 401 later. +_api_key_store: ApiKeyStore | None = None def _build_storage() -> StorageAdapter: @@ -108,6 +114,38 @@ def get_storage() -> StorageAdapter: return _storage +def get_api_key_store() -> ApiKeyStore | None: + """The configured key store, or None when this deployment has no keys.""" + return _api_key_store + + +async def verified_agent_id( + x_amp_agent_id: str | None = Header(default=None), + x_amp_api_key: str | None = Header(default=None), +) -> str | None: + """The caller's agent id, proven against the key store when one is configured. + + The single gate for every route that acts as an agent: identity resolution + and its proof live here rather than in six handlers, so a new route cannot + get one without the other. Returns None when no header was sent - the create + route falls back to the body's `created_by`, and every other route raises + `MISSING_AGENT_ID` itself. + """ + store = get_api_key_store() + if x_amp_agent_id is None: + # With keys configured, an unproven identity must not slip in through the + # body either: `create_memory` falls back to `identity.created_by`, which + # would let any caller create a cell as any agent. + if store is not None: + raise unauthenticated() + return None + if store is None: + return x_amp_agent_id + if x_amp_api_key is None or not store.verify(x_amp_agent_id, x_amp_api_key): + raise unauthenticated() + return x_amp_agent_id + + def get_lifecycle() -> LifecycleEngine: assert _lifecycle is not None, ( "LifecycleEngine not initialized — lifespan not started" @@ -122,12 +160,15 @@ def get_lifecycle() -> LifecycleEngine: @asynccontextmanager async def lifespan(app: FastAPI): - global _storage, _lifecycle, _lifecycle_settings + global _storage, _lifecycle, _lifecycle_settings, _api_key_store configure_logging() # Built here, not at import: a misconfigured backend must stop the server # from starting rather than surface as poor search results later. _storage = _build_storage() _lifecycle = LifecycleEngine(_storage) + # Same reason: a key store that cannot be read must stop the server, not + # fall back to trusting an unverified header. + _api_key_store = store_from_env() # Read settings here, not only at import: a test (or an embedder) can set # AMP_LIFECYCLE_* before starting the app and expect it to take effect. @@ -198,6 +239,13 @@ async def _amp_error_handler(request: Request, exc: AMPError) -> JSONResponse: "description": "X-AMP-Agent-ID header is required", } } +_UNAUTHENTICATED: dict[int | str, dict[str, Any]] = { + 401: { + "model": ErrorResponse, + "description": "X-AMP-API-Key is missing or not valid for this agent id " + "(only when AMP_API_KEYS_FILE is configured)", + } +} _ACCESS_DENIED: dict[int | str, dict[str, Any]] = { 403: { "model": ErrorResponse, @@ -235,6 +283,9 @@ async def spec() -> dict[str, Any]: "capabilities": { "mcp_compatible": False, "storage_backends": [get_storage().name], + # Open endpoint on purpose: a client can find out it needs a key + # before it makes a call that would fail with 401. + "api_keys_required": _api_key_store is not None, "embedding": get_storage().embedding, "max_cell_size_bytes": MAX_CELL_SIZE_BYTES, "retention_days": get_storage().retention_days, @@ -281,11 +332,11 @@ async def run_lifecycle_now( @router.post( "/memories", status_code=201, - responses=_MISSING_AGENT_ID | _CELL_TOO_LARGE, + responses=_UNAUTHENTICATED | _MISSING_AGENT_ID | _CELL_TOO_LARGE, ) async def create_memory( body: MemoryCellCreate, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> dict[str, Any]: # Determine effective agent_id: header takes priority, fallback to created_by agent_id = x_amp_agent_id or body.identity.created_by @@ -317,10 +368,13 @@ async def create_memory( return cell.model_dump(mode="json") -@router.get("/memories/{memory_id}", responses=_MISSING_AGENT_ID | _ACCESS_DENIED) +@router.get( + "/memories/{memory_id}", + responses=(_UNAUTHENTICATED | _MISSING_AGENT_ID | _ACCESS_DENIED), +) async def get_memory( memory_id: str, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> dict[str, Any]: if not x_amp_agent_id: raise missing_agent_id() @@ -347,13 +401,17 @@ async def get_memory( @router.patch( "/memories/{memory_id}", responses=( - _MISSING_AGENT_ID | _ACCESS_DENIED | _INVALID_TRANSITION | _CELL_TOO_LARGE + _UNAUTHENTICATED + | _MISSING_AGENT_ID + | _ACCESS_DENIED + | _INVALID_TRANSITION + | _CELL_TOO_LARGE ), ) async def update_memory( memory_id: str, body: MemoryCellUpdate, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> dict[str, Any]: if not x_amp_agent_id: raise missing_agent_id() @@ -381,11 +439,14 @@ async def update_memory( @router.delete( "/memories/{memory_id}", status_code=204, - responses=_MISSING_AGENT_ID | _ACCESS_DENIED | _INVALID_TRANSITION, + responses=_UNAUTHENTICATED + | _MISSING_AGENT_ID + | _ACCESS_DENIED + | _INVALID_TRANSITION, ) async def delete_memory( memory_id: str, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> Response: if not x_amp_agent_id: raise missing_agent_id() @@ -412,14 +473,14 @@ async def delete_memory( return Response(status_code=204) -@router.get("/memories", responses=_MISSING_AGENT_ID) -@router.get("/memories/query", responses=_MISSING_AGENT_ID) +@router.get("/memories", responses=_UNAUTHENTICATED | _MISSING_AGENT_ID) +@router.get("/memories/query", responses=_UNAUTHENTICATED | _MISSING_AGENT_ID) async def query_memories( owner_id: str | None = None, type: MemoryType | None = None, status: LifecycleStatus | None = None, limit: int = 20, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> dict[str, Any]: if not x_amp_agent_id: raise missing_agent_id() @@ -447,10 +508,10 @@ async def query_memories( # --------------------------------------------------------------------------- -@router.post("/memories/search", responses=_MISSING_AGENT_ID) +@router.post("/memories/search", responses=_UNAUTHENTICATED | _MISSING_AGENT_ID) async def search_memories( body: SearchRequest, - x_amp_agent_id: str | None = Header(default=None), + x_amp_agent_id: str | None = Depends(verified_agent_id), ) -> dict[str, Any]: if not x_amp_agent_id: raise missing_agent_id() diff --git a/server/tests/test_auth.py b/server/tests/test_auth.py new file mode 100644 index 0000000..d98f296 --- /dev/null +++ b/server/tests/test_auth.py @@ -0,0 +1,332 @@ +"""API keys: proof of identity, layered on the spec's identity header. + +RFC-AMP-001 §6.1 carries agent identity in the `X-AMP-Agent-ID` header and +defines no credential, so the header is an assertion rather than proof. These +tests pin both halves of the arrangement: + +- with no key store configured, the header is trusted exactly as the spec + describes, so the binding does not change for anyone who has not opted in; +- with a store configured, a caller must prove it owns the id it claims - and an + id supplied in a request body does not count, because that would be the same + unproven assertion wearing a different hat. +""" + +from __future__ import annotations + +import json +import uuid + +import pytest +from conftest import make_cell +from httpx import ASGITransport, AsyncClient + +from amp_server.auth import ApiKeyStore, digest, load_store, store_from_env + +_KEY = "agent-one-secret-key" +_OTHER_KEY = "agent-two-secret-key" +_HEADERS = {"X-AMP-Agent-ID": "agent-one"} + + +def _install_app(store: ApiKeyStore | None = None) -> None: + """Point the app at a fresh storage + engine, as the lifespan would.""" + import amp_server.main as main_mod + from amp_server.lifecycle import LifecycleEngine + from amp_server.storage.chroma import ChromaAdapter + + main_mod._storage = ChromaAdapter(collection_name=f"test_{uuid.uuid4().hex[:12]}") + main_mod._lifecycle = LifecycleEngine(main_mod._storage) + main_mod._api_key_store = store + + +def _store(monkeypatch, tmp_path, keys: dict[str, str]) -> ApiKeyStore: + """Write a key store file and install it on the app, as the lifespan would.""" + path = tmp_path / "api-keys.json" + path.write_text( + json.dumps({agent_id: digest(key) for agent_id, key in keys.items()}), + encoding="utf-8", + ) + monkeypatch.setenv("AMP_API_KEYS_FILE", str(path)) + + store = store_from_env() + assert store is not None + _install_app(store) + return store + + +@pytest.fixture(autouse=True) +def _no_keys_by_default(): + """Every test starts from "this deployment has no keys". + + Installed on the module, not read from the environment, so a stray + AMP_API_KEYS_FILE in a developer's shell cannot change what these assert. + """ + import amp_server.main as main_mod + + main_mod._api_key_store = None + yield + main_mod._api_key_store = None + + +async def _client() -> AsyncClient: + import amp_server.main as main_mod + + return AsyncClient( + transport=ASGITransport(app=main_mod.app), base_url="http://test" + ) + + +def _body() -> dict: + return { + "type": "semantic", + "content": {"text": "authenticated memory"}, + "identity": {"owner_id": "user-auth", "owner_type": "user"}, + } + + +# --------------------------------------------------------------------------- +# The store itself +# --------------------------------------------------------------------------- + + +def test_a_digest_is_sha256_and_never_the_key(): + value = digest(_KEY) + + assert value.startswith("sha256:") + assert len(value) == len("sha256:") + 64 + assert _KEY not in value + + +def test_verify_accepts_the_registered_key_only(): + store = ApiKeyStore(digests={"agent-one": digest(_KEY)}) + + assert store.verify("agent-one", _KEY) is True + assert store.verify("agent-one", _OTHER_KEY) is False + assert store.verify("agent-one", "") is False + + +def test_verify_rejects_an_unknown_agent_without_raising(): + """A key for an agent that does not exist is simply not valid.""" + store = ApiKeyStore(digests={"agent-one": digest(_KEY)}) + + assert store.verify("agent-nobody", _OTHER_KEY) is False + assert store.verify("agent-nobody", "") is False + + +def test_loading_reports_the_file_it_cannot_read(tmp_path): + missing = tmp_path / "nope.json" + + with pytest.raises(ValueError, match="does not exist"): + load_store(missing) + + +def test_loading_refuses_a_raw_key_in_the_store(tmp_path): + """The file holds digests; a pasted key would be a leak waiting to happen.""" + path = tmp_path / "api-keys.json" + path.write_text(json.dumps({"agent-one": _KEY}), encoding="utf-8") + + with pytest.raises(ValueError, match="digest, not a raw key"): + load_store(path) + + +def test_loading_refuses_a_malformed_store(tmp_path): + """Half a store is worse than none: the parsed entries would be enforced.""" + path = tmp_path / "api-keys.json" + + path.write_text("{not json", encoding="utf-8") + with pytest.raises(ValueError, match="not valid JSON"): + load_store(path) + + path.write_text(json.dumps(["agent-one"]), encoding="utf-8") + with pytest.raises(ValueError, match="must be a JSON object"): + load_store(path) + + path.write_text(json.dumps({"agent-one": "sha256:short"}), encoding="utf-8") + with pytest.raises(ValueError, match="not a sha256 digest"): + load_store(path) + + path.write_text(json.dumps({}), encoding="utf-8") + with pytest.raises(ValueError, match="contains no agents"): + load_store(path) + + +def test_no_file_means_no_keys(monkeypatch): + monkeypatch.delenv("AMP_API_KEYS_FILE", raising=False) + assert store_from_env() is None + + +# --------------------------------------------------------------------------- +# Without a store: the spec's binding, unchanged +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_header_alone_still_works_without_a_store(): + _install_app() + + async with await _client() as client: + created = await client.post("/amp/v1/memories", json=_body(), headers=_HEADERS) + + assert created.status_code == 201 + + +@pytest.mark.asyncio +async def test_without_a_store_a_body_supplied_identity_is_still_accepted(): + """The `created_by` fallback is the documented default, not a hole. + + It is only reachable when no key store is configured; the test below pins the + other half. + """ + _install_app() + body = _body() + body["identity"]["created_by"] = "agent-from-body" + + async with await _client() as client: + created = await client.post("/amp/v1/memories", json=body) + + assert created.status_code == 201 + assert created.json()["identity"]["created_by"] == "agent-from-body" + + +# --------------------------------------------------------------------------- +# With a store: proof required +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_request_without_a_key_is_refused(monkeypatch, tmp_path): + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + + async with await _client() as client: + response = await client.post("/amp/v1/memories", json=_body(), headers=_HEADERS) + + assert response.status_code == 401 + assert response.json()["error"]["code"] == "UNAUTHENTICATED" + + +@pytest.mark.asyncio +async def test_a_key_for_another_agent_is_refused(monkeypatch, tmp_path): + _store(monkeypatch, tmp_path, {"agent-one": _KEY, "agent-two": _OTHER_KEY}) + + async with await _client() as client: + response = await client.post( + "/amp/v1/memories", + json=_body(), + headers={**_HEADERS, "X-AMP-API-Key": _OTHER_KEY}, + ) + + assert response.status_code == 401 + + +@pytest.mark.asyncio +async def test_an_unknown_agent_answers_like_a_wrong_key(monkeypatch, tmp_path): + """Same response either way, so the route cannot enumerate agent ids.""" + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + + async with await _client() as client: + unknown = await client.post( + "/amp/v1/memories", + json=_body(), + headers={"X-AMP-Agent-ID": "agent-nobody", "X-AMP-API-Key": _KEY}, + ) + wrong = await client.post( + "/amp/v1/memories", + json=_body(), + headers={**_HEADERS, "X-AMP-API-Key": "not-the-key"}, + ) + + assert unknown.status_code == wrong.status_code == 401 + assert unknown.json() == wrong.json() + + +@pytest.mark.asyncio +async def test_the_right_key_gets_through(monkeypatch, tmp_path): + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + + async with await _client() as client: + created = await client.post( + "/amp/v1/memories", + json=_body(), + headers={**_HEADERS, "X-AMP-API-Key": _KEY}, + ) + fetched = await client.get( + f"/amp/v1/memories/{created.json()['id']}", + headers={**_HEADERS, "X-AMP-API-Key": _KEY}, + ) + + assert created.status_code == 201 + assert fetched.status_code == 200 + + +@pytest.mark.asyncio +async def test_a_body_supplied_identity_cannot_bypass_a_configured_store( + monkeypatch, tmp_path +): + """The fallback would otherwise let any caller create as any agent.""" + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + body = _body() + body["identity"]["created_by"] = "agent-impersonated" + + async with await _client() as client: + response = await client.post("/amp/v1/memories", json=body) + + assert response.status_code == 401 + assert response.json()["error"]["code"] == "UNAUTHENTICATED" + + +@pytest.mark.asyncio +async def test_health_and_spec_stay_open(monkeypatch, tmp_path): + """Operators and clients read these before they have a key.""" + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + + async with await _client() as client: + health = await client.get("/amp/v1/health") + spec = await client.get("/amp/v1/spec") + + assert health.status_code == 200 + assert spec.status_code == 200 + + +@pytest.mark.asyncio +async def test_every_agent_route_is_behind_the_same_gate(monkeypatch, tmp_path): + """A route that resolves an agent id without proving it would be the bug. + + Written as a sweep rather than one case per route so a route added later + without the dependency fails here instead of shipping unverified. + """ + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + cell = make_cell(created_by="agent-one", text="guarded") + import amp_server.main as main_mod + + await main_mod.get_storage().save(cell) + + calls = [ + ("post", "/amp/v1/memories", {"json": _body()}), + ("get", f"/amp/v1/memories/{cell.id}", {}), + ("patch", f"/amp/v1/memories/{cell.id}", {"json": {"content": {"text": "x"}}}), + ("delete", f"/amp/v1/memories/{cell.id}", {}), + ("get", "/amp/v1/memories?limit=5", {}), + ("get", "/amp/v1/memories/query?limit=5", {}), + ("post", "/amp/v1/memories/search", {"json": {"query": "guarded"}}), + ] + + async with await _client() as client: + for method, url, kwargs in calls: + response = await getattr(client, method)(url, headers=_HEADERS, **kwargs) + assert response.status_code == 401, ( + f"{method.upper()} {url} answered {response.status_code} without a key" + ) + + +@pytest.mark.asyncio +async def test_spec_says_whether_a_key_is_needed(monkeypatch, tmp_path): + """A client should learn it needs a key before a call fails with 401.""" + _install_app() + async with await _client() as client: + without = (await client.get("/amp/v1/spec")).json() + + _store(monkeypatch, tmp_path, {"agent-one": _KEY}) + async with await _client() as client: + with_keys = (await client.get("/amp/v1/spec")).json() + + assert without["capabilities"]["api_keys_required"] is False + assert with_keys["capabilities"]["api_keys_required"] is True diff --git a/spec/v0.1.0/openapi.json b/spec/v0.1.0/openapi.json index 0de68a1..aed3604 100644 --- a/spec/v0.1.0/openapi.json +++ b/spec/v0.1.0/openapi.json @@ -672,6 +672,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "responses": { @@ -728,6 +744,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "requestBody": { @@ -864,6 +896,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "responses": { @@ -922,6 +970,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "requestBody": { @@ -999,6 +1063,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "responses": { @@ -1075,6 +1155,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "responses": { @@ -1150,6 +1246,22 @@ ], "title": "X-Amp-Agent-Id" } + }, + { + "in": "header", + "name": "x-amp-api-key", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "X-Amp-Api-Key" + } } ], "requestBody": {