entity isolation and estimateTokenFix

packages/memory-service/src/graph/index.js

@@ -74,4 +74,17 @@ function getEpisodeIdsByEntities(entityIds) {
     ).all(...entityIds).map(r => r.episode_id);
 }
 
-module.exports = { getNeighborhood, getEntityNeighbors, getEpisodeIdsByEntities };
+//Entity IDs linked (via entity_episodes) to any episode in a given session
+//Scopes non-project entity search to the session's own entities under the "isolated chats" model.
+//Entities globally deduped
+function getEntityIdsBySession(sessionId){
+    const db = getDB();
+    return db.prepare(`
+        SELECT DISTINCT ee.entity_id
+        FROM entity_episodes ee
+        JOIN episodes e on e.id = ee.episode_id
+        WHERE e.session_id = ?
+    `).all(sessionId).map(r => r.entity_id);
+}
+
+module.exports = { getNeighborhood, getEntityNeighbors, getEpisodeIdsByEntities, getEntityIdsBySession };

packages/memory-service/src/index.js

@@ -162,6 +162,13 @@ app.get('/sessions/:id/episode-stats', (req, res) => {
   res.json(episodic.getSessionEpisodeStats(Number(req.params.id)));
 });
 
+//Entity IDs linked to this session's episodes: sesion-scoped entity search
+app.get('/sessions/:id/entity-ids', (req, res) => {
+  res.json({
+    entityIds: graph.getEntityIdsBySession(Number(req.params.id))
+  })
+})
+
 // Episodes newer than :afterId, chronological — the un-summarized tail.
 app.get('/sessions/:id/episodes/since/:afterId', (req, res) => {
   const episodes = episodic.getEpisodesSince(Number(req.params.id), Number(req.params.afterId));

packages/orchestration-service/src/chat/index.js

@@ -127,16 +127,21 @@ async function getSemanticEpisodes(
   }
 }
 
-async function getRelevantEntities(vector, projectId = null) {
+async function getRelevantEntities(vector, { projectId = null, sessionId } = {}) {
     if (!vector) return [];
     try {
-        const results = await qdrant.searchEntities(vector, { projectId });
-        logger.info(
-            '[orchestration] Entity search results:',
-            results.map((r) => ({ name: r.payload?.name, score: r.score })),
-        );
-        // Include the Qdrant point ID (== SQLite entity ID) for graph traversal
-        return results.map((r) => r.payload ? { id: r.id, ...r.payload } : null).filter(Boolean);
+        let allowedIds;
+        if (projectId === null || projectId === undefined) {
+            // Non-project chat is its own island — scope to entities linked to
+            // THIS session. No links yet ⇒ nothing to retrieve, and we return
+            // early so searchEntities is never called unfiltered.
+            allowedIds = await memory.getEntityIdsBySession(sessionId);
+            if (allowedIds.length === 0) return [];
+        }
+        const results = await qdrant.searchEntities(vector, { projectId, allowedIds });
+        logger.info('[orchestration] Entity search results:',
+            results.map(r => ({ name: r.payload?.name, score: r.score })));
+        return results.map(r => r.payload ? { id: r.id, ...r.payload } : null).filter(Boolean);
     } catch (err) {
         logger.debug('[orchestration] Entity search failed, continuing without:', err.message);
         return [];
@@ -302,7 +307,7 @@ async function assembleContext(externalId, userMessage) {
 
         [fusedWithScores, entityResults] = await Promise.all([
             getFusedEpisodes(userMessage, session, recentIds, projectSessionIds, queryVector, { semanticLimit, scoreThreshold, semanticWeight, keywordWeight }),
-            getRelevantEntities(queryVector, session.project_id ?? null),
+            getRelevantEntities(queryVector, { projectId: session.project_id ?? null, sessionId: session.id }),
         ]);
     } else {
         logger.debug('[orchestration] Trivial turn — skipping semantic/keyword/entity retrieval');

packages/orchestration-service/src/services/memory.js

@@ -216,6 +216,14 @@ async function getEpisodesByEntities(entityIds) {
     return res.json(); // { episodeIds: [...] }
 }
 
+async function getEntityIdsBySession(sessionId){
+    const res = await fetch(`${BASE_URL}/sessions/${sessionId}/entity-ids`);
+
+    if (!res.ok) throw new Error(`Entity-ids-by-session error: ${res.status}`);
+    const {entityIds} = await res.jaon();
+    return entityIds;
+}
+
 module.exports = {
     getSessionByExternalId,
     createSession,

packages/orchestration-service/src/services/qdrant.js

@@ -29,38 +29,28 @@ async function searchEpisodes( vector, {limit = ORCHESTRATION.RECENT_EPISODE_LIM
     return data.result;
 }
 
-async function searchEntities(vector, { limit = ORCHESTRATION.ENTITIES_LIMIT, scoreThreshold = ORCHESTRATION.ENTITIES_THRESHOLD, projectId = undefined } = {}) {
+async function searchEntities(vector, { limit = ORCHESTRATION.ENTITIES_LIMIT, scoreThreshold = ORCHESTRATION.ENTITIES_THRESHOLD, projectId, allowedIds } = {}) {
     const body = { vector, limit, score_threshold: scoreThreshold, with_payload: true };
 
-    // non-project chats must also be filters to the "no project" pool
     if (projectId !== null && projectId !== undefined) {
-        body.filter = {
-            must: [{ key: 'projectId', match: { value: projectId } }]
-        };
-    } else {
-        //entities from project sessions carry a projectId;
-        //without this branch, a non-project chat searches ALL entities and project knowledge leaks into the common pool.
-        //is_empty matches null AND missing payload keys, so pre-isolation-era entities are covered too
-        body.filter = {
-            must: [{ is_empty: {key: 'projectId'} }]
-        }
+        // Project chat: entities shared across the project's sessions.
+        body.filter = { must: [{ key: 'projectId', match: { value: projectId } }] };
+    } else if (allowedIds && allowedIds.length > 0) {
+        // Non-project chat: restrict to this session's own entities (Model 2).
+        body.filter = { must: [{ has_id: allowedIds }] };
     }
+    // No else: the caller returns early when a non-project session has no linked
+    // entities, so an unfiltered (leaky) search is never reached.
+
     const res = await fetch(
         `${BASE_URL}/collections/${COLLECTIONS.ENTITIES}/points/search`,
-        {
-            method: 'POST',
-            headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify(body),
-        }
+        { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body) }
     );
-
     if (!res.ok) {
-        const body = await res.text();
-        throw new Error(`Qdrant error: ${res.status} - ${body}`);
+        const text = await res.text();
+        throw new Error(`Qdrant error: ${res.status} - ${text}`);
     }
-        
-    const data = await res.json();
-    return data.result;
+    return (await res.json()).result;
 }
 
 module.exports = { searchEpisodes, searchEntities };