From 594d1863893281a721dbaa68a0796ea94f3554b9 Mon Sep 17 00:00:00 2001
From: whiterabbit1983 <asmatic075@gmail.com>
Date: Tue, 16 Apr 2024 01:53:55 +0000
Subject: [PATCH] fix: Raise agent not found error on assertion failure (#191)

---
 agents-api/agents_api/routers/agents/routers.py   |  2 +-
 agents-api/agents_api/routers/sessions/routers.py | 12 ++++++++++++
 agents-api/agents_api/routers/users/routers.py    |  2 +-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/agents-api/agents_api/routers/agents/routers.py b/agents-api/agents_api/routers/agents/routers.py
index fc1d91c4e..e8ce03a0b 100644
--- a/agents-api/agents_api/routers/agents/routers.py
+++ b/agents-api/agents_api/routers/agents/routers.py
@@ -133,7 +133,7 @@ async def update_agent(
             detail="Agent not found",
         )
     except QueryException as e:
-        if e.code == "transact::assertion_failure":
+        if e.code in ("transact::assertion_failure", "eval::assert_some_failure"):
             raise AgentNotFoundError(x_developer_id, agent_id)
 
         raise
diff --git a/agents-api/agents_api/routers/sessions/routers.py b/agents-api/agents_api/routers/sessions/routers.py
index 06821b6e0..59c841fed 100644
--- a/agents-api/agents_api/routers/sessions/routers.py
+++ b/agents-api/agents_api/routers/sessions/routers.py
@@ -161,6 +161,12 @@ async def update_session(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="Session not found",
         )
+    except QueryException as e:
+        # the code is not so informative now, but it may be a good solution in the future
+        if e.code in ("transact::assertion_failure", "eval::assert_some_failure"):
+            raise SessionNotFoundError(x_developer_id, session_id)
+
+        raise
 
 
 @router.patch("/sessions/{session_id}", tags=["sessions"])
@@ -186,6 +192,12 @@ async def patch_session(
             status_code=status.HTTP_404_NOT_FOUND,
             detail="Session not found",
         )
+    except QueryException as e:
+        # the code is not so informative now, but it may be a good solution in the future
+        if e.code in ("transact::assertion_failure", "eval::assert_some_failure"):
+            raise SessionNotFoundError(x_developer_id, session_id)
+
+        raise
 
 
 @router.get("/sessions/{session_id}/suggestions", tags=["sessions"])
diff --git a/agents-api/agents_api/routers/users/routers.py b/agents-api/agents_api/routers/users/routers.py
index 66ed4d710..5ae3c0fb3 100644
--- a/agents-api/agents_api/routers/users/routers.py
+++ b/agents-api/agents_api/routers/users/routers.py
@@ -108,7 +108,7 @@ async def update_user(
         )
     except QueryException as e:
         # the code is not so informative now, but it may be a good solution in the future
-        if e.code == "transact::assertion_failure":
+        if e.code in ("transact::assertion_failure", "eval::assert_some_failure"):
             raise UserNotFoundError(x_developer_id, user_id)
 
         raise