Merge pull request #146 from FSoft-AI4Code/fix/module-clustering-code-size-coverage

Fix module clustering size estimates and source file coverage

Author: anhnh2002

.gitignore

@@ -48,6 +48,7 @@ tests/*
 !tests/test_artifact_analyzer.py
 !tests/test_scala_analyzer.py
 !tests/test_agent_limits.py
+!tests/test_module_scope.py
 
 # Jupyter
 *.ipynb

codewiki/cli/adapters/doc_generator.py

@@ -278,6 +278,10 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
             super_group_modules,
         )
         from codewiki.src.config import FIRST_MODULE_TREE_FILENAME, MODULE_TREE_FILENAME
+        from codewiki.src.be.module_scope import (
+            get_module_token_count,
+            validate_module_file_coverage,
+        )
         from codewiki.src.utils import file_manager
 
         working_dir = str(self.output_dir.absolute())
@@ -298,19 +302,20 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
                     self.progress_tracker.update_stage(0.5, "Loaded cached module tree")
             else:
                 if self.verbose:
-                    clustering_tokens = get_clustering_input_token_count(leaf_nodes, components)
+                    clustering_tokens = get_module_token_count(components)
                     self.progress_tracker.update_stage(
                         0.3,
                         (
                             f"Preparing {len(leaf_nodes)} leaf nodes for clustering "
-                            f"({clustering_tokens} tokens, threshold "
+                            f"(code_tokens={clustering_tokens}, leaf_tokens="
+                            f"{get_clustering_input_token_count(leaf_nodes, components)}, threshold "
                             f"{backend_config.max_token_per_module})"
                         ),
                     )
                     if clustering_tokens <= backend_config.max_token_per_module:
                         self.progress_tracker.update_stage(
                             0.4,
-                            "Skipping LLM clustering; selected leaf nodes fit within the module token threshold",
+                            "Skipping LLM clustering; code size fits within the module token threshold",
                         )
                     else:
                         self.progress_tracker.update_stage(
@@ -342,6 +347,7 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
                 from codewiki.src.be.module_naming import dedupe_module_tree_names
 
                 module_tree = dedupe_module_tree_names(module_tree)
+                validate_module_file_coverage(module_tree, components)
                 file_manager.save_json(module_tree, first_module_tree_path)
                 file_manager.save_json(module_tree, module_tree_path)
 

codewiki/src/be/agent_tools/deps.py

@@ -20,3 +20,17 @@ class CodeWikiDeps:
     # (create / str_replace / insert / undo_edit) to a docs file whose
     # absolute, resolved path is not in this set. ``None`` = unrestricted.
     allowed_write_paths: set[str] | None = None
+    # Entry points belonging to this agent. Sizes expand these IDs to all
+    # components in their files; delegation may never claim a sibling's file.
+    scope_component_ids: list[str] | None = None
+
+    def module_component_ids(self) -> list[str]:
+        if self.scope_component_ids is not None:
+            return self.scope_component_ids
+        if not self.path_to_current_module:
+            return list(self.components)
+        level = self.module_tree
+        for name in self.path_to_current_module:
+            info = level[name]
+            level = info.get("children", {})
+        return info.get("components", [])

codewiki/src/be/agent_tools/generate_sub_module_documentations.py

@@ -12,8 +12,12 @@
     format_system_prompt,
     format_user_prompt,
 )
-from codewiki.src.be.utils import is_complex_module, count_tokens
-from codewiki.src.be.cluster_modules import format_potential_core_components
+from codewiki.src.be.utils import is_complex_module
+from codewiki.src.be.module_scope import (
+    component_files,
+    get_module_token_count,
+    prepare_submodule_partition,
+)
 
 import logging
 from codewiki.src.be.api_usage import usage_target
@@ -38,6 +42,8 @@ async def generate_sub_module_documentation(
 
     deps = ctx.deps
     previous_module_name = deps.current_module_name
+    if deps.current_depth >= deps.max_depth:
+        return "Maximum delegation depth reached; write this module directly."
 
     # Create fallback models from config
     fallback_models = create_fallback_models(deps.config)
@@ -64,10 +70,20 @@ async def generate_sub_module_documentation(
         if requested_name in name_map
     }
 
-    # add the sub-module to the module tree
     value = deps.module_tree
     for key in deps.path_to_current_module:
         value = value[key]["children"]
+    final_specs = prepare_submodule_partition(
+        final_specs, deps.module_component_ids(), deps.components, value
+    )
+    for requested, final_name in list(name_map.items()):
+        if final_name not in final_specs:
+            plan.skipped[requested] = "no unassigned files in the current module scope"
+            del name_map[requested]
+    if not final_specs:
+        return skipped_report(plan.skipped, parent_doc)
+
+    # Reserve every child's complete file scope before running any agent.
     for sub_module_name, core_component_ids in final_specs.items():
         value[sub_module_name] = {"components": core_component_ids, "children": {}}
 
@@ -78,14 +94,15 @@ async def generate_sub_module_documentation(
 
         logger.info(f"{indent}{arrow} Generating documentation for sub-module: {sub_module_name}")
 
-        num_tokens = count_tokens(
-            format_potential_core_components(core_component_ids, ctx.deps.components)[-1]
+        num_tokens = get_module_token_count(
+            deps.components, component_files(core_component_ids, deps.components)
         )
+        child_depth = deps.current_depth + 1
 
         if (
             is_complex_module(ctx.deps.components, core_component_ids)
-            and ctx.deps.current_depth < ctx.deps.max_depth
-            and num_tokens >= ctx.deps.config.max_token_per_leaf_module * ctx.deps.current_depth
+            and child_depth < ctx.deps.max_depth
+            and num_tokens >= ctx.deps.config.max_token_per_leaf_module * child_depth
         ):
             sub_agent = Agent(
                 model=fallback_models,
@@ -119,6 +136,8 @@ async def generate_sub_module_documentation(
                 tools=[read_code_components_tool, str_replace_editor_tool],
             )
 
+        previous_scope = deps.scope_component_ids
+        deps.scope_component_ids = core_component_ids
         deps.current_module_name = sub_module_name
         deps.path_to_current_module.append(sub_module_name)
         deps.current_depth += 1
@@ -143,6 +162,7 @@ async def generate_sub_module_documentation(
             deps.path_to_current_module.pop()
             deps.current_depth -= 1
             deps.current_module_name = previous_module_name
+            deps.scope_component_ids = previous_scope
 
     # restore the previous module name
     deps.current_module_name = previous_module_name

codewiki/src/be/caw_backend.py

@@ -31,14 +31,19 @@
 
 from codewiki.src.be.agent_tools.deps import CodeWikiDeps
 from codewiki.src.be.backend import AgentReply, LLMBackend
-from codewiki.src.be.cluster_modules import format_potential_core_components
+from codewiki.src.be.module_scope import (
+    complete_file_candidates,
+    component_files,
+    get_module_token_count,
+    validate_module_file_coverage,
+)
 from codewiki.src.be.dependency_analyzer.models.core import Node
 from codewiki.src.be.prompt_template import (
     format_leaf_system_prompt,
     format_system_prompt,
     format_user_prompt,
 )
-from codewiki.src.be.utils import count_tokens, is_complex_module, set_main_loop
+from codewiki.src.be.utils import is_complex_module, set_main_loop
 from codewiki.src.be.doc_layout import config_layout, find_doc
 from codewiki.src.config import MODULE_TREE_FILENAME, OVERVIEW_FILENAME, Config
 from codewiki.src.utils import file_manager
@@ -399,21 +404,18 @@ def _run_module_agent_sync(
             return module_tree
 
         custom_instructions = config.get_prompt_addition()
+        if not module_path:
+            core_component_ids = complete_file_candidates(core_component_ids, components)
 
-        # Mirror PydanticAIBackend's early-cut: a module is only worth
-        # delegating to sub-agents when it spans multiple files AND has enough
-        # content to justify the cost AND we still have recursion budget.
-        # Without this gate the caw path would give every multi-file sub-module
-        # the recursive SYSTEM_PROMPT + delegation tool and fan out one extra
-        # agent call per sub-spec even when a single leaf write would suffice.
-        # See generate_sub_module_documentation_tool for the pydantic-ai
-        # equivalent.
-        _, components_with_code = format_potential_core_components(core_component_ids, components)
-        num_tokens = count_tokens(components_with_code)
+        # Whole-repository agents may optionally delegate even below the leaf
+        # token budget. Nested modules retain the cost gate and depth limit.
+        num_tokens = get_module_token_count(
+            components, component_files(core_component_ids, components)
+        )
         can_delegate = (
             is_complex_module(components, core_component_ids)
             and start_depth < config.max_depth
-            and num_tokens >= config.max_token_per_leaf_module
+            and (not module_path or num_tokens >= config.max_token_per_leaf_module)
         )
         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}"
@@ -440,6 +442,7 @@ def _run_module_agent_sync(
             current_depth=start_depth,
             config=config,
             custom_instructions=custom_instructions,
+            scope_component_ids=core_component_ids,
         )
 
         toolkit = CawToolKit(deps=deps, backend=self, allow_subagent=can_delegate)
@@ -521,6 +524,7 @@ def _run_module_agent_sync(
                 traj.total_tool_calls,
                 (usage or {}).get("cost_usd", 0.0),
             )
+            validate_module_file_coverage(deps.module_tree, components)
             file_manager.save_json(deps.module_tree, module_tree_path)
             return deps.module_tree
         except Exception as e:

codewiki/src/be/caw_toolkit.py

@@ -30,6 +30,7 @@
 from codewiki.src.be.agent_tools.deps import CodeWikiDeps
 from codewiki.src.be.doc_layout import config_layout, module_doc_file
 from codewiki.src.be.module_naming import skipped_report, plan_sub_module_specs, sub_module_report
+from codewiki.src.be.module_scope import prepare_submodule_partition
 
 if TYPE_CHECKING:
     from codewiki.src.be.caw_backend import CawBackend
@@ -325,6 +326,17 @@ def _run_sub_modules(self, sub_module_specs: dict[str, list[str]]) -> str:
         value = deps.module_tree
         for key in deps.path_to_current_module:
             value = value[key]["children"]
+        final_specs = prepare_submodule_partition(
+            final_specs, deps.module_component_ids(), deps.components, value
+        )
+        for requested, final_name in list(name_map.items()):
+            if final_name not in final_specs:
+                plan.skipped[requested] = "no unassigned files in the current module scope"
+                del name_map[requested]
+        if not final_specs:
+            return skipped_report(
+                plan.skipped, module_doc_file(previous_module_name, parent_path, layout)
+            )
         for sub_name, core_ids in final_specs.items():
             value[sub_name] = {"components": core_ids, "children": {}}