diff --git a/src/context_engine/serve_http.py b/src/context_engine/serve_http.py index b2d8ea8..bc17b76 100644 --- a/src/context_engine/serve_http.py +++ b/src/context_engine/serve_http.py @@ -15,6 +15,7 @@ from context_engine.utils import project_storage_dir from context_engine.indexer.embedder import Embedder from context_engine.compression.compressor import Compressor +from context_engine.retrieval.retriever import HybridRetriever from context_engine.models import Chunk, ChunkType, GraphNode, GraphEdge, NodeType, EdgeType try: @@ -27,6 +28,9 @@ _MAX_REQUEST_BYTES = 10 * 1024 * 1024 # 10 MB — generous for bulk ingest, not unbounded +# Mirror the MCP server's guard (mcp_server._MAX_QUERY_CHARS) so a buggy or +# malicious client can't submit a multi-MB query string for embedding. +_MAX_QUERY_CHARS = 10_000 _LOOPBACK_HOSTS = {"127.0.0.1", "::1", "localhost"} @@ -35,6 +39,7 @@ def __init__(self, backend: LocalBackend, embedder: Embedder, compressor: Compre self.backend = backend self.embedder = embedder self.compressor = compressor + self.retriever = HybridRetriever(backend=backend, embedder=embedder) async def handle_vector_search(self, request: web.Request) -> web.Response: data = await _read_json(request) @@ -89,6 +94,55 @@ async def handle_delete_file(self, request: web.Request) -> web.Response: await self.backend.delete_by_file(file_path) return web.json_response({"ok": True}) + async def handle_search(self, request: web.Request) -> web.Response: + data = await _read_json(request) + query = (data.get("query") or "").strip() + if not query: + return web.json_response({"error": "query cannot be empty"}, status=400) + if len(query) > _MAX_QUERY_CHARS: + return web.json_response( + {"error": f"query too long (max {_MAX_QUERY_CHARS} characters)"}, + status=400, + ) + # NOTE: the 0.2 default below is hardcoded, so a project overriding + # `retrieval_confidence_threshold` in .context-engine.yaml is honoured by the MCP + # context_search handler but NOT here. Deliberate for v1 — wiring config through + # is a follow-up, not a widening of this PR. + # Validate + clamp: non-numeric input would otherwise raise ValueError and + # surface as a 400 "missing field" via the generic handler; clamp to the same + # ranges the MCP context_search handler uses. + try: + top_k = max(1, min(int(data.get("top_k", 10)), 100)) + confidence_threshold = max(0.0, min(float(data.get("confidence_threshold", 0.2)), 1.0)) + except (TypeError, ValueError): + return web.json_response( + {"error": "top_k must be an int and confidence_threshold a float"}, + status=400, + ) + # NOTE: unlike the MCP context_search handler, this endpoint does not call + # _record(), so queries via /search are not reflected in `cce savings`. + chunks = await self.retriever.retrieve( + query, + top_k=top_k, + confidence_threshold=confidence_threshold, + ) + return web.json_response({ + "results": [ + { + "id": c.id, + "file_path": c.file_path, + "start_line": c.start_line, + "end_line": c.end_line, + "content": c.content, + "chunk_type": c.chunk_type.value, + "language": c.language, + "confidence_score": c.confidence_score, + "metadata": c.metadata, + } + for c in chunks + ] + }) + async def handle_health(self, request: web.Request) -> web.Response: return web.json_response({"status": "ok"}) @@ -195,6 +249,7 @@ def create_app(backend, embedder, compressor, *, api_token: str | None = None) - middlewares=[_make_auth_middleware(api_token), _error_middleware], ) app.router.add_get("/health", handler.handle_health) + app.router.add_post("/search", handler.handle_search) app.router.add_post("/vector_search", handler.handle_vector_search) app.router.add_post("/fts_search", handler.handle_fts_search) app.router.add_post("/chunks_by_ids", handler.handle_chunks_by_ids) diff --git a/tests/test_serve_http_search.py b/tests/test_serve_http_search.py new file mode 100644 index 0000000..780aaf4 --- /dev/null +++ b/tests/test_serve_http_search.py @@ -0,0 +1,221 @@ +"""Tests for POST /search on the loopback HTTP server (PR #96). + +The endpoint is a thin wrapper around HybridRetriever, so what needs covering is +not the retrieval itself (tested elsewhere) but the input handling added in +review: empty query, over-length query, and non-numeric top_k / +confidence_threshold all return 400 rather than raising. + +The retriever is stubbed so these tests stay hermetic — no embedder, no backend, +no index. That also lets the happy-path test assert the clamped values actually +reach retrieve(), which is the part a smoke test would miss. +""" +from __future__ import annotations + +import json + +import pytest +from aiohttp import web +from aiohttp.test_utils import make_mocked_request + +from context_engine.models import Chunk, ChunkType +from context_engine.serve_http import _MAX_QUERY_CHARS, ContextEngineHTTP + + +class _StubRetriever: + """Records the arguments it was called with and returns one fixed chunk.""" + + def __init__(self, chunks=None): + self.chunks = chunks if chunks is not None else [] + self.calls: list[dict] = [] + + async def retrieve(self, query, top_k=10, confidence_threshold=0.2): + self.calls.append( + {"query": query, "top_k": top_k, "confidence_threshold": confidence_threshold} + ) + return self.chunks + + +def _chunk() -> Chunk: + return Chunk( + id="b1294739d28245a1", + content="def record_usage(model, provider, input_tokens, ...)", + chunk_type=ChunkType.FUNCTION, + file_path="memory/journal.py", + start_line=81, + end_line=86, + language="python", + metadata={"_distance": 0.774}, + confidence_score=0.878, + ) + + +def _server(chunks=None) -> ContextEngineHTTP: + """A ContextEngineHTTP with its retriever swapped for the stub. + + __init__ builds a real HybridRetriever from a backend and embedder, neither of + which these tests need, so the instance is created without running __init__ and + only the attribute handle_search touches is set. + """ + server = ContextEngineHTTP.__new__(ContextEngineHTTP) + server.retriever = _StubRetriever(chunks) + return server + + +def _request(payload) -> web.Request: + """A mocked POST /search carrying `payload` as its JSON body.""" + body = json.dumps(payload).encode() + req = make_mocked_request( + "POST", "/search", + headers={"Content-Type": "application/json"}, + payload=body, + ) + + async def _json(*args, **kwargs): + return json.loads(body) + + req.json = _json + return req + + +def _body(response: web.Response) -> dict: + return json.loads(response.text) + + +@pytest.mark.asyncio +async def test_search_happy_path_returns_results(): + """A valid query returns 200 and a serialised results list.""" + server = _server([_chunk()]) + resp = await server.handle_search(_request({"query": "cost tracking", "top_k": 5})) + + assert resp.status == 200 + results = _body(resp)["results"] + assert len(results) == 1 + + r = results[0] + assert r["id"] == "b1294739d28245a1" + assert r["file_path"] == "memory/journal.py" + assert r["start_line"] == 81 + assert r["end_line"] == 86 + assert r["chunk_type"] == "function" # the enum is serialised by .value + assert r["language"] == "python" + assert r["confidence_score"] == pytest.approx(0.878) + assert r["metadata"] == {"_distance": 0.774} + + # top_k must reach the retriever, not just be accepted and dropped. + assert server.retriever.calls == [ + {"query": "cost tracking", "top_k": 5, "confidence_threshold": 0.2} + ] + + +@pytest.mark.asyncio +async def test_search_empty_results_is_still_200(): + """No matches is a valid answer, not an error.""" + server = _server([]) + resp = await server.handle_search(_request({"query": "nothing matches this"})) + + assert resp.status == 200 + assert _body(resp)["results"] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("query", ["", " ", None]) +async def test_search_empty_query_returns_400(query): + """Empty, whitespace-only and missing queries are all rejected.""" + server = _server([_chunk()]) + resp = await server.handle_search(_request({"query": query})) + + assert resp.status == 400 + assert "empty" in _body(resp)["error"] + assert server.retriever.calls == [], "retriever must not be reached on a bad request" + + +@pytest.mark.asyncio +async def test_search_over_length_query_returns_400(): + """A query past _MAX_QUERY_CHARS is rejected before it is embedded.""" + server = _server([_chunk()]) + resp = await server.handle_search(_request({"query": "x" * (_MAX_QUERY_CHARS + 1)})) + + assert resp.status == 400 + assert "too long" in _body(resp)["error"] + assert server.retriever.calls == [] + + +@pytest.mark.asyncio +async def test_search_at_max_query_length_is_accepted(): + """The limit is inclusive — exactly _MAX_QUERY_CHARS must pass.""" + server = _server([]) + resp = await server.handle_search(_request({"query": "x" * _MAX_QUERY_CHARS})) + + assert resp.status == 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + {"query": "ok", "top_k": "abc"}, + {"query": "ok", "top_k": None}, + {"query": "ok", "top_k": [5]}, + {"query": "ok", "confidence_threshold": "high"}, + {"query": "ok", "confidence_threshold": {}}, + ], +) +async def test_search_non_numeric_params_return_400(payload): + """Non-numeric top_k / confidence_threshold return 400, never a 500.""" + server = _server([_chunk()]) + resp = await server.handle_search(_request(payload)) + + assert resp.status == 400 + assert "top_k" in _body(resp)["error"] + assert server.retriever.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "given, expected", + [(0, 1), (-5, 1), (1, 1), (100, 100), (101, 100), (10_000, 100)], +) +async def test_search_top_k_is_clamped(given, expected): + """top_k clamps to 1..100 rather than being rejected or passed through.""" + server = _server([]) + resp = await server.handle_search(_request({"query": "ok", "top_k": given})) + + assert resp.status == 200 + assert server.retriever.calls[0]["top_k"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "given, expected", + [(-1.0, 0.0), (0.0, 0.0), (0.5, 0.5), (1.0, 1.0), (2.5, 1.0)], +) +async def test_search_confidence_threshold_is_clamped(given, expected): + """confidence_threshold clamps to 0.0..1.0.""" + server = _server([]) + resp = await server.handle_search( + _request({"query": "ok", "confidence_threshold": given}) + ) + + assert resp.status == 200 + assert server.retriever.calls[0]["confidence_threshold"] == pytest.approx(expected) + + +@pytest.mark.asyncio +async def test_search_defaults_when_params_omitted(): + """Omitted params fall back to top_k=10, confidence_threshold=0.2.""" + server = _server([]) + resp = await server.handle_search(_request({"query": "ok"})) + + assert resp.status == 200 + assert server.retriever.calls[0]["top_k"] == 10 + assert server.retriever.calls[0]["confidence_threshold"] == pytest.approx(0.2) + + +@pytest.mark.asyncio +async def test_search_query_is_stripped(): + """Surrounding whitespace is trimmed before the query reaches the retriever.""" + server = _server([]) + resp = await server.handle_search(_request({"query": " cost tracking "})) + + assert resp.status == 200 + assert server.retriever.calls[0]["query"] == "cost tracking"