Skip to content

Commit b243e0a

Browse files
authored
Merge pull request #94 from FSoft-AI4Code/fix/module-prompt-input-too-large
Fix module/cluster prompts exceeding codex input cap; fix resume skipping missing module docs
2 parents 003b876 + 85b6f0a commit b243e0a

4 files changed

Lines changed: 184 additions & 112 deletions

File tree

‎codewiki/src/be/caw_backend.py‎

Lines changed: 24 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import logging
2424
import os
2525
import shutil
26-
from typing import Any, Dict, List
26+
from typing import Any
2727

2828
from caw import Agent as CawAgent
2929
from caw import ToolGroup
@@ -170,7 +170,7 @@ def _with_allowed_tools(cmd):
170170
allowed = ",".join(f"mcp__{s}" for s in servers)
171171
logger.info("Injected --allowedTools for MCP servers: %s", servers)
172172
return list(cmd) + ["--allowedTools", allowed]
173-
except Exception as e: # never break the spawn on a patch hiccup
173+
except Exception as e: # noqa: BLE001 — never break the spawn on a patch hiccup
174174
logger.warning("claude allowedTools patch skipped: %s", e)
175175
return cmd
176176

@@ -266,11 +266,11 @@ def complete(
266266
async def run_module_agent(
267267
self,
268268
module_name: str,
269-
components: Dict[str, Node],
270-
core_component_ids: List[str],
271-
module_path: List[str],
269+
components: dict[str, Node],
270+
core_component_ids: list[str],
271+
module_path: list[str],
272272
working_dir: str,
273-
) -> Dict[str, Any]:
273+
) -> dict[str, Any]:
274274
# caw.completion shells out to a subprocess and blocks the calling
275275
# thread. Push it off the event loop so the rest of the async
276276
# pipeline keeps moving.
@@ -293,13 +293,13 @@ async def run_module_agent(
293293
def _run_module_agent_sync(
294294
self,
295295
module_name: str,
296-
components: Dict[str, Node],
297-
core_component_ids: List[str],
298-
module_path: List[str],
296+
components: dict[str, Node],
297+
core_component_ids: list[str],
298+
module_path: list[str],
299299
working_dir: str,
300300
start_depth: int = 1,
301-
module_tree: Dict[str, Any] | None = None,
302-
) -> Dict[str, Any]:
301+
module_tree: dict[str, Any] | None = None,
302+
) -> dict[str, Any]:
303303
# ``start_depth`` lets the recursion preserve max_depth across nested
304304
# _run_module_agent_sync calls — each fresh deps object would otherwise
305305
# reset current_depth to 1 and silently bypass max_depth guards.
@@ -314,10 +314,15 @@ def _run_module_agent_sync(
314314
if module_tree is None:
315315
module_tree = file_manager.load_json(module_tree_path)
316316

317-
overview_docs_path = os.path.join(working_dir, OVERVIEW_FILENAME)
318-
if os.path.exists(overview_docs_path):
319-
logger.info("✓ Overview docs already exists at %s", overview_docs_path)
320-
return module_tree
317+
# overview.md is the root module's own doc (renamed from
318+
# {repo_name}.md), so its presence only proves the root is done —
319+
# nested modules must still be checked against their own doc file,
320+
# or a resume after a partial run silently skips every missing one.
321+
if not module_path:
322+
overview_docs_path = os.path.join(working_dir, OVERVIEW_FILENAME)
323+
if os.path.exists(overview_docs_path):
324+
logger.info("✓ Overview docs already exists at %s", overview_docs_path)
325+
return module_tree
321326
docs_path = os.path.join(working_dir, f"{module_name}.md")
322327
if os.path.exists(docs_path):
323328
logger.info("✓ Module docs already exists at %s", docs_path)
@@ -333,16 +338,16 @@ def _run_module_agent_sync(
333338
# agent call per sub-spec even when a single leaf write would suffice.
334339
# See generate_sub_module_documentation_tool for the pydantic-ai
335340
# equivalent.
336-
_, components_with_code = format_potential_core_components(
337-
core_component_ids, components
338-
)
341+
_, components_with_code = format_potential_core_components(core_component_ids, components)
339342
num_tokens = count_tokens(components_with_code)
340343
can_delegate = (
341344
is_complex_module(components, core_component_ids)
342345
and start_depth < config.max_depth
343346
and num_tokens >= config.max_token_per_leaf_module
344347
)
345-
logger.info(f"Module {module_name} can delegate: {can_delegate} - is_complex_module: {is_complex_module(components, core_component_ids)} - start_depth: {start_depth} - num_tokens: {num_tokens} - max_depth: {config.max_depth} - max_token_per_leaf_module: {config.max_token_per_leaf_module}")
348+
logger.info(
349+
f"Module {module_name} can delegate: {can_delegate} - is_complex_module: {is_complex_module(components, core_component_ids)} - start_depth: {start_depth} - num_tokens: {num_tokens} - max_depth: {config.max_depth} - max_token_per_leaf_module: {config.max_token_per_leaf_module}"
350+
)
346351

347352
if can_delegate:
348353
system_prompt = format_system_prompt(module_name, custom_instructions)

‎codewiki/src/be/documentation_generator.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -297,13 +297,10 @@ async def generate_parent_module_docs(
297297
module_tree_path = os.path.join(working_dir, MODULE_TREE_FILENAME)
298298
module_tree = file_manager.load_json(module_tree_path)
299299

300-
# check if overview docs already exists
301-
overview_docs_path = os.path.join(working_dir, OVERVIEW_FILENAME)
302-
if os.path.exists(overview_docs_path):
303-
logger.info(f"✓ Overview docs already exists at {overview_docs_path}")
304-
return module_tree
305-
306-
# check if parent docs already exists
300+
# check if parent docs already exists (for the root module the doc
301+
# path below resolves to overview.md, so a blanket "overview exists →
302+
# skip" check is not needed here and would mask missing parent docs
303+
# on resume)
307304
parent_docs_path = os.path.join(
308305
working_dir,
309306
f"{module_name if len(module_path) >= 1 else OVERVIEW_FILENAME.replace('.md', '')}.md",

‎codewiki/src/be/prompt_template.py‎

Lines changed: 142 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -257,10 +257,32 @@
257257
Reasoning at first, then return the list of relative paths in JSON format.
258258
"""
259259

260+
import logging
261+
from collections import defaultdict
260262
from typing import Any
261263

262264
from codewiki.src.utils import file_manager
263265

266+
logger = logging.getLogger(__name__)
267+
268+
# codex rejects any turn whose total input exceeds 1,048,576 characters
269+
# (input_too_large, code -32602 — server-side, not configurable). Cap the
270+
# user prompt below that, leaving headroom for the system prompt, tool
271+
# schemas and protocol overhead.
272+
MAX_USER_PROMPT_CHARS = 900_000
273+
274+
MODULE_TREE_TRIMMED_NOTE = (
275+
"NOTE: per-module component listings were omitted because the full module "
276+
"tree exceeds the model input limit. Module names and hierarchy are "
277+
"complete; read the referenced modules' documentation files or use your "
278+
"code-reading tools when you need component-level detail."
279+
)
280+
281+
CODE_TRUNCATED_NOTE = (
282+
"\n... [file contents truncated to fit the model input limit — use your "
283+
"file-reading tools to read the full files]"
284+
)
285+
264286
EXTENSION_TO_LANGUAGE = {
265287
".py": "python",
266288
".md": "markdown",
@@ -289,6 +311,54 @@
289311
}
290312

291313

314+
def _format_module_tree_str(
315+
module_tree: dict[str, Any],
316+
current_module_name: str | None = None,
317+
include_components: bool = True,
318+
) -> str:
319+
"""
320+
Render a module tree as an indented text outline.
321+
322+
With include_components=False only module names and hierarchy are
323+
emitted, which keeps the outline small enough for huge trees that would
324+
otherwise blow past MAX_USER_PROMPT_CHARS.
325+
"""
326+
lines: list[str] = []
327+
328+
def _walk(tree: dict[str, Any], indent: int = 0) -> None:
329+
for key, value in tree.items():
330+
if key == current_module_name:
331+
lines.append(f"{' ' * indent}{key} (current module)")
332+
else:
333+
lines.append(f"{' ' * indent}{key}")
334+
335+
if include_components:
336+
# Group components by file
337+
by_file = defaultdict(list)
338+
for c in value["components"]:
339+
if "::" in c:
340+
fpath, name = c.split("::", 1)
341+
by_file[fpath].append(name)
342+
else:
343+
by_file[""].append(c)
344+
for fpath, names in by_file.items():
345+
if fpath:
346+
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
347+
else:
348+
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
349+
350+
if (
351+
("children" in value)
352+
and isinstance(value["children"], dict)
353+
and len(value["children"]) > 0
354+
):
355+
lines.append(f"{' ' * (indent + 1)} Children:")
356+
_walk(value["children"], indent + 2)
357+
358+
_walk(module_tree, 0)
359+
return "\n".join(lines)
360+
361+
292362
def format_user_prompt(
293363
module_name: str,
294364
core_component_ids: list[str],
@@ -306,41 +376,7 @@ def format_user_prompt(
306376
Returns:
307377
Formatted user prompt string
308378
"""
309-
310-
# format module tree
311-
lines = []
312-
313-
def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
314-
for key, value in module_tree.items():
315-
if key == module_name:
316-
lines.append(f"{' ' * indent}{key} (current module)")
317-
else:
318-
lines.append(f"{' ' * indent}{key}")
319-
320-
# Group components by file
321-
from collections import defaultdict
322-
323-
by_file = defaultdict(list)
324-
for c in value["components"]:
325-
if "::" in c:
326-
fpath, name = c.split("::", 1)
327-
by_file[fpath].append(name)
328-
else:
329-
by_file[""].append(c)
330-
for fpath, names in by_file.items():
331-
if fpath:
332-
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
333-
else:
334-
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
335-
336-
if isinstance(value["children"], dict) and len(value["children"]) > 0:
337-
lines.append(f"{' ' * (indent + 1)} Children:")
338-
_format_module_tree(value["children"], indent + 2)
339-
340-
_format_module_tree(module_tree, 0)
341-
formatted_module_tree = "\n".join(lines)
342-
343-
# print(f"Formatted module tree:\n{formatted_module_tree}")
379+
formatted_module_tree = _format_module_tree_str(module_tree, module_name)
344380

345381
# Group core component IDs by their file path
346382
grouped_components: dict[str, list[str]] = {}
@@ -375,12 +411,56 @@ def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
375411

376412
core_component_codes += "```\n\n"
377413

378-
return USER_PROMPT.format(
414+
prompt = USER_PROMPT.format(
379415
module_name=module_name,
380416
formatted_core_component_codes=core_component_codes,
381417
module_tree=formatted_module_tree,
382418
)
383419

420+
if len(prompt) > MAX_USER_PROMPT_CHARS:
421+
full_len = len(prompt)
422+
formatted_module_tree = (
423+
MODULE_TREE_TRIMMED_NOTE
424+
+ "\n\n"
425+
+ _format_module_tree_str(module_tree, module_name, include_components=False)
426+
)
427+
prompt = USER_PROMPT.format(
428+
module_name=module_name,
429+
formatted_core_component_codes=core_component_codes,
430+
module_tree=formatted_module_tree,
431+
)
432+
logger.warning(
433+
"Module %s: user prompt (%d chars) exceeds %d; "
434+
"module tree trimmed to names only (%d chars)",
435+
module_name,
436+
full_len,
437+
MAX_USER_PROMPT_CHARS,
438+
len(prompt),
439+
)
440+
441+
if len(prompt) > MAX_USER_PROMPT_CHARS:
442+
# Even the slim tree was not enough — the inlined file contents
443+
# dominate. Truncation is recoverable: the agent has file-reading
444+
# tools (read_code_components / str_replace_editor).
445+
excess = len(prompt) - MAX_USER_PROMPT_CHARS + len(CODE_TRUNCATED_NOTE)
446+
core_component_codes = (
447+
core_component_codes[: max(0, len(core_component_codes) - excess)] + CODE_TRUNCATED_NOTE
448+
)
449+
prompt = USER_PROMPT.format(
450+
module_name=module_name,
451+
formatted_core_component_codes=core_component_codes,
452+
module_tree=formatted_module_tree,
453+
)
454+
logger.warning(
455+
"Module %s: user prompt still over %d chars after tree trim; "
456+
"truncated inlined file contents (now %d chars)",
457+
module_name,
458+
MAX_USER_PROMPT_CHARS,
459+
len(prompt),
460+
)
461+
462+
return prompt
463+
384464

385465
def format_cluster_prompt(
386466
potential_core_components: str,
@@ -393,53 +473,38 @@ def format_cluster_prompt(
393473
if module_tree is None:
394474
module_tree = {}
395475

396-
# format module tree
397-
lines = []
398-
399-
# print(f"Module tree:\n{json.dumps(module_tree, indent=2)}")
400-
401-
def _format_module_tree(module_tree: dict[str, any], indent: int = 0):
402-
for key, value in module_tree.items():
403-
if key == module_name:
404-
lines.append(f"{' ' * indent}{key} (current module)")
405-
else:
406-
lines.append(f"{' ' * indent}{key}")
407-
408-
# Group components by file
409-
from collections import defaultdict
410-
411-
by_file = defaultdict(list)
412-
for c in value["components"]:
413-
if "::" in c:
414-
fpath, name = c.split("::", 1)
415-
by_file[fpath].append(name)
416-
else:
417-
by_file[""].append(c)
418-
for fpath, names in by_file.items():
419-
if fpath:
420-
lines.append(f"{' ' * (indent + 1)} {fpath}: {', '.join(names)}")
421-
else:
422-
lines.append(f"{' ' * (indent + 1)} {', '.join(names)}")
423-
424-
if (
425-
("children" in value)
426-
and isinstance(value["children"], dict)
427-
and len(value["children"]) > 0
428-
):
429-
lines.append(f"{' ' * (indent + 1)} Children:")
430-
_format_module_tree(value["children"], indent + 2)
431-
432-
_format_module_tree(module_tree, 0)
433-
formatted_module_tree = "\n".join(lines)
434-
435476
if module_tree == {}:
436477
return CLUSTER_REPO_PROMPT.format(potential_core_components=potential_core_components)
437-
else:
438-
return CLUSTER_MODULE_PROMPT.format(
478+
479+
formatted_module_tree = _format_module_tree_str(module_tree, module_name)
480+
prompt = CLUSTER_MODULE_PROMPT.format(
481+
potential_core_components=potential_core_components,
482+
module_tree=formatted_module_tree,
483+
module_name=module_name,
484+
)
485+
486+
if len(prompt) > MAX_USER_PROMPT_CHARS:
487+
full_len = len(prompt)
488+
formatted_module_tree = (
489+
MODULE_TREE_TRIMMED_NOTE
490+
+ "\n\n"
491+
+ _format_module_tree_str(module_tree, module_name, include_components=False)
492+
)
493+
prompt = CLUSTER_MODULE_PROMPT.format(
439494
potential_core_components=potential_core_components,
440495
module_tree=formatted_module_tree,
441496
module_name=module_name,
442497
)
498+
logger.warning(
499+
"Module %s: cluster prompt (%d chars) exceeds %d; "
500+
"module tree trimmed to names only (%d chars)",
501+
module_name,
502+
full_len,
503+
MAX_USER_PROMPT_CHARS,
504+
len(prompt),
505+
)
506+
507+
return prompt
443508

444509

445510
def format_super_group_prompt(module_tree: dict[str, Any]) -> str:

0 commit comments

Comments
 (0)