This commit was merged in pull request #842.
This commit is contained in:
@@ -12,6 +12,8 @@ from collections.abc import Callable, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from lease_generation import initialize_runtime_generation
|
||||
|
||||
MAX_FRAME: Final = 64 * 1024
|
||||
BROKER_TIMEOUT_SECONDS: Final = 1.5
|
||||
CLAUDE_DANGEROUS_FLAG: Final = "--dangerously-skip-permissions"
|
||||
@@ -44,6 +46,7 @@ def main(
|
||||
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,
|
||||
initialize_generation: Callable[[Path, int], None] = initialize_runtime_generation,
|
||||
) -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
|
||||
@@ -83,6 +86,8 @@ def main(
|
||||
or any(character not in "0123456789abcdef" for character in session_id)
|
||||
):
|
||||
raise ValueError("registration refused")
|
||||
generation_file = socket_path.parent / f"generation-{session_id}.state"
|
||||
initialize_generation(generation_file, generation)
|
||||
except (KeyError, ValueError, OSError, json.JSONDecodeError):
|
||||
print("Mosaic lease broker registration failed; runtime launch denied.", file=sys.stderr)
|
||||
return 1
|
||||
@@ -90,6 +95,7 @@ def main(
|
||||
environment = dict(source_environment)
|
||||
environment["MOSAIC_LEASE_SESSION_ID"] = session_id
|
||||
environment["MOSAIC_RUNTIME_GENERATION"] = str(generation)
|
||||
environment["MOSAIC_LEASE_GENERATION_FILE"] = str(generation_file)
|
||||
environment["MOSAIC_LEASE_RUNTIME"] = arguments.runtime
|
||||
try:
|
||||
execute(command[0], command, environment)
|
||||
|
||||
Reference in New Issue
Block a user