band-aider-pub/tests/band_aider/test_story_decomposition_check.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

import json
from unittest.mock import patch

from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
from band_aider.ai_client.json_array_prompt_reducer import GeneralPromptReducer
from band_aider.ai_client.token_jacker import TokenJacker
from band_aider.ai_client.token_miser import TokenMiser
from band_aider.context.token_budget_state import TokenBudgetSettings
from tests.support.completion_clients import ScriptedCompletionClient, StaticCompletionClient
from band_aider.judgement.story_decomposition_check import PromptToStoriesCheck
from band_aider.judgement.prompt_formatter import MarkdownPromptFormatterBuilder


class _Client(CompletionClient):
    """Return a configured JSON decomposition."""

    def __init__(self, payload: dict[str, object]) -> None:
        self.payload = payload

    @property
    def inner(self):
        return None

    def complete(self, prompt: str, logger=None) -> str:
        return json.dumps(self.payload)

    @property
    def openai_configs(self) -> tuple[OpenAIConfig, ...]:
        return ()


def test_no_unrelated_work_retains_original_prompt_unchanged() -> None:
    """One coherent feature remains a single, verbatim story."""

    prompt = "Add login validation and its acceptance tests."
    check = PromptToStoriesCheck(prompt, _Client({"stories": [prompt]}), MarkdownPromptFormatterBuilder())

    assert check() == [prompt]


def test_related_requirements_remain_in_one_story() -> None:
    """The model is instructed to keep work sharing one goal in one story."""

    story = "Add login validation, display the error, and test the invalid-password path."
    client = ScriptedCompletionClient((json.dumps({"stories": [story]}),))
    check = PromptToStoriesCheck(story, client, MarkdownPromptFormatterBuilder())

    assert check() == [story]
    assert (
        "Cluster the human task into independent stories.\n"
        "Independent means semantically independent and directionally uncorrelated.\n"
        "Do not split work that appears to share the same goal.\n"
        "Preserve every requirement and constraint, and keep first-appearance order.\n"
        "If the work forms one story, return the original human task unchanged."
    ) in client.prompts[0]
    assert 'Return strict JSON only: {"stories": ["..."]}.' in client.prompts[0]
    assert f"Human task:\n{story}" in client.prompts[0]
    assert "technically independent" not in client.prompts[0]
    assert "shared release goal" not in client.prompts[0]


def test_unrelated_stories_keep_model_returned_first_appearance_order() -> None:
    """The check exposes independent stories in the requested order."""

    first = "Update the billing invoice export."
    second = "Add a dark-mode toggle to settings."
    check = PromptToStoriesCheck(f"{first}\n{second}", _Client({"stories": [first, second]}), MarkdownPromptFormatterBuilder())

    assert check() == [first, second]


def test_story_decomposition_uses_highest_tier_through_corrections() -> None:
    """A shared goal and its corrections stay on the highest tier beyond its allowance."""

    goal = "Update the codec and both callers as one connected story; preserve rejected writes."
    rejected = '{"stories":[""]}'
    normal = ScriptedCompletionClient(())
    middle = ScriptedCompletionClient(())
    highest = ScriptedCompletionClient((*(rejected for _ in range(15)), json.dumps({"stories": [goal]})))
    runtime = TokenJacker(
        normal, TokenBudgetSettings(max_prompt_tokens=100_000), escalation_clients=(middle, highest),
    )
    check = PromptToStoriesCheck(goal, runtime, MarkdownPromptFormatterBuilder())

    assert check() == [goal]
    assert normal.prompts == []
    assert middle.prompts == []
    assert len(highest.prompts) == 16
    assert all(goal in prompt for prompt in highest.prompts)
    assert all("Do not split work that appears to share the same goal." in prompt for prompt in highest.prompts)
    assert highest.prompts[-1].count(rejected) == 1


def test_story_decomposition_highest_tier_reads_complete_task_without_reduction() -> None:
    """The higher provider receives the complete connected task in one initial read."""

    goal = (
        "Update the shared codec, migrate both callers, preserve rejected writes, "
        "and verify the resulting format as one connected goal."
    )
    settings = TokenBudgetSettings(max_prompt_tokens=64, max_output_tokens=16)
    local = ScriptedCompletionClient(())
    normal = TokenMiser(local, settings, reducers=[GeneralPromptReducer()])
    highest = ScriptedCompletionClient((json.dumps({"stories": [goal]}),))
    runtime = TokenJacker(normal, settings, escalation_clients=(highest,))
    check = PromptToStoriesCheck(goal, runtime, MarkdownPromptFormatterBuilder())

    assert check() == [goal]
    assert local.prompts == []
    assert len(highest.prompts) == 1
    assert goal in highest.prompts[0]
    assert "Reduce this oversized prompt" not in highest.prompts[0]


def test_story_decomposition_single_provider_retains_normal_scrolling_reader() -> None:
    """Without a higher tier, independent goals retain the supplied scrolling reader."""

    first = "Change the help screen color palette."
    second = "Add a monthly invoice CSV export."
    goal = f"{first}\n{second}"
    response = json.dumps({"stories": [first, second]})
    local = StaticCompletionClient(response)
    normal = TokenMiser(
        local, TokenBudgetSettings(max_prompt_tokens=64, max_output_tokens=16),
        reducers=[GeneralPromptReducer()],
    )
    check = PromptToStoriesCheck(goal, normal, MarkdownPromptFormatterBuilder())

    with patch.object(local, "complete", wraps=local.complete) as reads:
        result = check()

    assert result == [first, second]
    assert reads.call_count > 1
    assert "Reduce this oversized prompt" in reads.call_args_list[0].args[0]
    assert reads.call_args_list[-1].args[0] == response


def test_story_decomposition_provider_failure_retains_entire_human_task() -> None:
    """An external provider failure preserves every requirement as one fallback story."""

    goal = "Update the shared codec and both callers; preserve rejected writes."
    normal = ScriptedCompletionClient(())
    highest = ScriptedCompletionClient(())
    runtime = TokenJacker(
        normal, TokenBudgetSettings(max_prompt_tokens=100_000), escalation_clients=(highest,),
    )
    check = PromptToStoriesCheck(goal, runtime, MarkdownPromptFormatterBuilder())

    with patch.object(highest, "complete", side_effect=ConnectionError("Provider unavailable")):
        result = check()

    assert result == [goal]
    assert check.stories == [goal]
    assert check.additional["stories"] == [goal]
    assert normal.prompts == []