main
py 510 lines 16.8 KB
Raw
1 import threading
2 from pathlib import Path
3 from types import SimpleNamespace
4
5 import pytest
6
7 import initialize
8 from agent import AgentConfig, AgentContext
9 from helpers import dirty_json, files, persist_chat, projects, subagents
10 from helpers import state_monitor_integration
11
12
13 def _prepare_project_tree(monkeypatch, tmp_path: Path) -> None:
14 monkeypatch.setattr(files, "_base_dir", str(tmp_path))
15 (tmp_path / "usr" / "projects").mkdir(parents=True, exist_ok=True)
16 (tmp_path / "usr" / "plugins").mkdir(parents=True, exist_ok=True)
17 (tmp_path / "plugins").mkdir(parents=True, exist_ok=True)
18
19
20 @pytest.mark.parametrize(
21 "destination_project", ["project-y", None], ids=["project", "global"]
22 )
23 def test_project_switch_resets_only_profiles_missing_from_the_new_scope(
24 monkeypatch, destination_project
25 ):
26 context_id = "ctx-project-profile-switch"
27 AgentContext.remove(context_id)
28 context = AgentContext(
29 config=AgentConfig(mcp_servers="", profile="project-only"),
30 id=context_id,
31 set_current=False,
32 )
33 monkeypatch.setattr(
34 projects,
35 "load_edit_project_data",
36 lambda name: {"title": name.title(), "color": ""},
37 )
38 monkeypatch.setattr(persist_chat, "save_tmp_chat", lambda _context: None)
39 monkeypatch.setattr(
40 subagents,
41 "get_agents_dict",
42 lambda project_name=None: {
43 "agent0": subagents.SubAgentListItem(name="agent0"),
44 **(
45 {
46 "project-only": subagents.SubAgentListItem(
47 name="project-only"
48 )
49 }
50 if project_name == "project-x"
51 else {}
52 ),
53 },
54 )
55 monkeypatch.setattr(
56 initialize,
57 "initialize_agent",
58 lambda override_settings=None: AgentConfig(
59 mcp_servers="",
60 profile=(override_settings or {}).get("agent_profile", "agent0"),
61 ),
62 )
63
64 try:
65 projects.activate_project(context_id, "project-x", mark_dirty=False)
66 assert context.config.profile == "project-only"
67
68 if destination_project:
69 projects.activate_project(
70 context_id, destination_project, mark_dirty=False
71 )
72 else:
73 projects.deactivate_project(context_id, mark_dirty=False)
74 assert context.config.profile == "agent0"
75 assert context.agent0.config.profile == "agent0"
76 finally:
77 AgentContext.remove(context_id)
78
79
80 def test_project_agent_availability_retains_project_only_profiles(
81 monkeypatch,
82 ) -> None:
83 monkeypatch.setattr(
84 subagents,
85 "get_agents_dict",
86 lambda project_name=None: {
87 "global": subagents.SubAgentListItem(name="global", enabled=True),
88 **(
89 {
90 "project-only": subagents.SubAgentListItem(
91 name="project-only", enabled=True
92 )
93 }
94 if project_name == "demo"
95 else {}
96 ),
97 },
98 )
99
100 assert projects._normalize_subagents(
101 {
102 "global": {"enabled": True},
103 "project-only": {"enabled": False},
104 "missing": {"enabled": False},
105 },
106 "demo",
107 ) == {"project-only": {"enabled": False}}
108
109
110 def test_project_profile_toggle_preserves_other_entries_and_refuses_bad_json(
111 monkeypatch,
112 tmp_path: Path,
113 ) -> None:
114 _prepare_project_tree(monkeypatch, tmp_path)
115 meta = tmp_path / "usr" / "projects" / "demo" / ".a0proj"
116 meta.mkdir(parents=True)
117 availability = meta / "agents.json"
118 monkeypatch.setattr(
119 subagents,
120 "get_agents_dict",
121 lambda _project=None: {
122 "default": subagents.SubAgentListItem(name="default", enabled=True),
123 "researcher": subagents.SubAgentListItem(
124 name="researcher", enabled=True
125 ),
126 },
127 )
128 availability.write_text(
129 '{"default":{"enabled":false}}',
130 encoding="utf-8",
131 )
132
133 projects.set_project_subagent_enabled("demo", "researcher", False)
134
135 assert dirty_json.parse(availability.read_text(encoding="utf-8")) == {
136 "default": {"enabled": False},
137 "researcher": {"enabled": False},
138 }
139 broken = b'{"default":'
140 availability.write_bytes(broken)
141
142 with pytest.raises(ValueError, match="Project agent availability"):
143 projects.set_project_subagent_enabled("demo", "researcher", True)
144
145 assert availability.read_bytes() == broken
146
147
148 def test_project_edit_ignores_stale_agent_availability(
149 monkeypatch,
150 tmp_path: Path,
151 ) -> None:
152 _prepare_project_tree(monkeypatch, tmp_path)
153 meta = tmp_path / "usr" / "projects" / "demo" / ".a0proj"
154 meta.mkdir(parents=True)
155 (meta / "project.json").write_text('{"title":"Demo"}', encoding="utf-8")
156 availability = meta / "agents.json"
157 original = b'{"default":{"enabled":false}}'
158 availability.write_bytes(original)
159 monkeypatch.setattr("helpers.git.get_repo_status", lambda _path: {})
160 monkeypatch.setattr(projects, "reactivate_project_in_chats", lambda _name: None)
161 extended: list[dict] = []
162 monkeypatch.setattr(
163 projects,
164 "save_project_extended_data",
165 lambda _name, data: extended.append(data),
166 )
167
168 loaded = projects.load_edit_project_data("demo")
169 projects.update_project(
170 "demo",
171 {
172 **loaded,
173 "title": "Renamed",
174 "subagents": {"default": {"enabled": True}},
175 },
176 )
177
178 assert "subagents" not in loaded
179 assert availability.read_bytes() == original
180 assert extended and all("subagents" not in data for data in extended)
181
182
183 def test_profile_reconciliation_uses_an_available_fallback(monkeypatch) -> None:
184 context_id = "ctx-profile-availability-fallback"
185 AgentContext.remove(context_id)
186 context = AgentContext(
187 config=AgentConfig(mcp_servers="", profile="disabled"),
188 id=context_id,
189 set_current=False,
190 )
191 monkeypatch.setattr(
192 subagents,
193 "get_available_agents_dict",
194 lambda _project_name: {
195 "researcher": subagents.SubAgentListItem(name="researcher")
196 },
197 )
198 monkeypatch.setattr(
199 initialize,
200 "initialize_agent",
201 lambda override_settings=None: AgentConfig(
202 mcp_servers="",
203 profile=(override_settings or {}).get("agent_profile", "default"),
204 ),
205 )
206
207 try:
208 assert projects.reconcile_agent_profile(context, None) is True
209 assert context.config.profile == "researcher"
210 assert context.agent0.config.profile == "researcher"
211 finally:
212 AgentContext.remove(context_id)
213
214
215 def test_context_lookup_reconciles_only_new_contexts(monkeypatch) -> None:
216 from helpers.context_utils import use_context
217
218 existing_id = "ctx-existing-profile"
219 created_id = "ctx-new-profile"
220 AgentContext.remove(existing_id)
221 AgentContext.remove(created_id)
222 existing = AgentContext(
223 config=AgentConfig(mcp_servers="", profile="default"),
224 id=existing_id,
225 set_current=False,
226 )
227 reconciled: list[str] = []
228 monkeypatch.setattr(
229 initialize,
230 "initialize_agent",
231 lambda: AgentConfig(mcp_servers="", profile="default"),
232 )
233 monkeypatch.setattr(
234 projects,
235 "reconcile_agent_profile",
236 lambda context, _project: reconciled.append(context.id),
237 )
238
239 try:
240 assert use_context(threading.RLock(), existing_id) is existing
241 assert reconciled == []
242
243 assert use_context(threading.RLock(), created_id).id == created_id
244 assert reconciled == [created_id]
245 finally:
246 AgentContext.remove(existing_id)
247 AgentContext.remove(created_id)
248
249
250 @pytest.mark.parametrize(
251 ("all_scopes", "expected"),
252 [
253 (False, ["global-changed"]),
254 (True, ["global-changed", "project-changed"]),
255 ],
256 )
257 def test_bulk_profile_reconciliation_persists_only_changed_chats(
258 monkeypatch, all_scopes: bool, expected: list[str]
259 ) -> None:
260 unchanged = SimpleNamespace(id="global-unchanged", project=None)
261 global_changed = SimpleNamespace(id="global-changed", project=None)
262 project_changed = SimpleNamespace(id="project-changed", project="demo")
263 saved: list[str] = []
264 dirty: list[str] = []
265 catalog_lookups: list[str | None] = []
266 monkeypatch.setattr(
267 AgentContext,
268 "all",
269 classmethod(
270 lambda _cls: [unchanged, global_changed, project_changed]
271 ),
272 )
273 monkeypatch.setattr(
274 projects, "get_context_project_name", lambda context: context.project
275 )
276 monkeypatch.setattr(
277 projects,
278 "reconcile_agent_profile",
279 lambda context, _project, _available: context is not unchanged,
280 )
281 monkeypatch.setattr(
282 subagents,
283 "get_available_agents_dict",
284 lambda project: catalog_lookups.append(project) or {},
285 )
286 monkeypatch.setattr(
287 persist_chat, "save_tmp_chat", lambda context: saved.append(context.id)
288 )
289 monkeypatch.setattr(
290 state_monitor_integration,
291 "mark_dirty_for_context",
292 lambda context_id, **_kwargs: dirty.append(context_id),
293 )
294
295 projects.reconcile_agent_profiles(None, all_scopes=all_scopes)
296
297 assert saved == expected
298 assert dirty == expected
299 assert catalog_lookups == ([None, "demo"] if all_scopes else [None])
300
301
302 def test_project_refresh_touches_only_matching_chats(monkeypatch) -> None:
303 contexts = [
304 SimpleNamespace(id="matching", get_data=lambda _key: "demo"),
305 SimpleNamespace(id="unrelated", get_data=lambda _key: "other"),
306 ]
307 calls: list[tuple] = []
308 monkeypatch.setattr(
309 AgentContext, "all", staticmethod(lambda: contexts)
310 )
311 monkeypatch.setattr(
312 projects,
313 "activate_project",
314 lambda context_id, name, *, mark_dirty: calls.append(
315 ("activate", context_id, name, mark_dirty)
316 ),
317 )
318 monkeypatch.setattr(
319 projects,
320 "deactivate_project",
321 lambda context_id, *, mark_dirty: calls.append(
322 ("deactivate", context_id, mark_dirty)
323 ),
324 )
325 monkeypatch.setattr(state_monitor_integration, "mark_dirty_all", lambda **_kwargs: None)
326
327 projects.reactivate_project_in_chats("demo")
328 projects.deactivate_project_in_chats("demo")
329
330 assert calls == [
331 ("activate", "matching", "demo", False),
332 ("deactivate", "matching", False),
333 ]
334
335
336 def test_project_include_agents_md_defaults_true_and_saves(monkeypatch, tmp_path):
337 _prepare_project_tree(monkeypatch, tmp_path)
338 meta = tmp_path / "usr" / "projects" / "demo" / ".a0proj"
339 meta.mkdir(parents=True)
340 (meta / "project.json").write_text('{"title": "Demo"}', encoding="utf-8")
341
342 data = projects.load_basic_project_data("demo")
343
344 assert data["include_agents_md"] is True
345
346 projects.save_project_header("demo", data)
347 saved = dirty_json.parse((meta / "project.json").read_text(encoding="utf-8"))
348
349 assert saved["include_agents_md"] is True
350
351
352 def test_project_mcp_servers_persist_in_project_meta(monkeypatch, tmp_path):
353 _prepare_project_tree(monkeypatch, tmp_path)
354 config = '{"mcpServers":{"demo":{"url":"https://example.com/mcp"}}}'
355
356 projects.create_project(
357 "demo",
358 {
359 "title": "Demo",
360 "mcp_servers": config,
361 },
362 )
363
364 assert projects.load_project_mcp_servers("demo") == config
365 assert projects.load_edit_project_data("demo")["mcp_servers"] == config
366
367 updated = '{"mcpServers":{"other":{"command":"uvx","args":["pkg"]}}}'
368 projects.save_project_mcp_servers("demo", updated)
369
370 assert projects.load_project_mcp_servers("demo") == updated
371
372
373 def test_project_mcp_servers_reject_path_names(monkeypatch, tmp_path):
374 _prepare_project_tree(monkeypatch, tmp_path)
375
376 for name in ("../escape", "nested/project", ".", "..", ""):
377 try:
378 projects.save_project_mcp_servers(name, '{"mcpServers":{}}')
379 except ValueError:
380 pass
381 else:
382 raise AssertionError(f"Expected invalid project name: {name!r}")
383
384
385 def test_project_creation_creates_skills_folder(monkeypatch, tmp_path):
386 _prepare_project_tree(monkeypatch, tmp_path)
387
388 projects.create_project("demo", {"title": "Demo"})
389
390 assert (tmp_path / "usr" / "projects" / "demo" / ".a0proj" / "skills").is_dir()
391
392
393 def test_project_load_repairs_missing_skills_folder(monkeypatch, tmp_path):
394 _prepare_project_tree(monkeypatch, tmp_path)
395 meta = tmp_path / "usr" / "projects" / "demo" / ".a0proj"
396 meta.mkdir(parents=True)
397 (meta / "project.json").write_text('{"title": "Demo"}', encoding="utf-8")
398
399 assert not (meta / "skills").exists()
400
401 projects.load_edit_project_data("demo")
402
403 assert (meta / "skills").is_dir()
404
405
406 def test_project_system_prompt_includes_root_agents_md_with_path(monkeypatch, tmp_path):
407 _prepare_project_tree(monkeypatch, tmp_path)
408 projects.create_project(
409 "demo",
410 {
411 "title": "Demo",
412 "instructions": "Main project rule.",
413 },
414 )
415 project_root = tmp_path / "usr" / "projects" / "demo"
416 (project_root / "AGENTS.md").write_text("Root AGENTS rule.", encoding="utf-8")
417 (
418 project_root / ".a0proj" / "instructions" / "extra.md"
419 ).write_text("Folder instruction rule.", encoding="utf-8")
420
421 prompt_vars = projects.build_system_prompt_vars("demo")
422 instructions = prompt_vars["project_instructions"]
423
424 assert "Main project rule." in instructions
425 assert instructions.count("## project instruction files") == 1
426 assert "## project instruction file\n" not in instructions
427 assert "### path: /a0/usr/projects/demo/AGENTS.md" in instructions
428 assert "Root AGENTS rule." in instructions
429 assert "### path: /a0/usr/projects/demo/.a0proj/instructions/extra.md" in instructions
430 assert "Folder instruction rule." in instructions
431
432
433 def test_project_system_prompt_prefers_agents_override_md(monkeypatch, tmp_path):
434 _prepare_project_tree(monkeypatch, tmp_path)
435 projects.create_project("demo", {"title": "Demo"})
436 project_root = tmp_path / "usr" / "projects" / "demo"
437 (project_root / "AGENTS.md").write_text("Standard rule.", encoding="utf-8")
438 (project_root / "AGENTS.override.md").write_text("Override rule.", encoding="utf-8")
439
440 instructions = projects.build_system_prompt_vars("demo")["project_instructions"]
441
442 assert "### path: /a0/usr/projects/demo/AGENTS.override.md" in instructions
443 assert "Override rule." in instructions
444 assert "Standard rule." not in instructions
445
446
447 def test_project_system_prompt_respects_disabled_agents_md(monkeypatch, tmp_path):
448 _prepare_project_tree(monkeypatch, tmp_path)
449 projects.create_project(
450 "demo",
451 {
452 "title": "Demo",
453 "include_agents_md": False,
454 },
455 )
456 project_root = tmp_path / "usr" / "projects" / "demo"
457 (project_root / "AGENTS.md").write_text("Root AGENTS rule.", encoding="utf-8")
458
459 prompt_vars = projects.build_system_prompt_vars("demo")
460
461 assert "Root AGENTS rule." not in prompt_vars["project_instructions"]
462 assert "AGENTS.md" not in prompt_vars["project_instructions"]
463
464
465 def test_agents_md_chain_walks_direct_path_only(monkeypatch, tmp_path):
466 _prepare_project_tree(monkeypatch, tmp_path)
467 root = tmp_path
468 (root / "AGENTS.md").write_text("root doc", encoding="utf-8")
469 target = root / "services" / "payments"
470 sibling = root / "services" / "auth"
471 target.mkdir(parents=True)
472 sibling.mkdir(parents=True)
473 (root / "services" / "AGENTS.md").write_text("services doc", encoding="utf-8")
474 (target / "AGENTS.md").write_text("payments doc", encoding="utf-8")
475 (sibling / "AGENTS.md").write_text("auth doc", encoding="utf-8")
476
477 chain = projects.get_agents_md_chain(str(root), str(target / "handler.py"))
478 contents = [content for _, content in chain]
479
480 assert contents == ["root doc", "services doc", "payments doc"]
481
482
483 def test_agents_md_protocol_excludes_project_root_and_keeps_subdir(
484 monkeypatch, tmp_path
485 ):
486 _prepare_project_tree(monkeypatch, tmp_path)
487 prompt_name = "agent.protocol.projects.agents_md.md"
488 prompt_source = Path(__file__).resolve().parents[1] / "prompts" / prompt_name
489 prompt_dir = tmp_path / "prompts"
490 prompt_dir.mkdir()
491 (prompt_dir / prompt_name).write_text(
492 prompt_source.read_text(encoding="utf-8"),
493 encoding="utf-8",
494 )
495 projects.create_project("demo", {"title": "Demo"})
496 (tmp_path / "AGENTS.md").write_text("framework doc", encoding="utf-8")
497 project_root = tmp_path / "usr" / "projects" / "demo"
498 (project_root / "AGENTS.md").write_text("project root doc", encoding="utf-8")
499 api_dir = project_root / "api"
500 api_dir.mkdir()
501 (api_dir / "AGENTS.md").write_text("api doc", encoding="utf-8")
502
503 protocol = projects.build_agents_md_protocol(
504 "demo",
505 target=str(api_dir / "handler.py"),
506 )
507
508 assert "framework doc" in protocol
509 assert "api doc" in protocol
510 assert "project root doc" not in protocol