Merge #4020 into stage-ml

static/sessions.js

@@ -3408,7 +3408,7 @@ function _applySessionListPayload(sessData, projData){
   _reconcileActiveSessionIdleStateFromList(serverSessions);
   _allSessions = _mergeOptimisticFirstTurnSessions(serverSessions);
   _syncSessionAttentionSoundState(_allSessions);
-  _clearLineageReportCache();
+  _pruneLineageReportCacheToVisibleSessions(_allSessions);
   _allProjects = projData.projects||[];
   _markPollingCompletionUnreadTransitions(_allSessions);
   const isStreaming = _allSessions.some(s => Boolean(s && s.is_streaming));
@@ -4247,6 +4247,20 @@ function _clearLineageReportCache(){
   _lineageReportCacheGeneration++;
 }
 
+function _pruneLineageReportCacheToVisibleSessions(sessions){
+  const visibleKeys=new Set();
+  for(const s of (Array.isArray(sessions)?sessions:[])){
+    const key=_sidebarLineageKeyForRow(s);
+    if(key) visibleKeys.add(key);
+  }
+  for(const key of Array.from(_lineageReportCache.keys())){
+    if(!visibleKeys.has(key)) _lineageReportCache.delete(key);
+  }
+  for(const key of Array.from(_lineageReportInflight.keys())){
+    if(!visibleKeys.has(key)) _lineageReportInflight.delete(key);
+  }
+}
+
 function _lineageReportCacheKey(s,lineageKey){
   return lineageKey||_sidebarLineageKeyForRow(s)||null;
 }
@@ -4291,14 +4305,16 @@ function _fetchLineageReportForRow(s,lineageKey){
   let request;
   request=api('/api/session/lineage/report?session_id='+encodeURIComponent(s.session_id))
     .then(report=>{
-      if(generation===_lineageReportCacheGeneration){
+      if(generation===_lineageReportCacheGeneration&&_lineageReportInflight.get(key)===request){
         _lineageReportCache.set(key,(report&&report.found!==false)?report:{error:true});
       }
       return report;
     })
     .catch(err=>{
       console.warn('lineage report',err);
-      if(generation===_lineageReportCacheGeneration) _lineageReportCache.set(key,{error:true});
+      if(generation===_lineageReportCacheGeneration&&_lineageReportInflight.get(key)===request){
+        _lineageReportCache.set(key,{error:true});
+      }
       return null;
     })
     .finally(()=>{
@@ -5231,6 +5247,9 @@ function renderSessionListFromCache(){
     const lineageReportKey=showLineageMetadata?_lineageReportCacheKey(s,lineageKey):null;
     const canExpandLineageSegments=showLineageMetadata&&Boolean(lineageKey&&segmentCount>1&&(lineageSegments.length>0||needsLineageReport||_lineageReportInflight.has(lineageReportKey)));
     const lineageSegmentsExpanded=canExpandLineageSegments&&_expandedLineageKeys.has(lineageKey);
+    if(lineageSegmentsExpanded&&needsLineageReport){
+      _fetchLineageReportForRow(s,lineageKey).then(()=>renderSessionListFromCache());
+    }
     if(segmentCount>0){
       const segmentCountEl=document.createElement('span');
       segmentCountEl.className='session-lineage-count'+(canExpandLineageSegments?' expandable':'');

tests/test_session_lineage_collapse.py

@@ -469,12 +469,14 @@ def test_lineage_segment_expansion_static_contract():
     assert "const _expandedLineageKeys = new Set();" in js
     assert "const _lineageReportCache = new Map();" in js
     assert "const _lineageReportInflight = new Map();" in js
+    assert "_pruneLineageReportCacheToVisibleSessions(_allSessions);" in js
     assert "session-lineage-count,.session-lineage-segments,.session-lineage-segment" in js
     assert "segmentCountEl.setAttribute('aria-expanded'" in js
     assert "_expandedLineageKeys.has(lineageKey)" in js
     assert "_expandedLineageKeys.add(lineageKey)" in js
     assert "_expandedLineageKeys.delete(lineageKey)" in js
     assert "_fetchLineageReportForRow(s,lineageKey).then" in js
+    assert js.count("_fetchLineageReportForRow(s,lineageKey).then(()=>renderSessionListFromCache());") == 2
     assert "'/api/session/lineage/report?session_id='" in js
     assert "encodeURIComponent(s.session_id)" in js
     assert "className='session-lineage-segments'" in js
@@ -617,6 +619,98 @@ eval(extractFunc('_fetchLineageReportForRow'));
     }
 
 
+def test_lineage_refresh_cache_prune_keeps_visible_keys_and_drops_missing_ones():
+    js = SESSIONS_JS_PATH.read_text(encoding="utf-8")
+    source = f"""
+const src = {js!r};
+function extractFunc(name) {{
+  const re = new RegExp('function\\\\s+' + name + '\\\\s*\\\\(');
+  const start = src.search(re);
+  if (start < 0) throw new Error(name + ' not found');
+  let i = src.indexOf('{{', start);
+  let depth = 1; i++;
+  while (depth > 0 && i < src.length) {{
+    if (src[i] === '{{') depth++;
+    else if (src[i] === '}}') depth--;
+    i++;
+  }}
+  return src.slice(start, i);
+}}
+const _lineageReportCache = new Map();
+const _lineageReportInflight = new Map();
+eval(extractFunc('_sidebarLineageKeyForRow'));
+eval(extractFunc('_pruneLineageReportCacheToVisibleSessions'));
+const visibleRequest = Promise.resolve({{found:true}});
+const staleRequest = Promise.resolve({{found:true}});
+_lineageReportCache.set('root', {{segments:[{{session_id:'root'}}]}});
+_lineageReportCache.set('stale', {{segments:[{{session_id:'stale'}}]}});
+_lineageReportInflight.set('root', visibleRequest);
+_lineageReportInflight.set('stale', staleRequest);
+_pruneLineageReportCacheToVisibleSessions([
+  {{session_id:'tip', _lineage_key:'root'}},
+  {{session_id:'child', parent_session_id:'root'}},
+]);
+console.log(JSON.stringify({{
+  cacheKeys:Array.from(_lineageReportCache.keys()),
+  inflightKeys:Array.from(_lineageReportInflight.keys()),
+}}));
+"""
+    assert json.loads(_run_node(source)) == {
+        "cacheKeys": ["root"],
+        "inflightKeys": ["root"],
+    }
+
+
+def test_pruned_lineage_inflight_request_cannot_repopulate_cache():
+    js = SESSIONS_JS_PATH.read_text(encoding="utf-8")
+    source = f"""
+const src = {js!r};
+function extractFunc(name) {{
+  const re = new RegExp('function\\\\s+' + name + '\\\\s*\\\\(');
+  const start = src.search(re);
+  if (start < 0) throw new Error(name + ' not found');
+  let i = src.indexOf('{{', start);
+  let depth = 1; i++;
+  while (depth > 0 && i < src.length) {{
+    if (src[i] === '{{') depth++;
+    else if (src[i] === '}}') depth--;
+    i++;
+  }}
+  return src.slice(start, i);
+}}
+const _lineageReportCache = new Map();
+const _lineageReportInflight = new Map();
+let _lineageReportCacheGeneration = 0;
+let resolveApi;
+function api(path) {{
+  return new Promise(resolve => {{
+    resolveApi = () => resolve({{found:true, path, segments:[{{session_id:'stale'}}]}});
+  }});
+}}
+eval(extractFunc('_sidebarLineageKeyForRow'));
+eval(extractFunc('_lineageReportCacheKey'));
+eval(extractFunc('_pruneLineageReportCacheToVisibleSessions'));
+eval(extractFunc('_fetchLineageReportForRow'));
+(async()=>{{
+  const staleRow = {{session_id:'stale-tip', _lineage_key:'stale'}};
+  const request = _fetchLineageReportForRow(staleRow, 'stale');
+  _pruneLineageReportCacheToVisibleSessions([{{session_id:'tip', _lineage_key:'root'}}]);
+  resolveApi();
+  await request;
+  console.log(JSON.stringify({{
+    staleCached:_lineageReportCache.has('stale'),
+    staleInflight:_lineageReportInflight.has('stale'),
+    visibleCached:_lineageReportCache.has('root'),
+  }}));
+}})().catch(err=>{{console.error(err); process.exit(1);}});
+"""
+    assert json.loads(_run_node(source)) == {
+        "staleCached": False,
+        "staleInflight": False,
+        "visibleCached": False,
+    }
+
+
 def test_active_hidden_lineage_segment_auto_expands_parent():
     js = SESSIONS_JS_PATH.read_text(encoding="utf-8")
     source = f"""