propagate agent_id where needed

Author: flaviusburca

surogates/api/routes/sessions.py

@@ -80,9 +80,12 @@ def _get_session_store(request: Request) -> SessionStore:
 
 
 async def _get_session_for_tenant(
-    store: SessionStore, session_id: UUID, tenant: TenantContext
+    request: Request,
+    session_id: UUID,
+    tenant: TenantContext,
 ) -> Session:
-    """Fetch a session and verify it belongs to the tenant's org."""
+    """Fetch a session and verify it belongs to the tenant's org and this agent."""
+    store = _get_session_store(request)
     try:
         session = await store.get_session(session_id)
     except SessionNotFoundError:
@@ -91,7 +94,8 @@ async def _get_session_for_tenant(
             detail=f"Session {session_id} not found.",
         )
 
-    if session.org_id != tenant.org_id:
+    agent_id = request.app.state.settings.agent_id
+    if session.org_id != tenant.org_id or session.agent_id != agent_id:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=f"Session {session_id} not found.",
@@ -170,7 +174,7 @@ async def send_message(
 ) -> SendMessageResponse:
     """Send a user message to a session, triggering agent processing."""
     store = _get_session_store(request)
-    session = await _get_session_for_tenant(store, session_id, tenant)
+    session = await _get_session_for_tenant(request, session_id, tenant)
 
     if session.status not in ("active", "idle", "failed", "paused"):
         raise HTTPException(
@@ -226,8 +230,7 @@ async def confirm_disclosure(
     enforcement is enabled.  Typically called by the frontend after
     showing the AI disclosure notice to the user.
     """
-    store = _get_session_store(request)
-    await _get_session_for_tenant(store, session_id, tenant)
+    await _get_session_for_tenant(request, session_id, tenant)
 
     governance = getattr(request.app.state, "governance_gate", None)
     if governance is not None:
@@ -241,8 +244,7 @@ async def get_session(
     tenant: TenantContext = Depends(get_current_tenant),
 ) -> Session:
     """Retrieve metadata for a single session."""
-    store = _get_session_store(request)
-    return await _get_session_for_tenant(store, session_id, tenant)
+    return await _get_session_for_tenant(request, session_id, tenant)
 
 
 @router.get("/sessions", response_model=ListSessionsResponse)
@@ -252,8 +254,9 @@ async def list_sessions(
     limit: int = 50,
     offset: int = 0,
 ) -> ListSessionsResponse:
-    """List the authenticated user's sessions (paginated, newest first)."""
+    """List the authenticated user's sessions for this agent, newest first."""
     store = _get_session_store(request)
+    settings = request.app.state.settings
 
     if limit < 1:
         limit = 1
@@ -265,6 +268,7 @@ async def list_sessions(
     sessions = await store.list_sessions(
         org_id=tenant.org_id,
         user_id=tenant.user_id,
+        agent_id=settings.agent_id,
         limit=limit,
         offset=offset,
     )
@@ -285,7 +289,7 @@ async def pause_session(
 ) -> Session:
     """Pause an active session."""
     store = _get_session_store(request)
-    session = await _get_session_for_tenant(store, session_id, tenant)
+    session = await _get_session_for_tenant(request, session_id, tenant)
 
     if session.status not in ("active", "processing", "paused"):
         raise HTTPException(
@@ -319,7 +323,7 @@ async def resume_session(
 ) -> Session:
     """Resume a paused session."""
     store = _get_session_store(request)
-    session = await _get_session_for_tenant(store, session_id, tenant)
+    session = await _get_session_for_tenant(request, session_id, tenant)
 
     if session.status != "paused":
         raise HTTPException(
@@ -348,7 +352,7 @@ async def delete_session(
 ) -> None:
     """Archive (soft-delete) a session and delete its workspace storage."""
     store = _get_session_store(request)
-    await _get_session_for_tenant(store, session_id, tenant)
+    await _get_session_for_tenant(request, session_id, tenant)
 
     await store.update_session_status(session_id, "archived")
 

surogates/channels/identity.py

@@ -113,6 +113,7 @@ async def get_or_create_channel_session(
         result = await db.execute(
             select(SessionRow)
             .where(SessionRow.user_id == user_id)
+            .where(SessionRow.agent_id == agent_id)
             .where(SessionRow.channel == channel)
             .where(SessionRow.status.in_(["active", "processing", "paused"]))
             .where(

surogates/harness/tool_exec.py

@@ -569,6 +569,7 @@ async def execute_single_tool(
                 tool_name,
                 tool_args,
                 session_id=str(session.id),
+                agent_id=session.agent_id,
                 tenant=tenant,
                 session_store=store,
                 redis=redis,

surogates/jobs/reset_idle_sessions.py

@@ -517,6 +517,7 @@ async def reset_idle_sessions(dry_run: bool = False) -> int:
     daily_at_hour = reset_cfg.at_hour if reset_cfg.mode in ("daily", "both") else None
     idle_sessions = await session_store.find_idle_sessions(
         idle_minutes=reset_cfg.idle_minutes,
+        agent_id=settings.agent_id,
         daily_at_hour=daily_at_hour,
         mode=reset_cfg.mode,
     )

surogates/orchestrator/worker.py

@@ -167,6 +167,17 @@ async def run_worker(settings: Settings) -> None:
         )
     configured_org_id = UUID(settings.org_id)
 
+    # Resolve the agent_id from config (required).  Sessions belong to an
+    # agent; a worker refuses to process sessions that belong to a
+    # different agent.
+    if not settings.agent_id:
+        raise RuntimeError(
+            "SUROGATES_AGENT_ID is not set. Each worker instance serves "
+            "exactly one agent. Set agent_id in config.yaml or "
+            "SUROGATES_AGENT_ID env var."
+        )
+    configured_agent_id = settings.agent_id
+
     # 7. Harness factory -- creates a fully-wired AgentHarness for a given session.
     async def harness_factory(session_id: UUID) -> AgentHarness:
         """Build an AgentHarness with all dependencies injected.
@@ -176,6 +187,15 @@ async def harness_factory(session_id: UUID) -> AgentHarness:
         # Load session to get user_id.
         session = await session_store.get_session(session_id)
 
+        # Refuse to process sessions that belong to a different agent —
+        # defence-in-depth in case a foreign session id leaks into this
+        # worker's queue.
+        if session.agent_id != configured_agent_id:
+            raise RuntimeError(
+                f"session {session_id} belongs to agent {session.agent_id!r}, "
+                f"this worker serves agent {configured_agent_id!r}"
+            )
+
         # Load org + user from DB.
         from sqlalchemy import select as sa_select
         from surogates.db.models import Org, User

surogates/session/store.py

@@ -136,17 +136,19 @@ async def list_sessions(
         self,
         org_id: UUID,
         user_id: UUID,
+        agent_id: str,
         *,
         limit: int = 50,
         offset: int = 0,
     ) -> list[Session]:
-        """Return sessions for a user within an org, newest first."""
+        """Return sessions for a user within an org, scoped to one agent, newest first."""
         async with self._sf() as db:
             result = await db.execute(
                 select(SessionRow)
                 .where(
                     SessionRow.org_id == org_id,
                     SessionRow.user_id == user_id,
+                    SessionRow.agent_id == agent_id,
                     SessionRow.status != "archived",
                 )
                 .order_by(SessionRow.created_at.desc())
@@ -508,6 +510,7 @@ async def get_pending_events(self, session_id: UUID) -> list[Event]:
     async def find_idle_sessions(
         self,
         idle_minutes: int,
+        agent_id: str,
         *,
         daily_at_hour: int | None = None,
         mode: str = "idle",
@@ -557,13 +560,18 @@ async def find_idle_sessions(
             LEFT JOIN session_leases l
                 ON l.session_id = s.id AND l.expires_at > now()
             WHERE s.status IN ('active', 'idle')
+              AND s.agent_id = :agent_id
               AND l.session_id IS NULL
               AND {where_clause}
             ORDER BY s.updated_at ASC
             LIMIT :lim
         """
 
-        params: dict[str, Any] = {"idle_minutes": idle_minutes, "lim": limit}
+        params: dict[str, Any] = {
+            "idle_minutes": idle_minutes,
+            "agent_id": agent_id,
+            "lim": limit,
+        }
         if daily_at_hour is not None:
             params["at_hour"] = daily_at_hour
 

surogates/tools/builtin/session_search.py

@@ -266,6 +266,7 @@ async def _list_recent_sessions(
     session_store: Any,
     org_id: UUID,
     user_id: UUID,
+    agent_id: str,
     limit: int,
     current_session_id: UUID | None = None,
 ) -> str:
@@ -274,6 +275,7 @@ async def _list_recent_sessions(
         sessions = await session_store.list_sessions(
             org_id=org_id,
             user_id=user_id,
+            agent_id=agent_id,
             limit=limit + 5,  # fetch extra to skip current
         )
 
@@ -355,6 +357,7 @@ async def session_search(
     session_store: Any = None,
     org_id: UUID | None = None,
     user_id: UUID | None = None,
+    agent_id: str = "",
     current_session_id: UUID | None = None,
     auxiliary_fn: Any | None = None,
 ) -> str:
@@ -386,13 +389,21 @@ async def session_search(
             "error": "Tenant context (org_id/user_id) not available.",
         })
 
+    # Search stays within the current session's agent — cross-agent history is
+    # not addressable.
+    if not agent_id:
+        return json.dumps({
+            "success": False,
+            "error": "agent_id is required to scope search to this agent.",
+        })
+
     limit = min(limit, 5)  # Cap at 5 sessions to avoid excessive LLM calls
 
     # Recent sessions mode: when query is empty, return metadata for recent sessions.
     # No LLM calls -- just DB queries for titles, previews, timestamps.
     if not query or not query.strip():
         return await _list_recent_sessions(
-            session_store, org_id, user_id, limit, current_session_id,
+            session_store, org_id, user_id, agent_id, limit, current_session_id,
         )
 
     query = query.strip()
@@ -447,6 +458,7 @@ async def session_search(
                     JOIN sessions s ON s.id = e.session_id
                     WHERE s.org_id = :org_id
                       AND s.user_id = :user_id
+                      AND s.agent_id = :agent_id
                       AND s.status != 'archived'
                       AND to_tsvector('english', COALESCE(e.data->>'content', '') || ' ' || COALESCE(e.data->>'result', ''))
                           @@ plainto_tsquery('english', :query)
@@ -458,6 +470,7 @@ async def session_search(
                     "query": query,
                     "org_id": org_id,
                     "user_id": user_id,
+                    "agent_id": agent_id,
                     "limit": 50,  # Get more matches to find unique sessions
                 }
 
@@ -740,6 +753,7 @@ async def _session_search_handler(
     tenant = kwargs.get("tenant", {})
     org_id = tenant.get("org_id") if isinstance(tenant, dict) else getattr(tenant, "org_id", None)
     user_id = tenant.get("user_id") if isinstance(tenant, dict) else getattr(tenant, "user_id", None)
+    agent_id = kwargs.get("agent_id", "")
     current_session_id = kwargs.get("session_id")
     auxiliary_fn = kwargs.get("auxiliary_fn")
 
@@ -758,6 +772,7 @@ async def _session_search_handler(
         session_store=store,
         org_id=org_id,
         user_id=user_id,
+        agent_id=agent_id,
         current_session_id=current_session_id,
         auxiliary_fn=auxiliary_fn,
     )

tests/integration/test_session_store.py

@@ -97,17 +97,19 @@ async def test_list_sessions(session_store, session_factory):
         ids.append(s.id)
 
     # Fetch all
-    sessions = await session_store.list_sessions(org_id, user_id)
+    sessions = await session_store.list_sessions(org_id, user_id, "test-agent")
     returned_ids = {s.id for s in sessions}
     assert all(sid in returned_ids for sid in ids)
 
     # Pagination: limit=2
-    page1 = await session_store.list_sessions(org_id, user_id, limit=2)
+    page1 = await session_store.list_sessions(
+        org_id, user_id, "test-agent", limit=2
+    )
     assert len(page1) == 2
 
     # Pagination: offset=2
     page2 = await session_store.list_sessions(
-        org_id, user_id, limit=10, offset=2
+        org_id, user_id, "test-agent", limit=10, offset=2
     )
     assert len(page2) >= 1