tywrap 0.7.0 → 0.9.0

package/runtime/python_bridge.py

@@ -1,9 +1,14 @@
 #!/usr/bin/env python3
 """
-Reference tywrap Python bridge server (subprocess + HTTP transports).
+Reference tywrap Python subprocess bridge.
+
+This module implements only the stdin/stdout subprocess loop for the Node/Bun/
+Deno transport. `HttpBridge` connects to an HTTP server provided by the
+operator; tywrap does not ship that server or an HTTP transport in this module.
+The operator is responsible for securing any network exposure.
 
 This module owns the I/O concerns and process identity for the Node/Bun/Deno
-subprocess transport and the HTTP transport:
+subprocess transport:
   - the stdin/stdout JSONL request/response loop (main())
   - env-var size guards (TYWRAP_CODEC_MAX_BYTES / TYWRAP_REQUEST_MAX_BYTES)
   - TYWRAP_CODEC_FALLBACK=json marker mode and TYWRAP_TORCH_ALLOW_COPY
@@ -14,7 +19,7 @@ subprocess transport and the HTTP transport:
     the explicit byte-size limit error message)
 
 SECURITY / TRUST MODEL: the bridge imports the requested module and getattrs the
-requested function/class/method, so call/instantiate/call_method are an arbitrary
+requested function/class method, so call is an arbitrary
 import+getattr+call surface. Two guards (implemented in tywrap_bridge_core) bound
 that surface:
   * TYWRAP_ALLOWED_MODULES (comma/space separated): when set, only those modules
@@ -37,9 +42,8 @@ import json
 import os
 import importlib  # noqa: F401  (re-exported for compat / used by handlers via core)
 
-from safe_codec import BridgeCodec, CodecError
-
 import tywrap_bridge_core as core
+from frame_codec import FrameError, Reassembler, encode_frames
 
 # Re-export the shared protocol/serialization surface so existing importers of
 # python_bridge keep working after the extraction (codex-flagged: runtime/ ships).
@@ -47,8 +51,9 @@ from tywrap_bridge_core import (  # noqa: F401
     PROTOCOL,
     PROTOCOL_VERSION,
     CODEC_VERSION,
+    BridgeCodec,
+    CodecError,
     ProtocolError,
-    InstanceHandleError,
     ImportNotAllowedError,
     AttributeNotAllowedError,
     import_allowed_module,
@@ -89,8 +94,6 @@ except (OSError, ValueError, TypeError, AttributeError) as exc:
     except (OSError, ValueError):
         pass
 
-instances = {}
-
 FALLBACK_JSON = os.environ.get('TYWRAP_CODEC_FALLBACK', '').lower() == 'json'
 TORCH_ALLOW_COPY = os.environ.get('TYWRAP_TORCH_ALLOW_COPY', '').lower() in ('1', 'true', 'yes')
 BRIDGE_NAME = 'python-subprocess'
@@ -221,6 +224,24 @@ def get_request_max_bytes():
 REQUEST_MAX_BYTES = get_request_max_bytes()
 
 
+def get_transport_frame_bytes():
+    """Read the private frame ceiling argument supplied by SubprocessTransport."""
+    if "--tywrap-max-frame-bytes" not in sys.argv:
+        return 100 * 1024 * 1024
+    try:
+        index = sys.argv.index("--tywrap-max-frame-bytes")
+        value = int(sys.argv[index + 1])
+    except (ValueError, IndexError) as exc:
+        raise CodecConfigError("--tywrap-max-frame-bytes must be a positive integer") from exc
+    if value <= 0:
+        raise CodecConfigError("--tywrap-max-frame-bytes must be a positive integer")
+    return value
+
+
+MAX_FRAME_BYTES = get_transport_frame_bytes()
+
+
+
 def serialize(obj):
     """
     Backward-compatible result serializer (subprocess identity).
