feat: implement JSON-RPC caching layer and backend method handlers for model specification and orchestration

Author: apstenku123

cppmega_mlx/data/nanochat_pipeline/tokenized_enriched.py

@@ -9,6 +9,10 @@
 
 from cppmega_mlx.data.nanochat_pipeline.tokenized_enriched_schema import (
     PLATFORM_IDS_COLUMN,
+    CHANGED_CHUNK_IDS_COLUMN,
+    CHANGED_CHUNK_SPANS_COLUMN,
+    EDIT_OP_PER_TOKEN_COLUMN,
+    HUNK_ID_PER_TOKEN_COLUMN,
     TOKEN_AST_DEPTH_COLUMN,
     TOKEN_AST_NODE_TYPE_COLUMN,
     TOKEN_CALL_EDGES_COLUMN,
@@ -25,6 +29,8 @@
     TOKEN_SYMBOL_IDS_COLUMN,
     TOKEN_TYPE_EDGES_COLUMN,
     TOKEN_TYPE_REFS_COLUMN,
+    TOKEN_CHANGE_MASK_POST_COLUMN,
+    TOKEN_CHANGE_MASK_PRE_COLUMN,
 )
 
 _KIND_STR_TO_INT = {
@@ -277,6 +283,44 @@ def _build_token_chunk_layout(
     }
 
 
+def _tokenize_optional_char_field(
+    doc: dict[str, Any],
+    keys: tuple[str, ...],
+    token_spans: list[tuple[int, int]],
+) -> list[int]:
+    for key in keys:
+        values = doc.get(key, [])
+        if values:
+            return _chars_to_tokens_structure_ids(values, "", token_spans)
+    return []
+
+
+def _changed_chunk_metadata(
+    token_chunk_starts: list[int],
+    token_chunk_ends: list[int],
+    change_masks: tuple[list[int], ...],
+) -> tuple[list[int], list[dict[str, int]]]:
+    if not token_chunk_starts or not token_chunk_ends:
+        return [], []
+    changed_ids: list[int] = []
+    changed_spans: list[dict[str, int]] = []
+    for chunk_idx, (start, end) in enumerate(zip(token_chunk_starts, token_chunk_ends)):
+        start_i = max(int(start), 0)
+        end_i = max(int(end), start_i)
+        changed = False
+        for mask in change_masks:
+            if not mask:
+                continue
+            bounded_end = min(end_i, len(mask))
+            if start_i < bounded_end and any(int(value) != 0 for value in mask[start_i:bounded_end]):
+                changed = True
+                break
+        if changed:
+            changed_ids.append(chunk_idx)
+            changed_spans.append({"start": start_i, "end": end_i})
+    return changed_ids, changed_spans
+
+
 def _platform_ids_from_doc(doc: dict[str, Any]) -> list[int]:
     platform_info = doc.get("platform_info")
     if not platform_info:
@@ -344,7 +388,28 @@ def materialize_tokenized_enriched_batch(
             len(token_ids),
         )
 
-        row = {
+        token_change_mask_pre = _tokenize_optional_char_field(
+            doc,
+            ("change_mask_pre", TOKEN_CHANGE_MASK_PRE_COLUMN),
+            token_spans,
+        )
+        token_change_mask_post = _tokenize_optional_char_field(
+            doc,
+            ("change_mask_post", "token_change_mask", TOKEN_CHANGE_MASK_POST_COLUMN),
+            token_spans,
+        )
+        hunk_id_per_token = _tokenize_optional_char_field(
+            doc,
+            ("hunk_id_per_char", HUNK_ID_PER_TOKEN_COLUMN),
+            token_spans,
+        )
+        edit_op_per_token = _tokenize_optional_char_field(
+            doc,
+            ("edit_op_per_char", "token_edit_op", EDIT_OP_PER_TOKEN_COLUMN),
+            token_spans,
+        )
+
+        row: dict[str, Any] = {
             TOKEN_IDS_COLUMN: token_ids,
             PLATFORM_IDS_COLUMN: _platform_ids_from_doc(doc),
             TOKEN_STRUCTURE_IDS_COLUMN: token_structure_ids,
@@ -384,8 +449,19 @@ def materialize_tokenized_enriched_batch(
                 "",
                 token_spans,
             ),
+            TOKEN_CHANGE_MASK_PRE_COLUMN: token_change_mask_pre,
+            TOKEN_CHANGE_MASK_POST_COLUMN: token_change_mask_post,
+            HUNK_ID_PER_TOKEN_COLUMN: hunk_id_per_token,
+            EDIT_OP_PER_TOKEN_COLUMN: edit_op_per_token,
         }
         row.update(cast(dict[str, list[int]], _build_token_chunk_layout(doc, tok_chunks, len(token_ids))))
+        changed_chunk_ids, changed_chunk_spans = _changed_chunk_metadata(
+            cast(list[int], row[TOKEN_CHUNK_STARTS_COLUMN]),
+            cast(list[int], row[TOKEN_CHUNK_ENDS_COLUMN]),
+            (token_change_mask_pre, token_change_mask_post),
+        )
+        row[CHANGED_CHUNK_IDS_COLUMN] = changed_chunk_ids
+        row[CHANGED_CHUNK_SPANS_COLUMN] = changed_chunk_spans
         out.append(row)
 
     return out

cppmega_mlx/runtime/path_c_fusion.py

@@ -2380,6 +2380,8 @@ def _tilelang_optimizer_for(
     capture = io.StringIO()
     with redirect_stdout(capture), redirect_stderr(capture):
         from tilelang.engine.fusion import (
+            FusionBlockDescriptor,
+            FusionBlockRegistry,
             FusionOptimizer,
             FusionScheduleRegistry,
         )
@@ -2408,14 +2410,25 @@ def _tilelang_optimizer_for(
         require_single_kernel=True,
         enable_z3_sync_async=region.z3_sync.enabled,
     )
-    for node in region.nodes:
-        optimizer.add_node(
-            node.name,
-            op=node.op_name,
-            inputs=node.inputs,
-            outputs=node.outputs,
-            attrs=_tilelang_node_attrs(node),
+    block_registry = FusionBlockRegistry(
+        tuple(
+            FusionBlockDescriptor(
+                op=node.op_name,
+                inputs=node.inputs,
+                outputs=node.outputs,
+                attrs=_tilelang_node_attrs(node),
+                aliases=(f"{node.name}:{node.op_name}",),
+            )
+            for node in region.nodes
         )
+    )
+    optimizer.add_blocks(
+        tuple(
+            {"name": node.name, "kind": f"{node.name}:{node.op_name}"}
+            for node in region.nodes
+        ),
+        block_registry=block_registry,
+    )
     workspace_edge_buffers = _declared_workspace_edge_buffers(schedule_template)
     node_by_name = {node.name: node for node in region.nodes}
     for edge in region.edges:

cppmega_v4/jsonrpc/cache.py

@@ -0,0 +1,115 @@
+"""LRU cache keyed on canonical sha256 of the spec.
+
+Cache key: SHA-256 of ``canonicalize(spec)`` — sorted keys, no
+whitespace, NaN/Inf rejected. Node positions are stripped via
+:func:`strip_layout` before canonicalisation so that dragging nodes
+around in the canvas doesn't bust the cache.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections import OrderedDict
+from threading import Lock
+from typing import Any, Mapping
+
+
+DEFAULT_CAPACITY: int = 50
+
+
+_LAYOUT_KEYS: frozenset[str] = frozenset({"x", "y", "position", "layout"})
+
+
+def strip_layout(payload: Any) -> Any:
+    """Recursively drop layout-only fields before hashing.
+
+    Visual canvas positions never influence backend semantics; stripping
+    them lets the cache survive drag operations without rebuilding.
+    """
+    if isinstance(payload, Mapping):
+        return {
+            k: strip_layout(v)
+            for k, v in payload.items()
+            if k not in _LAYOUT_KEYS
+        }
+    if isinstance(payload, list):
+        return [strip_layout(v) for v in payload]
+    if isinstance(payload, tuple):
+        return [strip_layout(v) for v in payload]
+    return payload
+
+
+def canonical_json(payload: Any) -> str:
+    """Produce a deterministic JSON string for cache keying."""
+    return json.dumps(
+        strip_layout(payload),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+        allow_nan=False,
+    )
+
+
+def canonical_sha256(payload: Any) -> str:
+    """Stable cache key — SHA-256 over :func:`canonical_json`."""
+    return hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest()
+
+
+class LRUCache:
+    """Bounded LRU cache for backend RPC results.
+
+    Thread-safe via a single Lock — RPC backends often share one cache
+    across event-loop tasks. Capacity defaults to 50 per VBPlan §5.1.
+    """
+
+    def __init__(self, capacity: int = DEFAULT_CAPACITY) -> None:
+        if capacity < 1:
+            raise ValueError(f"capacity must be ≥ 1, got {capacity}")
+        self._capacity = capacity
+        self._data: OrderedDict[str, Any] = OrderedDict()
+        self._lock = Lock()
+        self._hits = 0
+        self._misses = 0
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._data)
+
+    def get(self, key: str) -> Any | None:
+        with self._lock:
+            if key not in self._data:
+                self._misses += 1
+                return None
+            self._data.move_to_end(key)
+            self._hits += 1
+            return self._data[key]
+
+    def set(self, key: str, value: Any) -> None:
+        with self._lock:
+            if key in self._data:
+                self._data.move_to_end(key)
+                self._data[key] = value
+                return
+            self._data[key] = value
+            if len(self._data) > self._capacity:
+                self._data.popitem(last=False)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._data.clear()
+            self._hits = 0
+            self._misses = 0
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return {
+                "size": len(self._data),
+                "capacity": self._capacity,
+                "hits": self._hits,
+                "misses": self._misses,
+            }