administrator/DocForge2
1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions Pages

Make task context page packing logarithmic

Browse source
This commit is contained in:
Andraxion 2026-07-29 08:51:22 -04:00
parent 9a48233983
commit eb9355b003
2 changed files with 130 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

37
src/docforge/mcp_server.py
View file

@ -1170,24 +1170,39 @@ class DocForgeService:
"pagination": pagination,
}
for kind, item in items[position:]:
if consumed >= selected_limit:
break
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 += 1
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
response_limited = True
if consumed == 0:
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)
@ -1202,7 +1217,7 @@ class DocForgeService:
}
)
consumed = 1
break
response_limited = True
return page_result()
def _page_context_result(

90
tests/test_mcp_server.py
View file

@ -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))