diff --git a/.gitignore b/.gitignore
index 0b89baef..9390c427 100644
--- a/.gitignore
+++ b/.gitignore
@@ -49,6 +49,7 @@ tests/*
!tests/test_scala_analyzer.py
!tests/test_agent_limits.py
!tests/test_module_scope.py
+!tests/test_min_module_size.py
!tests/test_doc_write_guard.py
!tests/test_overview_generation.py
diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py
index 46933d0d..e514eabe 100644
--- a/codewiki/cli/adapters/doc_generator.py
+++ b/codewiki/cli/adapters/doc_generator.py
@@ -20,7 +20,7 @@
# Import backend modules
from codewiki.src.be.documentation_generator import DocumentationGenerator
from codewiki.src.be.doc_layout import list_doc_files
-from codewiki.src.config import DEFAULT_LAYOUT
+from codewiki.src.config import DEFAULT_LAYOUT, DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
from codewiki.src.config import Config as BackendConfig
from codewiki.src.config import set_cli_context
@@ -154,8 +154,11 @@ def generate(self) -> DocumentationJob:
google_project=self.config.get("google_project", ""),
google_location=self.config.get("google_location", ""),
max_tokens=self.config.get("max_tokens", 32768),
+ min_token_per_module=self.config.get("min_token_per_module", 1000),
max_token_per_module=self.config.get("max_token_per_module", 36369),
- max_token_per_leaf_module=self.config.get("max_token_per_leaf_module", 16000),
+ max_token_per_leaf_module=self.config.get(
+ "max_token_per_leaf_module", DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
+ ),
max_leaf_nodes_per_cluster=self.config.get("max_leaf_nodes_per_cluster", 600),
max_depth=self.config.get("max_depth", 2),
request_limit=self.config.get("request_limit", 100),
@@ -341,7 +344,12 @@ async def _run_backend_generation(self, backend_config: BackendConfig):
# Artifact nodes the clustering LLM dropped get a fixed module
# so build/CI/config coverage does not depend on the LLM.
if getattr(backend_config, "artifacts_enabled", True):
- module_tree = ensure_artifact_module(module_tree, leaf_nodes, components)
+ module_tree = ensure_artifact_module(
+ module_tree,
+ leaf_nodes,
+ components,
+ min_tokens=backend_config.min_token_per_module,
+ )
# Only freshly clustered trees are deduped: renaming a cached
# key whose .md already exists would orphan the doc.
from codewiki.src.be.module_naming import dedupe_module_tree_names
diff --git a/codewiki/cli/commands/config.py b/codewiki/cli/commands/config.py
index 19d5c92d..4d7134f1 100644
--- a/codewiki/cli/commands/config.py
+++ b/codewiki/cli/commands/config.py
@@ -11,6 +11,7 @@
from codewiki.cli.config_manager import ConfigManager
from codewiki.cli.models.config import AgentInstructions
from codewiki.src.language import normalize_language
+from codewiki.src.config import DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
from codewiki.cli.utils.errors import ConfigurationError, handle_error, EXIT_CONFIG_ERROR
from codewiki.cli.utils.validation import (
validate_url,
@@ -145,13 +146,21 @@ def pricing_remove(model, provider):
@click.option("--cluster-model", type=str, help="Model for module clustering (recommend top-tier)")
@click.option("--fallback-model", type=str, help="Fallback model for documentation generation")
@click.option("--max-tokens", type=int, help="Maximum tokens for LLM response (default: 32768)")
+@click.option(
+ "--min-token-per-module",
+ type=click.IntRange(min=1),
+ default=None,
+ help="Minimum code tokens per standalone module (default: 1000)",
+)
@click.option(
"--max-token-per-module",
type=int,
help="Maximum tokens per module for clustering (default: 36369)",
)
@click.option(
- "--max-token-per-leaf-module", type=int, help="Maximum tokens per leaf module (default: 16000)"
+ "--max-token-per-leaf-module",
+ type=int,
+ help=f"Base code-token threshold to allow delegation (default: {DEFAULT_MAX_TOKEN_PER_LEAF_MODULE})",
)
@click.option(
"--max-depth", type=int, help="Maximum depth for hierarchical decomposition (default: 2)"
@@ -223,6 +232,7 @@ def config_set(
max_token_per_module: Optional[int],
max_token_per_leaf_module: Optional[int],
max_depth: Optional[int],
+ min_token_per_module: Optional[int] = None,
request_limit: Optional[int] = None,
agent_retries: Optional[int] = None,
provider: Optional[str] = None,
@@ -303,6 +313,7 @@ def config_set(
cluster_model,
fallback_model,
max_tokens,
+ min_token_per_module,
max_token_per_module,
max_token_per_leaf_module,
max_depth,
@@ -369,6 +380,9 @@ def config_set(
raise ConfigurationError("max_tokens must be a positive integer")
validated_data["max_tokens"] = max_tokens
+ if min_token_per_module is not None:
+ validated_data["min_token_per_module"] = min_token_per_module
+
if max_token_per_module is not None:
if max_token_per_module < 1:
raise ConfigurationError("max_token_per_module must be a positive integer")
@@ -425,6 +439,7 @@ def config_set(
cluster_model=validated_data.get("cluster_model"),
fallback_model=validated_data.get("fallback_model"),
max_tokens=validated_data.get("max_tokens"),
+ min_token_per_module=validated_data.get("min_token_per_module"),
max_token_per_module=validated_data.get("max_token_per_module"),
max_token_per_leaf_module=validated_data.get("max_token_per_leaf_module"),
max_depth=validated_data.get("max_depth"),
@@ -476,6 +491,9 @@ def config_set(
if max_tokens:
click.secho(f"✓ Max tokens: {max_tokens}", fg="green")
+ if min_token_per_module is not None:
+ click.secho(f"✓ Min token per module: {min_token_per_module}", fg="green")
+
if max_token_per_module:
click.secho(f"✓ Max token per module: {max_token_per_module}", fg="green")
@@ -570,8 +588,11 @@ def config_show(output_json: bool):
"fallback_model": config.fallback_model if config else "glm-4p5",
"default_output": config.default_output if config else "docs",
"max_tokens": config.max_tokens if config else 32768,
+ "min_token_per_module": config.min_token_per_module if config else 1000,
"max_token_per_module": config.max_token_per_module if config else 36369,
- "max_token_per_leaf_module": config.max_token_per_leaf_module if config else 16000,
+ "max_token_per_leaf_module": config.max_token_per_leaf_module
+ if config
+ else DEFAULT_MAX_TOKEN_PER_LEAF_MODULE,
"max_depth": config.max_depth if config else 2,
"request_limit": config.request_limit if config else 100,
"agent_retries": config.agent_retries if config else 3,
@@ -644,6 +665,7 @@ def config_show(output_json: bool):
click.secho("Token Settings", fg="cyan", bold=True)
if config:
click.echo(f" Max Tokens: {config.max_tokens}")
+ click.echo(f" Min Token/Module: {config.min_token_per_module}")
click.echo(f" Max Token/Module: {config.max_token_per_module}")
click.echo(f" Max Token/Leaf Module: {config.max_token_per_leaf_module}")
click.echo(f" Prompt Caching: {config.prompt_caching}")
diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py
index b9002da4..c9578786 100644
--- a/codewiki/cli/commands/generate.py
+++ b/codewiki/cli/commands/generate.py
@@ -302,6 +302,12 @@ def _find_affected(tree, parent_names=None):
default=None,
help="Maximum tokens for LLM response (overrides config)",
)
+@click.option(
+ "--min-token-per-module",
+ type=click.IntRange(min=1),
+ default=None,
+ help="Minimum code tokens per standalone module (default: 1000)",
+)
@click.option(
"--max-token-per-module",
type=int,
@@ -312,7 +318,7 @@ def _find_affected(tree, parent_names=None):
"--max-token-per-leaf-module",
type=int,
default=None,
- help="Maximum tokens per leaf module (overrides config)",
+ help="Base code-token threshold to allow delegation (overrides config)",
)
@click.option(
"--max-depth",
@@ -469,6 +475,7 @@ def generate_command(
request_limit: int | None,
agent_retries: int | None,
prompt_caching: bool | None,
+ min_token_per_module: int | None = None,
main_model: str | None = None,
cluster_model: str | None = None,
fallback_model: str | None = None,
@@ -745,6 +752,9 @@ def generate_command(
)
logger.debug(f"Main model: {main_model or config.main_model}")
logger.debug(f"Cluster model: {cluster_model or config.cluster_model}")
+ logger.debug(
+ f"Min token/module: {min_token_per_module if min_token_per_module is not None else getattr(config, 'min_token_per_module', 1000)}"
+ )
logger.debug(f"Max tokens: {effective_max_tokens}")
logger.debug(f"Max token/module: {effective_max_token_per_module}")
logger.debug(f"Max token/leaf module: {effective_max_token_per_leaf}")
@@ -796,6 +806,9 @@ def generate_command(
"agent_instructions": agent_instructions_dict,
# Max token settings (runtime overrides take precedence)
"max_tokens": max_tokens if max_tokens is not None else config.max_tokens,
+ "min_token_per_module": min_token_per_module
+ if min_token_per_module is not None
+ else getattr(config, "min_token_per_module", 1000),
"max_token_per_module": max_token_per_module
if max_token_per_module is not None
else config.max_token_per_module,
diff --git a/codewiki/cli/config_manager.py b/codewiki/cli/config_manager.py
index 59d48625..6fc5760c 100644
--- a/codewiki/cli/config_manager.py
+++ b/codewiki/cli/config_manager.py
@@ -147,6 +147,7 @@ def save(
use_gitignore: Optional[bool] = None,
prompt_caching: Optional[bool] = None,
pricing: Optional[dict] = None,
+ min_token_per_module: Optional[int] = None,
):
"""
Save configuration to file and keyring.
@@ -159,6 +160,7 @@ def save(
fallback_model: Fallback model
default_output: Default output directory
max_tokens: Maximum tokens for LLM response
+ min_token_per_module: Minimum code tokens for a standalone module
max_token_per_module: Maximum tokens per module for clustering
max_token_per_leaf_module: Maximum tokens per leaf module
max_depth: Maximum depth for hierarchical decomposition
@@ -220,6 +222,14 @@ def save(
self._config.default_output = default_output
if max_tokens is not None:
self._config.max_tokens = max_tokens
+ if min_token_per_module is not None:
+ from codewiki.src.config import validate_min_module_tokens
+
+ try:
+ validate_min_module_tokens(min_token_per_module)
+ except ValueError as exc:
+ raise ConfigurationError(str(exc)) from exc
+ self._config.min_token_per_module = min_token_per_module
if max_token_per_module is not None:
self._config.max_token_per_module = max_token_per_module
if max_token_per_leaf_module is not None:
diff --git a/codewiki/cli/models/config.py b/codewiki/cli/models/config.py
index 6539e4c1..07c7982b 100644
--- a/codewiki/cli/models/config.py
+++ b/codewiki/cli/models/config.py
@@ -9,6 +9,11 @@
from dataclasses import dataclass, field
from codewiki.src.language import format_language_directive
+from codewiki.src.config import (
+ DEFAULT_MAX_TOKEN_PER_LEAF_MODULE,
+ DEFAULT_MIN_TOKEN_PER_MODULE,
+ validate_min_module_tokens,
+)
from codewiki.cli.utils.validation import (
validate_model_name,
validate_url,
@@ -154,8 +159,9 @@ class Configuration:
google_project: Google Cloud project for google-vertex
google_location: Vertex AI location for google-vertex
max_tokens: Maximum tokens for LLM response (default: 32768)
+ min_token_per_module: Minimum code tokens per standalone module (default: 1000)
max_token_per_module: Maximum tokens per module for clustering (default: 36369)
- max_token_per_leaf_module: Maximum tokens per leaf module (default: 16000)
+ max_token_per_leaf_module: Base token threshold to allow delegation (default: 8000)
max_depth: Maximum depth for hierarchical decomposition (default: 2)
request_limit: Maximum model requests per agent run (default: 100)
agent_retries: Retries for a failing tool call (default: 3)
@@ -176,8 +182,9 @@ class Configuration:
google_project: str = ""
google_location: str = ""
max_tokens: int = 32768
+ min_token_per_module: int = DEFAULT_MIN_TOKEN_PER_MODULE
max_token_per_module: int = 36369
- max_token_per_leaf_module: int = 16000
+ max_token_per_leaf_module: int = DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
max_depth: int = 2
request_limit: int = 100
agent_retries: int = 3
@@ -192,6 +199,7 @@ def __post_init__(self):
try:
validate_pricing(self.pricing)
+ validate_min_module_tokens(self.min_token_per_module)
except ValueError as exc:
raise ConfigurationError(str(exc)) from exc
@@ -207,6 +215,8 @@ def validate(self):
"""
from codewiki.src.be.backend import is_caw_provider, needs_base_url
+ validate_min_module_tokens(self.min_token_per_module)
+
if is_caw_provider(self.provider):
validate_model_name(self.main_model)
return
@@ -230,6 +240,7 @@ def to_dict(self) -> dict:
"google_project": self.google_project,
"google_location": self.google_location,
"max_tokens": self.max_tokens,
+ "min_token_per_module": self.min_token_per_module,
"max_token_per_module": self.max_token_per_module,
"max_token_per_leaf_module": self.max_token_per_leaf_module,
"max_depth": self.max_depth,
@@ -272,8 +283,11 @@ def from_dict(cls, data: dict) -> "Configuration":
google_project=data.get("google_project", ""),
google_location=data.get("google_location", ""),
max_tokens=data.get("max_tokens", 32768),
+ min_token_per_module=data.get("min_token_per_module", DEFAULT_MIN_TOKEN_PER_MODULE),
max_token_per_module=data.get("max_token_per_module", 36369),
- max_token_per_leaf_module=data.get("max_token_per_leaf_module", 16000),
+ max_token_per_leaf_module=data.get(
+ "max_token_per_leaf_module", DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
+ ),
max_depth=data.get("max_depth", 2),
request_limit=data.get("request_limit", 100),
agent_retries=data.get("agent_retries", 3),
@@ -342,6 +356,7 @@ def to_backend_config(
google_project=self.google_project,
google_location=self.google_location,
max_tokens=self.max_tokens,
+ min_token_per_module=self.min_token_per_module,
max_token_per_module=self.max_token_per_module,
max_token_per_leaf_module=self.max_token_per_leaf_module,
max_depth=self.max_depth,
diff --git a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py
index 29ff31d9..71903d85 100644
--- a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py
+++ b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py
@@ -13,11 +13,10 @@
format_system_prompt,
format_user_prompt,
)
-from codewiki.src.be.utils import is_complex_module
+from codewiki.src.config import DEFAULT_MIN_TOKEN_PER_MODULE
from codewiki.src.be.module_scope import (
- component_files,
- get_module_token_count,
- prepare_submodule_partition,
+ can_split_module,
+ prepare_sized_submodule_partition,
)
import logging
@@ -38,16 +37,30 @@ async def generate_sub_module_documentation(
Sub-module names must be unique across the whole wiki. A name already used by another module is
prefixed with the current module name; a name that is already documented (plain or prefixed)
is SKIPPED, not regenerated. The tool result reports the final file names actually saved and
- every skipped request.
+ every skipped request. Children below the configured minimum code-token size are
+ merged into a nearby sibling. Partitions with fewer than two viable children are
+ rejected; document the module directly instead.
"""
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)
+ if not can_split_module(
+ deps.components,
+ deps.module_component_ids(),
+ deps.config,
+ deps.current_depth,
+ max_depth=deps.max_depth,
+ ):
+ return skipped_report(
+ {
+ name: "module cannot delegate: file count, token threshold or depth limit; "
+ "write this module directly"
+ for name in sub_module_specs
+ },
+ module_doc_file(
+ previous_module_name, deps.path_to_current_module, config_layout(deps.config)
+ ),
+ )
# Resolve name collisions against the module tree and files already on disk
# before touching the tree (issue #76): names are unique across the wiki.
@@ -74,16 +87,23 @@ async def generate_sub_module_documentation(
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
+ final_specs, rejected = prepare_sized_submodule_partition(
+ final_specs,
+ deps.module_component_ids(),
+ deps.components,
+ value,
+ getattr(deps.config, "min_token_per_module", DEFAULT_MIN_TOKEN_PER_MODULE),
+ deps.config.max_token_per_module,
)
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"
+ plan.skipped[requested] = rejected[final_name]
del name_map[requested]
if not final_specs:
return skipped_report(plan.skipped, parent_doc)
+ fallback_models = create_fallback_models(deps.config)
+
# 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": {}}
@@ -95,15 +115,14 @@ async def generate_sub_module_documentation(
logger.info(f"{indent}{arrow} Generating documentation for sub-module: {sub_module_name}")
- 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 child_depth < ctx.deps.max_depth
- and num_tokens >= ctx.deps.config.max_token_per_leaf_module * child_depth
+ if can_split_module(
+ deps.components,
+ core_component_ids,
+ deps.config,
+ child_depth,
+ max_depth=deps.max_depth,
):
sub_agent = Agent(
model=fallback_models,
diff --git a/codewiki/src/be/caw_backend.py b/codewiki/src/be/caw_backend.py
index 47075811..b8e5de4a 100644
--- a/codewiki/src/be/caw_backend.py
+++ b/codewiki/src/be/caw_backend.py
@@ -33,9 +33,8 @@
from codewiki.src.be.agent_tools.doc_guard import DocumentationWriteError, doc_written
from codewiki.src.be.backend import AgentReply, LLMBackend
from codewiki.src.be.module_scope import (
+ can_split_module,
complete_file_candidates,
- component_files,
- get_module_token_count,
validate_module_file_coverage,
)
from codewiki.src.be.dependency_analyzer.models.core import Node
@@ -44,7 +43,7 @@
format_system_prompt,
format_user_prompt,
)
-from codewiki.src.be.utils import is_complex_module, set_main_loop
+from codewiki.src.be.utils import set_main_loop
from codewiki.src.be.doc_layout import config_layout, find_doc, module_doc_file
from codewiki.src.be.overview_support import reusable_overview, root_recovery_prompt
from codewiki.src.config import MODULE_TREE_FILENAME, OVERVIEW_FILENAME, Config
@@ -409,19 +408,8 @@ def _run_module_agent_sync(
if not module_path:
core_component_ids = complete_file_candidates(core_component_ids, components)
- # 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 (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}"
- )
+ can_delegate = can_split_module(components, core_component_ids, config, start_depth)
+ logger.info("Module %s can delegate: %s (depth %d)", module_name, can_delegate, start_depth)
if can_delegate:
system_prompt = format_system_prompt(
diff --git a/codewiki/src/be/caw_toolkit.py b/codewiki/src/be/caw_toolkit.py
index 19f06b9e..040585bb 100644
--- a/codewiki/src/be/caw_toolkit.py
+++ b/codewiki/src/be/caw_toolkit.py
@@ -30,7 +30,8 @@
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
+from codewiki.src.be.module_scope import can_split_module, prepare_sized_submodule_partition
+from codewiki.src.config import DEFAULT_MIN_TOKEN_PER_MODULE
if TYPE_CHECKING:
from codewiki.src.be.caw_backend import CawBackend
@@ -260,7 +261,9 @@ async def str_replace_editor(
"Sub-module names must be unique across the whole wiki. A name already used by another "
"module is prefixed with the current module name; a name that is already documented "
"(plain or prefixed) is SKIPPED, not regenerated. The result reports the files actually "
- "saved and every skipped request."
+ "saved and every skipped request. Children below the configured minimum code-token "
+ "size are merged into a nearby sibling. Partitions with fewer than two viable "
+ "children are rejected; document the module directly instead."
)
)
async def generate_sub_module_documentation(
@@ -298,6 +301,23 @@ async def generate_sub_module_documentation(
def _run_sub_modules(self, sub_module_specs: dict[str, list[str]]) -> str:
deps = self._deps
previous_module_name = deps.current_module_name
+ if not can_split_module(
+ deps.components,
+ deps.module_component_ids(),
+ deps.config,
+ deps.current_depth,
+ max_depth=deps.max_depth,
+ ):
+ return skipped_report(
+ {
+ name: "module cannot delegate: file count, token threshold or depth limit; "
+ "write this module directly"
+ for name in sub_module_specs
+ },
+ module_doc_file(
+ previous_module_name, deps.path_to_current_module, config_layout(deps.config)
+ ),
+ )
# Resolve name collisions against the module tree and files already on
# disk before touching the tree (issue #76): names are unique across the wiki.
@@ -326,12 +346,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
+ final_specs, rejected = prepare_sized_submodule_partition(
+ final_specs,
+ deps.module_component_ids(),
+ deps.components,
+ value,
+ getattr(deps.config, "min_token_per_module", DEFAULT_MIN_TOKEN_PER_MODULE),
+ deps.config.max_token_per_module,
)
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"
+ plan.skipped[requested] = rejected[final_name]
del name_map[requested]
if not final_specs:
return skipped_report(
diff --git a/codewiki/src/be/cluster_modules.py b/codewiki/src/be/cluster_modules.py
index 0a31d25b..12659baf 100644
--- a/codewiki/src/be/cluster_modules.py
+++ b/codewiki/src/be/cluster_modules.py
@@ -10,10 +10,12 @@
from codewiki.src.be.module_naming import resolve_unique_name, sanitize_module_name
from codewiki.src.be.module_scope import (
complete_file_candidates,
+ coalesce_small_modules,
component_files,
file_token_counts,
get_module_token_count,
normalize_module_partition,
+ nearest_module,
validate_module_file_coverage,
)
from codewiki.src.be.prompt_template import format_cluster_prompt, format_super_group_prompt
@@ -21,6 +23,7 @@
from codewiki.src.config import (
DEFAULT_MAX_LEAF_NODES_PER_CLUSTER,
DEFAULT_MIN_MODULES_FOR_SUPER_GROUPING,
+ DEFAULT_MIN_TOKEN_PER_MODULE,
Config,
)
@@ -271,9 +274,12 @@ def _batch_fallback_name(
def _structural_fallback(
- candidates: list[str], components: dict[str, Node], threshold: int
+ candidates: list[str],
+ components: dict[str, Node],
+ threshold: int,
+ min_tokens: int = DEFAULT_MIN_TOKEN_PER_MODULE,
) -> dict[str, Any]:
- """Pack path-adjacent whole files by code size when the LLM cannot split."""
+ """Pack whole files by code size, or return {} when no split remains."""
by_file = defaultdict(list)
for cid in candidates:
by_file[components[cid].relative_path].append(cid)
@@ -298,7 +304,8 @@ def _structural_fallback(
"components": group,
"children": {},
}
- return tree
+ tree, _ = coalesce_small_modules(tree, components, min_tokens, threshold)
+ return tree if len(tree) > 1 else {}
def cluster_modules(
@@ -344,6 +351,7 @@ def cluster_modules(
leaf_files = component_files(leaves, components)
coverage = len(leaf_files) / len(scope_files) if scope_files else 1.0
threshold = config.max_token_per_module
+ minimum = config.min_token_per_module
module_label = current_module_name or "repository"
logger.info(
@@ -370,10 +378,11 @@ def cluster_modules(
len(candidates),
)
- if input_tokens <= threshold:
+ if input_tokens <= threshold or input_tokens < 2 * minimum:
logger.info(
"Skipping LLM module clustering for %s because %d code tokens fit within the "
- "%d-token threshold; using whole-module documentation mode.",
+ "%d-token threshold or cannot form two minimum-sized modules; "
+ "using whole-module documentation mode.",
module_label,
input_tokens,
threshold,
@@ -459,7 +468,16 @@ def cluster_modules(
module_label,
len(module_tree),
)
- module_tree = _structural_fallback(candidates, components, threshold)
+ module_tree = _structural_fallback(candidates, components, threshold, minimum)
+
+ module_tree, _ = coalesce_small_modules(module_tree, components, minimum, threshold)
+ if len(module_tree) < 2:
+ logger.info(
+ "Keeping %s intact: fewer than two modules remain after applying the %d-token minimum.",
+ module_label,
+ minimum,
+ )
+ return {}
logger.info(
"Module clustering for %s produced %d top-level modules.",
@@ -699,14 +717,16 @@ def ensure_artifact_module(
leaf_nodes: list[str],
components: dict[str, Node],
min_share: float = ARTIFACT_MIN_SHARE,
+ min_tokens: int = DEFAULT_MIN_TOKEN_PER_MODULE,
) -> dict[str, Any]:
"""Guarantee that artifact leaf nodes are documented.
Clustering is an LLM call and may drop or scatter artifact nodes despite
the prompt. If fewer than ``min_share`` of the artifact leaf nodes landed
in ``module_tree``, add a fixed top-level module holding every unassigned
- artifact node. Returns ``module_tree`` unchanged in whole-repository mode
- (empty tree: one agent documents all leaf nodes anyway).
+ artifact node, or merge a bundle below ``min_tokens`` into an existing
+ leaf. Returns ``module_tree`` unchanged in whole-repository mode (empty
+ tree: one agent documents all leaf nodes anyway).
"""
if not module_tree:
return module_tree
@@ -737,6 +757,53 @@ def _order(node_id: str):
return (rank, node.relative_path, node_id)
unassigned.sort(key=_order)
+ # Resolve existing file owners first: an artifact entry point may share
+ # a file with a source component that the tree already lists.
+ leaves: dict[str, dict[str, Any]] = {}
+ ancestors: dict[str, list[dict[str, Any]]] = {}
+
+ def collect(level: dict[str, Any], parents: list[dict[str, Any]], prefix: str) -> None:
+ for name, info in level.items():
+ key = f"{prefix}/{name}"
+ if info.get("children"):
+ collect(info["children"], parents + [info], key)
+ else:
+ leaves[key] = info
+ ancestors[key] = parents
+
+ def attach(target: str, ids: list[str]) -> None:
+ # Ancestors aggregate their descendants' scope, including on resume.
+ for info in [*ancestors[target], leaves[target]]:
+ info["components"] = sorted(set(info["components"]) | set(ids))
+ info["path"] = _common_path_prefix(
+ sorted(component_files(info["components"], components))
+ )
+
+ collect(module_tree, [], "")
+ scopes = {
+ name: component_files(info["components"], components) for name, info in leaves.items()
+ }
+ owners = {path: name for name, paths in scopes.items() for path in paths}
+ new_files = []
+ for cid in unassigned:
+ owner = owners.get(components[cid].relative_path)
+ if owner is not None:
+ attach(owner, [cid])
+ else:
+ new_files.append(cid)
+ unassigned = new_files
+ if not unassigned:
+ return module_tree
+
+ unassigned_files = component_files(unassigned, components)
+ tokens = get_module_token_count(components, unassigned_files)
+ if tokens < min_tokens:
+ sizes = file_token_counts(components)
+ loads = {name: sum(sizes[path] for path in paths) for name, paths in scopes.items()}
+ target = nearest_module(unassigned_files, scopes, loads)
+ attach(target, unassigned)
+ logger.info("Merged %d artifact tokens into existing leaf '%s'", tokens, target)
+ return module_tree
name = resolve_unique_name(ARTIFACT_MODULE_NAME, None, collect_module_tree_names(module_tree))
module_tree[name] = {
"path": _common_path_prefix([components[n].relative_path for n in unassigned]) or ".",
diff --git a/codewiki/src/be/documentation_generator.py b/codewiki/src/be/documentation_generator.py
index 84e0a6b8..b920d10a 100644
--- a/codewiki/src/be/documentation_generator.py
+++ b/codewiki/src/be/documentation_generator.py
@@ -533,7 +533,12 @@ async def run(self) -> None:
# Artifact nodes the clustering LLM dropped get a fixed module
# so build/CI/config coverage does not depend on the LLM.
if getattr(self.config, "artifacts_enabled", True):
- module_tree = ensure_artifact_module(module_tree, leaf_nodes, components)
+ module_tree = ensure_artifact_module(
+ module_tree,
+ leaf_nodes,
+ components,
+ min_tokens=self.config.min_token_per_module,
+ )
# Only freshly clustered trees are deduped: renaming a cached
# key whose .md already exists would orphan the doc.
module_tree = dedupe_module_tree_names(module_tree)
diff --git a/codewiki/src/be/module_naming.py b/codewiki/src/be/module_naming.py
index f64a8f5f..11387536 100644
--- a/codewiki/src/be/module_naming.py
+++ b/codewiki/src/be/module_naming.py
@@ -209,7 +209,9 @@ def skipped_report(skipped: dict[str, str], parent_doc: str) -> str:
items = ", ".join(f"'{name}' ({reason})" for name, reason in skipped.items())
return (
f"Skipped sub-modules: {items}. Do NOT call generate_sub_module_documentation again "
- f"for these; link the existing pages from `{parent_doc}` instead."
+ f"for these. Link existing pages only when reported as already documented; "
+ f"document undelegated files directly in `{parent_doc}`. "
+ "For merged requests, use the surviving module page reported above."
)
diff --git a/codewiki/src/be/module_scope.py b/codewiki/src/be/module_scope.py
index 5f6df990..7ee5c662 100644
--- a/codewiki/src/be/module_scope.py
+++ b/codewiki/src/be/module_scope.py
@@ -15,6 +15,7 @@
from codewiki.src.be.dependency_analyzer.models.core import Node
from codewiki.src.be.utils import count_tokens
+from codewiki.src.config import DEFAULT_MIN_TOKEN_PER_MODULE, Config
logger = logging.getLogger(__name__)
@@ -39,6 +40,104 @@ def get_module_token_count(components: dict[str, Node], scope_files: set[str] |
)
+def can_split_module(
+ components: dict[str, Node],
+ ids: list[str],
+ config: Config,
+ depth: int,
+ *,
+ max_depth: int | None = None,
+) -> bool:
+ """Use the same cost and depth gate for every documentation agent."""
+ files = component_files(ids, components)
+ limit = config.max_depth if max_depth is None else max_depth
+ minimum = getattr(config, "min_token_per_module", DEFAULT_MIN_TOKEN_PER_MODULE)
+ return (
+ len(files) > 1
+ and depth < limit
+ and get_module_token_count(components, files)
+ >= max(2 * minimum, config.max_token_per_leaf_module * depth)
+ )
+
+
+def _directory_proximity(left: str, right: str) -> int:
+ depth = 0
+ for a, b in zip(left.split("/")[:-1], right.split("/")[:-1]):
+ if a != b:
+ break
+ depth += 1
+ return depth
+
+
+def nearest_module(
+ files: set[str],
+ groups: dict[str, set[str]],
+ loads: dict[str, int],
+ *,
+ tokens: int = 0,
+ max_tokens: int | None = None,
+) -> str:
+ """Prefer a fitting sibling, then directory proximity, load and name."""
+ return min(
+ groups,
+ key=lambda name: (
+ max_tokens is not None and loads[name] + tokens > max_tokens,
+ -max(
+ (_directory_proximity(left, right) for left in files for right in groups[name]),
+ default=0,
+ ),
+ loads[name],
+ name,
+ ),
+ )
+
+
+def coalesce_small_modules(
+ tree: dict[str, Any],
+ components: dict[str, Node],
+ min_tokens: int,
+ max_tokens: int | None = None,
+) -> tuple[dict[str, Any], dict[str, str]]:
+ """Coalesce a flat file partition before children or pages are created.
+
+ Returns the partition and reasons for removed names. A singleton is left
+ for the caller to reject as a non-split. No input dictionaries are mutated.
+ """
+ result = {name: {**info, "components": list(info["components"])} for name, info in tree.items()}
+ sizes = file_token_counts(components)
+ files = {name: component_files(info["components"], components) for name, info in result.items()}
+ loads = {name: sum(sizes[path] for path in paths) for name, paths in files.items()}
+ merged: dict[str, tuple[int, str]] = {}
+ while len(result) > 1:
+ donor = min(result, key=lambda name: (loads[name], name))
+ if loads[donor] >= min_tokens:
+ break
+ target = nearest_module(
+ files[donor],
+ {name: paths for name, paths in files.items() if name != donor},
+ loads,
+ tokens=loads[donor],
+ max_tokens=max_tokens,
+ )
+ merged[donor] = (loads[donor], target)
+ files[target].update(files.pop(donor))
+ loads[target] += loads.pop(donor)
+ result[target]["components"] = sorted(
+ set(result[target]["components"]) | set(result.pop(donor)["components"]),
+ key=lambda cid: (components[cid].relative_path, cid),
+ )
+ result[target]["path"] = posixpath.commonpath(sorted(files[target]))
+ logger.info(
+ "Merged undersized module '%s' into '%s' (minimum %d tokens)", donor, target, min_tokens
+ )
+ reasons = {}
+ for donor, (tokens, target) in merged.items():
+ while target in merged:
+ target = merged[target][1]
+ reasons[donor] = f"{tokens} code tokens below minimum {min_tokens}; merged into '{target}'"
+ return result, reasons
+
+
def complete_file_candidates(
ids: list[str], components: dict[str, Node], scope_files: set[str] | None = None
) -> list[str]:
@@ -194,3 +293,37 @@ def prepare_submodule_partition(
{name: {"components": ids} for name, ids in specs.items()}, candidates, components
)
return {name: info["components"] for name, info in normalized.items()}
+
+
+def prepare_sized_submodule_partition(
+ specs: dict[str, list[str]],
+ parent_ids: list[str],
+ components: dict[str, Node],
+ existing_children: dict[str, Any],
+ min_tokens: int,
+ max_tokens: int | None = None,
+) -> tuple[dict[str, list[str]], dict[str, str]]:
+ """Plan a meaningful delegation, including explanations before any mutation."""
+ partition = prepare_submodule_partition(specs, parent_ids, components, existing_children)
+ tree, reasons = coalesce_small_modules(
+ {name: {"components": ids} for name, ids in partition.items()},
+ components,
+ min_tokens,
+ max_tokens,
+ )
+ for name in specs.keys() - partition.keys():
+ reasons[name] = "no unassigned files in the current module scope"
+ if len(tree) < 2:
+ for name in partition:
+ tokens = get_module_token_count(
+ components, component_files(partition[name], components)
+ )
+ detail = (
+ f"{tokens} code tokens below minimum {min_tokens}; " if tokens < min_tokens else ""
+ )
+ reasons[name] = (
+ detail + "partition has fewer than two viable children after applying the "
+ f"{min_tokens}-token minimum; no child pages created, write these files directly in the parent"
+ )
+ return {}, reasons
+ return {name: info["components"] for name, info in tree.items()}, reasons
diff --git a/codewiki/src/be/prompt_template.py b/codewiki/src/be/prompt_template.py
index 8e65fcb0..db16a55d 100644
--- a/codewiki/src/be/prompt_template.py
+++ b/codewiki/src/be/prompt_template.py
@@ -42,6 +42,7 @@
1. Analyze the provided code components and module structure, explore the not given dependencies between the components if needed
2. Create the main `{doc_file}` file with overview and architecture in working directory
3. Use `generate_sub_module_documentation` to generate detailed sub-modules documentation for COMPLEX modules which at least have more than 1 code file and are able to clearly split into sub-modules. Sub-module names must be unique across the whole wiki ({uniqueness_reason}) — prefer names prefixed with the current module name, e.g. `{module_name}_search`
+ The tool enforces minimum module sizes and rejects partitions with fewer than two viable children. Small groups may be merged into siblings; use only the final pages reported by the tool. If delegation is rejected, document the code directly in the current page.
4. Include relevant Mermaid diagrams throughout the documentation
5. After all sub-modules are documented, adjust `{doc_file}` with ONLY ONE STEP to ensure all generated files including sub-modules documentation are properly cross-refered, using the final file names reported by `generate_sub_module_documentation`
diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py
index 7d5bed95..df0c04a1 100644
--- a/codewiki/src/be/pydantic_ai_backend.py
+++ b/codewiki/src/be/pydantic_ai_backend.py
@@ -28,14 +28,17 @@
from codewiki.src.be.backend import AgentReply, LLMBackend, UsageLedger, usage_to_dict
from codewiki.src.be.dependency_analyzer.models.core import Node
from codewiki.src.be.llm_services import call_llm, create_fallback_models, pop_last_usage
-from codewiki.src.be.module_scope import complete_file_candidates, validate_module_file_coverage
+from codewiki.src.be.module_scope import (
+ can_split_module,
+ complete_file_candidates,
+ validate_module_file_coverage,
+)
from codewiki.src.be.overview_support import reusable_overview, root_recovery_prompt
from codewiki.src.be.prompt_template import (
format_leaf_system_prompt,
format_system_prompt,
format_user_prompt,
)
-from codewiki.src.be.utils import is_complex_module
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
@@ -149,7 +152,7 @@ async def run_module_agent(
if not module_path:
core_component_ids = complete_file_candidates(core_component_ids, components)
- if is_complex_module(components, core_component_ids) and config.max_depth > 1:
+ if can_split_module(components, core_component_ids, config, 1):
agent = Agent(
self._fallback_models,
name=module_name,
diff --git a/codewiki/src/config.py b/codewiki/src/config.py
index 338a062b..e5a32d1e 100644
--- a/codewiki/src/config.py
+++ b/codewiki/src/config.py
@@ -25,8 +25,9 @@
MAX_DEPTH = 2
# Default max token settings
DEFAULT_MAX_TOKENS = 32_768
+DEFAULT_MIN_TOKEN_PER_MODULE = 1_000
DEFAULT_MAX_TOKEN_PER_MODULE = 36_369
-DEFAULT_MAX_TOKEN_PER_LEAF_MODULE = 4_000
+DEFAULT_MAX_TOKEN_PER_LEAF_MODULE = 8_000
# pydantic-ai caps each agent run at 50 model requests by default; complex
# modules that read many components and delegate sub-modules can need more.
DEFAULT_REQUEST_LIMIT = 100
@@ -76,6 +77,11 @@ def is_cli_context() -> bool:
ATLAS_CLOUD_BASE_URL = "https://api.atlascloud.ai/v1"
+def validate_min_module_tokens(value: int) -> None:
+ if type(value) is not int or value < 1:
+ raise ValueError("min_token_per_module must be a positive integer")
+
+
@dataclass
class Config:
"""Configuration class for CodeWiki."""
@@ -105,6 +111,7 @@ class Config:
google_location: str = ""
# Max token settings
max_tokens: int = DEFAULT_MAX_TOKENS
+ min_token_per_module: int = DEFAULT_MIN_TOKEN_PER_MODULE
max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE
max_token_per_leaf_module: int = DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
min_modules_for_super_grouping: int = DEFAULT_MIN_MODULES_FOR_SUPER_GROUPING
@@ -130,6 +137,9 @@ class Config:
# Docs layout (LAYOUT_HIERARCHICAL or LAYOUT_FLAT)
layout: str = DEFAULT_LAYOUT
+ def __post_init__(self):
+ validate_min_module_tokens(self.min_token_per_module)
+
@property
def artifact_exclude(self) -> list[str] | None:
"""Extra patterns excluded from artifact analysis (from agent instructions)."""
@@ -229,6 +239,9 @@ def from_args(cls, args: argparse.Namespace) -> "Config":
main_model=MAIN_MODEL,
cluster_model=CLUSTER_MODEL,
fallback_model=FALLBACK_MODEL_1,
+ min_token_per_module=getattr(
+ args, "min_token_per_module", DEFAULT_MIN_TOKEN_PER_MODULE
+ ),
use_gitignore=getattr(args, "use_gitignore", True),
layout=getattr(args, "layout", DEFAULT_LAYOUT),
)
@@ -265,6 +278,7 @@ def from_cli(
with_prose: bool = False,
layout: str = DEFAULT_LAYOUT,
pricing: dict | None = None,
+ min_token_per_module: int = DEFAULT_MIN_TOKEN_PER_MODULE,
) -> "Config":
"""
Create configuration for CLI context.
@@ -285,6 +299,7 @@ def from_cli(
google_project: Google Cloud project for the google-vertex provider
google_location: Vertex AI location for the google-vertex provider
max_tokens: Maximum tokens for LLM response
+ min_token_per_module: Minimum code tokens for a standalone module
max_token_per_module: Maximum tokens per module for clustering
max_token_per_leaf_module: Maximum tokens per leaf module
min_modules_for_super_grouping: Super-group the top level into
@@ -330,6 +345,7 @@ def from_cli(
google_project=google_project,
google_location=google_location,
max_tokens=max_tokens,
+ min_token_per_module=min_token_per_module,
max_token_per_module=max_token_per_module,
max_token_per_leaf_module=max_token_per_leaf_module,
min_modules_for_super_grouping=min_modules_for_super_grouping,
diff --git a/guides/cli-reference.md b/guides/cli-reference.md
index 3c499237..03284f3f 100644
--- a/guides/cli-reference.md
+++ b/guides/cli-reference.md
@@ -138,10 +138,23 @@ Override the stored limits for one run.
| Flag | Stored default | Meaning |
| --- | --- | --- |
| `--max-tokens N` | `32768` | Maximum output tokens per LLM response |
+| `--min-token-per-module N` | `1000` | Minimum code tokens per standalone module; smaller groups merge into nearby siblings |
| `--max-token-per-module N` | `36369` | Input-token threshold that triggers module clustering |
-| `--max-token-per-leaf-module N` | `16000` | Input-token threshold for leaf modules |
+| `--max-token-per-leaf-module N` | `8000` | Base code-token threshold to allow an agent to delegate |
| `--max-depth N` | `2` | Maximum depth of the hierarchical decomposition |
+Module sizes include all indexed components in the files a module owns, using
+the GPT-4 tokenizer. Both agent backends apply the minimum to new child pages
+and reject a partition that collapses to one child. Delegation requires at
+least `max(2 * min_token_per_module, max_token_per_leaf_module * depth)` code
+tokens, multiple files, and a depth below `max_depth`. Delegation depth starts
+at 1, so the default token gates are 8,000, 16,000, 24,000, and so on; reaching
+the gate allows delegation without requiring it. A small repository stays
+on one overview page. Whole files remain indivisible, so an oversized file
+can exceed the upper threshold, including nearby small files merged into it.
+Cached module trees are retained; generate into a fresh output directory to
+apply these rules throughout an existing tree.
+
### Examples
```bash
@@ -178,7 +191,7 @@ keys you pass are changed. Provider examples are in [Providers](providers.md).
| `--azure-deployment NAME` | Azure OpenAI only |
| `--google-project ID` | `google-vertex` only. Defaults to `GOOGLE_CLOUD_PROJECT` |
| `--google-location LOC` | `google-vertex` only, e.g. `global`, `eu`, `us`. Defaults to `GOOGLE_CLOUD_LOCATION`, then `global` |
-| `--max-tokens N`, `--max-token-per-module N`, `--max-token-per-leaf-module N`, `--max-depth N` | Stored token limits, see the table above |
+| `--max-tokens N`, `--min-token-per-module N`, `--max-token-per-module N`, `--max-token-per-leaf-module N`, `--max-depth N` | Stored token limits, see the table above |
| `--use-gitignore` / `--no-gitignore` | Stored default for Git ignore handling |
| `--prompt-caching` / `--no-prompt-caching` | Stored default for prompt caching |
diff --git a/tests/test_agent_limits.py b/tests/test_agent_limits.py
index 21c004df..cf084b37 100644
--- a/tests/test_agent_limits.py
+++ b/tests/test_agent_limits.py
@@ -106,22 +106,31 @@ async def run(self, prompt, deps, **kwargs):
relative_path="b.py",
source_code=source.read_text(),
)
+ other = component.model_copy(
+ update={"id": "c.py::fc", "file_path": "c.py", "relative_path": "c.py"}
+ )
deps = CodeWikiDeps(
absolute_docs_path=str(tmp_path),
absolute_repo_path=str(tmp_path),
registry={},
- components={component.id: component},
+ components={component.id: component, other.id: other},
path_to_current_module=[],
current_module_name="root",
module_tree={},
max_depth=2,
current_depth=1,
- config=SimpleNamespace(max_token_per_leaf_module=4000, agent_retries=7, request_limit=321),
+ config=SimpleNamespace(
+ min_token_per_module=1,
+ max_token_per_module=36369,
+ max_token_per_leaf_module=1,
+ agent_retries=7,
+ request_limit=321,
+ ),
custom_instructions=None,
)
asyncio.run(
sub_mod.generate_sub_module_documentation(
- SimpleNamespace(deps=deps), {"billing": ["b.py::fb"]}
+ SimpleNamespace(deps=deps), {"billing": ["b.py::fb"], "other": [other.id]}
)
)
assert seen["init"]["retries"] == 7
diff --git a/tests/test_artifact_analyzer.py b/tests/test_artifact_analyzer.py
index b9009cd8..6ba521a2 100644
--- a/tests/test_artifact_analyzer.py
+++ b/tests/test_artifact_analyzer.py
@@ -439,19 +439,19 @@ def test_ensure_artifact_module() -> None:
art_ids = [i for i in leaves if i.startswith("ci/")]
# 1/5 assigned -> module inserted with the other 4
tree = {"Core": {"path": "src", "components": ["src/a.py::A", art_ids[0]], "children": {}}}
- out = ensure_artifact_module(tree, leaves, comps)
+ out = ensure_artifact_module(tree, leaves, comps, min_tokens=1)
assert ARTIFACT_MODULE_NAME in out
mod = out[ARTIFACT_MODULE_NAME]
assert set(mod["components"]) == set(art_ids[1:])
assert mod["children"] == {} and mod["path"]
# 4/5 assigned -> unchanged
tree = {"Core": {"path": "src", "components": ["src/a.py::A", *art_ids[:4]], "children": {}}}
- assert ARTIFACT_MODULE_NAME not in ensure_artifact_module(tree, leaves, comps)
+ assert ARTIFACT_MODULE_NAME not in ensure_artifact_module(tree, leaves, comps, min_tokens=1)
# whole-repo mode -> unchanged
assert ensure_artifact_module({}, leaves, comps) == {}
# name collision -> unique variant
tree = {ARTIFACT_MODULE_NAME: {"path": "", "components": ["src/a.py::A"], "children": {}}}
- out = ensure_artifact_module(tree, leaves, comps)
+ out = ensure_artifact_module(tree, leaves, comps, min_tokens=1)
inserted = [k for k in out if k != ARTIFACT_MODULE_NAME]
assert len(inserted) == 1 and inserted[0].startswith(ARTIFACT_MODULE_NAME)
diff --git a/tests/test_cluster_partitioning.py b/tests/test_cluster_partitioning.py
index cdd435d9..1d85ceb0 100644
--- a/tests/test_cluster_partitioning.py
+++ b/tests/test_cluster_partitioning.py
@@ -50,6 +50,7 @@ def _config(**overrides) -> Config:
llm_api_key="key",
main_model="m",
cluster_model="m",
+ min_token_per_module=1,
)
return Config.from_cli(**kwargs, **overrides)
diff --git a/tests/test_doc_write_guard.py b/tests/test_doc_write_guard.py
index 77fb69b5..fb070596 100644
--- a/tests/test_doc_write_guard.py
+++ b/tests/test_doc_write_guard.py
@@ -262,10 +262,16 @@ def test_subagent_failure_restores_dependencies(tmp_path, monkeypatch):
source_code="class A: pass",
)
deps = _deps(tmp_path)
- deps.components = {node.id: node}
- deps.scope_component_ids = [node.id]
+ other = node.model_copy(update={"id": "b.py::B", "file_path": "b.py", "relative_path": "b.py"})
+ deps.components = {node.id: node, other.id: other}
+ deps.scope_component_ids = [node.id, other.id]
deps.config = SimpleNamespace(
- agent_retries=1, request_limit=10, max_token_per_leaf_module=100, layout="flat"
+ agent_retries=1,
+ request_limit=10,
+ max_token_per_leaf_module=1,
+ layout="flat",
+ min_token_per_module=1,
+ max_token_per_module=36369,
)
original_scope = deps.scope_component_ids
@@ -281,7 +287,9 @@ async def fail(*args, **kwargs):
monkeypatch.setattr(mod, "run_until_doc_written", fail)
with pytest.raises(DocumentationWriteError):
asyncio.run(
- mod.generate_sub_module_documentation(SimpleNamespace(deps=deps), {"child": [node.id]})
+ mod.generate_sub_module_documentation(
+ SimpleNamespace(deps=deps), {"child": [node.id], "other": [other.id]}
+ )
)
assert deps.current_module_name == "repo" and deps.current_depth == 1
assert deps.path_to_current_module == []
diff --git a/tests/test_min_module_size.py b/tests/test_min_module_size.py
new file mode 100644
index 00000000..b7345245
--- /dev/null
+++ b/tests/test_min_module_size.py
@@ -0,0 +1,412 @@
+"""Minimum sizes use full file scopes and are enforced before agents run."""
+
+from __future__ import annotations
+
+import asyncio
+from copy import deepcopy
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+from click.testing import CliRunner
+
+from codewiki.src.be.agent_tools.deps import CodeWikiDeps
+from codewiki.src.be.cluster_modules import (
+ _structural_fallback,
+ cluster_modules,
+ ensure_artifact_module,
+)
+from codewiki.src.be.dependency_analyzer.models.core import Node
+from codewiki.src.be.module_scope import (
+ can_split_module,
+ coalesce_small_modules,
+ component_files,
+ get_module_token_count,
+ prepare_sized_submodule_partition,
+ validate_module_file_coverage,
+)
+from codewiki.src.be.utils import count_tokens
+from codewiki.src.config import Config
+
+
+def node(path, tokens):
+ source = " x" * tokens
+ assert count_tokens(source) == tokens
+ cid = f"{path}::entry"
+ return Node(
+ id=cid,
+ name="entry",
+ component_type="function",
+ file_path=path,
+ relative_path=path,
+ source_code=source,
+ )
+
+
+def graph(**sizes):
+ return {n.id: n for path, tokens in sizes.items() for n in [node(path, tokens)]}
+
+
+def config(tmp_path, **overrides):
+ return Config.from_cli(
+ repo_path=str(tmp_path),
+ output_dir=str(tmp_path),
+ llm_base_url="http://localhost",
+ llm_api_key="test",
+ main_model="test",
+ cluster_model="test",
+ **overrides,
+ )
+
+
+def partition(components):
+ return {cid.split("::")[0]: {"components": [cid], "children": {}} for cid in components}
+
+
+@pytest.mark.parametrize(
+ "tokens,depth,expected",
+ [
+ (1999, 1, False),
+ (2000, 1, True),
+ (3999, 2, False),
+ (4000, 2, True),
+ (12000, 3, False),
+ ],
+)
+def test_shared_gate_respects_minimum_leaf_budget_and_depth(tmp_path, tokens, depth, expected):
+ components = graph(**{"a.py": tokens // 2, "b.py": tokens - tokens // 2})
+ cfg = config(tmp_path, max_depth=3, max_token_per_leaf_module=2000)
+ assert can_split_module(components, list(components), cfg, depth) is expected
+ one_file = {
+ cid: n.model_copy(update={"relative_path": "one.py"}) for cid, n in components.items()
+ }
+ assert not can_split_module(one_file, list(one_file), cfg, depth)
+
+
+@pytest.mark.parametrize("extension", ["c", "cpp", "cs", "go", "py", "ts"])
+def test_size_includes_unlisted_methods_in_owned_files(tmp_path, extension):
+ components = graph(**{f"a.{extension}": 10, f"b.{extension}": 1000})
+ entry_ids = list(components)
+ method = node(f"a.{extension}", 990).model_copy(update={"id": f"a.{extension}::method"})
+ components[method.id] = method
+ cfg = config(tmp_path, max_token_per_leaf_module=1)
+ assert can_split_module(components, entry_ids, cfg, 1)
+ specs, rejected = prepare_sized_submodule_partition(
+ {"a": [entry_ids[0]], "b": [entry_ids[1]]},
+ entry_ids,
+ components,
+ {},
+ 1000,
+ )
+ assert set(specs) == {"a", "b"} and rejected == {}
+
+
+@pytest.mark.parametrize("small,expected_count", [(999, 2), (1000, 3)])
+def test_coalesce_boundary_preserves_all_files_and_input(small, expected_count):
+ components = graph(**{"src/a.c": small, "src/b.c": 1000, "other/c.c": 1000})
+ original = partition(components)
+ before = deepcopy(original)
+ tree, rejected = coalesce_small_modules(original, components, 1000)
+ assert len(tree) == expected_count
+ assert original == before
+ assert all(
+ get_module_token_count(components, component_files(v["components"], components)) >= 1000
+ for v in tree.values()
+ )
+ validate_module_file_coverage(tree, components)
+ if small == 999:
+ assert "999 code tokens below minimum 1000" in rejected["src/a.c"]
+ assert "src/b.c" in rejected["src/a.c"]
+ reverse, _ = coalesce_small_modules(dict(reversed(list(original.items()))), components, 1000)
+ assert reverse == tree
+
+
+def test_coalesce_prefers_fitting_sibling_and_reports_final_destination():
+ components = graph(**{"src/tiny.c": 100, "src/huge.c": 4000, "other/room.c": 1000})
+ tree, rejected = coalesce_small_modules(partition(components), components, 1000, 4000)
+ assert "other/room.c" in rejected["src/tiny.c"]
+ validate_module_file_coverage(tree, components)
+ smalls = graph(**{"a.c": 100, "b.c": 200, "c.c": 800})
+ tree, rejected = coalesce_small_modules(partition(smalls), smalls, 1000)
+ assert set(tree) == {"c.c"}
+ assert all("merged into 'c.c'" in reason for reason in rejected.values())
+
+
+@pytest.mark.parametrize("mode", ["llm", "fallback"])
+def test_scanner_header_does_not_create_an_identical_child(tmp_path, mode):
+ components = graph(**{"src/scanners.c": 70645, "src/scanners.h": 879})
+ cfg = config(tmp_path)
+ calls = []
+
+ def complete(prompt):
+ calls.append(prompt)
+ return (
+ f"{partition(components)!r}"
+ if mode == "llm"
+ else None
+ )
+
+ tree = cluster_modules(list(components), components, cfg, completer=complete)
+ assert tree == {} and len(calls) == 1
+ fallback = _structural_fallback(list(components), components, cfg.max_token_per_module)
+ assert fallback == {}
+ packed, _ = coalesce_small_modules(partition(components), components, 1000)
+ assert len(packed) == 1
+ validate_module_file_coverage(packed, components)
+
+
+def test_clustering_coalesces_before_recursion_and_restores_coverage(tmp_path):
+ components = graph(**{"src/a.c": 999, "src/b.c": 2000, "other/c.c": 2000})
+ cfg = config(tmp_path, max_token_per_module=4000)
+ calls = []
+
+ def complete(prompt):
+ calls.append(prompt)
+ return f"{partition(components)!r}"
+
+ tree = cluster_modules(list(components), components, cfg, completer=complete)
+ assert len(tree) == 2 and len(calls) == 1
+ assert all(not v["children"] for v in tree.values())
+ validate_module_file_coverage(tree, components)
+
+
+def test_all_small_clusters_can_combine_into_two_viable_modules(tmp_path):
+ components = graph(**{f"src/{i}.c": 500 for i in range(4)})
+ cfg = config(tmp_path, max_token_per_module=1000)
+ tree = cluster_modules(
+ list(components),
+ components,
+ cfg,
+ completer=lambda p: f"{partition(components)!r}",
+ )
+ assert len(tree) == 2
+ validate_module_file_coverage(tree, components)
+
+
+def test_parent_too_small_for_two_children_skips_clustering(tmp_path):
+ components = graph(**{"a.c": 700, "b.c": 700})
+ calls = []
+ assert (
+ cluster_modules(
+ list(components),
+ components,
+ config(tmp_path, max_token_per_module=1000),
+ completer=lambda p: calls.append(p),
+ )
+ == {}
+ )
+ assert calls == []
+
+
+@pytest.fixture(params=["pydantic", "caw"])
+def delegate(request, tmp_path, monkeypatch):
+ runs = []
+ if request.param == "pydantic":
+ from codewiki.src.be.agent_tools import generate_sub_module_documentations as mod
+
+ class FakeAgent:
+ def __init__(self, *args, **kwargs):
+ pass
+
+ async def write(agent, prompt, *, deps, **kwargs):
+ runs.append((deps.current_module_name, list(deps.module_component_ids())))
+ Path(deps.absolute_docs_path, f"{deps.current_module_name}.md").write_text(
+ "# Module\n\nDocumentation."
+ )
+
+ monkeypatch.setattr(mod, "Agent", FakeAgent)
+ monkeypatch.setattr(mod, "create_fallback_models", lambda c: None)
+ monkeypatch.setattr(mod, "run_until_doc_written", write)
+
+ def call(deps, specs):
+ return asyncio.run(
+ mod.generate_sub_module_documentation(SimpleNamespace(deps=deps), specs)
+ )
+ else:
+ pytest.importorskip("caw")
+ from codewiki.src.be.caw_toolkit import CawToolKit
+
+ class Backend:
+ def _run_module_agent_sync(self, **kwargs):
+ runs.append((kwargs["module_name"], list(kwargs["core_component_ids"])))
+ Path(tmp_path, f"{kwargs['module_name']}.md").write_text(
+ "# Module\n\nDocumentation."
+ )
+
+ def call(deps, specs):
+ return CawToolKit(deps, Backend(), True)._run_sub_modules(specs)
+
+ return call, runs
+
+
+def deps_for(tmp_path, components):
+ return CodeWikiDeps(
+ absolute_docs_path=str(tmp_path),
+ absolute_repo_path=str(tmp_path),
+ registry={},
+ components=components,
+ path_to_current_module=[],
+ current_module_name="repo",
+ module_tree={},
+ max_depth=3,
+ current_depth=1,
+ config=config(tmp_path, max_depth=3, max_token_per_leaf_module=4000),
+ custom_instructions="",
+ )
+
+
+@pytest.mark.parametrize("kind", ["all", "almost", "single_file", "small_parent", "depth"])
+def test_tools_reject_without_mutating_or_running_agents(tmp_path, delegate, kind):
+ components = graph(**{"a.c": 5000, "b.c": 100})
+ deps = deps_for(tmp_path, components)
+ ids = list(components)
+ specs = {"all": ids} if kind == "all" else {"large": ids[:1], "tiny": ids[1:]}
+ if kind == "single_file":
+ deps.components = {ids[0]: components[ids[0]]}
+ specs = {"one": ids[:1]}
+ elif kind == "small_parent":
+ deps.components = graph(**{"a.c": 240, "b.c": 66})
+ elif kind == "depth":
+ deps.current_depth = deps.max_depth
+ before = deepcopy(deps.module_tree)
+ call, runs = delegate
+ report = call(deps, specs)
+ assert deps.module_tree == before and runs == []
+ assert list(tmp_path.iterdir()) == []
+ assert "Skipped" in report and "directly" in report
+ assert deps.current_module_name == "repo" and deps.path_to_current_module == []
+
+
+def test_tools_merge_tiny_child_into_nearby_sibling_and_restore_scope(tmp_path, delegate):
+ components = graph(**{"src/tiny.c": 999, "src/near.c": 2000, "other/far.c": 2000})
+ deps = deps_for(tmp_path, components)
+ call, runs = delegate
+ report = call(
+ deps,
+ {
+ "tiny": ["src/tiny.c::entry"],
+ "near": ["src/near.c::entry"],
+ "far": ["other/far.c::entry"],
+ },
+ )
+ assert {name for name, ids in runs} == {"near", "far"}
+ assert not (tmp_path / "tiny.md").exists()
+ assert "999 code tokens below minimum 1000" in report and "merged into 'near'" in report
+ assert "link the existing pages" not in report
+ assert deps.current_depth == 1 and deps.scope_component_ids is None
+ validate_module_file_coverage(deps.module_tree, components)
+ again = call(deps, {"tiny": list(components)})
+ assert len(runs) == 2 and "no unassigned files" in again
+
+
+def test_small_artifact_bundle_merges_into_leaf_and_updates_ancestors():
+ components = graph(**{"src/core.c": 2000, "ci/task.yml": 50})
+ artifact = components["ci/task.yml::entry"]
+ artifact.component_type = "artifact"
+ artifact.artifact_class = "ci"
+ tree = {
+ "Core": {
+ "components": ["src/core.c::entry"],
+ "children": {"Implementation": {"components": ["src/core.c::entry"], "children": {}}},
+ }
+ }
+ result = ensure_artifact_module(tree, list(components), components)
+ assert set(result) == {"Core"}
+ assert artifact.id in result["Core"]["components"]
+ assert artifact.id in result["Core"]["children"]["Implementation"]["components"]
+ validate_module_file_coverage(result, components)
+
+
+@pytest.mark.parametrize("value", [0, -1, True, 1.5, "1000"])
+def test_invalid_minimum_is_rejected_by_config(tmp_path, value):
+ from codewiki.cli.models.config import Configuration
+ from codewiki.cli.utils.errors import ConfigurationError
+
+ with pytest.raises(ValueError, match="positive integer"):
+ config(tmp_path, min_token_per_module=value)
+ with pytest.raises(ConfigurationError, match="positive integer"):
+ Configuration.from_dict({"min_token_per_module": value})
+
+
+def test_default_and_custom_minimum_reach_backend(tmp_path):
+ from codewiki.cli.models.config import Configuration
+
+ for data, expected in [({}, 1000), ({"min_token_per_module": 2000}, 2000)]:
+ cfg = Configuration.from_dict(data)
+ assert cfg.min_token_per_module == expected
+ round_trip = Configuration.from_dict(cfg.to_dict())
+ assert (
+ round_trip.to_backend_config(".", str(tmp_path), "test").min_token_per_module
+ == expected
+ )
+
+
+def test_cli_persists_and_shows_minimum(tmp_path, monkeypatch):
+ import json
+ import codewiki.cli.config_manager as cm
+ from codewiki.cli.commands.config import config_set, config_show
+ from codewiki.cli.commands.generate import generate_command
+
+ monkeypatch.setattr(cm, "CONFIG_DIR", tmp_path)
+ monkeypatch.setattr(cm, "CONFIG_FILE", tmp_path / "config.json")
+ monkeypatch.setenv("CODEWIKI_NO_KEYRING", "1")
+ runner = CliRunner()
+ result = runner.invoke(config_set, ["--min-token-per-module", "2000"])
+ assert result.exit_code == 0, result.output
+ assert json.loads((tmp_path / "config.json").read_text())["min_token_per_module"] == 2000
+ assert "2000" in runner.invoke(config_show).output
+ for command in (config_set, generate_command):
+ result = runner.invoke(command, ["--min-token-per-module", "0"])
+ assert result.exit_code != 0 and "0" in result.output
+ assert "--min-token-per-module" in runner.invoke(command, ["--help"]).output
+
+
+def test_artifact_entry_points_follow_their_existing_file_owner():
+ components = graph(**{"src/core.c": 2000, "src/other.c": 2000})
+ artifact = node("src/core.c", 50).model_copy(
+ update={
+ "id": "src/core.c::artifact",
+ "component_type": "artifact",
+ "artifact_class": "config",
+ }
+ )
+ components[artifact.id] = artifact
+ tree = {
+ "Core": {"components": ["src/core.c::entry"], "children": {}},
+ "Other": {"components": ["src/other.c::entry"], "children": {}},
+ }
+ result = ensure_artifact_module(tree, list(components), components)
+ assert set(result) == {"Core", "Other"}
+ assert artifact.id in result["Core"]["components"]
+ validate_module_file_coverage(result, components)
+
+
+@pytest.mark.parametrize(
+ "tokens,depth,expected",
+ [
+ (7999, 1, False),
+ (8000, 1, True),
+ (15999, 2, False),
+ (16000, 2, True),
+ ],
+)
+def test_default_delegation_gate_is_eight_thousand(tmp_path, tokens, depth, expected):
+ from codewiki.cli.models.config import Configuration
+ from codewiki.src.config import DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
+
+ cfg = config(tmp_path, max_depth=3)
+ cli_cfg = Configuration.from_dict({})
+ assert (
+ cfg.max_token_per_leaf_module
+ == cli_cfg.max_token_per_leaf_module
+ == DEFAULT_MAX_TOKEN_PER_LEAF_MODULE
+ == 8000
+ )
+ assert cli_cfg.to_backend_config(".", str(tmp_path), "test").max_token_per_leaf_module == 8000
+ components = graph(**{"a.py": tokens // 2, "b.py": tokens - tokens // 2})
+ assert can_split_module(components, list(components), cfg, depth) is expected
+ # Existing explicit settings continue to override the new default.
+ assert (
+ Configuration.from_dict({"max_token_per_leaf_module": 4000}).max_token_per_leaf_module
+ == 4000
+ )
diff --git a/tests/test_model_flags.py b/tests/test_model_flags.py
index 40aeec37..456bd54e 100644
--- a/tests/test_model_flags.py
+++ b/tests/test_model_flags.py
@@ -24,6 +24,7 @@ def captured(tmp_path, monkeypatch):
aws_region="us-east-1",
agent_instructions=None,
max_tokens=1,
+ min_token_per_module=2000,
max_token_per_module=1,
max_token_per_leaf_module=1,
max_depth=1,
@@ -76,3 +77,12 @@ def test_flags_override_saved_models(captured):
"c",
"f",
)
+
+
+@pytest.mark.parametrize("override,expected", [(None, 2000), (1500, 1500)])
+def test_minimum_module_size_override(captured, override, expected):
+ args = ["-o", "out"]
+ if override is not None:
+ args += ["--min-token-per-module", str(override)]
+ CliRunner().invoke(gen.generate_command, args)
+ assert captured["min_token_per_module"] == expected
diff --git a/tests/test_module_scope.py b/tests/test_module_scope.py
index 71211118..4a618c2f 100644
--- a/tests/test_module_scope.py
+++ b/tests/test_module_scope.py
@@ -48,6 +48,8 @@ def _config(repo: Path, out: Path, threshold: int = 36369) -> Config:
cluster_model="test",
max_token_per_module=threshold,
artifacts_enabled=False,
+ min_token_per_module=1,
+ max_token_per_leaf_module=1,
)
@@ -341,8 +343,10 @@ def test_delegation_covers_parent_without_stealing_siblings():
@pytest.mark.parametrize("backend_kind", ["pydantic", "caw"])
@pytest.mark.parametrize("max_depth", [1, 2])
-def test_small_whole_repo_agent_has_all_files_and_optional_delegate(
- tmp_path, monkeypatch, backend_kind, max_depth
+@pytest.mark.parametrize("file_tokens", [1, 1000, 2000])
+@pytest.mark.parametrize("leaf_budget", [1, 4000])
+def test_whole_repo_agent_has_all_files_and_shared_delegate_gate(
+ tmp_path, monkeypatch, backend_kind, max_depth, file_tokens, leaf_budget
):
from codewiki.src.be.backend import UsageLedger
@@ -354,12 +358,14 @@ def test_small_whole_repo_agent_has_all_files_and_optional_delegate(
for path, name in (("types.go", "Buffer"), ("methods.go", "Buffer.Read")):
file = repo / path
file.write_text("package tiny\n")
- node = _node(f"{path}::{name}", "small")
+ node = _node(f"{path}::{name}", " x" * file_tokens)
node.file_path = str(file)
components[node.id] = node
(out / "module_tree.json").write_text("{}")
config = _config(repo, out)
config.max_depth = max_depth
+ config.min_token_per_module = 1000
+ config.max_token_per_leaf_module = leaf_budget
captured = {}
if backend_kind == "pydantic":
@@ -405,7 +411,7 @@ def completion(self, prompt):
backend._config = config
asyncio.run(backend.run_module_agent("tiny", components, ["types.go::Buffer"], [], str(out)))
- assert captured["can_delegate"] == (max_depth > 1)
+ assert captured["can_delegate"] == (max_depth > 1 and 2 * file_tokens >= max(2000, leaf_budget))
assert component_files(captured["ids"], components) == {"types.go", "methods.go"}
assert json.loads((out / "module_tree.json").read_text()) == {}
@@ -434,7 +440,7 @@ def test_delegate_normalizes_files_before_running_subagents(tmp_path, monkeypatc
custom_instructions="",
)
seen = []
- specs = {"One": [ids[0]], "Two": [ids[1], ids[3]]}
+ specs = {"One": [ids[0]], "Two": [ids[1], ids[2], ids[3]]}
if backend_kind == "pydantic":
class FakeAgent:
diff --git a/tests/test_sub_module_dedupe.py b/tests/test_sub_module_dedupe.py
index 4e1eb7c6..32adf8a4 100644
--- a/tests/test_sub_module_dedupe.py
+++ b/tests/test_sub_module_dedupe.py
@@ -53,7 +53,13 @@ def _deps(tmp_path) -> CodeWikiDeps:
module_tree={},
max_depth=2,
current_depth=1,
- config=SimpleNamespace(max_token_per_leaf_module=4000, agent_retries=3, request_limit=100),
+ config=SimpleNamespace(
+ min_token_per_module=1,
+ max_token_per_module=36369,
+ max_token_per_leaf_module=1,
+ agent_retries=3,
+ request_limit=100,
+ ),
custom_instructions="",
)
@@ -95,8 +101,16 @@ def test_mixed_batch_generates_only_new_names(tmp_path):
deps = _deps(tmp_path)
(tmp_path / "approval.md").write_text("old")
(tmp_path / "overview_approval.md").write_text("old")
- report = _call(deps, {"approval": ["a.py::fa"], "billing": ["b.py::fb"]})
+ deps.components["c.py::fc"] = _node("c.py::fc")
+ report = _call(
+ deps,
+ {
+ "approval": ["a.py::fa"],
+ "billing": ["b.py::fb"],
+ "shipping": ["c.py::fc"],
+ },
+ )
assert (tmp_path / "billing.md").exists()
- assert FakeAgent.runs == ["billing"]
+ assert FakeAgent.runs == ["billing", "shipping"]
assert "billing.md" in report and "'approval'" in report and "Skipped" in report
- assert set(deps.module_tree) == {"billing"}
+ assert set(deps.module_tree) == {"billing", "shipping"}