from agentbyte.termination.token_usage import TokenUsageTermination
from agentbyte.messages import AssistantMessage, Usage


class TestTokenUsageTermination:
    def test_token_usage(self):
        term = TokenUsageTermination(max_tokens=100)

        # Message with 40 tokens
        msg1 = AssistantMessage(
            content="test",
            source="agent",
            usage=Usage(tokens_input=20, tokens_output=20),
        )
        assert term.check([msg1]) is None
        assert term.token_count == 40

        # Message with 70 tokens (reaches 110 total)
        msg2 = AssistantMessage(
            content="test",
            source="agent",
            usage=Usage(tokens_input=35, tokens_output=35),
        )
        stop = term.check([msg2])
        assert stop is not None
        assert "Token budget exceeded" in stop.content
        assert term.token_count == 110

    def test_reset(self):
        term = TokenUsageTermination(max_tokens=10)
        msg = AssistantMessage(
            content="t", source="a", usage=Usage(tokens_input=10, tokens_output=10)
        )
        term.check([msg])
        assert term.is_met()

        term.reset()
        assert not term.is_met()
        assert term.token_count == 0
