Secure companion access and migrate to MCP v2 - #1
Open
flujo-app wants to merge 4 commits into
Open
Conversation
📊 Coverage ReportOverall Coverage: 62% Diff: origin/main...HEAD
Summary
Line-by-lineView line-by-line diff coveragesrc/kilntainers/auth.pyLines 66-75 66 parsed.scheme,
67 parsed.hostname.lower(),
68 parsed.port or (443 if parsed.scheme == "https" else 80),
69 )
! 70 except ValueError:
! 71 return None
72
73
74 def _matches(supplied: str, expected: str | None) -> bool:
75 """Compare bytes so malformed non-ASCII credentials cannot raise TypeError."""Lines 155-163 155 def _allowed_host(self, scope: Scope, headers: Headers) -> bool:
156 """Bind optional anonymous MCP to configured addresses, not DNS rebinding."""
157 hosts = headers.getlist("host")
158 if len(hosts) != 1:
! 159 return False
160 scheme = "https" if scope.get("scheme") == "https" else "http"
161 return _origin(f"{scheme}://{hosts[0]}") in self.allowed_origins
162
163 async def _reject(src/kilntainers/cli.pyLines 276-284 276 if server_config.output_limit < 1:
277 _startup_error("--output-limit must be at least 1 byte.")
278
279 if not server_config.computer_id:
! 280 _startup_error(
281 "COMPUTER_ID is required (example: COMPUTER_ID=agent-workstation)."
282 )
283 try:
284 validate_computer_id(server_config.computer_id)Lines 347-355 347 host = server_config.host
348 if host in {"127.0.0.1", "localhost", "::1", "0.0.0.0", "::"}:
349 hosts = ("127.0.0.1", "localhost", "[::1]")
350 else:
! 351 hosts = (f"[{host}]" if ":" in host else host,)
352 origins = tuple(f"http://{item}:{server_config.port}" for item in hosts)
353 app.add_middleware(
354 BearerTokenMiddleware, # ty: ignore[invalid-argument-type]
355 token=server_config.auth_token,Lines 363-373 363 return app
364
365
366 async def _run_http(mcp, server_config: ServerConfig) -> None:
! 367 import uvicorn
368
! 369 server = uvicorn.Server(
370 uvicorn.Config(
371 _protected_http_app(mcp, server_config),
372 host=server_config.host,
373 port=server_config.port,Lines 375-385 375 access_log=False,
376 timeout_graceful_shutdown=30,
377 )
378 )
! 379 await server.serve()
! 380 if mcp.cleanup_failed is True:
! 381 raise BackendError("Computer application cleanup did not complete.")
382
383
384 async def _run_stdio_with_dashboard(mcp, server_config: ServerConfig) -> None:
385 """Serve stdio MCP and a standalone loopback dashboard on one event loop."""Lines 384-392 384 async def _run_stdio_with_dashboard(mcp, server_config: ServerConfig) -> None:
385 """Serve stdio MCP and a standalone loopback dashboard on one event loop."""
386 import uvicorn
387
! 388 app = _protected_http_app(mcp, server_config)
389 uvicorn_config = uvicorn.Config(
390 app,
391 host="127.0.0.1",
392 port=server_config.port,Lines 399-407 399 try:
400 while not dashboard_server.started:
401 if dashboard_task.done():
402 await dashboard_task
! 403 raise RuntimeError(
404 "The standalone dashboard server stopped during startup."
405 )
406 await asyncio.sleep(0.01)
407 await mcp.run_stdio_async()Lines 407-416 407 await mcp.run_stdio_async()
408 finally:
409 dashboard_server.should_exit = True
410 await dashboard_task
! 411 if mcp.cleanup_failed is True:
! 412 raise BackendError("Computer application cleanup did not complete.")
413
414
415 def main() -> None:
416 """CLI entry point. Parses args, configures, and runs the server.Lines 447-456 447 with interruptible_stdin(server_config.transport == "stdio") as stop_stdin:
448 # Route termination through asyncio cancellation and awaited lifespan cleanup.
449 # A forced successful exit could hide failed or unfinished cleanup.
450 def _handle_sigterm(signum: int, frame: object) -> None:
! 451 stop_stdin()
! 452 signal.raise_signal(signal.SIGINT)
453
454 async def _run_stdio() -> None:
455 # asyncio.Runner has installed its cancellation handler by this point.
456 # Wake the file reader, then preserve that normal cancellation path.Lines 456-470 456 # Wake the file reader, then preserve that normal cancellation path.
457 runner_handler = signal.getsignal(signal.SIGINT)
458
459 def _handle_sigint(signum: int, frame: FrameType | None) -> None:
! 460 stop_stdin()
! 461 if callable(runner_handler):
! 462 cast(Callable[[int, FrameType | None], Any], runner_handler)(
463 signum, frame
464 )
465 else:
! 466 signal.default_int_handler(signum, frame)
467
468 signal.signal(signal.SIGINT, _handle_sigint)
469 try:
470 await _run_stdio_with_dashboard(mcp, server_config)Lines 472-480 472 signal.signal(signal.SIGINT, runner_handler)
473
474 handlers = {signal.SIGTERM: _handle_sigterm}
475 if hasattr(signal, "SIGBREAK"):
! 476 handlers[signal.SIGBREAK] = _handle_sigterm
477 previous = {
478 sig: signal.signal(sig, handler) for sig, handler in handlers.items()
479 }
480 try:Lines 480-488 480 try:
481 if server_config.transport == "stdio":
482 asyncio.run(_run_stdio())
483 else:
! 484 asyncio.run(_run_http(mcp, server_config))
485 except KeyboardInterrupt:
486 pass # Lifespan cleanup completed during asyncio runner shutdown.
487 finally:
488 for sig, handler in previous.items():src/kilntainers/server.pyLines 977-990 977
978
979 async def _notify_catalog_changed(ctx: Context[SessionContext, Any]) -> None:
980 """Use the appropriate public notification channel for the request's era."""
! 981 if ctx.protocol_version in MODERN_PROTOCOL_VERSIONS:
! 982 await ctx.notify_tools_changed()
! 983 await ctx.notify_resources_changed()
984 else:
! 985 await ctx.session.send_tool_list_changed()
! 986 await ctx.session.send_resource_list_changed()
987
988
989 def _activity_snapshot(
990 phase: str,Lines 1281-1289 1281 if payload.get("desktop_url"):
1282 payload["desktop_url"] = public_desktop_url(
1283 registry.peek(config.computer_id)
1284 )
! 1285 result = _result(payload, is_error=bool(result.is_error))
1286 publish_activity("result", "terminal_execute", arguments, payload)
1287 return result
1288
1289 mcp.add_tool(Lines 1340-1350 1340 is_error=True,
1341 )
1342 capabilities_changed = False
1343 if desktop_capability_sync is not None:
! 1344 capabilities_changed = desktop_capability_sync(sandbox.desktop_environment)
1345 if capabilities_changed and ctx is not None:
! 1346 await _notify_catalog_changed(ctx)
1347 sync_dashboard_resource_meta(sandbox.desktop_url)
1348 return _result(
1349 {
1350 "operation": "idle",Lines 1401-1411 1401 is_error=True,
1402 )
1403 capabilities_changed = False
1404 if sandbox is not None and desktop_capability_sync is not None:
! 1405 capabilities_changed = desktop_capability_sync(sandbox.desktop_environment)
1406 if capabilities_changed and ctx is not None:
! 1407 await _notify_catalog_changed(ctx)
1408 return _result(
1409 {
1410 "operation": "idle",
1411 "computer_id": config.computer_id,Lines 1439-1447 1439 ) -> CallToolResult:
1440 """Plug in or unplug the virtual computer's real network connection."""
1441 session = _session_from_context(ctx)
1442 if session is None:
! 1443 return _result(
1444 {"error": "Internal error: no context provided"}, is_error=True
1445 )
1446 try:
1447 async with runtime_switch_lock:Lines 1472-1480 1472 ) -> CallToolResult:
1473 """Switch the same computer between virtual and real desktops."""
1474 session = _session_from_context(ctx)
1475 if session is None:
! 1476 return _result(
1477 {"error": "Internal error: no context provided"}, is_error=True
1478 )
1479 try:
1480 async with runtime_switch_lock:Lines 1485-1493 1485 )
1486 if desktop_capability_sync is not None:
1487 desktop_capability_sync(sandbox.desktop_environment)
1488 if ctx is not None:
! 1489 await _notify_catalog_changed(ctx)
1490 return _result(
1491 {
1492 "operation": "idle",
1493 "computer_id": config.computer_id,src/kilntainers/stdio_input.pyLines 20-52 20 class _InputBuffer(io.RawIOBase):
21 def __init__(
22 self, messages: queue.Queue[bytes | OSError | None], stopped: threading.Event
23 ):
! 24 super().__init__()
! 25 self.messages = messages
! 26 self.stopped = stopped
! 27 self.pending = b""
28
29 def readable(self) -> bool:
! 30 return True
31
32 def readinto(self, buffer) -> int:
! 33 while not self.pending:
! 34 if self.stopped.is_set():
! 35 return 0
! 36 try:
! 37 data = self.messages.get(timeout=0.05)
! 38 except queue.Empty:
! 39 continue
! 40 if data is None:
! 41 return 0
! 42 if isinstance(data, OSError):
! 43 raise data
! 44 self.pending = data
! 45 size = min(len(buffer), len(self.pending))
! 46 buffer[:size] = self.pending[:size]
! 47 self.pending = self.pending[size:]
! 48 return size
49
50
51 @contextmanager
52 def interruptible_stdin(enabled: bool) -> Iterator[Callable[[], None]]:Lines 53-62 53 """Divert child stdin from MCP and permit SIGTERM without client EOF."""
54 stopped = threading.Event()
55 original = sys.stdin
56 if not enabled:
! 57 yield stopped.set
! 58 return
59 try:
60 fd = original.fileno()
61 except (AttributeError, OSError, io.UnsupportedOperation):
62 yield stopped.setLines 60-119 60 fd = original.fileno()
61 except (AttributeError, OSError, io.UnsupportedOperation):
62 yield stopped.set
63 return
! 64 if fd != 0:
! 65 yield stopped.set
! 66 return
67
! 68 restore_fd = os.dup(fd)
! 69 read_fd = os.dup(fd)
! 70 messages: queue.Queue[bytes | OSError | None] = queue.Queue(maxsize=2)
71
! 72 def send(data: bytes | OSError | None) -> bool:
! 73 while not stopped.is_set():
! 74 try:
! 75 messages.put(data, timeout=0.05)
! 76 return True
! 77 except queue.Full:
! 78 pass
! 79 return False
80
! 81 def pump() -> None:
! 82 try:
! 83 with os.fdopen(read_fd, "rb", buffering=0) as wire:
! 84 while not stopped.is_set():
! 85 data = wire.read(65536)
! 86 if not send(data or None) or not data:
! 87 break
! 88 except OSError as error:
! 89 send(error)
90
! 91 thread = threading.Thread(target=pump, name="mcp-stdin-pump", daemon=True)
! 92 raw = _InputBuffer(messages, stopped)
! 93 stream = io.TextIOWrapper(
94 io.BufferedReader(raw), encoding="utf-8", errors="replace"
95 )
! 96 started = False
! 97 try:
! 98 null_fd = os.open(os.devnull, os.O_RDONLY)
! 99 try:
! 100 os.dup2(null_fd, fd)
! 101 if sys.platform == "win32":
! 102 rebind_std_handle_to_fd(fd)
103 finally:
! 104 os.close(null_fd)
! 105 sys.stdin = stream
! 106 thread.start()
! 107 started = True
! 108 yield stopped.set
109 finally:
! 110 stopped.set()
! 111 sys.stdin = original
! 112 os.dup2(restore_fd, fd)
! 113 if sys.platform == "win32":
! 114 rebind_std_handle_to_fd(fd)
! 115 os.close(restore_fd)
! 116 if not started:
! 117 os.close(read_fd)
! 118 stream.close()
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Companion endpoints previously exposed recorded tool arguments without the configured MCP bearer token, and desktop WebSockets relied on loopback IP alone. This change authenticates sensitive HTTP and WebSocket routes, validates browser Origin/Host, gives the dashboard a separate process-scoped capability, and removes raw commands, file contents, output and capability URLs from activity history.
The server now uses Python MCP SDK 2.1.1 through public APIs and serves MCP 2026-07-28 plus supported legacy clients. The standalone browser uses the TypeScript v2 client. Computer identity and persisted files remain independent of protocol connections. The unused
--session-timeoutoption is removed with a migration note; audio relay URLs now preserve their route and authorization.Shutdown now awaits shielded, bounded cleanup instead of a five-second forced-success watchdog. A public stdin wrapper permits SIGTERM/SIGINT/SIGBREAK while the client keeps its input pipe open, and HTTP cleanup failures produce nonzero exit. Permanent computers still survive shutdown; no destructive MCP tool was added.
Validation:
npm run checkpasses: browser typecheck/build, release metadata, Python lint/types, 383 tests, and wheel/sdist builds; 48 provider/integration tests are excluded from that credential-free gate.24e9b135cb8a9bbaa7f520492d456c119f74cf6f. Release gate, native Linux/Windows signal/stdin tests and Coverage Report are green. Built-wheel Docker acceptance proves both-era persistence/reconnect, SIGTERM with input still open, permanent survival and awaited explicit backend deletion. Live Fly and interactive production desktop acceptance were not run.