fix(#829): gate claudex and close branch coverage gaps
All checks were successful
ci/woodpecker/pr/ci Pipeline was successful

This commit is contained in:
ms-lead-reviewer
2026-07-17 23:10:58 -05:00
parent 046896c612
commit 6a3505dfbc
10 changed files with 253 additions and 40 deletions

View File

@@ -8,6 +8,7 @@ import json
import os
import socket
import sys
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import Final
@@ -36,11 +37,17 @@ def broker_request(socket_path: Path, request: dict[str, object]) -> dict[str, o
return value
def main() -> int:
def main(
argv: Sequence[str] | None = None,
*,
environ: Mapping[str, str] | None = None,
request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
execute: Callable[[str, list[str], dict[str, str]], object] = os.execvpe,
) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
parser.add_argument("command", nargs=argparse.REMAINDER)
arguments = parser.parse_args()
arguments = parser.parse_args(argv)
command = arguments.command
if command and command[0] == "--":
command = command[1:]
@@ -48,12 +55,13 @@ def main() -> int:
print("lease-gated runtime command is required", file=sys.stderr)
return 64
source_environment = os.environ if environ is None else environ
try:
socket_path = Path(os.environ["MOSAIC_LEASE_BROKER_SOCKET"])
generation = int(os.environ.get("MOSAIC_RUNTIME_GENERATION", "1"))
socket_path = Path(source_environment["MOSAIC_LEASE_BROKER_SOCKET"])
generation = int(source_environment.get("MOSAIC_RUNTIME_GENERATION", "1"))
if generation < 0:
raise ValueError("invalid generation")
reply = broker_request(
reply = request(
socket_path,
{
"action": "register_anchor",
@@ -72,15 +80,16 @@ def main() -> int:
print("Mosaic lease broker registration failed; runtime launch denied.", file=sys.stderr)
return 1
environment = dict(os.environ)
environment = dict(source_environment)
environment["MOSAIC_LEASE_SESSION_ID"] = session_id
environment["MOSAIC_RUNTIME_GENERATION"] = str(generation)
environment["MOSAIC_LEASE_RUNTIME"] = arguments.runtime
try:
os.execvpe(command[0], command, environment)
execute(command[0], command, environment)
except OSError:
print("Mosaic lease-gated runtime exec failed.", file=sys.stderr)
return 1
return 0
if __name__ == "__main__":

View File

@@ -8,8 +8,9 @@ import json
import os
import socket
import sys
from collections.abc import Callable, Mapping, Sequence
from pathlib import Path
from typing import Final
from typing import BinaryIO, Final
MAX_FRAME: Final = 64 * 1024
BROKER_TIMEOUT_SECONDS: Final = 1.5
@@ -20,8 +21,9 @@ def deny(code: str) -> int:
return 2
def read_tool_name() -> str:
raw = sys.stdin.buffer.read(MAX_FRAME + 1)
def read_tool_name(stream: BinaryIO | None = None) -> str:
source = sys.stdin.buffer if stream is None else stream
raw = source.read(MAX_FRAME + 1)
if len(raw) > MAX_FRAME:
raise ValueError("INVALID_GATE_INPUT")
value = json.loads(raw)
@@ -56,19 +58,26 @@ def broker_request(socket_path: Path, request: dict[str, object]) -> dict[str, o
return value
def main() -> int:
def main(
argv: Sequence[str] | None = None,
*,
environ: Mapping[str, str] | None = None,
stream: BinaryIO | None = None,
request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
arguments = parser.parse_args()
arguments = parser.parse_args(argv)
source_environment = os.environ if environ is None else environ
try:
tool_name = read_tool_name()
socket_value = os.environ["MOSAIC_LEASE_BROKER_SOCKET"]
session_id = os.environ["MOSAIC_LEASE_SESSION_ID"]
generation = int(os.environ["MOSAIC_RUNTIME_GENERATION"])
tool_name = read_tool_name(stream)
socket_value = source_environment["MOSAIC_LEASE_BROKER_SOCKET"]
session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
generation = int(source_environment["MOSAIC_RUNTIME_GENERATION"])
if generation < 0:
raise ValueError("INVALID_GENERATION")
reply = broker_request(
reply = request(
Path(socket_value),
{
"action": "authorize_tool",