From eb9355b00394d539f521039d41b1a3b8c6fe8602 Mon Sep 17 00:00:00 2001 From: Andraxion Date: Wed, 29 Jul 2026 08:51:22 -0400 Subject: [PATCH] Make task context page packing logarithmic --- src/docforge/mcp_server.py | 67 +++++++++++++++++----------- tests/test_mcp_server.py | 90 +++++++++++++++++++++++++++++++++++++- 2 files changed, 130 insertions(+), 27 deletions(-) diff --git a/src/docforge/mcp_server.py b/src/docforge/mcp_server.py index 833dc1d..d90e7b2 100644 --- a/src/docforge/mcp_server.py +++ b/src/docforge/mcp_server.py @@ -1170,39 +1170,54 @@ class DocForgeService: "pagination": pagination, } - for kind, item in items[position:]: - if consumed >= selected_limit: - break - destination = page_evidence if kind == "evidence" else page_omissions - destination.append(item) - consumed += 1 + candidates = items[position : position + selected_limit] + + def populate(candidate_count: int) -> None: + nonlocal consumed, response_limited + page_evidence.clear() + page_omissions.clear() + for kind, item in candidates[:candidate_count]: + destination = page_evidence if kind == "evidence" else page_omissions + destination.append(item) + consumed = candidate_count + response_limited = candidate_count < len(candidates) + + def fits(candidate_count: int) -> bool: + populate(candidate_count) decorated = { **page_result(), "server_version": SERVER_VERSION, "content_warning": CONTENT_WARNING, "staleness": "current", } - if self._encoded_length(decorated) <= maximum: - continue - destination.pop() - consumed -= 1 + return self._encoded_length(decorated) <= maximum + + lower = 0 + upper = len(candidates) + while lower < upper: + midpoint = (lower + upper + 1) // 2 + if fits(midpoint): + lower = midpoint + else: + upper = midpoint - 1 + populate(lower) + if lower == 0 and candidates: + _, item = candidates[0] + subject = "unknown" + if isinstance(item, Mapping): + item_payload = cast(Mapping[str, object], item) + candidate = item_payload.get("node_id") or item_payload.get("subject") + if isinstance(candidate, str) and candidate: + subject = candidate[:256] + page_omissions.append( + { + "code": "response_limit", + "subject": subject, + "detail_hash": canonical_hash(cast(object, item)), + } + ) + consumed = 1 response_limited = True - if consumed == 0: - subject = "unknown" - if isinstance(item, Mapping): - item_payload = cast(Mapping[str, object], item) - candidate = item_payload.get("node_id") or item_payload.get("subject") - if isinstance(candidate, str) and candidate: - subject = candidate[:256] - page_omissions.append( - { - "code": "response_limit", - "subject": subject, - "detail_hash": canonical_hash(cast(object, item)), - } - ) - consumed = 1 - break return page_result() def _page_context_result( diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 287af17..8c65dd2 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -26,11 +26,12 @@ from docforge.mcp_server import ( CONTENT_WARNING, PROPOSAL_TOOLS, READ_TOOLS, + SERVER_VERSION, DocForgeService, _create_bound_server, create_server, ) -from docforge.project import Project +from docforge.project import Project, project_root_fingerprint from docforge.viewer_manager import ViewerManager ROOT = Path(__file__).resolve().parents[1] @@ -416,6 +417,93 @@ class DocForgeMcpTests(unittest.IsolatedAsyncioTestCase): 8_000, ) + def test_dense_task_context_page_packing_is_logarithmic(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = self.copy_fixture("alpha", Path(directory)) + descriptor = root / ".docforge" / "project.toml" + descriptor.write_text( + descriptor.read_text(encoding="utf-8") + .replace("max_results = 20", "max_results = 1000") + .replace( + "max_context_tokens = 2000", + "max_context_tokens = 2000\nmax_tool_output_chars = 20000", + ), + encoding="utf-8", + ) + project = Project.open(root) + service = DocForgeService(project, capability_mode_name="read") + capsule = { + "schema_version": 1, + "plan": { + "effective_policy_hash": "1" * 64, + "request_hash": "2" * 64, + "plan_hash": "3" * 64, + }, + "generation": {"index_schema_version": 3}, + "evidence": [ + { + "node_id": f"node.{index:04d}", + "content": "bounded evidence " * 40, + } + for index in range(1_000) + ], + "gaps": [], + "omissions": [], + "collection_hash": "4" * 64, + "capsule_hash": "5" * 64, + "state": "complete", + "summary": {"evidence_count": 1_000}, + } + result = { + "status": "ok", + "project_id": project.descriptor.project_id, + "project_root_fingerprint": project_root_fingerprint(project.descriptor.root), + "adapter": project.descriptor.adapter, + "revision": "test-revision", + "source_hash": "6" * 64, + "capsule": capsule, + } + + with mock.patch.object( + service, + "_encoded_length", + wraps=service._encoded_length, + ) as encoded_length: + page = service._page_task_context_result( + result, + selected_limit=1_000, + cursor=None, + ) + + pagination = page["pagination"] + self.assertGreater(pagination["returned_count"], 0) + self.assertLess(pagination["returned_count"], 1_000) + returned_count = pagination["returned_count"] + self.assertEqual( + capsule["evidence"][:returned_count], + page["capsule"]["evidence"], + ) + self.assertEqual([], page["capsule"]["omissions"]) + self.assertEqual( + pagination["next_cursor"], + page["capsule"]["pagination"]["next_cursor"], + ) + self.assertEqual( + returned_count, + page["capsule"]["summary"]["page_item_count"], + ) + self.assertLessEqual(encoded_length.call_count, 11) + decorated = { + **page, + "server_version": SERVER_VERSION, + "content_warning": CONTENT_WARNING, + "staleness": "current", + } + self.assertLessEqual( + len(json.dumps(decorated, sort_keys=True, separators=(",", ":"))), + 20_000, + ) + def test_task_context_default_page_clamps_to_small_project_limit(self) -> None: with tempfile.TemporaryDirectory() as directory: root = self.copy_fixture("alpha", Path(directory))