187 lines
8.6 KiB
Python
187 lines
8.6 KiB
Python
import json
|
|
from typing import cast
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from band_aider.ai_client.stabilization_escalation import StabilizationEscalation
|
|
from band_aider.ai_client.token_jacker import TokenJacker
|
|
from band_aider.context import token_budget_state
|
|
from band_aider.context.token_budget_state import TokenBudgetSettings
|
|
from band_aider.judgement.contextual_file_renderer import ContextualFileRenderer
|
|
from band_aider.judgement.file_manager import FileManager
|
|
from band_aider.judgement.source_discovery_strategy import SearchableBeat, SourceDiscoveryStrategy
|
|
from band_aider.judgement.story_decomposition_check import PromptToStoriesCheck
|
|
from band_aider.strategies import LanguageStrategySet
|
|
from tests.support.completion_clients import ScriptedCompletionClient
|
|
from tests.support.memory_repo_exec import MemoryRepoExec
|
|
|
|
from band_aider.context.class_meta import ClassMeta
|
|
from band_aider.context.cot_task_generator import CotTaskGenerator, CotTaskRecord
|
|
from band_aider.judgement.prompt_formatter import MarkdownPromptFormatterBuilder
|
|
|
|
|
|
def _candidate() -> ClassMeta:
|
|
"""Return one source root with concrete ownership evidence."""
|
|
|
|
return ClassMeta(
|
|
location="band_aider/context/memento.py",
|
|
summary="Owns Memento lifecycle state and snapshot behavior.",
|
|
relevance=[
|
|
"Memento lifecycle transitions and managed context behavior are owned here.",
|
|
],
|
|
relevant_symbols={
|
|
"Memento.__enter__": "Activates the Memento lifecycle.",
|
|
"Memento.commit": "Commits the active Memento transaction.",
|
|
},
|
|
)
|
|
|
|
|
|
def _generator() -> CotTaskGenerator:
|
|
"""Return one prompt-only generator fixture."""
|
|
|
|
return CotTaskGenerator(
|
|
client=object(),
|
|
prompt_formatter_builder=MarkdownPromptFormatterBuilder(),
|
|
)
|
|
|
|
|
|
def test_broad_prompt_preserves_file_applicable_contract_clauses() -> None:
|
|
"""Broad COT selection must retain relevant contract instead of asking for a label."""
|
|
|
|
prompt = _generator()._prompt(
|
|
"Keep Memento construction inert; context entry captures once; callers use managed scope; "
|
|
"all lifecycle and caller migration clauses are one atomic story.",
|
|
[_candidate()],
|
|
"",
|
|
)
|
|
|
|
assert "independently executable file-owned COT tasks" in prompt
|
|
assert "retain every original-contract clause that applies to the selected file" in prompt
|
|
assert "Do not summarize the task into a label" in prompt
|
|
assert "Do not restate broader upstream wording" not in prompt
|
|
assert prompt.index("ORIGINAL RECURSIVE-PASS CONTRACT") < prompt.index(
|
|
"CANDIDATE PRODUCTION FILES"
|
|
)
|
|
|
|
|
|
def test_refinement_prompt_preserves_authority_and_requires_a_file_owned_mini_plan() -> None:
|
|
"""Refinement must enrich the original contract with current-file evidence."""
|
|
|
|
candidate = _candidate()
|
|
broad = CotTaskRecord(
|
|
task="Update Memento lifecycle behavior owned by this file.",
|
|
location=candidate.location,
|
|
root=candidate,
|
|
)
|
|
prompt = _generator()._refinement_prompt(
|
|
"Keep Memento construction inert and migrate all callers through managed contexts.",
|
|
broad,
|
|
candidate,
|
|
"class Memento:\n def __enter__(self): ...",
|
|
"No dependency reference cards resolved.",
|
|
[
|
|
{"kind": "method", "name": "Memento.__enter__"},
|
|
{"kind": "method", "name": "Memento.commit"},
|
|
],
|
|
"",
|
|
"",
|
|
)
|
|
|
|
assert "The original contract is authoritative" in prompt
|
|
assert "must not weaken or discard an applicable original requirement" in prompt
|
|
assert "FILE-OWNED CONTRACT" in prompt
|
|
assert "CURRENT SOURCE EVIDENCE" in prompt
|
|
assert "REQUIRED CHANGES" in prompt
|
|
assert "COORDINATION / GUARDRAILS" in prompt
|
|
assert "VERIFICATION" in prompt
|
|
assert "numbered, concrete steps" in prompt
|
|
assert prompt.index("ORIGINAL RECURSIVE-PASS CONTRACT") < prompt.index(
|
|
"BROAD FILE ASSIGNMENT"
|
|
)
|
|
assert prompt.index("BROAD FILE ASSIGNMENT") < prompt.index(
|
|
"CURRENT FILE EVIDENCE"
|
|
)
|
|
|
|
|
|
def test_connected_story_keeps_nine_paged_roots_with_cloud_available(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Highest-tier story clustering retains nine roots through normal paged assignment."""
|
|
|
|
monkeypatch.setattr(token_budget_state, "_settings", TokenBudgetSettings(max_prompt_tokens=128))
|
|
goal = "Migrate the codec and all eight callers together; preserve rejected writes before acceptance."
|
|
candidates = list(map(lambda index: ClassMeta(
|
|
f"src/component_{index}.ext", "Owns one part of the connected migration.", [goal], {},
|
|
), range(9)))
|
|
replies = tuple(map(lambda candidate: json.dumps({"cot_tasks": [{
|
|
"location": candidate.location, "task": f"Migrate {candidate.location} under the shared completion gate.",
|
|
}]}), candidates))
|
|
normal = ScriptedCompletionClient(replies)
|
|
cloud = ScriptedCompletionClient((json.dumps({"stories": [goal]}),))
|
|
runtime = TokenJacker(normal, TokenBudgetSettings(max_prompt_tokens=100_000), escalation_clients=(cloud,))
|
|
builder = MarkdownPromptFormatterBuilder()
|
|
|
|
stories = PromptToStoriesCheck(goal, runtime, builder)()
|
|
records = CotTaskGenerator(runtime, prompt_formatter_builder=builder).generate(stories[0], candidates)
|
|
|
|
assert stories == [goal]
|
|
assert len(normal.prompts) == 9
|
|
assert all(goal in prompt for prompt in normal.prompts)
|
|
assert tuple(map(lambda record: record.root, records)) == tuple(candidates)
|
|
# Candidate paging retains each owner in its separate original assignment.
|
|
for candidate, prompt in zip(candidates, normal.prompts, strict=True):
|
|
assert candidate.location in prompt
|
|
assert len(cloud.prompts) == 1
|
|
assert goal in cloud.prompts[0]
|
|
assert "Do not split work that appears to share the same goal." in cloud.prompts[0]
|
|
|
|
|
|
def test_cot_refinement_uses_highest_tier_and_keeps_renderer_on_normal_client() -> None:
|
|
"""Normal assignment feeds strongest refinement without changing authority or source evidence."""
|
|
|
|
repo = MemoryRepoExec()
|
|
source = "def work():\n return 1\n"
|
|
repo.write_text(repo.root / "src/work.py", source)
|
|
candidate = ClassMeta("src/work.py", "Work", ["Owns the work function."], {"work": "Update the scoped attempt."})
|
|
goal = "Keep inner and outer attempt completion separate."
|
|
broad = "Update work while retaining separate nested attempt scopes."
|
|
refined = "Update work to complete the inner attempt before validating the enclosing attempt."
|
|
builder = MarkdownPromptFormatterBuilder()
|
|
normal = ScriptedCompletionClient((json.dumps({"cot_tasks": [{"task": broad, "location": candidate.location}]}),))
|
|
middle = ScriptedCompletionClient(())
|
|
wrong = json.dumps({"task": "Update work using invented_symbol.", "location": candidate.location, "symbols": ["invented_symbol"]})
|
|
highest = ScriptedCompletionClient((
|
|
*(wrong for _ in range(15)),
|
|
json.dumps({"task": refined, "location": candidate.location, "symbols": ["work"]}),
|
|
))
|
|
runtime = StabilizationEscalation(
|
|
(normal, middle, highest), check_name="PlanningRuntime", prompt_formatter_builder=builder,
|
|
)
|
|
manager = cast(FileManager, Mock(
|
|
spec=FileManager, repo_root=repo.root, repo_exec=repo,
|
|
src_root=repo.root / "src", test_root=repo.root / "tests",
|
|
dependency_strategy=None, src_classes=[candidate],
|
|
))
|
|
discovery = Mock(spec=SourceDiscoveryStrategy)
|
|
discovery.searchable_beats.return_value = [SearchableBeat(
|
|
kind="function", text="work", attributes=(("name", "work"),), structural_depth=0,
|
|
)]
|
|
renderer = Mock(spec=ContextualFileRenderer)
|
|
renderer.read.return_value = source
|
|
renderer.get_reference_card_evidence.return_value = "No dependency reference cards resolved."
|
|
|
|
with (
|
|
patch.object(LanguageStrategySet, "get", return_value=LanguageStrategySet({SourceDiscoveryStrategy: discovery})),
|
|
patch.object(ContextualFileRenderer, "build", return_value=renderer) as build,
|
|
):
|
|
records = CotTaskGenerator(runtime, file_manager=manager, prompt_formatter_builder=builder).generate(goal, [candidate])
|
|
|
|
assert len(normal.prompts) == 1
|
|
assert "COT TASK GENERATION" in normal.prompts[0]
|
|
assert middle.prompts == []
|
|
assert len(highest.prompts) == 16
|
|
assert all(goal in prompt and broad in prompt and source in prompt for prompt in highest.prompts)
|
|
assert highest.prompts[-1].count(wrong) == 1
|
|
assert "unsupported symbols: invented_symbol" in highest.prompts[-1]
|
|
assert build.call_args.kwargs["client"] is runtime
|
|
assert build.call_args.kwargs["prompt"] == broad
|
|
assert records == [CotTaskRecord(refined, candidate.location, candidate, ("work",))]
|