diff --git a/solvent/chat.py b/solvent/chat.py index fb1651a..ce9de74 100644 --- a/solvent/chat.py +++ b/solvent/chat.py @@ -21,6 +21,15 @@ _BUDGET_RE = re.compile(r"(?:budget|pay|spend)\s*[:\$]?\s*\$?(\d+(?:\.\d{1,2})?)", re.I) +def _new_chat_job_id(agent: Solvent) -> str: + """Generate a server-owned chat job ID that does not overwrite an existing job.""" + for _ in range(10): + job_id = "T" + uuid.uuid4().hex[:8] + if not agent.t.get_job(job_id): + return job_id + return "T" + uuid.uuid4().hex + + def _make_executor(agent: Solvent, session_id: str, live_search: bool): ctx = tools.ToolContext() @@ -65,7 +74,7 @@ def run(name: str, args: dict) -> str: topic = args.get("topic", "") budget = int(args.get("budget_cents", 0)) email = args.get("customer_email", "client@example.com") - job_id = args.get("job_id") or ("T" + uuid.uuid4().hex[:8]) + job_id = _new_chat_job_id(agent) job = { "id": job_id, "topic": topic, @@ -77,7 +86,8 @@ def run(name: str, args: dict) -> str: "context": f"Commissioned via chat session {session_id}", } result = agent.enqueue_job(job) - agent.t.update_chat_session(session_id, notify_job_id=job_id, pending_job_json="") + if result.get("stage") != "declined" and not result.get("error"): + agent.t.update_chat_session(session_id, notify_job_id=job_id, pending_job_json="") if result.get("url"): return json.dumps({ "job_id": job_id, diff --git a/tests/test_chat_tools.py b/tests/test_chat_tools.py index a0f076c..8041a93 100644 --- a/tests/test_chat_tools.py +++ b/tests/test_chat_tools.py @@ -1,3 +1,4 @@ +import os import tempfile import unittest from pathlib import Path @@ -12,6 +13,8 @@ class TestChatTools(unittest.TestCase): def setUp(self): + self._env = patch.dict(os.environ, {"SOLVENT_DELIVERY_SECRET": "x" * 32}, clear=False) + self._env.start() self.tmp = tempfile.TemporaryDirectory() self.db = Path(self.tmp.name) / "t.db" self.t = Treasury(path=self.db) @@ -24,6 +27,7 @@ def setUp(self): def tearDown(self): self.tmp.cleanup() + self._env.stop() def test_format_job_notification(self): msg = format_job_notification({"stage": "delivered", "job_id": "J1", "url": "https://x"}) @@ -74,6 +78,42 @@ def counting(self, name, args): self.assertLessEqual(calls["n"], 5) self.assertIn("set", reply.lower()) + @patch.object(nemotron, "complete") + def test_submit_brief_ignores_tool_supplied_job_id(self, mock_complete): + self.t.upsert_job( + "VICTIM1", + "awaiting_payment", + topic="Victim topic", + budget_cents=7500, + customer_email="victim@example.com", + job_payload_json={"id": "VICTIM1", "topic": "Victim topic"}, + ) + mock_complete.side_effect = [ + ( + '{"name": "submit_brief", "arguments": ' + '{"job_id": "VICTIM1", "topic": "Attacker topic", ' + '"budget_cents": 5000, "customer_email": "attacker@example.com"}}', + {}, + ), + ("Done.", {}), + ] + + handle_message( + self.session["id"], + "Submit a brief", + agent=self.agent, + memory=self.memory, + ) + + victim = self.t.get_job("VICTIM1") + self.assertEqual(victim["topic"], "Victim topic") + self.assertEqual(victim["customer_email"], "victim@example.com") + session = self.t.get_chat_session(self.session["id"]) + self.assertNotEqual(session.get("notify_job_id"), "VICTIM1") + created = [j for j in self.t.list_jobs() if j["id"] != "VICTIM1"] + self.assertEqual(len(created), 1) + self.assertEqual(created[0]["topic"], "Attacker topic") + @patch.object(nemotron, "complete") def test_submit_brief_via_tool(self, mock_complete): mock_complete.side_effect = [