chore: standardize session not found error (#3031)

# What does this PR do?
1. Creates a new `SessionNotFoundError` class
2. Implements the new class where appropriate 

Relates to #2379

Signed-off-by: Nathan Weinberg <nweinber@redhat.com>

llama_stack/apis/common/errors.py

@@ -50,3 +50,11 @@ class ToolGroupNotFoundError(ValueError):
             f"Tool group '{toolgroup_name}' not found. Use client.toolgroups.list() to list available tool groups."
         )
         super().__init__(message)
+
+
+class SessionNotFoundError(ValueError):
+    """raised when Llama Stack cannot find a referenced session or access is denied"""
+
+    def __init__(self, session_name: str) -> None:
+        message = f"Session '{session_name}' not found or access denied."
+        super().__init__(message)

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -44,6 +44,7 @@ from llama_stack.apis.common.content_types import (
     ToolCallDelta,
     ToolCallParseStatus,
 )
+from llama_stack.apis.common.errors import SessionNotFoundError
 from llama_stack.apis.inference import (
     ChatCompletionResponseEventType,
     CompletionMessage,
@@ -214,7 +215,7 @@ class ChatAgent(ShieldRunnerMixin):
         is_resume = isinstance(request, AgentTurnResumeRequest)
         session_info = await self.storage.get_session_info(request.session_id)
         if session_info is None:
-            raise ValueError(f"Session {request.session_id} not found")
+            raise SessionNotFoundError(request.session_id)
 
         turns = await self.storage.get_session_turns(request.session_id)
         if is_resume and len(turns) == 0:

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -230,8 +230,6 @@ class MetaReferenceAgentsImpl(Agents):
         agent = await self._get_agent_impl(agent_id)
 
         session_info = await agent.storage.get_session_info(session_id)
-        if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
         turns = await agent.storage.get_session_turns(session_id)
         if turn_ids:
             turns = [turn for turn in turns if turn.turn_id in turn_ids]
@@ -244,9 +242,6 @@ class MetaReferenceAgentsImpl(Agents):
 
     async def delete_agents_session(self, agent_id: str, session_id: str) -> None:
         agent = await self._get_agent_impl(agent_id)
-        session_info = await agent.storage.get_session_info(session_id)
-        if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
 
         # Delete turns first, then the session
         await agent.storage.delete_session_turns(session_id)

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -10,6 +10,7 @@ import uuid
 from datetime import UTC, datetime
 
 from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn
+from llama_stack.apis.common.errors import SessionNotFoundError
 from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
 from llama_stack.core.access_control.datatypes import AccessRule
 from llama_stack.core.datatypes import User
@@ -61,12 +62,12 @@ class AgentPersistence:
         )
         return session_id
 
-    async def get_session_info(self, session_id: str) -> AgentSessionInfo | None:
+    async def get_session_info(self, session_id: str) -> AgentSessionInfo:
         value = await self.kvstore.get(
             key=f"session:{self.agent_id}:{session_id}",
         )
         if not value:
-            return None
+            raise SessionNotFoundError(session_id)
 
         session_info = AgentSessionInfo(**json.loads(value))
 
@@ -95,7 +96,7 @@ class AgentPersistence:
     async def add_vector_db_to_session(self, session_id: str, vector_db_id: str):
         session_info = await self.get_session_if_accessible(session_id)
         if session_info is None:
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         session_info.vector_db_id = vector_db_id
         await self.kvstore.set(
@@ -105,7 +106,7 @@ class AgentPersistence:
 
     async def add_turn_to_session(self, session_id: str, turn: Turn):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}",
@@ -114,7 +115,7 @@ class AgentPersistence:
 
     async def get_session_turns(self, session_id: str) -> list[Turn]:
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         values = await self.kvstore.values_in_range(
             start_key=f"session:{self.agent_id}:{session_id}:",
@@ -137,7 +138,7 @@ class AgentPersistence:
 
     async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None:
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         value = await self.kvstore.get(
             key=f"session:{self.agent_id}:{session_id}:{turn_id}",
@@ -148,7 +149,7 @@ class AgentPersistence:
 
     async def set_in_progress_tool_call_step(self, session_id: str, turn_id: str, step: ToolExecutionStep):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
@@ -166,7 +167,7 @@ class AgentPersistence:
 
     async def set_num_infer_iters_in_turn(self, session_id: str, turn_id: str, num_infer_iters: int):
         if not await self.get_session_if_accessible(session_id):
-            raise ValueError(f"Session {session_id} not found or access denied")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.set(
             key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
@@ -218,6 +219,6 @@ class AgentPersistence:
         """
         session_info = await self.get_session_info(session_id)
         if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
+            raise SessionNotFoundError(session_id)
 
         await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}")