fix(#829): gate claudex and close branch coverage gaps
All checks were successful
ci/woodpecker/pr/ci Pipeline was successful
All checks were successful
ci/woodpecker/pr/ci Pipeline was successful
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user