diff --git a/Beam/Broker/Backend/Lean.lean b/Beam/Broker/Backend/Lean.lean index 4b55be08..b70ac470 100644 --- a/Beam/Broker/Backend/Lean.lean +++ b/Beam/Broker/Backend/Lean.lean @@ -21,18 +21,13 @@ open Lean.Lsp namespace Beam.Broker.Backend.Lean -private def pluginPath (config : BrokerConfig) : IO System.FilePath := do - match config.leanPlugin? with - | some path => Beam.resolveExistingPath path - | none => throw <| IO.userError "missing Beam daemon --lean-plugin configuration" - def command (config : BrokerConfig) : IO (String × Array String × Array (String × Option String)) := do - let some cmd := config.leanCmd? - | throw <| IO.userError "missing Beam daemon --lean-cmd configuration" - let plugin := ← pluginPath config - let lakeEnv ← leanServerLakeEnv config.root config.leanCmd? config.leanLakeHelper? + let some leanConfig := config.lean? + | throw <| IO.userError "Lean backend is not configured" + let plugin ← Beam.resolveExistingPath leanConfig.plugin + let lakeEnv ← leanServerLakeEnv config.root (some leanConfig.command) leanConfig.lakeHelper? pure ( - cmd, + leanConfig.command, #["--server"] ++ lakeEnv.moreServerArgs ++ #[s!"--plugin={plugin}", "-Dexperimental.module=true"], lakeEnv.env) diff --git a/Beam/Broker/Backend/Rocq.lean b/Beam/Broker/Backend/Rocq.lean index 30aaaa3f..13bb2426 100644 --- a/Beam/Broker/Backend/Rocq.lean +++ b/Beam/Broker/Backend/Rocq.lean @@ -15,13 +15,10 @@ open Lean.Lsp namespace Beam.Broker.Backend.Rocq -private def lspPath (config : BrokerConfig) : IO String := do - match config.rocqCmd? with - | some path => pure path - | none => throw <| IO.userError "missing Beam daemon --rocq-cmd configuration" - def command (config : BrokerConfig) : IO (String × Array String) := do - pure ((← lspPath config), #[]) + let some rocqConfig := config.rocq? + | throw <| IO.userError "Rocq backend is not configured" + pure (rocqConfig.command, #[]) def initializeParams (root : System.FilePath) : Json := let rootUri := System.Uri.pathToUri root diff --git a/Beam/Broker/Config.lean b/Beam/Broker/Config.lean index 0f1efba1..6e36f35d 100644 --- a/Beam/Broker/Config.lean +++ b/Beam/Broker/Config.lean @@ -8,12 +8,49 @@ import Lean namespace Beam.Broker +/-- Complete process configuration for one Lean backend. -/ +structure LeanBackendConfig where + command : String + plugin : System.FilePath + lakeHelper? : Option System.FilePath := none + deriving BEq, Inhabited, Repr + +/-- Complete process configuration for one Rocq backend. -/ +structure RocqBackendConfig where + command : String + deriving BEq, Inhabited, Repr + +/-- +Runtime configuration for one broker workspace. + +The optional backends permit an intentionally backend-less standalone bootstrap workspace. Once a +backend is present, its required process configuration is complete by construction. +-/ structure BrokerConfig where root : System.FilePath - leanCmd? : Option String := none - leanPlugin? : Option System.FilePath := none - leanLakeHelper? : Option System.FilePath := none - rocqCmd? : Option String := none - deriving Inhabited, Repr + lean? : Option LeanBackendConfig := none + rocq? : Option RocqBackendConfig := none + deriving BEq, Inhabited, Repr + +namespace BrokerConfig + +/-- Assemble the typed runtime model from optional fields at a process boundary. -/ +def ofOptions + (root : System.FilePath) + (leanCommand? : Option String) + (leanPlugin? : Option System.FilePath) + (rocqCommand? : Option String := none) : Except String BrokerConfig := do + let lean? ← + match leanCommand?, leanPlugin? with + | none, none => pure none + | some command, some plugin => pure <| some { command, plugin } + | _, _ => throw "Lean backend configuration requires a command and plugin together" + pure { + root + lean? + rocq? := rocqCommand?.map fun command => { command } + } + +end BrokerConfig end Beam.Broker diff --git a/Beam/Broker/Protocol.lean b/Beam/Broker/Protocol.lean index a6c7d462..08eaccd7 100644 --- a/Beam/Broker/Protocol.lean +++ b/Beam/Broker/Protocol.lean @@ -300,11 +300,14 @@ structure ReleaseRequest where path : String handle : Handle +structure InitLeanBackendConfig where + command : String + plugin : String + structure InitWorkspaceRequest where workspaceMode? : Option Beam.Workspace.InitMode := none root : String - leanCmd? : Option String := none - leanPlugin? : Option String := none + lean? : Option InitLeanBackendConfig := none rocqCmd? : Option String := none /-- The fields owned by exactly one broker operation. -/ @@ -535,8 +538,12 @@ private def RequestPayload.jsonFields : RequestPayload → List (String × Json) | .initWorkspace request => optionalJsonField "workspaceMode" request.workspaceMode? ++ [("root", toJson request.root)] ++ - optionalJsonField "leanCmd" request.leanCmd? ++ - optionalJsonField "leanPlugin" request.leanPlugin? ++ + (match request.lean? with + | some lean => [ + ("leanCmd", toJson lean.command), + ("leanPlugin", toJson lean.plugin) + ] + | none => []) ++ optionalJsonField "rocqCmd" request.rocqCmd? | .stats | .listWorkspaces | .dropWorkspace | .shutdown => [] @@ -739,11 +746,18 @@ instance : FromJson Request where handle } | .initWorkspace => + let leanCmd? ← optionalField? (α := String) j "leanCmd" + let leanPlugin? ← optionalField? (α := String) j "leanPlugin" + let lean? ← + match leanCmd?, leanPlugin? with + | none, none => pure none + | some command, some plugin => pure <| some { command, plugin } + | some _, none => throw "'leanCmd' requires 'leanPlugin'" + | none, some _ => throw "'leanPlugin' requires 'leanCmd'" pure <| .initWorkspace { workspaceMode? := ← optionalField? (α := Beam.Workspace.InitMode) j "workspaceMode" root := ← requiredField j "root" - leanCmd? := ← optionalField? (α := String) j "leanCmd" - leanPlugin? := ← optionalField? (α := String) j "leanPlugin" + lean? rocqCmd? := ← optionalField? (α := String) j "rocqCmd" } | .listWorkspaces => pure .listWorkspaces diff --git a/Beam/Broker/Server.lean b/Beam/Broker/Server.lean index 1ff325a8..1390f102 100644 --- a/Beam/Broker/Server.lean +++ b/Beam/Broker/Server.lean @@ -74,6 +74,7 @@ structure BackendState where session? : Option Session := none structure WorkspaceState where + generation : Nat config : BrokerConfig nextFileSnapshotSeq : Nat := 1 lean : BackendState := {} @@ -84,6 +85,7 @@ structure WorkspaceState where structure State where bootstrapConfig : BrokerConfig startMonoNanos : Nat := 0 + nextWorkspaceGeneration : Nat := 2 workspaces : Std.TreeMap WorkspaceId WorkspaceState := {} streamSink? : Option (StreamMessage → IO Unit) := none currentClientRequestId? : Option String := none @@ -267,7 +269,9 @@ private def sessionExited (session : Session) : IO Bool := do catch _ => pure true -private def mkWorkspaceState (config : BrokerConfig) : WorkspaceState := { config } +private def mkWorkspaceState + (config : BrokerConfig) + (generation : Nat) : WorkspaceState := { config, generation } private def mkInitialState (config : BrokerConfig) @@ -275,7 +279,7 @@ private def mkInitialState (startMonoNanos : Nat) : State := { bootstrapConfig := config startMonoNanos - workspaces := Std.TreeMap.empty.insert workspaceId (mkWorkspaceState config) + workspaces := Std.TreeMap.empty.insert workspaceId (mkWorkspaceState config 1) } private def validWorkspaceId (workspaceId : WorkspaceId) : Bool := @@ -290,6 +294,17 @@ private def setWorkspace (workspace : WorkspaceState) : State := { state with workspaces := state.workspaces.insert workspaceId workspace } +private def setFreshWorkspace + (state : State) + (workspaceId : WorkspaceId) + (config : BrokerConfig) : State := + let generation := state.nextWorkspaceGeneration + { + state with + nextWorkspaceGeneration := generation + 1 + workspaces := state.workspaces.insert workspaceId (mkWorkspaceState config generation) + } + private def getBackendState (workspace : WorkspaceState) (backend : Backend) : BackendState := match backend with | .lean => workspace.lean @@ -755,11 +770,14 @@ An immutable view of a source file used to synchronize the LSP session. The file contents and metadata are computed before the broker state mutex is held. This keeps potentially slow filesystem work out of the critical section. +The workspace generation is captured before that read and checked before the +snapshot is applied, so a concurrent workspace reset cannot retarget it. For request handlers that can race with each other, `readSeq` is reserved while holding the mutex and is later used by `DocumentState.syncFileDecision` to ignore stale snapshots that completed after a newer read was already applied. -/ private structure FileSyncSnapshot where + workspaceGeneration : Nat path : System.FilePath uri : DocumentUri text : String @@ -774,6 +792,7 @@ private structure SyncedFileSnapshot where private def readFileSyncSnapshot (root path : System.FilePath) (backend : Backend) + (workspaceGeneration : Nat) (readSeq : Nat := 0) : IO FileSyncSnapshot := do let path ← resolvePath root path let text ← IO.FS.readFile path @@ -781,6 +800,7 @@ private def readFileSyncSnapshot let uri := sessionUri path let moduleName? := DocumentState.trackedModuleName? root path backend pure { + workspaceGeneration path uri text @@ -793,6 +813,40 @@ private def readFileSyncSnapshot } } +private def workspaceForSnapshot + (workspaceId : WorkspaceId) + (snapshot : FileSyncSnapshot) : M (Except ResponseFailure WorkspaceState) := do + let state ← get + match getWorkspace? state workspaceId with + | some workspace => + if workspace.generation == snapshot.workspaceGeneration then + pure (.ok workspace) + else + pure <| .error <| responseFailureFor .contentModified <| + s!"workspace '{workspaceId}' changed while the request source file was being read; retry the request" + | none => + pure <| .error <| responseFailureFor .contentModified <| + s!"workspace '{workspaceId}' was removed while the request source file was being read; retry after initializing it" + +private def withWorkspaceForSnapshot + (workspaceId : WorkspaceId) + (snapshot : FileSyncSnapshot) + (act : WorkspaceState → M (Except ResponseFailure α)) : + M (Except ResponseFailure α) := do + match ← workspaceForSnapshot workspaceId snapshot with + | .ok workspace => act workspace + | .error failure => pure (.error failure) + +private def withSessionForSnapshot + (workspaceId : WorkspaceId) + (backend : Backend) + (snapshot : FileSyncSnapshot) + (act : Session → M (Except ResponseFailure α)) : + M (Except ResponseFailure α) := + withWorkspaceForSnapshot workspaceId snapshot fun _ => do + let session ← ensureSession workspaceId backend + act session + private def syncFileSnapshotDetailed (session : Session) (snapshot : FileSyncSnapshot) : IO SyncedFileSnapshot := do @@ -1081,13 +1135,6 @@ def ServerRuntime.create closeDone := ← IO.Promise.new } -private def brokerConfigSame (left right : BrokerConfig) : Bool := - left.root == right.root && - left.leanCmd? == right.leanCmd? && - left.leanPlugin? == right.leanPlugin? && - left.leanLakeHelper? == right.leanLakeHelper? && - left.rocqCmd? == right.rocqCmd? - private def detachBackendSession (backend : BackendState) : BackendState × Option Session := match backend.session? with @@ -1226,15 +1273,14 @@ private def initWorkspaceTransition s!"workspace root {config.root} is already owned by workspace '{otherId}'" } else let (_, detachedSessions) := detachWorkspaceSessions current - let replacement := mkWorkspaceState config { - state := setWorkspace state workspaceId replacement + state := setFreshWorkspace state workspaceId config result := .ok <| workspaceInitResult workspaceId config.root mode false true (some current.config.root) detachedSessions } - else if brokerConfigSame current.config config then + else if current.config == config then { state, result := .ok <| workspaceInitResult workspaceId current.config.root mode true false } else @@ -1250,7 +1296,7 @@ private def initWorkspaceTransition s!"workspace root {config.root} is already owned by workspace '{otherId}'" } else { - state := setWorkspace state workspaceId (mkWorkspaceState config) + state := setFreshWorkspace state workspaceId config result := .ok <| workspaceInitResult workspaceId config.root mode false false } @@ -1543,6 +1589,9 @@ private def startSyncedDocumentRequest (emitDiagnostic? : Option (StreamDiagnostic → IO Unit) := none) (cancelRef? : Option (IO.Ref Bool) := none) : M (Except ResponseFailure StartedSyncedRequest) := do + match ← workspaceForSnapshot session.workspaceId snapshot with + | .error failure => return .error failure + | .ok _ => pure () let session ← syncFileSnapshot session snapshot let uri := snapshot.uri let docState ← requireDocState session uri @@ -1574,6 +1623,24 @@ private def startSyncedDocumentRequest pending } +private def startSyncedWorkspaceRequest + (workspaceId : WorkspaceId) + (backend : Backend) + (snapshot : FileSyncSnapshot) + (method : String) + (mkParams : DocumentUri → DocState → Json) + (trackedFor : DocumentUri → DocState → Option (DocumentUri × Nat)) + (expectedVersion? : Option Nat := none) + (clientRequestId? : Option String := none) + (emitProgress? : Option (SyncFileProgress → IO Unit) := none) + (diagnosticScope : DiagnosticScope := .errors) + (emitDiagnostic? : Option (StreamDiagnostic → IO Unit) := none) + (cancelRef? : Option (IO.Ref Bool) := none) : + M (Except ResponseFailure StartedSyncedRequest) := + withSessionForSnapshot workspaceId backend snapshot fun session => + startSyncedDocumentRequest session snapshot method mkParams trackedFor expectedVersion? + clientRequestId? emitProgress? diagnosticScope emitDiagnostic? cancelRef? + private def awaitSyncedDocumentRequest (server : ServerRuntime) (started : StartedSyncedRequest) @@ -1583,24 +1650,33 @@ private def awaitSyncedDocumentRequest if started.tracked.isSome then withFailureProgress pending.progress? <| liftHandlerIO <| mergeFileProgressIfCurrent server started.session started.uri pending.progress? + withFailureProgress pending.progress? <| + withCurrentMatchingSession server started.session fun _ => pure () pure pending private def readRequestSyncSnapshot (server : ServerRuntime) (req : BackendWorkspaceRequest) - (path : System.FilePath) : IO FileSyncSnapshot := do - let (root, readSeq) ← server.withState do - let workspace ← requireWorkspace req.workspaceId - let readSeq := workspace.nextFileSnapshotSeq - let workspace := { workspace with nextFileSnapshotSeq := readSeq + 1 } - modify fun state => setWorkspace state req.workspaceId workspace - pure (workspace.config.root, readSeq) + (path : System.FilePath) : IO (Except ResponseFailure FileSyncSnapshot) := do + let readContext? ← server.withState do + let state ← get + match getWorkspace? state req.workspaceId with + | none => pure none + | some workspace => + let readSeq := workspace.nextFileSnapshotSeq + let workspace := { workspace with nextFileSnapshotSeq := readSeq + 1 } + set <| setWorkspace state req.workspaceId workspace + pure <| some (workspace.config.root, workspace.generation, readSeq) + let some (root, workspaceGeneration, readSeq) := readContext? + | return .error <| responseFailureFor .contentModified <| + s!"workspace '{req.workspaceId}' was removed before the request source file could be read; retry after initializing it" -- Reserve the ordering token under the mutex, then do the slow file IO -- outside it. - readFileSyncSnapshot root path req.backend (readSeq := readSeq) + pure <| .ok <| ← readFileSyncSnapshot root path req.backend workspaceGeneration (readSeq := readSeq) private structure StartedTrackedBarrier where session : Session + leanConfig? : Option LeanBackendConfig uri : DocumentUri version : Nat textHash : UInt64 @@ -1618,38 +1694,43 @@ private def startTrackedDiagnosticsBarrierIO (emitProgress? : Option (SyncFileProgress → IO Unit) := none) (emitDiagnostic? : Option (StreamDiagnostic → IO Unit) := none) (cancelRef? : Option (IO.Ref Bool) := none) : - IO StartedTrackedBarrier := do - let snapshot ← readRequestSyncSnapshot server req path + IO (Except ResponseFailure StartedTrackedBarrier) := do + let snapshot ← + match ← readRequestSyncSnapshot server req path with + | .ok snapshot => pure snapshot + | .error failure => return .error failure server.withState do - let session ← ensureSession req.workspaceId req.backend - let synced ← syncFileSnapshotDetailed session snapshot - let session := synced.session - let uri := synced.uri - let docState ← requireDocState session uri - let tracked := trackedDocumentVersion uri docState - let params := toJson (WaitForDiagnosticsParams.mk uri docState.version) - let method ← IO.ofExcept <| diagnosticsBarrierMethod session.backend - let (session, pending) ← - startRequestJsonTrackedDetailed session method params - (clientRequestId? := req.clientRequestId?) - (tracked := tracked) - (initialProgress? := docState.fileProgress?) - (emitProgress? := emitProgress?) - (diagnosticScope := diagnosticScope) - (emitDiagnostic? := emitDiagnostic?) - (cancelRef? := cancelRef?) - updateSession session - pure { - session - uri - version := synced.version - textHash := docState.textHash - textTraceHash := docState.textTraceHash - textMTime := docState.textMTime - changed := synced.changed - priorProgress? := docState.fileProgress? - pending - } + withWorkspaceForSnapshot req.workspaceId snapshot fun workspace => do + let session ← ensureSession req.workspaceId req.backend + let synced ← syncFileSnapshotDetailed session snapshot + let session := synced.session + let uri := synced.uri + let docState ← requireDocState session uri + let tracked := trackedDocumentVersion uri docState + let params := toJson (WaitForDiagnosticsParams.mk uri docState.version) + let method ← IO.ofExcept <| diagnosticsBarrierMethod session.backend + let (session, pending) ← + startRequestJsonTrackedDetailed session method params + (clientRequestId? := req.clientRequestId?) + (tracked := tracked) + (initialProgress? := docState.fileProgress?) + (emitProgress? := emitProgress?) + (diagnosticScope := diagnosticScope) + (emitDiagnostic? := emitDiagnostic?) + (cancelRef? := cancelRef?) + updateSession session + pure <| .ok { + session + leanConfig? := workspace.config.lean? + uri + version := synced.version + textHash := docState.textHash + textTraceHash := docState.textTraceHash + textMTime := docState.textMTime + changed := synced.changed + priorProgress? := docState.fileProgress? + pending + } private def finalizeSavedDoc (server : ServerRuntime) @@ -1792,11 +1873,10 @@ private def saveOleanCore (emitDiagnostic? : Option (StreamDiagnostic → IO Unit) := none) : HandlerM SaveOleanCompleted := do liftFailureIO <| ensureRequestNotCancelled cancelRef? - let started ← liftHandlerIO <| startTrackedDiagnosticsBarrierIO server req path diagnosticScope + let started ← liftFailureIO <| startTrackedDiagnosticsBarrierIO server req path diagnosticScope emitProgress? emitDiagnostic? (cancelRef? := cancelRef?) - let (leanCmd?, lakeHelper?) ← liftHandlerIO <| server.withState do - let workspace ← requireWorkspace req.workspaceId - pure (workspace.config.leanCmd?, workspace.config.leanLakeHelper?) + let some leanConfig := started.leanConfig? + | throw <| responseFailureFor .invalidParams "Lean backend is not configured" liftHandlerIO <| propagatePendingCancellation started.session cancelRef? let barrier ← awaitWaitForDiagnosticsBarrier s!"save_olean sync barrier clientRequestId={optionLabel req.clientRequestId?} uri={started.uri} version={started.version}" @@ -1824,7 +1904,8 @@ private def saveOleanCore barrierOutcome.completionDiagnostics barrierProgress? let spec ← withFailureProgress barrierProgress? <| liftBrokerFailureIO <| mkLeanSaveSpec started.session.root path - { hash := started.textTraceHash, mtime := started.textMTime } leanCmd? lakeHelper? + { hash := started.textTraceHash, mtime := started.textMTime } + (some leanConfig.command) leanConfig.lakeHelper? let syncResult := mkSyncFileResult spec.relPath started.version currentDiagnostics saveReadiness withFailureProgress barrierProgress? <| @@ -1941,7 +2022,7 @@ private def handleSyncFileOp let path := System.FilePath.mk request.path liftFailureIO <| ensureRequestNotCancelled cancelRef? let diagnosticScope := request.diagnosticScope?.getD .errors - let started ← liftHandlerIO <| startTrackedDiagnosticsBarrierIO server req path diagnosticScope + let started ← liftFailureIO <| startTrackedDiagnosticsBarrierIO server req path diagnosticScope emitProgress? emitDiagnostic? (cancelRef? := cancelRef?) liftHandlerIO <| traceBroker s!"sync_file await barrier clientRequestId={optionLabel req.clientRequestId?} uri={started.uri} version={started.version}" @@ -2019,12 +2100,12 @@ private def handleUpdateFileOp HandlerM Response := do let path := System.FilePath.mk request.path liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| readRequestSyncSnapshot server req path - let updated ← liftHandlerIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - let synced ← syncFileSnapshotDetailed session snapshot - updateSession synced.session - pure synced + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req path + let updated ← liftFailureIO <| server.withState do + withSessionForSnapshot req.workspaceId req.backend snapshot fun session => do + let synced ← syncFileSnapshotDetailed session snapshot + updateSession synced.session + pure (.ok synced) pure <| Response.success (toJson ({ version := updated.version changed := updated.changed @@ -2073,10 +2154,9 @@ private def handleRunAtOp let method ← requestMethod <| runAtMethod request.backend let path := System.FilePath.mk request.path liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| readRequestSyncSnapshot server req path + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req path let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun uri _ => Json.mkObj <| [ ("textDocument", toJson ({ uri := uri, version? := some request.version : VersionedTextDocumentIdentifier })) , ("position", toJson ({ line := request.line, character := request.character : Lsp.Position })) @@ -2117,11 +2197,10 @@ private def handlePositionLspOp (emitProgress? : Option (SyncFileProgress → IO Unit) := none) : HandlerM Response := do liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun uri _ => positionLspParams request uri extraFields) (trackedLeanDocumentVersion req.backend) (expectedVersion? := some request.version) @@ -2187,11 +2266,10 @@ private def handleDocumentSymbolsOp HandlerM Response := do let method ← requestMethod <| documentSymbolsMethod request.backend liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun uri _ => Json.mkObj [ ("textDocument", toJson ({ uri := uri : TextDocumentIdentifier })) ]) @@ -2221,6 +2299,7 @@ private def handleWorkspaceSymbolsOp pure (session, pending) liftHandlerIO <| propagatePendingCancellation session cancelRef? let result ← awaitPending pending + withCurrentMatchingSession server session fun _ => pure () pure <| Response.success result.result private def codeActionResolveSourceUri @@ -2244,15 +2323,14 @@ private def handleCodeActionResolveOp HandlerM Response := do let method ← requestMethod <| codeActionResolveMethod request.backend liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let sourceUri ← requestArg <| codeActionResolveSourceUri request.codeAction if sourceUri != snapshot.uri then throw <| responseFailureFor .invalidParams s!"codeAction.data targets {sourceUri}, not requested document {snapshot.uri}" let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun _uri _docState => toJson request.codeAction) (trackedLeanDocumentVersion req.backend) (expectedVersion? := some request.version) @@ -2293,12 +2371,11 @@ private def handleGoalsOp throw <| responseFailureFor .invalidParams "lean goals does not accept speculative text; use lean-beam run-at for execution" liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend let position : Lsp.Position := { line := request.line, character := request.character } - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun uri docState => match req.backend with | .lean => @@ -2340,11 +2417,10 @@ private def handleTodoOp start := { line := request.line, character := request.character } «end» := { line := request.endLine, character := request.endCharacter } } - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do - let session ← ensureSession req.workspaceId req.backend - startSyncedDocumentRequest session snapshot method + startSyncedWorkspaceRequest req.workspaceId req.backend snapshot method (fun uri _docState => Json.mkObj <| [ ("textDocument", toJson ({ uri := uri, version? := some request.version : VersionedTextDocumentIdentifier })) , ("range", toJson range) @@ -2372,7 +2448,7 @@ private def handleRunWithOp HandlerM Response := do let method ← requestMethod <| runWithMethod request.handle.backend liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do match ← resolveCurrentHandle req.workspaceId request.handle with @@ -2409,7 +2485,7 @@ private def handleReleaseOp HandlerM Response := do let method ← requestMethod <| releaseMethod request.handle.backend liftFailureIO <| ensureRequestNotCancelled cancelRef? - let snapshot ← liftHandlerIO <| + let snapshot ← liftFailureIO <| readRequestSyncSnapshot server req (System.FilePath.mk request.path) let started ← liftFailureIO <| server.withState do match ← resolveCurrentHandle req.workspaceId request.handle with @@ -2436,12 +2512,17 @@ private def initWorkspaceConfigFromRequest resolveRoot (System.FilePath.mk request.root) catch e => return .error (responseFailureFor .invalidParams e.toString) - let leanPlugin? ← - try - request.leanPlugin?.mapM (fun path => Beam.resolveExistingPath <| System.FilePath.mk path) - catch e => - return .error (responseFailureFor .invalidParams e.toString) - if request.leanCmd?.isNone && leanPlugin?.isNone && request.rocqCmd?.isNone then + let lean? ← + match request.lean? with + | none => pure none + | some lean => + let plugin ← + try + Beam.resolveExistingPath <| System.FilePath.mk lean.plugin + catch e => + return .error (responseFailureFor .invalidParams e.toString) + pure <| some ({ command := lean.command, plugin } : LeanBackendConfig) + if lean?.isNone && request.rocqCmd?.isNone then let bootstrapConfig ← server.withState do let state ← get pure state.bootstrapConfig @@ -2449,9 +2530,8 @@ private def initWorkspaceConfigFromRequest return .ok bootstrapConfig pure <| .ok { root - leanCmd? := request.leanCmd? - leanPlugin? := leanPlugin? - rocqCmd? := request.rocqCmd? + lean? + rocq? := request.rocqCmd?.map fun command => { command } } private def handleRequestIO @@ -2993,13 +3073,10 @@ def main (args : List String) : IO Unit := do else pure <| .standalone daemonIdentity? let root ← Beam.resolveExistingPath <| System.FilePath.mk root - let leanPlugin? ← opts.leanPlugin?.mapM (fun path => Beam.resolveExistingPath <| System.FilePath.mk path) - let config : BrokerConfig := { - root := root - leanCmd? := opts.leanCmd? - leanPlugin? := leanPlugin? - rocqCmd? := opts.rocqCmd? - } + let leanPlugin? ← opts.leanPlugin?.mapM fun path => + Beam.resolveExistingPath <| System.FilePath.mk path + let config ← IO.ofExcept <| BrokerConfig.ofOptions root opts.leanCmd? leanPlugin? + (rocqCommand? := opts.rocqCmd?) withDaemonResources opts config workspaceId mode root fun resources => do match mode with | .wrapper identity _ => emitWrapperReady resources.transport identity diff --git a/Beam/Mcp/Runtime.lean b/Beam/Mcp/Runtime.lean index 4777a58b..0588909e 100644 --- a/Beam/Mcp/Runtime.lean +++ b/Beam/Mcp/Runtime.lean @@ -111,9 +111,11 @@ def mkBrokerConfig (opts : Options) (root : System.FilePath) : IO (Except RpcErr | .ok runtime => pure <| .ok { root := root - leanCmd? := some runtime.leanCmd - leanPlugin? := some runtime.leanPlugin - leanLakeHelper? := some runtime.leanLakeHelper + lean? := some { + command := runtime.leanCmd + plugin := runtime.leanPlugin + lakeHelper? := some runtime.leanLakeHelper + } } end Beam.Mcp.Runtime diff --git a/tests/lean/BeamTest/Broker/McpProtocolTest.lean b/tests/lean/BeamTest/Broker/McpProtocolTest.lean index 53560f7d..3848ccc0 100644 --- a/tests/lean/BeamTest/Broker/McpProtocolTest.lean +++ b/tests/lean/BeamTest/Broker/McpProtocolTest.lean @@ -561,7 +561,7 @@ private def checkRuntimeSetupErrors : IO Unit := do throw <| IO.userError s!"explicit MCP runtime setup failed: {err.message}" | .ok config => require "explicit MCP runtime should infer its sibling target Lake helper" - config.leanLakeHelper?.isSome + (config.lean?.bind (·.lakeHelper?)).isSome for (label, options) in #[ ("Lean command only", ({ beamCli? := some "unused", leanCmd? := some "lean" } : Beam.Mcp.Runtime.Options)), diff --git a/tests/lean/BeamTest/Broker/ProtocolTest.lean b/tests/lean/BeamTest/Broker/ProtocolTest.lean index 1f76bd3c..73d0e241 100644 --- a/tests/lean/BeamTest/Broker/ProtocolTest.lean +++ b/tests/lean/BeamTest/Broker/ProtocolTest.lean @@ -813,6 +813,36 @@ private def checkWorkspaceRoutingFields : IO Unit := do require s!"minimal {op.key} request changed its flat JSON wire shape" (toJson decoded == requestJson) + let initWithLean : Request := { + payload := .initWorkspace { + root := "/workspace" + lean? := some { command := "lean", plugin := "/beam/libbeam.so" } + rocqCmd? := some "rocq" + } + } + let initWithLeanJson := toJson initWithLean + requireJsonString "typed workspace init" "leanCmd" "lean" initWithLeanJson + requireJsonString "typed workspace init" "leanPlugin" "/beam/libbeam.so" initWithLeanJson + requireJsonString "typed workspace init" "rocqCmd" "rocq" initWithLeanJson + let decodedInitWithLean ← expectOk "typed workspace init round trip" <| + fromJson? (α := Request) initWithLeanJson + require "typed workspace init should preserve its flat wire shape" + (toJson decodedInitWithLean == initWithLeanJson) + + for (label, field, value) in #[ + ("command-only workspace init", "leanCmd", "lean"), + ("plugin-only workspace init", "leanPlugin", "/beam/libbeam.so") + ] do + match fromJson? (α := Request) <| Json.mkObj [ + ("op", toJson "init_workspace"), + ("root", toJson "/workspace"), + (field, toJson value) + ] with + | .ok _ => throw <| IO.userError s!"{label}: partial Lean configuration decoded" + | .error err => + require s!"{label}: error should name both coupled fields" + (err.contains "leanCmd" && err.contains "leanPlugin") + let unscopedReq := Request.stats require "missing workspace id remains unscoped" unscopedReq.resolvedWorkspaceId?.isNone requireFieldAbsent "stats request serialization" "backend" (toJson unscopedReq) @@ -1123,10 +1153,11 @@ private def checkLifecycleTeardownReleasesStateMutex runtime.state.atomically do let state ← get let targetWorkspace : WorkspaceState := { + generation := 1 config := targetConfig lean := { nextEpoch := 2, session? := some session } } - let observerWorkspace : WorkspaceState := { config := observerConfig } + let observerWorkspace : WorkspaceState := { generation := 2, config := observerConfig } let workspaces := state.workspaces.insert targetId targetWorkspace let workspaces := workspaces.insert observerId observerWorkspace set { state with workspaces } @@ -1173,6 +1204,201 @@ private def checkLifecycleTeardownConcurrency : IO Unit := do checkLifecycleTeardownReleasesStateMutex .reset checkLifecycleTeardownReleasesStateMutex .drop +private def checkWorkspaceSnapshotResetIsolation : IO Unit := do + let nonce ← IO.monoNanosNow + let workspaceId := "snapshot-reset" + let oldRoot := System.FilePath.mk s!"/tmp/beam-snapshot-old-{nonce}" + let newRoot := System.FilePath.mk s!"/tmp/beam-snapshot-new-{nonce}" + let fifo := oldRoot / "Slow.lean" + let writerReady := oldRoot / "writer-ready" + let releaseWriter := oldRoot / "release-writer" + IO.FS.createDirAll oldRoot + IO.FS.createDirAll newRoot + let mkfifo ← IO.Process.output { cmd := "mkfifo", args := #[fifo.toString] } + unless mkfifo.exitCode == 0 do + throw <| IO.userError s!"failed to create snapshot-reset FIFO: {mkfifo.stderr}" + let runtime ← ServerRuntime.create ({ root := oldRoot } : BrokerConfig) workspaceId + let requestTask ← IO.asTask (prio := Task.Priority.dedicated) <| + runtime.dispatchRequest { + payload := .updateFile { path := "Slow.lean" } + workspaceId? := some workspaceId + } + let writer ← IO.Process.spawn { + cmd := "python3" + args := #[ + "-c", + "import pathlib, sys, time\nwith open(sys.argv[1], 'w') as stream:\n pathlib.Path(sys.argv[2]).write_text('ready')\n while not pathlib.Path(sys.argv[3]).exists(): time.sleep(0.01)\n stream.write('def slowSnapshot : Nat := 1\\n')", + fifo.toString, + writerReady.toString, + releaseWriter.toString + ] + } + try + unless ← waitForPath writerReady do + throw <| IO.userError "snapshot-reset writer did not rendezvous with the broker read" + let oldGeneration ← runtime.state.atomically do + let state ← get + let some workspace := state.workspaces.get? workspaceId + | throw <| IO.userError "snapshot-reset workspace disappeared before reset" + pure workspace.generation + match ← runtime.initWorkspaceWithConfig workspaceId { root := newRoot } (some .reset) with + | .error failure => + throw <| IO.userError s!"snapshot-reset workspace reset failed: {failure.error.message}" + | .ok result => + require "snapshot-reset should invalidate the previous workspace generation" + result.invalidatedHandles + let newGeneration ← runtime.state.atomically do + let state ← get + let some workspace := state.workspaces.get? workspaceId + | throw <| IO.userError "snapshot-reset workspace disappeared after reset" + pure workspace.generation + require "workspace reset should allocate a fresh generation" (newGeneration != oldGeneration) + IO.FS.writeFile releaseWriter "release" + let response ← IO.ofExcept <| ← IO.wait requestTask + require "an old-generation source snapshot should be rejected as stale" + (response.error?.any fun err => + err.code == "contentModified" && err.message.contains "changed while") + let leanSessionActive ← runtime.state.atomically do + let state ← get + pure <| state.workspaces.get? workspaceId |>.bind (fun workspace => workspace.lean.session?) + |>.isSome + require "a stale snapshot should not start a backend in the replacement workspace" + !leanSessionActive + finally + if !(← releaseWriter.pathExists) then + IO.FS.writeFile releaseWriter "release" + if (← writer.tryWait).isNone then + writer.kill + discard <| writer.wait + runtime.close + if ← oldRoot.pathExists then + IO.FS.removeDirAll oldRoot + if ← newRoot.pathExists then + IO.FS.removeDirAll newRoot + +private def pendingOnlySession + (workspaceId : WorkspaceId) + (root exit : System.FilePath) : IO Session := do + let proc ← IO.Process.spawn { + toStdioConfig := brokerStdio + cmd := "python3" + args := #[ + "-c", + "import pathlib, sys, time\nwhile not pathlib.Path(sys.argv[1]).exists(): time.sleep(0.01)", + exit.toString + ] + } + let pending ← Std.Mutex.new ({} : Std.TreeMap Lean.JsonRpc.RequestID PendingRequest) + let stderrCapture ← startBackendStderrCapture proc.stderr + pure { + workspaceId + backend := .lean + root + epoch := 1 + sessionToken := s!"pending-only-{workspaceId}" + proc + stdin := IO.FS.Stream.ofHandle proc.stdin + stdout := IO.FS.Stream.ofHandle proc.stdout + stderrCapture + pending + } + +private partial def takePendingRequests + (store : PendingRequestStore) + (count : Nat) + (tries : Nat := 200) : IO (Array PendingRequest) := do + if (← PendingRequestStore.snapshot store).size >= count then + PendingRequestStore.clear store + else if tries == 0 then + throw <| IO.userError s!"timed out waiting for {count} pending backend request(s)" + else + IO.sleep 10 + takePendingRequests store count (tries - 1) + +private partial def waitForWorkspaceRoot + (runtime : ServerRuntime) + (workspaceId : WorkspaceId) + (expected : System.FilePath) + (tries : Nat := 200) : IO Bool := do + if (← runtime.workspaceRoot? workspaceId) == some expected then + pure true + else if tries == 0 then + pure false + else + IO.sleep 10 + waitForWorkspaceRoot runtime workspaceId expected (tries - 1) + +private def checkCompletedRequestResetIsolation : IO Unit := do + let nonce ← IO.monoNanosNow + let workspaceId := s!"completed-reset-{nonce}" + let oldRoot := System.FilePath.mk s!"/tmp/beam-completed-reset-old-{nonce}" + let newRoot := System.FilePath.mk s!"/tmp/beam-completed-reset-new-{nonce}" + let exit := oldRoot / "exit-backend" + IO.FS.createDirAll oldRoot + IO.FS.createDirAll newRoot + IO.FS.writeFile (oldRoot / "Demo.lean") "def demo : Nat := 1\n" + let config : BrokerConfig := { root := oldRoot } + let runtime ← ServerRuntime.create config workspaceId + let session ← pendingOnlySession workspaceId oldRoot exit + runtime.state.atomically do + let state ← get + let some workspace := state.workspaces.get? workspaceId + | throw <| IO.userError "completed-reset workspace disappeared" + let workspace := { workspace with + lean := { nextEpoch := 2, session? := some session } + } + set { state with workspaces := state.workspaces.insert workspaceId workspace } + let documentTask ← IO.asTask (prio := Task.Priority.dedicated) <| runtime.dispatchRequest { + payload := .runAt { + path := "Demo.lean" + version := 1 + line := 0 + character := 0 + text := "rfl" + } + workspaceId? := some workspaceId + } + let symbolsTask ← IO.asTask (prio := Task.Priority.dedicated) <| runtime.dispatchRequest { + payload := .workspaceSymbols { query := "demo" } + workspaceId? := some workspaceId + } + try + let requests ← takePendingRequests session.pending 2 + let resetTask ← IO.asTask (prio := Task.Priority.dedicated) <| + runtime.initWorkspaceWithConfig workspaceId { root := newRoot } (some .reset) + unless ← waitForWorkspaceRoot runtime workspaceId newRoot do + throw <| IO.userError "completed-reset workspace reset did not commit" + for request in requests do + PendingRequest.resolveResponse request (Json.mkObj []) + for (label, task) in [("document request", documentTask), ("workspace symbols", symbolsTask)] do + let response ← IO.ofExcept <| ← IO.wait task + require s!"{label}: an old-session result should be rejected after reset" + (response.error?.any fun err => err.code == "workerExited") + let shutdownRequests ← takePendingRequests session.pending 1 + for request in shutdownRequests do + PendingRequest.resolveResponse request Json.null + IO.FS.writeFile exit "exit" + match ← IO.ofExcept <| ← IO.wait resetTask with + | .ok result => + require "completed reset should invalidate the old session" result.invalidatedHandles + | .error failure => + throw <| IO.userError s!"completed workspace reset failed: {failure.error.message}" + finally + if !(← exit.pathExists) then + IO.FS.writeFile exit "exit" + try + runtime.close + catch _ => + pure () + try + session.proc.kill + catch _ => + pure () + if ← oldRoot.pathExists then + IO.FS.removeDirAll oldRoot + if ← newRoot.pathExists then + IO.FS.removeDirAll newRoot + private partial def waitForCancellation (cancelRef : IO.Ref Bool) (tries : Nat := 100) : IO Unit := do @@ -1218,6 +1444,32 @@ private def checkSessionCloseAdmission : IO Unit := do require "closed admission should leave no active request" ((← ActiveRequestRegistry.count runtime.activeRequests) == 0) +private def checkBrokerConfigBoundary : IO Unit := do + let root := System.FilePath.mk "/workspace" + let plugin := System.FilePath.mk "/beam/libbeam.so" + let config ← + match BrokerConfig.ofOptions root (some "lean") (some plugin) (some "rocq") with + | .ok config => pure config + | .error err => throw <| IO.userError s!"complete broker config failed: {err}" + match config.lean?, config.rocq? with + | some leanConfig, some rocqConfig => + require "broker config should preserve the Lean command" (leanConfig.command == "lean") + require "broker config should preserve the Lean plugin" (leanConfig.plugin == plugin) + require "broker CLI boundary should not invent a Lake helper" leanConfig.lakeHelper?.isNone + require "broker config should preserve the Rocq command" (rocqConfig.command == "rocq") + | _, _ => throw <| IO.userError "complete broker config lost a configured backend" + + let partialConfigs : Array (String × Option String × Option System.FilePath) := #[ + ("command only", some "lean", none), + ("plugin only", none, some plugin) + ] + for (label, command?, plugin?) in partialConfigs do + match BrokerConfig.ofOptions root command? plugin? with + | .ok _ => throw <| IO.userError s!"partial broker config '{label}' was accepted" + | .error err => + require s!"partial broker config '{label}' should explain the coupled fields" + (err.contains "command and plugin together") + private def checkWrapperDaemonAuthorization : IO Unit := do let base := System.FilePath.mk s!"/tmp/beam-wrapper-daemon-authorization-{← IO.monoNanosNow}" let rootPath := base / "workspace" @@ -1276,7 +1528,10 @@ def main : IO Unit := do checkWorkspaceRoutingFields checkWorkspaceLifecycleProtocol checkLifecycleTeardownConcurrency + checkWorkspaceSnapshotResetIsolation + checkCompletedRequestResetIsolation checkSessionCloseAdmission + checkBrokerConfigBoundary checkWrapperDaemonAuthorization end BeamTest.Broker.ProtocolTest diff --git a/tests/lean/BeamTest/Broker/SmokeTest.lean b/tests/lean/BeamTest/Broker/SmokeTest.lean index f8cbadc1..52297a65 100644 --- a/tests/lean/BeamTest/Broker/SmokeTest.lean +++ b/tests/lean/BeamTest/Broker/SmokeTest.lean @@ -827,8 +827,7 @@ private def runWorkspaceLifecycleSmoke let initResp ← runClient endpoint { payload := .initWorkspace { root := otherRoot.toString - leanCmd? := some leanCmd - leanPlugin? := some plugin.toString + lean? := some { command := leanCmd, plugin := plugin.toString } } workspaceId? := some workspaceId } @@ -839,8 +838,7 @@ private def runWorkspaceLifecycleSmoke let duplicateRoot ← runClient endpoint { payload := .initWorkspace { root := otherRoot.toString - leanCmd? := some leanCmd - leanPlugin? := some plugin.toString + lean? := some { command := leanCmd, plugin := plugin.toString } } workspaceId? := some "duplicate" } @@ -877,8 +875,7 @@ private def runWorkspaceLifecycleSmoke payload := .initWorkspace { root := otherRoot.toString workspaceMode? := some .reset - leanCmd? := some leanCmd - leanPlugin? := some plugin.toString + lean? := some { command := leanCmd, plugin := plugin.toString } } workspaceId? := some workspaceId } diff --git a/tests/lean/BeamTest/Broker/StreamDedupTest.lean b/tests/lean/BeamTest/Broker/StreamDedupTest.lean index 85ba311a..f6b6c7b3 100644 --- a/tests/lean/BeamTest/Broker/StreamDedupTest.lean +++ b/tests/lean/BeamTest/Broker/StreamDedupTest.lean @@ -159,6 +159,7 @@ private def fakeServerWithLeanSession (session : Beam.Broker.Session) : IO Beam.Broker.ServerRuntime := do let config : Beam.Broker.BrokerConfig := { root } let workspace : Beam.Broker.WorkspaceState := { + generation := 1 config lean := { nextEpoch := 1, session? := some session } }