@@ -262,7 +283,6 @@ def handle_meta():
     real pid and bridge='python-subprocess'.
     """
     return core.build_meta(
-        instances,
         bridge=BRIDGE_NAME,
         pid=os.getpid(),
         python_version=sys.version.split()[0],
@@ -270,7 +290,7 @@ def handle_meta():
     )
 
 
-def dispatch_request(msg):
+def dispatch_request(msg, *, has_envelope_markers=True):
     """
     Dispatch a validated request to the correct handler (subprocess identity).
 
@@ -280,7 +300,6 @@ def dispatch_request(msg):
     """
     out = core.dispatch_request(
         msg,
-        instances,
         bridge=BRIDGE_NAME,
         pid=os.getpid(),
         force_json_markers=FALLBACK_JSON,
@@ -289,6 +308,7 @@ def dispatch_request(msg):
         torch_allow_copy=TORCH_ALLOW_COPY,
         allowed_modules=ALLOWED_MODULES,
         allow_private_attrs=ALLOW_PRIVATE_ATTRS,
+        has_envelope_markers=has_envelope_markers,
     )
     return out['id'], out['result']
 
@@ -305,10 +325,11 @@ def encode_response(out):
     except CodecError as exc:
         # Convert CodecError to ValueError for consistent error handling
         raise ValueError(str(exc)) from exc
-    payload_bytes = len(payload.encode('utf-8'))
+    payload_utf8 = payload.encode('utf-8')
+    payload_bytes = len(payload_utf8)
     if CODEC_MAX_BYTES is not None and payload_bytes > CODEC_MAX_BYTES:
         raise PayloadTooLargeError(payload_bytes, CODEC_MAX_BYTES)
-    return payload
+    return payload, payload_utf8
 
 
 def write_payload(payload: str) -> bool:
@@ -326,55 +347,193 @@ def write_payload(payload: str) -> bool:
         return False
 
 
-def main():
-    for line in sys.stdin:
-        line = line.strip()
-        if not line:
-            continue
-        mid = None
-        out = None
+def write_response(payload: str, payload_utf8: bytes, response_id) -> bool:
+    """
+    Write a fully-encoded JSONL response, fragmenting it into ``tywrap-frame/1``
+    frames when the payload exceeds the fixed per-frame ceiling.
+
+    Why: a single response can exceed the JS-side JSONL line ceiling
+    (``maxLineLength``). The packaged subprocess transport always enables
+    framing, so an oversize response is split into frames, each written as its
+    own JSONL line and flushed one at a time so the OS pipe provides
+    backpressure. The TS reassembler rebuilds the logical response before the
+    codec sees it. Small responses keep the single-line fast path.
+
+    Frames require an integer correlation id. A response with a non-integer id
+    (only reachable on a malformed request whose error envelope carries id=None)
+    is never large enough to chunk, so it is always written as a single line.
+
+    Returns False if the parent's stdin/our stdout closed mid-write (BrokenPipe),
+    so the caller can exit the loop cleanly.
+    """
+    payload_bytes = len(payload_utf8)
+    if payload_bytes <= MAX_FRAME_BYTES:
+        return write_payload(payload)
+
+    if not isinstance(response_id, int) or isinstance(response_id, bool):
+        # Cannot correlate frames without an integer id; emit as one line. This
+        # only happens for tiny malformed-request error envelopes (id=None),
+        # which never exceed a sane frame ceiling.
+        return write_payload(payload)
+
+    frames = encode_frames(
+        payload,
+        id=response_id,
+        stream='response',
+        max_frame_bytes=MAX_FRAME_BYTES,
+        total_bytes=payload_bytes,
+    )
+    for frame in frames:
+        # One frame per JSONL line; flush per frame so the pipe backpressures
+        # and the TS reader can interleave reassembly with the write.
+        if not write_payload(_response_codec.encode(frame)):
+            return False
+    return True
+
+
+# Reassembles `tywrap-frame/1` REQUEST frames (W5) back into a single logical
+# request line. Framing is always accepted and restricted to the 'request'
+# stream. No reassembly-bytes cap here: the request size is
+# already bounded by the TS codec's maxPayloadBytes on the sending side, and
+# TYWRAP_REQUEST_MAX_BYTES is enforced on the complete payload after reassembly
+# (process_request_line), preserving the W5 post-reassembly semantics.
+_request_reassembler = Reassembler(expected_stream='request')
+
+
+def _accept_request_frame(frame):
+    """Accept one request frame, abandoning a different older serialized burst."""
+    frame_id = frame.get('id')
+    if frame.get('seq') == 0 and not _request_reassembler.is_pending(frame_id):
+        # The stdin loop is deliberately blocking/serial: an abandoned partial
+        # buffer persists while the worker is idle, then is freed by the next
+        # logical request (or by process restart). A background TTL timer would
+        # add concurrency solely to reclaim memory during an otherwise idle loop.
+        _request_reassembler.clear_pending()
+    return _request_reassembler.accept(frame)
+
+
+def _try_parse_frame_line(line):
+    """
+    Return a frame dict if ``line`` is a ``tywrap-frame/1`` envelope, else None.
+
+    Why: a request frame is a JSON object carrying ``__tywrap_frame__``. We only
+    treat a line as a frame when it parses as such an
+    object; anything else (including invalid JSON) falls through to the normal
+    single-line request path, which reports the JSON error exactly as before.
+    Structural validity (protocol, seq/total ranges, etc.) is enforced by the
+    Reassembler, not here.
+    """
+    try:
+        parsed = json.loads(line)
+    except json.JSONDecodeError:
+        return None
+    if isinstance(parsed, dict) and '__tywrap_frame__' in parsed:
+        return parsed
+    return None
+
+
+def process_request_line(line):
+    """
+    Process one complete logical request line and write its response.
+
+    ``line`` is the full logical JSON request: either a single line read from
+    stdin, or the payload reassembled from ``tywrap-frame/1`` request frames.
+    TYWRAP_REQUEST_MAX_BYTES is enforced on this complete payload (so for a
+    chunked request the limit applies to the REASSEMBLED size, not per frame).
+
+    Returns True to keep the loop running, or False if the parent's stdin / our
+    stdout closed mid-write (BrokenPipe), so main() can exit cleanly.
+    """
+    mid = None
+    out = None
+    try:
+        if REQUEST_MAX_BYTES is not None:
+            payload_bytes = len(line.encode('utf-8'))
+            if payload_bytes > REQUEST_MAX_BYTES:
+                raise RequestTooLargeError(payload_bytes, REQUEST_MAX_BYTES)
+        msg = json.loads(line)
+        if isinstance(msg, dict):
+            req_id = msg.get('id')
+            if isinstance(req_id, int):
+                # Why: preserve request ids even when handlers raise.
+                mid = req_id
         try:
-            if REQUEST_MAX_BYTES is not None:
-                payload_bytes = len(line.encode('utf-8'))
-                if payload_bytes > REQUEST_MAX_BYTES:
-                    raise RequestTooLargeError(payload_bytes, REQUEST_MAX_BYTES)
-            msg = json.loads(line)
-            if isinstance(msg, dict):
-                req_id = msg.get('id')
-                if isinstance(req_id, int):
-                    # Why: preserve request ids even when handlers raise.
-                    mid = req_id
-            try:
-                mid, result = dispatch_request(msg)
-                out = {'id': mid, 'protocol': PROTOCOL, 'result': result}
-            except ProtocolError as e:
-                emit_protocol_diagnostic(str(e))
-                out = build_error_payload(mid, e, include_traceback=False)
-            except Exception as e:  # noqa: BLE001
-                # Why: ensure any handler error becomes a protocol-compliant response.
-                out = build_error_payload(mid, e, include_traceback=True)
-        except RequestTooLargeError as e:
+            has_envelope_markers = '__tywrap' in line or '__type__' in line
+            mid, result = dispatch_request(msg, has_envelope_markers=has_envelope_markers)
+            out = {'id': mid, 'protocol': PROTOCOL, 'result': result}
+        except ProtocolError as e:
             emit_protocol_diagnostic(str(e))
             out = build_error_payload(mid, e, include_traceback=False)
-        except json.JSONDecodeError as e:
-            emit_protocol_diagnostic(f'Invalid JSON: {e}')
-            out = build_error_payload(mid, e, include_traceback=False)
         except Exception as e:  # noqa: BLE001
-            # Why: catch malformed input without breaking the JSONL protocol.
-            out = build_error_payload(mid, e, include_traceback=False)
+            # Why: ensure any handler error becomes a protocol-compliant response.
+            out = build_error_payload(mid, e, include_traceback=True)
+    except RequestTooLargeError as e:
+        emit_protocol_diagnostic(str(e))
+        out = build_error_payload(mid, e, include_traceback=False)
+    except json.JSONDecodeError as e:
+        emit_protocol_diagnostic(f'Invalid JSON: {e}')
+        out = build_error_payload(mid, e, include_traceback=False)
+    except Exception as e:  # noqa: BLE001
+        # Why: catch malformed input without breaking the JSONL protocol.
+        out = build_error_payload(mid, e, include_traceback=False)
 
+    try:
+        payload, payload_utf8 = encode_response(out)
+        # Correlate frames by the response id when chunking; out always
+        # carries the request id (or None for a malformed-request envelope).
+        response_id = out.get('id') if isinstance(out, dict) else None
+        if not write_response(payload, payload_utf8, response_id):
+            return False
+    except Exception as e:  # noqa: BLE001
+        # Why: fallback error keeps responses well-formed even if serialization fails.
+        err_out = build_error_payload(mid, e, include_traceback=False)
         try:
-            payload = encode_response(out)
-            if not write_payload(payload):
-                return
-        except Exception as e:  # noqa: BLE001
-            # Why: fallback error keeps responses well-formed even if serialization fails.
-            err_out = build_error_payload(mid, e, include_traceback=False)
+            # Error envelopes are tiny; write as a single line regardless of
+            # chunking so a serialization failure never recurses into framing.
+            if not write_payload(json.dumps(err_out)):
+                return False
+        except Exception:
+            return False
+    return True
+
+
+def main():
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+
+        # A line may be a `tywrap-frame/1` REQUEST frame. Reassemble per id; only
+        # the completed logical request is handed to process_request_line (which
+        # then enforces the request-size guard on the reassembled payload).
+        # Non-frame lines are normal requests.
+        frame = _try_parse_frame_line(line)
+        if frame is not None:
+            # JS serializes request frame bursts. A seq=0 frame for a different
+            # id proves any older partial stream was abandoned; a duplicate
+            # seq=0 for the same id must still fail loud in the reassembler.
             try:
-                if not write_payload(json.dumps(err_out)):
-                    return
-            except Exception:
+                reassembled = _accept_request_frame(frame)
+            except FrameError as exc:
+                # A framing-protocol violation desyncs the request stream and
+                # there is no correlatable response to write (the id may be
+                # malformed). Fail LOUD on stderr and stop the loop so the
+                # transport restarts the bridge rather than silently
+                # mis-parsing subsequent frames.
+                emit_protocol_diagnostic(f'Request frame error: {exc}')
+                return
+            if reassembled is None:
+                # More frames needed for this id; await the rest.
+                continue
+            if not process_request_line(reassembled):
                 return
+            continue
+
+        # Same idle-retention tradeoff as _accept_request_frame: a partial burst
+        # is retained only until the next request or process restart.
+        _request_reassembler.clear_pending()
+        if not process_request_line(line):
+            return
 
 
 if __name__ == '__main__':