main
py 349 lines 11 KB
Raw
1 from types import SimpleNamespace
2
3 import pytest
4
5 from agent import AgentContextType
6 from plugins._chat_naming.commands import rename_command
7 from plugins._chat_naming.extensions.python.monologue_end import _60_rename_chat as rename_chat
8 from plugins._chat_naming.helpers import naming
9
10
11 pytestmark = pytest.mark.asyncio
12
13
14 class _Message:
15 def __init__(self, content, *, ai=False, sequence=0):
16 self.content = content
17 self.ai = ai
18 self.sequence = sequence
19
20
21 class _History:
22 def __init__(self, messages):
23 self._messages = messages
24
25 def all_messages(self):
26 return list(self._messages)
27
28
29 class _Agent:
30 def __init__(self, messages, *, name="", response="Generated Name"):
31 self.context = SimpleNamespace(
32 id="ctx-naming",
33 name=name,
34 type=AgentContextType.USER,
35 )
36 self.context.agent0 = self
37 self.history = _History(messages)
38 self.config = SimpleNamespace(profile="agent0")
39 self._response = response
40 self.utility_calls = []
41
42 def read_prompt(self, name, **kwargs):
43 return f"{name}:{kwargs}"
44
45 async def call_utility_model(self, **kwargs):
46 self.utility_calls.append(kwargs)
47 return self._response
48
49
50 async def test_user_message_selection_excludes_assistant_work_and_tool_results():
51 agent = _Agent(
52 [
53 _Message({"user_message": "Plan a launch"}, sequence=1),
54 _Message({"tool_name": "search", "tool_result": "internal work"}, sequence=2),
55 _Message("assistant response", ai=True, sequence=3),
56 _Message({"user_intervention": "Okay, do it"}, sequence=4),
57 ]
58 )
59
60 assert naming.get_user_messages(agent) == ["Plan a launch", "Okay, do it"]
61 assert naming.latest_user_sequence(agent) == 4
62
63
64 async def test_once_mode_uses_first_message_and_does_not_override_a_name(monkeypatch):
65 scheduled = []
66 agent = _Agent(
67 [
68 _Message({"user_message": "First request"}, sequence=1),
69 _Message({"user_message": "Second request"}, sequence=2),
70 ]
71 )
72 monkeypatch.setattr(
73 rename_chat.naming,
74 "get_config",
75 lambda _agent: {"automatic_naming": True, "automatic_naming_mode": "once"},
76 )
77 monkeypatch.setattr(rename_chat.asyncio, "create_task", lambda coro: scheduled.append(coro))
78
79 await rename_chat.RenameChat(agent=agent).execute()
80 assert len(scheduled) == 1
81
82 captured = {}
83
84 async def generate(_agent, **kwargs):
85 captured.update(kwargs)
86 return "First Request"
87
88 saved = []
89 monkeypatch.setattr(rename_chat.naming, "generate_name", generate)
90 monkeypatch.setattr(
91 rename_chat.naming,
92 "save_context_name",
93 lambda _agent, name: saved.append(name),
94 )
95 monkeypatch.setattr(rename_chat.AgentContext, "get", lambda _id: agent.context)
96 await scheduled.pop()
97
98 assert captured["user_messages"] == ["First request"]
99 assert saved == ["First Request"]
100
101 agent.context.name = "Manual Name"
102 await rename_chat.RenameChat(agent=agent).execute()
103 assert scheduled == []
104
105
106 async def test_always_mode_passes_recent_user_context(monkeypatch):
107 messages = [
108 _Message({"user_message": f"Message {number}"}, sequence=number)
109 for number in range(1, 7)
110 ]
111 agent = _Agent(messages, name="Existing Name")
112 scheduled = []
113 monkeypatch.setattr(
114 rename_chat.naming,
115 "get_config",
116 lambda _agent: {"automatic_naming": True, "automatic_naming_mode": "always"},
117 )
118 monkeypatch.setattr(rename_chat.asyncio, "create_task", lambda coro: scheduled.append(coro))
119
120 await rename_chat.RenameChat(agent=agent).execute()
121 captured = {}
122
123 async def generate(_agent, **kwargs):
124 captured.update(kwargs)
125 return "Current Topic"
126
127 monkeypatch.setattr(rename_chat.naming, "generate_name", generate)
128 monkeypatch.setattr(rename_chat.naming, "save_context_name", lambda *_args: None)
129 monkeypatch.setattr(rename_chat.AgentContext, "get", lambda _id: agent.context)
130 await scheduled.pop()
131
132 assert captured["user_messages"] == [
133 "Message 3",
134 "Message 4",
135 "Message 5",
136 "Message 6",
137 ]
138 assert captured["current_name"] == "Existing Name"
139
140
141 async def test_rename_failure_sends_scoped_utility_model_notification(monkeypatch):
142 sent = []
143 agent = _Agent([_Message({"user_message": "Plan a launch"}, sequence=1)])
144
145 async def fail(*_args, **_kwargs):
146 raise RuntimeError("offline")
147
148 monkeypatch.setattr(rename_chat.naming, "generate_name", fail)
149 monkeypatch.setattr(
150 rename_chat.NotificationManager,
151 "send_notification",
152 lambda **kwargs: sent.append(kwargs),
153 )
154
155 await rename_chat.RenameChat(agent=agent).change_name(
156 messages=["Plan a launch"],
157 request_sequence=1,
158 only_if_unnamed=True,
159 )
160
161 assert len(sent) == 1
162 assert sent[0]["type"] == rename_chat.NotificationType.ERROR
163 assert sent[0]["title"] == "Chat Naming Failed"
164 assert sent[0]["id"] == "chat_naming_failed_ctx-naming"
165
166
167 async def test_generated_name_is_normalized_and_bounded(monkeypatch):
168 agent = _Agent(
169 [_Message({"user_message": "Plan the release"}, sequence=1)],
170 response=' "A very long generated release planning title that should be shortened" ',
171 )
172 monkeypatch.setattr(
173 "plugins._model_config.helpers.model_config.get_utility_model_config",
174 lambda _agent: {"ctx_length": 1000},
175 )
176
177 result = await naming.generate_name(agent)
178
179 assert len(result) <= naming.GENERATED_NAME_LIMIT
180 assert result.endswith("...")
181 assert agent.utility_calls[0]["background"] is True
182
183
184 async def test_naming_input_stays_within_utility_context_budget(monkeypatch):
185 agent = _Agent(
186 [
187 _Message({"user_message": "older context " * 2000}, sequence=1),
188 _Message({"user_message": "LATEST_CONTEXT " * 2000}, sequence=2),
189 ],
190 response="Budgeted Name",
191 )
192 monkeypatch.setattr(
193 "plugins._model_config.helpers.model_config.get_utility_model_config",
194 lambda _agent: {"ctx_length": 300},
195 )
196
197 await naming.generate_name(agent)
198
199 call = agent.utility_calls[0]
200 estimated_tokens = naming.tokens.approximate_tokens(
201 call["system"]
202 ) + naming.tokens.approximate_tokens(call["message"])
203 assert estimated_tokens <= int(300 * naming.UTILITY_CONTEXT_INPUT_RATIO)
204 assert "LATEST_CONTEXT" in call["message"]
205 assert "older context" not in call["message"]
206
207
208 async def test_manual_api_generates_and_saves_with_target_chat_agent(monkeypatch):
209 from plugins._chat_naming.api import chat_name
210
211 agent = _Agent([_Message({"user_message": "Name this chat"}, sequence=1)])
212 monkeypatch.setattr(chat_name.AgentContext, "get", lambda _id: agent.context)
213
214 async def generate(target_agent, **kwargs):
215 assert target_agent is agent
216 assert kwargs["current_name"] == ""
217 return "Generated Chat"
218
219 saved = []
220 monkeypatch.setattr(chat_name.naming, "generate_name", generate)
221 monkeypatch.setattr(
222 chat_name.naming,
223 "save_context_name",
224 lambda target_agent, name: saved.append((target_agent, name)),
225 )
226 handler = object.__new__(chat_name.ChatName)
227
228 generated = await handler.process(
229 {"action": "generate", "kind": "chat", "item_id": agent.context.id},
230 None,
231 )
232 renamed = await handler.process(
233 {
234 "action": "save",
235 "kind": "chat",
236 "item_id": agent.context.id,
237 "name": " Manual Chat ",
238 },
239 None,
240 )
241
242 assert generated == {"ok": True, "name": "Generated Chat"}
243 assert renamed == {"ok": True, "name": "Manual Chat"}
244 assert saved == [(agent, "Manual Chat")]
245
246
247 async def test_rename_command_saves_a_custom_name(monkeypatch):
248 agent = _Agent([_Message({"user_message": "Plan a launch"}, sequence=1)])
249 saved = []
250 monkeypatch.setattr(
251 rename_command.naming,
252 "save_context_name",
253 lambda target_agent, name: saved.append((target_agent, name)) or name,
254 )
255
256 result = await rename_command.run(
257 {
258 "invocation": {"raw_arguments": "New Chat Name"},
259 "context": {"agent": agent},
260 }
261 )
262
263 assert saved == [(agent, "New Chat Name")]
264 assert result == {
265 "text": "",
266 "effects": [
267 {
268 "type": "toast",
269 "message": 'Chat renamed to "New Chat Name".',
270 "level": "success",
271 }
272 ],
273 }
274
275
276 async def test_rename_auto_generates_and_saves_a_name(monkeypatch):
277 agent = _Agent([_Message({"user_message": "Plan a launch"}, sequence=1)])
278 generated = []
279 saved = []
280
281 async def generate(target_agent):
282 generated.append(target_agent)
283 return "Launch Plan"
284
285 monkeypatch.setattr(rename_command.naming, "generate_name", generate)
286 monkeypatch.setattr(
287 rename_command.naming,
288 "save_context_name",
289 lambda target_agent, name: saved.append((target_agent, name)) or name,
290 )
291
292 result = await rename_command.run(
293 {
294 "invocation": {"raw_arguments": "auto"},
295 "context": {"agent": agent},
296 }
297 )
298
299 assert generated == [agent]
300 assert saved == [(agent, "Launch Plan")]
301 assert result["effects"][0]["message"] == 'Chat renamed to "Launch Plan".'
302
303
304 async def test_manual_task_rename_updates_scheduler_and_context(monkeypatch):
305 from plugins._chat_naming.api import chat_name
306
307 agent = _Agent([_Message({"user_message": "Run a report"}, sequence=1)], name="Old Task")
308
309 class _Scheduler:
310 def __init__(self):
311 self.updated = []
312
313 async def reload(self):
314 return None
315
316 async def update_task(self, item_id, **kwargs):
317 self.updated.append((item_id, kwargs))
318 return SimpleNamespace(name=kwargs["name"])
319
320 scheduler = _Scheduler()
321 saved = []
322 dirty = []
323 monkeypatch.setattr(chat_name.AgentContext, "get", lambda _id: agent.context)
324 monkeypatch.setattr(chat_name.TaskScheduler, "get", lambda: scheduler)
325 monkeypatch.setattr(chat_name, "save_tmp_chat", lambda context: saved.append(context.name))
326 monkeypatch.setattr(chat_name, "mark_dirty_all", lambda *, reason: dirty.append(reason))
327
328 response = await object.__new__(chat_name.ChatName).process(
329 {
330 "action": "save",
331 "kind": "task",
332 "item_id": agent.context.id,
333 "name": "Daily Report",
334 },
335 None,
336 )
337
338 assert response == {"ok": True, "name": "Daily Report"}
339 assert scheduler.updated == [(agent.context.id, {"name": "Daily Report"})]
340 assert agent.context.name == "Daily Report"
341 assert saved == ["Daily Report"]
342 assert dirty == ["plugins._chat_naming.save_task_name"]
343
344
345 async def test_chat_naming_endpoints_keep_default_auth_and_csrf_protection():
346 from plugins._chat_naming.api.chat_name import ChatName
347
348 assert ChatName.requires_auth() is True
349 assert ChatName.requires_csrf() is True