[crazy-robotaxi] Add session replace_prompt passthrough + presenter extra_key_handlers

apps/crazy_robotaxi/crazy_robotaxi/hud_presenter.py

@@ -402,6 +402,7 @@ def __init__(
         scene_options: tuple[Any, ...],
         control_assets: Any,
         wheel: Any | None,
+        extra_key_handlers: dict[str, Callable[[], None]] | None = None,
     ) -> None:
         try:
             import slangpy as spy
@@ -436,6 +437,12 @@ def __init__(
         self._taxi_enclosure_segments_world = np.empty((0, 2, 3), dtype=np.float32)
         self._taxi_name_buffer = ""
         self._last_taxi_session_state: str | None = None
+        # Composition-root key extensions (e.g. live-edit abilities): keysym
+        # -> zero-arg callback, fired on discrete key press. Registered
+        # before ``_build_key_codes`` runs so the keysyms resolve to codes.
+        self._extra_key_handlers = _validate_extra_key_handlers(
+            extra_key_handlers, reserved=_RESERVED_HUD_KEYSYMS
+        )
 
         # Late-imports of helpers we need at runtime; ``demo`` imports
         # this module via the presenter factory, so direct top-level
@@ -2864,6 +2871,8 @@ def _build_key_codes(self) -> dict[str, Any]:
                 f"digit{character}",
                 f"num_{character}",
             )
+        for keysym in self._extra_key_handlers:
+            key_codes.setdefault(keysym, _lookup_key(spy.KeyCode, keysym))
         return key_codes
 
     def _taxi_name_character_for_key(self, key: Any) -> str | None:
@@ -2964,6 +2973,15 @@ def _on_keyboard_event(self, event: Any) -> None:
             self._keyboard.request_reset()
         elif self._key_matches(key, "x"):
             self.exit_scene()
+        else:
+            self._dispatch_extra_key(key)
+
+    def _dispatch_extra_key(self, key: Any) -> None:
+        """Fire the composition-root handler bound to this keysym, if any."""
+        for keysym, handler in self._extra_key_handlers.items():
+            if self._key_matches(key, keysym):
+                handler()
+                return
 
     def _expire_pending_drive_releases(self) -> None:
         """Commit any debounced release whose grace window has passed.
@@ -3398,6 +3416,52 @@ def configure_taxi_enclosure(self, segments_world: np.ndarray) -> None:
 # -- Module-level helpers ---------------------------------------------
 
 
+_RESERVED_HUD_KEYSYMS = frozenset(
+    {
+        "escape",
+        "f11",
+        "w",
+        "a",
+        "s",
+        "d",
+        "r",
+        "x",
+        "space",
+        "up",
+        "down",
+        "left",
+        "right",
+        "key1",
+        "key2",
+        "key3",
+        "backspace",
+        "enter",
+        "minus",
+        "underscore",
+    }
+)
+"""Keysyms the presenter dispatches itself; extra handlers may not shadow them."""
+
+
+def _validate_extra_key_handlers(
+    handlers: dict[str, Callable[[], None]] | None,
+    *,
+    reserved: frozenset[str],
+) -> dict[str, Callable[[], None]]:
+    """Reject handler keysyms that collide with the presenter's own keys.
+
+    A shadowed registration would never fire (the built-in branch wins),
+    so failing loudly at construction beats a silently dead key.
+    """
+    validated = dict(handlers or {})
+    conflicts = sorted(keysym for keysym in validated if keysym in reserved)
+    if conflicts:
+        raise ValueError(
+            f"extra_key_handlers may not rebind reserved keys: {conflicts}"
+        )
+    return validated
+
+
 def _lookup_key(key_enum: Any, *names: str) -> Any:
     for name in names:
         value = getattr(key_enum, name, None)

apps/crazy_robotaxi/crazy_robotaxi/streaming_presenter.py

@@ -103,6 +103,33 @@
     "3": "physx",
 }
 
+_RESERVED_BROWSER_KEYS = frozenset(
+    set(_BROWSER_KEY_TO_DRIVE_KEYSYM) | set(_BROWSER_KEY_TO_VIEW_MODE) | {"r", "R"}
+)
+"""Browser keys ``_apply_control`` dispatches itself; extras may not shadow them."""
+
+
+def _validate_extra_key_handlers(
+    handlers: dict[str, Callable[[], None]] | None,
+) -> dict[str, Callable[[], None]]:
+    """Reject handler keys that collide with the presenter's own bindings.
+
+    A shadowed registration would never fire (``_apply_control`` returns
+    before the extra-handler lookup), so failing loudly at construction
+    beats a silently dead key. Single-character keys are normalized to
+    lowercase because the browser posts ``e.key`` verbatim (``K`` when
+    Shift is held) while handlers are looked up case-insensitively.
+    """
+    validated: dict[str, Callable[[], None]] = {}
+    for key, handler in (handlers or {}).items():
+        validated[key.lower() if len(key) == 1 else key] = handler
+    conflicts = sorted(key for key in validated if key in _RESERVED_BROWSER_KEYS)
+    if conflicts:
+        raise ValueError(
+            f"extra_key_handlers may not rebind reserved keys: {conflicts}"
+        )
+    return validated
+
 
 class _KeyboardDriveSink:
     """In-process duck-typed ``ControlClient`` writing to ``KeyboardState``.
