fix: bind broker requests and diagnostics to source snapshots

Beam/Broker/DocumentState.lean

@@ -63,6 +63,7 @@ inductive SyncFileAction where
 structure SyncFileDecision where
   action : SyncFileAction
   version : Nat
+  nextVersion : Nat
   docs : Docs
 
 structure VersionMarkResult where
@@ -103,7 +104,8 @@ private def docStateOfSnapshot (version : Nat) (snapshot : FileSnapshot) : DocSt
 
 /--
 Decide how a freshly read file snapshot should update the broker's LSP document
-mirror.
+mirror. `nextVersion` is the session-wide allocator, retained when documents close.
+Only opens and source changes consume a revision; unchanged and superseded reads preserve it.
 
 Request handlers reserve nonzero `readSeq` values before reading the filesystem.
 If an older read finishes after a newer read has already been applied, the older
@@ -114,26 +116,30 @@ newer LSP document versions.
 def syncFileDecision
     (docs : Docs)
     (uri : DocumentUri)
-    (snapshot : FileSnapshot) : SyncFileDecision :=
+    (snapshot : FileSnapshot)
+    (nextVersion : Nat) : SyncFileDecision :=
   match docs.get? uri with
   | none =>
-      let version := 1
+      let version := nextVersion
       {
         action := .open
         version
+        nextVersion := version + 1
         docs := docs.insert uri (docStateOfSnapshot version snapshot)
       }
   | some docState =>
       if snapshot.readSeq != 0 && snapshot.readSeq < docState.syncSnapshotSeq then
         {
           action := .unchanged
           version := docState.version
+          nextVersion
           docs
         }
       else if docState.textHash == snapshot.textHash then
         {
           action := .unchanged
           version := docState.version
+          nextVersion
           docs := docs.insert uri {
             docState with
             textTraceHash := snapshot.textTraceHash
@@ -143,10 +149,11 @@ def syncFileDecision
           }
         }
       else
-        let version := docState.version + 1
+        let version := nextVersion
         {
           action := .change
           version
+          nextVersion := version + 1
           docs := docs.insert uri {
             (docStateOfSnapshot version snapshot) with
             checkpointedVersion? := none

Beam/Broker/OpenDocs.lean

@@ -17,6 +17,7 @@ namespace OpenDocs
 
 structure SessionView where
   root : System.FilePath
+  sessionToken : String
   docs : DocumentState.Docs := {}
 
 inductive DiskStatus where
@@ -48,6 +49,7 @@ def docDiskStatus (path : System.FilePath) (docState : DocState) : IO DiskStatus
 
 def docJson
     (root : System.FilePath)
+    (sessionToken : String)
     (uri : DocumentUri)
     (docState : DocState) : IO Json := do
   let path? := System.Uri.fileUriToPath? uri
@@ -65,7 +67,7 @@ def docJson
   pure <| Json.mkObj <|
     [
       ("uri", toJson uri),
-      ("version", toJson docState.version),
+      ("snapshot", toJson ({ session := sessionToken, revision := docState.version } : SnapshotRef)),
       ("diskStatus", toJson status),
       ("checkpointed", toJson checkpointed)
     ] ++
@@ -84,7 +86,7 @@ def sessionJson (session? : Option SessionView) : IO Json := do
       ]
   | some session =>
       let files ← session.docs.toList.mapM fun (uri, docState) =>
-        docJson session.root uri docState
+        docJson session.root session.sessionToken uri docState
       pure <| Json.mkObj [
         ("active", toJson true),
         ("files", Json.arr files.toArray)

Beam/Broker/Pending.lean

@@ -206,39 +206,45 @@ private def observeSyncFileProgress
       progress?
 
 private def trackedPublishDiagnosticsParam?
-    (trackedUri? : Option DocumentUri)
+    (tracked? : Option (DocumentUri × Nat))
     (diagnosticParam : PublishDiagnosticsParams) : Option PublishDiagnosticsParams :=
-  match trackedUri? with
-  | some uri =>
+  match tracked? with
+  | some (uri, version) =>
       let diagnosticParam := normalizePublishDiagnostics diagnosticParam
-      if diagnosticParam.uri == uri then
+      if diagnosticParam.uri == uri &&
+          diagnosticParam.version?.all (· == Int.ofNat version) then
         some diagnosticParam
       else
         none
   | none =>
       none
 
-private def diagnosticStreamKey (diagnostic : Diagnostic) : String :=
-  (toJson diagnostic).compress
+private def diagnosticStreamKey
+    (snapshot? : Option SnapshotRef) (diagnostic : Diagnostic) : String :=
+  (toJson (snapshot?, diagnostic)).compress
 
 private def emitNewTrackedDiagnostics
     (root : System.FilePath)
+    (sessionToken : String)
     (seen : Std.TreeSet String compare)
     (diagnosticParam : PublishDiagnosticsParams)
     (diagnosticScope : DiagnosticScope)
     (emitDiagnostic? : Option (StreamDiagnostic → IO Unit) := none) :
     IO (Std.TreeSet String compare) := do
+  let snapshot? := diagnosticParam.version?.bind fun version =>
+    if version > 0 then some { session := sessionToken, revision := version.toNat }
+    else none
   let mut seen := seen
   let diagnostics := filterSyncDiagnostics diagnosticScope diagnosticParam.diagnostics
   for diagnostic in diagnostics do
-    let key := diagnosticStreamKey diagnostic
+    let key := diagnosticStreamKey snapshot? diagnostic
     if !seen.contains key then
       seen := seen.insert key
       match emitDiagnostic? with
       | some emitDiagnostic =>
           try
             emitDiagnostic <|
-              streamDiagnosticOfDiagnostic root diagnosticParam.uri diagnosticParam.version? diagnostic
+              streamDiagnosticOfDiagnostic root diagnosticParam.uri snapshot? diagnostic
           catch _ =>
             pure ()
       | none =>
@@ -264,30 +270,20 @@ def observeProgress
 
 def observePublishDiagnostics
     (root : System.FilePath)
+    (sessionToken : String)
     (pending : PendingRequest)
     (diagnosticParam : PublishDiagnosticsParams) : IO Unit := do
-  match trackedPublishDiagnosticsParam? (pending.tracked?.map Prod.fst) diagnosticParam with
+  match trackedPublishDiagnosticsParam? pending.tracked? diagnosticParam with
   | none =>
       pure ()
   | some diagnosticParam =>
       pending.diagnosticsSeenRef.set true
       pending.diagnosticsRef.set diagnosticParam.diagnostics
       let seen ← pending.seenDiagnosticKeysRef.get
       let seen ←
-        emitNewTrackedDiagnostics root seen diagnosticParam pending.diagnosticScope pending.emitDiagnostic?
+        emitNewTrackedDiagnostics root sessionToken seen diagnosticParam pending.diagnosticScope pending.emitDiagnostic?
       pending.seenDiagnosticKeysRef.set seen
 
-def observeDiagnostics
-    [ToJson α]
-    (root : System.FilePath)
-    (pending : PendingRequest)
-    (param : α) : IO Unit := do
-  match fromJson? (toJson param) with
-  | .ok (diagnosticParam : PublishDiagnosticsParams) =>
-      observePublishDiagnostics root pending diagnosticParam
-  | .error _ =>
-      pure ()
-
 end PendingRequest
 
 namespace PendingRequestStore