Newer
Older
navi-1 / tests / unit / core / test_context_builder.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy on 12 May 5 KB Clarify knowledge persistence prompts
"""Unit tests for ContextBuilder."""

import pytest

from navi.core.context_builder import ContextBuilder
from navi.llm.base import Message
from tests.conftest_factory import make_profile, make_profile_registry


class FakeMcpManager:
    def __init__(self):
        self.calls = []

    def get_instructions(self, server_names=None):
        self.calls.append(server_names)
        if not server_names:
            return {}
        return {name: f"{name} instructions" for name in server_names}


class TestBuildSystemPrompt:
    def test_includes_persona(self):
        import navi.config as _config

        _config.settings.navi_persona = "You are Navi."
        _config.settings.navi_persona_file = ""
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test")
        prompt = builder.build_system_prompt(profile)
        assert "You are Navi." in prompt
        assert profile.system_prompt in prompt

    def test_includes_other_profiles(self):
        reg = make_profile_registry()
        builder = ContextBuilder(profile_registry=reg)
        profile = reg.get("secretary")
        prompt = builder.build_system_prompt(profile)
        assert "## Available profiles" in prompt
        assert "developer" in prompt

    def test_cache_returns_same_object(self):
        reg = make_profile_registry()
        builder = ContextBuilder(profile_registry=reg)
        profile = reg.get("secretary")
        p1 = builder.build_system_prompt(profile)
        p2 = builder.build_system_prompt(profile)
        assert p1 is p2  # cached

    def test_invalidate_cache(self):
        reg = make_profile_registry()
        builder = ContextBuilder(profile_registry=reg)
        profile = reg.get("secretary")
        p1 = builder.build_system_prompt(profile)
        builder.invalidate_system_prompt_cache(profile.id)
        p2 = builder.build_system_prompt(profile)
        assert p1 == p2
        assert p1 is not p2  # cache was invalidated


class TestBuildGoalAnchor:
    def test_includes_original_request(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        msg = builder._build_goal_anchor("sess-1", "Write tests")
        assert msg.role == "system"
        assert "Original request: Write tests" in msg.content
        assert "Stay on track" in msg.content


class TestBuild:
    def test_puts_system_first(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test")
        context = [Message(role="user", content="hi")]
        result = builder.build(context, profile, mem=None)
        assert result[0].role == "system"

    def test_injects_memory(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test")
        mem = Message(role="system", content="I remember you.")
        context = [Message(role="user", content="hi")]
        result = builder.build(context, profile, mem=mem)
        assert result[1] == mem

    def test_strips_existing_system(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test")
        context = [
            Message(role="system", content="old"),
            Message(role="user", content="hi"),
        ]
        result = builder.build(context, profile, mem=None)
        system_msgs = [m for m in result if m.role == "system"]
        assert len(system_msgs) == 1  # only the new system prompt

    def test_iteration_budget_injection(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test", iteration_budget_enabled=True)
        context = [Message(role="user", content="hi")]
        result = builder.build(context, profile, mem=None, iteration=7, max_iterations=10)
        last = result[-1]
        assert last.role == "system"
        assert "Iteration 8/10" in last.content
        assert "2 iteration(s) after this one" in last.content

    def test_critical_urgency(self):
        builder = ContextBuilder(profile_registry=make_profile_registry())
        profile = make_profile("test", iteration_budget_enabled=True)
        context = [Message(role="user", content="hi")]
        result = builder.build(context, profile, mem=None, iteration=9, max_iterations=10)
        assert "CRITICAL" in result[-1].content

    def test_does_not_inject_mcp_for_profiles_without_mcp_servers(self):
        mcp = FakeMcpManager()
        builder = ContextBuilder(profile_registry=make_profile_registry(), mcp_manager=mcp)
        profile = make_profile("test", mcp_servers={})
        context = [Message(role="user", content="hi")]

        result = builder.build(context, profile, mem=None)

        assert not any("MCP servers" in (m.content or "") for m in result)
        assert mcp.calls == []

    def test_injects_only_profile_mcp_server_instructions(self):
        mcp = FakeMcpManager()
        builder = ContextBuilder(profile_registry=make_profile_registry(), mcp_manager=mcp)
        profile = make_profile("test", mcp_servers={"gnexus-book": ["read"]})
        context = [Message(role="user", content="hi")]

        result = builder.build(context, profile, mem=None)

        assert any("gnexus-book instructions" in (m.content or "") for m in result)
        assert mcp.calls == [{"gnexus-book"}]