From f5c68899ac3f0262de172b6a8ca81acac1c72d9b Mon Sep 17 00:00:00 2001 From: Emilio Jesus Gallego Arias Date: Thu, 3 Sep 2026 22:56:52 +0200 Subject: [PATCH] refactor: tighten internal protocol types --- Beam/Broker/Errors.lean | 31 ++++---- Beam/Broker/LakeHelper.lean | 60 +++++++++----- Beam/Broker/LakeHelperMain.lean | 17 ++-- Beam/Mcp/Protocol.lean | 4 +- Beam/Mcp/Server.lean | 4 +- .../lean/BeamTest/Broker/McpProtocolTest.lean | 4 +- tests/lean/BeamTest/Broker/ProtocolTest.lean | 79 +++++++++++++++++-- 7 files changed, 146 insertions(+), 53 deletions(-) diff --git a/Beam/Broker/Errors.lean b/Beam/Broker/Errors.lean index cdf163a4..9235e3ae 100644 --- a/Beam/Broker/Errors.lean +++ b/Beam/Broker/Errors.lean @@ -34,26 +34,27 @@ def BrokerFailureCode.name : BrokerFailureCode → String | .saveTargetNotModule => saveTargetNotModuleCode | .internalError => "internalError" -def BrokerFailureCode.all : Array BrokerFailureCode := #[ - .invalidParams, - .requestCancelled, - .contentModified, - .workerExited, - .syncBarrierIncomplete, - .saveTraceStale, - .saveUnsupportedSetup, - .saveTargetNotModule, - .internalError -] - -def BrokerFailureCode.ofName? (name : String) : Option BrokerFailureCode := - BrokerFailureCode.all.find? fun code => code.name == name +instance : ToJson BrokerFailureCode where + toJson code := toJson code.name + +instance : FromJson BrokerFailureCode where + fromJson? + | .str "invalidParams" => pure .invalidParams + | .str "requestCancelled" => pure .requestCancelled + | .str "contentModified" => pure .contentModified + | .str "workerExited" => pure .workerExited + | .str "syncBarrierIncomplete" => pure .syncBarrierIncomplete + | .str "saveTraceStale" => pure .saveTraceStale + | .str "saveUnsupportedSetup" => pure .saveUnsupportedSetup + | .str "saveTargetNotModule" => pure .saveTargetNotModule + | .str "internalError" => pure .internalError + | json => throw s!"expected Beam broker failure code, got {json.compress}" structure BrokerFailure where code : BrokerFailureCode message : String := "" data? : Option Json := none - deriving Inhabited + deriving Inhabited, FromJson, ToJson def BrokerFailure.toResponseFailure (failure : BrokerFailure) : ResponseFailure := { diff --git a/Beam/Broker/LakeHelper.lean b/Beam/Broker/LakeHelper.lean index 2983f66b..67ba8b22 100644 --- a/Beam/Broker/LakeHelper.lean +++ b/Beam/Broker/LakeHelper.lean @@ -51,6 +51,43 @@ structure LakeHelperSaveSpec extends LakeHelperWriteTraceRequest where structure LakeHelperAck where deriving FromJson, ToJson +private def requireOnlyResponseFields + (allowed : Array String) : Json → Except String Unit + | .obj fields => + let unexpected := fields.foldl (init := #[]) fun unexpected field _ => + if allowed.contains field then unexpected else unexpected.push field + unless unexpected.isEmpty do + throw s!"target Lake helper response accepts no undeclared fields: {String.intercalate ", " unexpected.toList}" + | other => throw s!"target Lake helper response must be an object, got {other.compress}" + +namespace LakeHelperResponse + +/-- Encode one typed result on the private target-built Lake helper boundary. -/ +def encode [ToJson α] : Except BrokerFailure α → Json + | .ok result => Json.mkObj [ + ("ok", toJson true), + ("result", toJson result) + ] + | .error failure => Json.mkObj [ + ("ok", toJson false), + ("error", toJson failure) + ] + +/-- Decode one typed result from the private target-built Lake helper boundary. -/ +def decode [FromJson α] (json : Json) : Except String (Except BrokerFailure α) := do + requireOnlyResponseFields #["ok", "result", "error"] json + let ok ← json.getObjValAs? Bool "ok" + if ok then + if (json.getObjVal? "error").isOk then + throw "target Lake helper response with ok=true must not include 'error'" + pure <| .ok (← json.getObjValAs? α "result") + else + if (json.getObjVal? "result").isOk then + throw "target Lake helper response with ok=false must not include 'result'" + pure <| .error (← json.getObjValAs? BrokerFailure "error") + +end LakeHelperResponse + inductive LakeHelperOperation where | serverEnv | prepareSave @@ -96,33 +133,16 @@ private def runLakeHelperRequest [ToJson α] [FromJson β] helperOutputSummary out.stdout out.stderr } let response ← - match Json.parse out.stdout >>= fromJson? (α := Response) with + match Json.parse out.stdout >>= LakeHelperResponse.decode (α := β) with | .ok response => pure response | .error err => return .error { code := .internalError message := - s!"target Lake helper '{operation.key}' returned invalid JSON: {err}: " ++ + s!"target Lake helper '{operation.key}' returned an invalid response: {err}: " ++ helperOutputSummary out.stdout out.stderr } - match response with - | .successResult result _ => - match fromJson? result with - | .ok decoded => pure <| .ok decoded - | .error err => - pure <| .error { - code := .internalError - message := s!"target Lake helper '{operation.key}' returned an invalid result: {err}" - } - | .errorResult failure => - let error := failure.error - let some code := BrokerFailureCode.ofName? error.code - | return .error { - code := .internalError - message := - s!"target Lake helper '{operation.key}' returned unknown error code '{error.code}'" - } - pure <| .error { code, message := error.message, data? := error.data? } + pure response /-- Ask the target-built helper for the Lean server environment. -/ def runLakeHelperServerEnv diff --git a/Beam/Broker/LakeHelperMain.lean b/Beam/Broker/LakeHelperMain.lean index 600ebfc3..f4c663b7 100644 --- a/Beam/Broker/LakeHelperMain.lean +++ b/Beam/Broker/LakeHelperMain.lean @@ -16,25 +16,28 @@ private def readRequest [FromJson α] : IO α := do let json ← IO.ofExcept <| Json.parse input IO.ofExcept <| fromJson? json -private def writeResponse (response : Response) : IO Unit := do - IO.println (toJson response).compress +private def writeResponse [ToJson α] (response : Except BrokerFailure α) : IO Unit := do + IO.println (LakeHelperResponse.encode response).compress + +private def writeFailure (failure : BrokerFailure) : IO Unit := + writeResponse (α := Json) <| .error failure private def runServerEnv : IO Unit := do let request : LakeHelperEnvRequest ← readRequest let serverEnv ← leanServerLakeEnv (System.FilePath.mk request.root) (some request.leanCmd) - writeResponse <| Response.success (toJson serverEnv) + writeResponse <| .ok serverEnv private def runPrepareSave : IO Unit := do let request : LakeHelperSaveRequest ← readRequest match ← lakeHelperSaveSpec request with - | .ok spec => writeResponse <| Response.success (toJson spec) - | .error failure => writeResponse failure.toResponse + | .ok spec => writeResponse <| .ok spec + | .error failure => writeFailure failure private def runWriteSaveTrace : IO Unit := do let request : LakeHelperWriteTraceRequest ← readRequest lakeHelperWriteLeanSaveTrace request - writeResponse <| Response.success (toJson ({} : LakeHelperAck)) + writeResponse <| .ok ({} : LakeHelperAck) def main (args : List String) : IO Unit := do try @@ -47,7 +50,7 @@ def main (args : List String) : IO Unit := do | none => throw <| IO.userError "invalid target Lake helper operation" | _ => throw <| IO.userError "invalid target Lake helper operation" catch error => - writeResponse <| BrokerFailure.toResponse { + writeFailure { code := .internalError message := error.toString } diff --git a/Beam/Mcp/Protocol.lean b/Beam/Mcp/Protocol.lean index 45d6ce84..4b453fee 100644 --- a/Beam/Mcp/Protocol.lean +++ b/Beam/Mcp/Protocol.lean @@ -426,7 +426,7 @@ instance : ToJson ToolStatusState where toJson state := toJson state.key structure ToolStatus where - requestId : Json + requestId : RequestId tool : String state : ToolStatusState message : String @@ -437,7 +437,7 @@ instance : ToJson ToolStatus where toJson status := Json.mkObj <| [ - ("requestId", status.requestId), + ("requestId", status.requestId.json), ("tool", toJson status.tool), ("state", toJson status.state), ("message", toJson status.message) diff --git a/Beam/Mcp/Server.lean b/Beam/Mcp/Server.lean index 6da93008..b839dc52 100644 --- a/Beam/Mcp/Server.lean +++ b/Beam/Mcp/Server.lean @@ -279,7 +279,7 @@ private structure RequestStatusState where emitted : Bool := false private structure RequestStatusEmitter where - requestId : Json + requestId : RequestId tool : ToolName path? : Option String state : Std.Mutex RequestStatusState @@ -332,7 +332,7 @@ private def RequestStatusEmitter.create let timer ← Std.Internal.UV.Timer.mk delayMs.toUInt64 false try let emitter : RequestStatusEmitter := { - requestId := requestId.json + requestId tool path? state := ← Std.Mutex.new {} diff --git a/tests/lean/BeamTest/Broker/McpProtocolTest.lean b/tests/lean/BeamTest/Broker/McpProtocolTest.lean index 3848ccc0..ab872b09 100644 --- a/tests/lean/BeamTest/Broker/McpProtocolTest.lean +++ b/tests/lean/BeamTest/Broker/McpProtocolTest.lean @@ -723,8 +723,10 @@ private def checkProgressProtocol : IO Unit := do { setupDiagnostic with message := "ordinary Lean information" } == .diagnostic + let statusRequestId ← expectOk "decode status request id" <| + Beam.Mcp.RequestId.fromJson? (toJson (42 : Nat)) let statusNotification := Beam.Mcp.toolStatusNotification { - requestId := toJson (42 : Nat) + requestId := statusRequestId tool := "lean_sync" state := .preparingDependencies message := setupMessage diff --git a/tests/lean/BeamTest/Broker/ProtocolTest.lean b/tests/lean/BeamTest/Broker/ProtocolTest.lean index 9f9d9396..8d53923d 100644 --- a/tests/lean/BeamTest/Broker/ProtocolTest.lean +++ b/tests/lean/BeamTest/Broker/ProtocolTest.lean @@ -507,12 +507,6 @@ private def checkSyncFileResultDecode : IO Unit := do syncFileResultJson 7 incompleteReadiness private def checkFailureResponseConversions : IO Unit := do - for code in BrokerFailureCode.all do - require s!"broker failure code '{code.name}' should round-trip from its wire name" - (BrokerFailureCode.ofName? code.name == some code) - require "unknown broker failure code should remain unknown" - (BrokerFailureCode.ofName? "not-a-broker-failure" |>.isNone) - let data := Json.mkObj [("uri", toJson "file:///A.lean")] let failure : BrokerFailure := { code := .contentModified @@ -542,6 +536,78 @@ private def checkFailureResponseConversions : IO Unit := do require "response failure preserves progress metadata" (failure.fileProgress? == some progress) +private def checkLakeHelperResponseProtocol : IO Unit := do + let failureCodes : Array BrokerFailureCode := #[ + .invalidParams, + .requestCancelled, + .contentModified, + .workerExited, + .syncBarrierIncomplete, + .saveTraceStale, + .saveUnsupportedSetup, + .saveTargetNotModule, + .internalError + ] + for code in failureCodes do + require s!"broker failure code '{code.name}' should round-trip" <| + (fromJson? (α := BrokerFailureCode) (toJson code)).toOption == some code + + let success : Except BrokerFailure LakeHelperAck := .ok {} + let successJson := LakeHelperResponse.encode success + requireJsonBool "Lake helper success" "ok" true successJson + discard <| requireObjVal "Lake helper success" "result" successJson + requireFieldAbsent "Lake helper success" "error" successJson + match LakeHelperResponse.decode (α := LakeHelperAck) successJson with + | .ok (.ok _) => pure () + | .ok (.error failure) => + throw <| IO.userError s!"Lake helper success decoded as failure: {failure.message}" + | .error err => + throw <| IO.userError s!"Lake helper success failed to round-trip: {err}" + + let data := Json.mkObj [("path", toJson "Demo.lean")] + let failure : BrokerFailure := { + code := .saveTraceStale + message := "save trace is stale" + data? := some data + } + let response : Except BrokerFailure LakeHelperAck := .error failure + let failureJson := LakeHelperResponse.encode response + requireJsonBool "Lake helper failure" "ok" false failureJson + requireFieldAbsent "Lake helper failure" "result" failureJson + let failurePayload ← requireObjVal "Lake helper failure" "error" failureJson + requireJsonString "Lake helper failure payload" "code" failure.code.name failurePayload + requireJsonString "Lake helper failure payload" "message" failure.message failurePayload + let failureData ← requireObjVal "Lake helper failure payload" "data" failurePayload + require "Lake helper failure payload preserves data" (failureData == data) + match LakeHelperResponse.decode (α := LakeHelperAck) failureJson with + | .ok (.error decoded) => + require "Lake helper failure preserves its typed code" (decoded.code == failure.code) + require "Lake helper failure preserves its message" (decoded.message == failure.message) + require "Lake helper failure preserves its data" (decoded.data? == failure.data?) + | .ok (.ok _) => + throw <| IO.userError "Lake helper failure decoded as success" + | .error err => + throw <| IO.userError s!"Lake helper failure failed to round-trip: {err}" + + let expectLakeHelperDecodeFailure (label : String) (json : Json) : IO Unit := do + match LakeHelperResponse.decode (α := LakeHelperAck) json with + | .ok _ => throw <| IO.userError s!"{label}: expected decode failure" + | .error _ => pure () + expectLakeHelperDecodeFailure "Lake helper response rejects unknown failure codes" <| Json.mkObj [ + ("ok", toJson false), + ("error", Json.mkObj [ + ("code", toJson "not-a-broker-failure"), + ("message", toJson "unknown") + ]) + ] + expectLakeHelperDecodeFailure "Lake helper response rejects mixed success and failure payloads" <| Json.mkObj [ + ("ok", toJson true), + ("result", toJson ({} : LakeHelperAck)), + ("error", toJson failure) + ] + expectLakeHelperDecodeFailure "Lake helper response rejects undeclared fields" <| + (LakeHelperResponse.encode success).setObjVal! "legacy" Json.null + private def checkTypedLakeSaveTraceFailure : IO Unit := do let missingHelper := System.FilePath.mk s!"/tmp/beam-missing-lake-helper-{← IO.monoNanosNow}" @@ -1651,6 +1717,7 @@ def main : IO Unit := do checkOrderedJsonPretty checkSyncFileResultDecode checkFailureResponseConversions + checkLakeHelperResponseProtocol checkTypedLakeSaveTraceFailure checkDocumentVersionMismatchErrorData checkReadinessBoundary