@@ -812,9 +839,15 @@ def __init__(
         jpeg_quality: int = 85,
         scenes: tuple[dict[str, object], ...] = (),
         thumbnails: dict[str, bytes] | None = None,
+        extra_key_handlers: dict[str, Callable[[], None]] | None = None,
     ) -> None:
         self._raster = raster
         self._keyboard = keyboard
+        # Composition-root key extensions (e.g. live-edit abilities): browser
+        # keysym (lowercase for letters) -> zero-arg callback, fired on
+        # keydown. Mirrors ``SlangPyHudPresenter``'s ``extra_key_handlers``
+        # so both transports expose the same hook.
+        self._extra_key_handlers = _validate_extra_key_handlers(extra_key_handlers)
         self._visual_flare = CollisionVisualFlare()
         self._taxi_enabled = False
         self._bev_config: BevConfig | None = None
@@ -1248,6 +1281,15 @@ def _apply_control(self, key: str, down: bool) -> None:
             # holding the key doesn't trigger a cascade of resets.
             if key in ("r", "R"):
                 self._keyboard.request_reset()
+                return
+            # Composition-root key extensions; single characters were
+            # normalized to lowercase at registration, so fold case here to
+            # keep Shift-modified presses bound to the same handler.
+            handler = self._extra_key_handlers.get(
+                key.lower() if len(key) == 1 else key
+            )
+            if handler is not None:
+                handler()
 
     def _state_snapshot(self) -> dict[str, object]:
         """Return a JSON-serializable vehicle and taxi telemetry snapshot.

apps/crazy_robotaxi/tests/test_presenter.py

@@ -17,6 +17,10 @@
 from crazy_robotaxi.hud_presenter import (
     SlangPyHudPresenter as CrazyRobotaxiHudPresenter,
 )
+from crazy_robotaxi.hud_presenter import (
+    _RESERVED_HUD_KEYSYMS,
+    _validate_extra_key_handlers,
+)
 from crazy_robotaxi.hud_presenter import (
     _build_bev_panel_image as _build_taxi_bev_panel_image,
 )
@@ -1493,3 +1497,32 @@ def render(_status: object) -> None:
     presenter.wait_while_preloading(lambda: True)
 
     assert renders == 0
+
+
+def test_hud_extra_key_handler_fires_on_matching_key_only() -> None:
+    presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter)
+    fired: list[str] = []
+    presenter._extra_key_handlers = {"k": lambda: fired.append("k")}
+    presenter._key_codes = {"k": 42}
+
+    presenter._dispatch_extra_key(42)
+    presenter._dispatch_extra_key(7)
+
+    assert fired == ["k"]
+
+
+def test_hud_build_key_codes_resolves_extra_handler_keysyms() -> None:
+    presenter = CrazyRobotaxiHudPresenter.__new__(CrazyRobotaxiHudPresenter)
+    presenter._extra_key_handlers = {"k": lambda: None}
+    presenter._spy = SimpleNamespace(KeyCode=SimpleNamespace(k=101))
+
+    key_codes = presenter._build_key_codes()
+
+    assert key_codes["k"] == 101
+
+
+def test_hud_extra_key_handlers_reject_reserved_keys() -> None:
+    with pytest.raises(ValueError, match="reserved"):
+        _validate_extra_key_handlers(
+            {"r": lambda: None}, reserved=_RESERVED_HUD_KEYSYMS
+        )

apps/crazy_robotaxi/tests/test_streaming_presenter_realtime.py

@@ -6,6 +6,7 @@
 import threading
 
 import numpy as np
+import pytest
 from crazy_robotaxi.game import TaxiGameSnapshot
 from crazy_robotaxi.input import (
     CrazyRobotaxiKeyboardState,
@@ -15,6 +16,7 @@
     MJPEGStreamingPresenter,
     _as_rgb_host_uint8,
     _publish_if_open,
+    _validate_extra_key_handlers,
     _wait_for_bus_frame,
 )
 from omnidreams_game_engine.camera import FThetaCameraModel
@@ -196,3 +198,21 @@ def test_streaming_state_snapshot_keeps_upstream_shape_outside_taxi() -> None:
         "steer_rad": 0.25,
         "yaw_rad": 0.5,
     }
+
+
+def test_streaming_extra_key_handler_fires_on_keydown_case_insensitively() -> None:
+    presenter = MJPEGStreamingPresenter.__new__(MJPEGStreamingPresenter)
+    presenter._keyboard = CrazyRobotaxiKeyboardState()
+    fired: list[str] = []
+    presenter._extra_key_handlers = {"k": lambda: fired.append("k")}
+
+    presenter._apply_control("k", True)
+    presenter._apply_control("K", True)  # Shift held: browser posts uppercase
+    presenter._apply_control("k", False)  # keyup must not re-fire
+
+    assert fired == ["k", "k"]
+
+
+def test_streaming_extra_key_handlers_reject_reserved_browser_keys() -> None:
+    with pytest.raises(ValueError, match="reserved"):
+        _validate_extra_key_handlers({"R": lambda: None})

apps/crazy_robotaxi/tests/test_world_model_adapter.py

@@ -313,6 +313,76 @@ def test_session_uses_flashdreams_pipeline_for_rollout() -> None:
     assert fake_pipeline.finalize_calls == [(0, "cache"), (1, "cache")]
 
 
+class _FakeReplaceTextPipeline(_FakePipeline):
+    def __init__(self) -> None:
+        super().__init__()
+        self.replace_text_calls: list[dict[str, object]] = []
+
+    def replace_text(
+        self, cache: object, text: list[list[str]], **kwargs: object
+    ) -> None:
+        self.replace_text_calls.append({"cache": cache, "text": text, **kwargs})
+
+
+def test_replace_prompt_flushes_pending_finalize_then_swaps() -> None:
+    fake_pipeline = _FakeReplaceTextPipeline()
+    session = FlashdreamsWorldModelSession(
+        _manifest(),
+        pipeline_factory=lambda manifest, profile: fake_pipeline,
+    )
+    session.warmup_model()
+    initial_rgb = np.zeros((2, 3, 3), dtype=np.uint8)
+    first_conditions = [np.zeros((2, 3, 3), dtype=np.uint8) for _ in range(5)]
+    next_conditions = [np.zeros((2, 3, 3), dtype=np.uint8) for _ in range(8)]
+    session.start(initial_rgb, first_conditions, "demo prompt")
+
+    session.replace_prompt("night city", guidance_scale=2.0, guidance_chunks=3)
+
+    # The deferred chunk-0 finalize must run under the OLD text, before the
+    # swap; deferring it into the next continue_generation would re-commit
+    # the chunk under the new prompt.
+    assert fake_pipeline.finalize_calls == [(0, "cache")]
+    assert fake_pipeline.replace_text_calls == [
+        {
+            "cache": "cache",
+            "text": [["night city"]],
+            "guidance_scale": 2.0,
+            "guidance_chunks": 3,
+        }
+    ]
+
+    session.continue_generation(next_conditions)
+    # The flushed finalize is not repeated by continue_generation.
+    assert fake_pipeline.finalize_calls == [(0, "cache")]
+    assert fake_pipeline.generate_calls[1]["autoregressive_index"] == 1
+
+
+def test_replace_prompt_before_start_raises() -> None:
+    session = FlashdreamsWorldModelSession(
+        _manifest(),
+        pipeline_factory=lambda manifest, profile: _FakeReplaceTextPipeline(),
+    )
+    session.warmup_model()
+
+    with pytest.raises(RuntimeError, match="start\\(\\) must be called"):
+        session.replace_prompt("night city")
+
+
+def test_replace_prompt_requires_pipeline_replace_text() -> None:
+    fake_pipeline = _FakePipeline()
+    session = FlashdreamsWorldModelSession(
+        _manifest(),
+        pipeline_factory=lambda manifest, profile: fake_pipeline,
+    )
+    session.warmup_model()
+    initial_rgb = np.zeros((2, 3, 3), dtype=np.uint8)
+    first_conditions = [np.zeros((2, 3, 3), dtype=np.uint8) for _ in range(5)]
+    session.start(initial_rgb, first_conditions, "demo prompt")
+
+    with pytest.raises(RuntimeError, match="replace_text"):
+        session.replace_prompt("night city")
+
+
 def test_session_postprocesses_local_frames_and_supports_live_toggle(
     monkeypatch,
 ) -> None:

apps/omnidreams_game_engine/omnidreams_game_engine/world_model/flashdreams_adapter.py

@@ -684,6 +684,48 @@ def continue_generation(self, condition_frames: list[object]) -> list[object]:
             )
         return model_frames
 
+    def replace_prompt(
+        self,
+        prompt: str,
+        *,
+        guidance_scale: float = 1.0,
+        guidance_chunks: int = 0,
+    ) -> None:
+        """Hot-swap the rollout's prompt at the current chunk boundary.
+
+        Public passthrough to the pipeline's ``replace_text`` so callers
+        (e.g. live-edit abilities) never reach into the private cache. Call
+        between :meth:`start` / :meth:`continue_generation` calls; the swap
+        applies from the next generated chunk onward.
+
+        The chunk finalize this session defers into the next
+        ``continue_generation`` is flushed first: finalize must run under
+        the OLD text, otherwise the previous chunk's KV history is
+        re-committed under the new prompt (an implicit recache).
+
+        Args:
+            prompt: Replacement prompt for subsequent chunks.
+            guidance_scale: Optional edit strength forwarded to
+                ``replace_text``; 1.0 is a plain swap.
+            guidance_chunks: Number of upcoming chunks to guide, forwarded
+                to ``replace_text``.
+        """
+        if self._cache is None:
+            raise RuntimeError("start() must be called before replace_prompt()")
+        replace_text = getattr(self.pipeline, "replace_text", None)
+        if not callable(replace_text):
+            raise RuntimeError("replace_prompt requires flashdreams replace_text().")
+        with torch.no_grad():
+            if self._pending_finalization_index is not None:
+                self.pipeline.finalize(self._pending_finalization_index, self._cache)
+                self._pending_finalization_index = None
+            replace_text(
+                self._cache,
+                [[prompt]],
+                guidance_scale=guidance_scale,
+                guidance_chunks=guidance_chunks,
+            )
+
     def reset(self, *, clear_precomputed_embeddings: bool = False) -> None:
         self._close_postprocess_stream()
         self._cache = None