| 1 | import sys |
| 2 | from pathlib import Path |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 7 | if str(PROJECT_ROOT) not in sys.path: |
| 8 | sys.path.insert(0, str(PROJECT_ROOT)) |
| 9 | |
| 10 | from agent import AgentConfig, AgentContext, AgentContextType |
| 11 | from helpers import runtime, tokens |
| 12 | |
| 13 | |
| 14 | def _iter_prompt_files(): |
| 15 | yield from (PROJECT_ROOT / "prompts").rglob("*.md") |
| 16 | yield from (PROJECT_ROOT / "agents" / "agent0" / "prompts").rglob("*.md") |
| 17 | yield from (PROJECT_ROOT / "knowledge" / "main").rglob("*.md") |
| 18 | for prompts_dir in (PROJECT_ROOT / "plugins").glob("*/prompts"): |
| 19 | yield from prompts_dir.rglob("*.md") |
| 20 | |
| 21 | |
| 22 | async def _build_system_text(profile: str = "agent0", rendered: bool = False) -> str: |
| 23 | old_args = dict(runtime.args) |
| 24 | runtime.args.clear() |
| 25 | runtime.args["dockerized"] = "true" |
| 26 | |
| 27 | ctx = AgentContext( |
| 28 | config=AgentConfig( |
| 29 | profile=profile, |
| 30 | knowledge_subdirs=["custom", "default"], |
| 31 | mcp_servers='{"mcpServers": {}}', |
| 32 | ), |
| 33 | type=AgentContextType.USER, |
| 34 | set_current=False, |
| 35 | ) |
| 36 | try: |
| 37 | if rendered: |
| 38 | prompt = await ctx.agent0.prepare_prompt(ctx.agent0.loop_data) |
| 39 | return str(prompt[0].content) |
| 40 | system = await ctx.agent0.get_system_prompt(ctx.agent0.loop_data) |
| 41 | return "\n\n".join(system) |
| 42 | finally: |
| 43 | AgentContext.remove(ctx.id) |
| 44 | runtime.args.clear() |
| 45 | runtime.args.update(old_args) |
| 46 | |
| 47 | |
| 48 | @pytest.mark.asyncio |
| 49 | async def test_default_agent0_prompt_budget_and_guardrails(): |
| 50 | system_text = await _build_system_text() |
| 51 | rendered_system_text = await _build_system_text(rendered=True) |
| 52 | communication_prompt = ( |
| 53 | PROJECT_ROOT / "prompts" / "agent.system.main.communication.md" |
| 54 | ).read_text(encoding="utf-8") |
| 55 | |
| 56 | # The default prompt now intentionally includes the compact always-on tool |
| 57 | # surface plus skill metadata. Keep the guardrail close to the observed |
| 58 | # budget so prompt creep remains visible without pretending this surface is |
| 59 | # a tiny single-tool prompt. |
| 60 | assert tokens.approximate_tokens(system_text) <= 10000 |
| 61 | assert "`tool_name` must be one listed tool name" in system_text |
| 62 | assert "- tool_args: key value pairs tool arguments" in system_text |
| 63 | assert '"tool_name": "call_subordinate"' in system_text |
| 64 | assert '"tool_name": "parallel"' in system_text |
| 65 | assert "Each `tool_calls` item is a normal tool request object" in system_text |
| 66 | assert '"reset": true' in system_text |
| 67 | assert '"tool_name": "text_editor"' in system_text |
| 68 | assert '"action": "read"' in system_text |
| 69 | assert '"tool_name": "code_execution_tool"' in system_text |
| 70 | assert '"tool_name": "memory_load"' in system_text |
| 71 | assert "informative but tight" in system_text |
| 72 | assert "Your actual output starts with `{` and ends with `}`" in system_text |
| 73 | assert "~~~json" in communication_prompt |
| 74 | assert "~~~json" not in rendered_system_text |
| 75 | assert "```json" not in rendered_system_text |
| 76 | assert "# code_execution_remote tool" not in system_text |
| 77 | assert "# text_editor_remote tool" not in system_text |
| 78 | assert "### computer_use_remote" not in system_text |
| 79 | assert '"tool_name": "code_execution_remote"' not in system_text |
| 80 | assert '"tool_name": "text_editor_remote"' not in system_text |
| 81 | assert '"tool_name": "computer_use_remote"' not in system_text |
| 82 | assert "Computer Use enablement is scoped to the current CLI session" not in system_text |
| 83 | |
| 84 | |
| 85 | @pytest.mark.asyncio |
| 86 | @pytest.mark.parametrize( |
| 87 | "profile", ["agent0", "default", "developer", "researcher", "tiny-local"] |
| 88 | ) |
| 89 | async def test_rendered_profiles_strip_json_fences(profile: str): |
| 90 | system_text = await _build_system_text(profile, rendered=True) |
| 91 | |
| 92 | assert "~~~json" not in system_text |
| 93 | assert "```json" not in system_text |
| 94 | |
| 95 | if profile == "researcher": |
| 96 | assert "~~~python" in system_text |
| 97 | |
| 98 | |
| 99 | def test_remove_code_fences_can_target_json_only(): |
| 100 | from helpers import files |
| 101 | |
| 102 | prompt = """Before |
| 103 | ~~~json |
| 104 | {"tool_name":"response","tool_args":{"text":"done"}} |
| 105 | ~~~ |
| 106 | ~~~python |
| 107 | print("keep me fenced") |
| 108 | ~~~ |
| 109 | After |
| 110 | """ |
| 111 | |
| 112 | rendered = files.remove_code_fences(prompt, language="json") |
| 113 | |
| 114 | assert "~~~json" not in rendered |
| 115 | assert '{"tool_name":"response"' in rendered |
| 116 | assert '~~~python\nprint("keep me fenced")\n~~~' in rendered |
| 117 | |
| 118 | |
| 119 | @pytest.mark.asyncio |
| 120 | async def test_tiny_local_profile_prompt_is_action_first_json_contract(): |
| 121 | system_text = await _build_system_text("tiny-local") |
| 122 | communication_prompt = ( |
| 123 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "agent.system.main.communication.md" |
| 124 | ).read_text(encoding="utf-8") |
| 125 | code_prompt = ( |
| 126 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "agent.system.tool.code_exe.md" |
| 127 | ).read_text(encoding="utf-8") |
| 128 | response_prompt = ( |
| 129 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "agent.system.tool.response.md" |
| 130 | ).read_text(encoding="utf-8") |
| 131 | repeat_prompt = ( |
| 132 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "fw.msg_repeat.md" |
| 133 | ).read_text(encoding="utf-8") |
| 134 | text_editor_prompt = ( |
| 135 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "agent.system.tool.text_editor.md" |
| 136 | ).read_text(encoding="utf-8") |
| 137 | solving_prompt = ( |
| 138 | PROJECT_ROOT / "agents" / "tiny-local" / "prompts" / "agent.system.main.solving.md" |
| 139 | ).read_text(encoding="utf-8") |
| 140 | |
| 141 | assert "You are Agent Zero. Act on the user's behalf." in system_text |
| 142 | assert "Your visible assistant message must be exactly one valid JSON object." in system_text |
| 143 | assert 'Use exactly these top-level fields: `"tool_name"` and `"tool_args"`.' in system_text |
| 144 | assert 'For a final user-facing answer, use the `response` tool.' in system_text |
| 145 | assert "Use `response` only when the work is complete, blocked, or the user is only acknowledging completed work." in system_text |
| 146 | assert "If the user says \"proceed\", \"continue\", \"go ahead\", \"do it\", \"excellent proceed\"" in system_text |
| 147 | assert "Do not explain what command the user could run manually." in system_text |
| 148 | assert "output a corrected JSON tool request immediately" in system_text |
| 149 | assert "do not resend the same JSON" in system_text |
| 150 | assert "## Tiny Local Output Rule" in system_text |
| 151 | assert "~~~json" not in communication_prompt |
| 152 | assert "~~~json" not in code_prompt |
| 153 | assert "~~~json" not in response_prompt |
| 154 | assert "~~~json" not in text_editor_prompt |
| 155 | assert "No JSON in markdown fences" not in communication_prompt |
| 156 | assert "thoughts: array thoughts before execution" not in communication_prompt |
| 157 | assert "headline: short headline summary" not in communication_prompt |
| 158 | assert "explain each step in thoughts" not in solving_prompt |
| 159 | assert "Continuation words" in solving_prompt |
| 160 | assert "Do not respond by saying you will begin, continue, start, proceed, or investigate." in solving_prompt |
| 161 | assert "Do not use this tool for \"proceed\", \"continue\", \"go ahead\"" in response_prompt |
| 162 | assert "Your repeated JSON was recorded, but it did not execute another tool." in repeat_prompt |
| 163 | assert "replace it with the next real tool call" in repeat_prompt |
| 164 | assert "do not repeat the same status response or exact tool request" in solving_prompt |
| 165 | assert "do not repeat the same exact tool call" in solving_prompt |
| 166 | assert '"open_in_canvas":true' in text_editor_prompt |
| 167 | assert "do not repeat the same tool call" in text_editor_prompt |
| 168 | assert '"headline"' not in code_prompt |
| 169 | assert '"headline"' not in response_prompt |
| 170 | assert '"headline"' not in text_editor_prompt |
| 171 | |
| 172 | |
| 173 | def test_tiny_local_profile_is_discoverable(): |
| 174 | from helpers import subagents |
| 175 | |
| 176 | profiles = { |
| 177 | str(item.get("key") or ""): str(item.get("label") or "") |
| 178 | for item in subagents.get_all_agents_list() |
| 179 | } |
| 180 | |
| 181 | assert profiles["tiny-local"] == "Tiny Local" |
| 182 | |
| 183 | |
| 184 | def test_removed_small_profile_and_prompt_text_generic(): |
| 185 | removed_profile = "a0" + "_" + "small" |
| 186 | |
| 187 | assert not (PROJECT_ROOT / "agents" / removed_profile).exists() |
| 188 | assert not ( |
| 189 | PROJECT_ROOT / "knowledge" / "main" / f"{removed_profile}_tool_call_examples.md" |
| 190 | ).exists() |
| 191 | assert not (PROJECT_ROOT / "knowledge" / "main" / "tool_call_reference_examples.md").exists() |
| 192 | |
| 193 | for path in _iter_prompt_files(): |
| 194 | assert removed_profile not in path.read_text(encoding="utf-8") |
| 195 | |
| 196 | |
| 197 | def test_prompt_token_estimate_omits_embedded_image_data_urls(): |
| 198 | embedded_png = "data:image/png;base64," + ("ABCDabcd0123+/==" * 20_000) |
| 199 | prompt_text = f"user: please inspect this screenshot {embedded_png}" |
| 200 | |
| 201 | sanitized = tokens.sanitize_embedded_image_data_urls(prompt_text) |
| 202 | |
| 203 | assert "ABCDabcd0123+/==" not in sanitized |
| 204 | assert "data:image/png;base64," in sanitized |
| 205 | assert tokens.EMBEDDED_IMAGE_DATA_PLACEHOLDER in sanitized |
| 206 | assert tokens.approximate_prompt_tokens(prompt_text) < 100 |
| 207 | assert tokens.approximate_prompt_tokens(prompt_text) < tokens.approximate_tokens(prompt_text) / 100 |