refactor: tighten internal protocol types

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 :=
   {

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

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
     }

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)

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 {}

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

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