Artur Mukhamadiev
e65fbf4b67
FastAPI app, upstream Confluence/model adapters, authoritative history, rootless container lifecycle, artifact storage and downloads, fake peers under backend/dev, tests under tests/backend. Root pytest.ini deselects the live marker by default; requirements gain the backend dependencies.
317 lines
13 KiB
Python
317 lines
13 KiB
Python
"""NDJSON bridge transport, framing, and protocol state machine (Revision 1)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import enum
|
|
import json
|
|
import re
|
|
from typing import Any, Dict, Optional, Set
|
|
|
|
from backend.errors import AppError
|
|
|
|
PROTOCOL_VERSION = 1
|
|
ID_REGEX = re.compile(r"^[A-Za-z0-9_-]{1,128}$")
|
|
|
|
# Wire limits (excluding trailing newline)
|
|
ENVELOPE_ALLOWANCE = 64 * 1024 # 64 KiB
|
|
MAX_TOOL_REQUEST_PAYLOAD = 16 * 1024 * 1024 # 16 MiB
|
|
MAX_MODEL_REQUEST_PAYLOAD = 128 * 1024 * 1024 # 128 MiB
|
|
MAX_COLLECTION_START_FRAME = 6 * 128 * 1024 * 1024 + ENVELOPE_ALLOWANCE
|
|
GLOBAL_MAX_FRAME_BYTES = MAX_COLLECTION_START_FRAME
|
|
|
|
MAX_INFLIGHT_REMOTE_CALLS = 4
|
|
|
|
|
|
class BridgeState(enum.Enum):
|
|
INIT = "INIT"
|
|
RUNNING = "RUNNING"
|
|
COLLECTING = "COLLECTING"
|
|
COMPLETE = "COMPLETE"
|
|
FAILED = "FAILED"
|
|
|
|
|
|
class NDJSONProtocolError(AppError):
|
|
def __init__(self, message: str):
|
|
super().__init__("execution_failed", message, status_code=500)
|
|
|
|
|
|
class BridgeTransport:
|
|
"""Async NDJSON reader and writer over attached container streams."""
|
|
|
|
def __init__(
|
|
self,
|
|
reader: asyncio.StreamReader,
|
|
writer: asyncio.StreamWriter,
|
|
max_frame_bytes: int = GLOBAL_MAX_FRAME_BYTES,
|
|
):
|
|
self.reader = reader
|
|
self.writer = writer
|
|
self.max_frame_bytes = max_frame_bytes
|
|
self._buffer = bytearray()
|
|
self.state = BridgeState.INIT
|
|
self._seen_ids: Set[str] = set()
|
|
self._inflight_calls: Set[str] = set() # set of request IDs awaiting response
|
|
self._b_counter = 0
|
|
|
|
def _next_backend_id(self) -> str:
|
|
self._b_counter += 1
|
|
return f"b_{self._b_counter}"
|
|
|
|
def validate_id(self, msg_id: str) -> None:
|
|
if not isinstance(msg_id, str) or not ID_REGEX.match(msg_id):
|
|
raise NDJSONProtocolError(f"Invalid message ID: '{msg_id}'")
|
|
if msg_id in self._seen_ids:
|
|
raise NDJSONProtocolError(f"Duplicate message ID: '{msg_id}'")
|
|
self._seen_ids.add(msg_id)
|
|
|
|
async def read_message_bounded(self) -> Dict[str, Any]:
|
|
"""Bounded frame read using readline with strict length check and type-specific limits."""
|
|
try:
|
|
if not hasattr(self.reader, "read"):
|
|
line = await self.reader.readline() # Test stream compatibility
|
|
else:
|
|
while True:
|
|
newline = self._buffer.find(b"\n")
|
|
if newline >= 0:
|
|
if newline > self.max_frame_bytes:
|
|
raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")
|
|
line = bytes(self._buffer[:newline + 1])
|
|
del self._buffer[:newline + 1]
|
|
break
|
|
if len(self._buffer) > self.max_frame_bytes:
|
|
raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")
|
|
chunk = await self.reader.read(min(64 * 1024, self.max_frame_bytes + 1 - len(self._buffer)))
|
|
if not chunk:
|
|
line = bytes(self._buffer)
|
|
self._buffer.clear()
|
|
break
|
|
self._buffer.extend(chunk)
|
|
except NDJSONProtocolError:
|
|
raise
|
|
except Exception as exc:
|
|
raise NDJSONProtocolError("Error reading bridge stream") from exc
|
|
|
|
if not line:
|
|
if self.state == BridgeState.COMPLETE:
|
|
return {"type": "eof"}
|
|
raise NDJSONProtocolError(f"Premature EOF in state {self.state.value}")
|
|
|
|
if not line.endswith(b"\n"):
|
|
raise NDJSONProtocolError("Bridge frame missing trailing newline")
|
|
|
|
raw_frame = line[:-1]
|
|
if len(raw_frame) > self.max_frame_bytes:
|
|
raise NDJSONProtocolError("Bridge frame exceeded maximum allowed bytes")
|
|
|
|
try:
|
|
msg = await asyncio.to_thread(json.loads, raw_frame)
|
|
except Exception as e:
|
|
raise NDJSONProtocolError("Malformed JSON in bridge frame") from e
|
|
if not isinstance(msg, dict) or set(msg) != {"v", "type", "id", "payload"}:
|
|
raise NDJSONProtocolError("Invalid bridge envelope fields")
|
|
if type(msg["v"]) is not int or msg["v"] != PROTOCOL_VERSION:
|
|
raise NDJSONProtocolError("Unsupported protocol version")
|
|
msg_type = msg["type"]
|
|
self.validate_id(msg["id"])
|
|
if not msg["id"].startswith("a_"):
|
|
raise NDJSONProtocolError("Agent message IDs must use a_ prefix")
|
|
schemas = {
|
|
"tool_request": {"tool": str, "parameters": dict},
|
|
"model_request": {"messages": list, "tools": list},
|
|
"collection_start": {"markdown": str, "warnings": list},
|
|
"artifact_begin": {"transfer_id": str, "name": str, "size_bytes": int},
|
|
"artifact_chunk": {"transfer_id": str, "index": int, "data_base64": str},
|
|
"artifact_end": {"transfer_id": str, "size_bytes": int, "chunks": int},
|
|
"complete": {"accepted_transfer_count": int},
|
|
"error": {"code": str, "message": str},
|
|
}
|
|
if not isinstance(msg_type, str) or msg_type not in schemas:
|
|
raise NDJSONProtocolError("Unknown bridge type")
|
|
payload = msg["payload"]
|
|
schema = schemas[msg_type]
|
|
if not isinstance(payload, dict) or set(payload) != set(schema) or any(type(payload[k]) is not t for k, t in schema.items()):
|
|
raise NDJSONProtocolError("Invalid bridge payload fields")
|
|
if any(t is int and payload[k] < 0 for k, t in schema.items()):
|
|
raise NDJSONProtocolError("Bridge counts must be nonnegative integers")
|
|
if "transfer_id" in payload and not ID_REGEX.fullmatch(payload["transfer_id"]):
|
|
raise NDJSONProtocolError("Invalid artifact transfer ID")
|
|
payload_limit = MAX_TOOL_REQUEST_PAYLOAD if msg_type == "tool_request" else MAX_MODEL_REQUEST_PAYLOAD if msg_type == "model_request" else None
|
|
if payload_limit is not None:
|
|
payload_bytes = await asyncio.to_thread(lambda: len(json.dumps(payload).encode("utf-8")))
|
|
if payload_bytes > payload_limit or len(raw_frame) > payload_limit + ENVELOPE_ALLOWANCE:
|
|
raise NDJSONProtocolError("Bridge payload exceeds type-specific byte limit")
|
|
elif msg_type == "collection_start":
|
|
metadata = {**msg, "payload": {**payload, "markdown": ""}}
|
|
metadata_bytes = await asyncio.to_thread(lambda: len(json.dumps(metadata).encode("utf-8")))
|
|
if metadata_bytes > ENVELOPE_ALLOWANCE:
|
|
raise NDJSONProtocolError("Collection metadata exceeds envelope allowance")
|
|
if len(payload["markdown"].encode("utf-8")) > 128 * 1024 * 1024:
|
|
raise NDJSONProtocolError("Collection markdown exceeds 128 MiB limit")
|
|
elif len(raw_frame) > ENVELOPE_ALLOWANCE + 6 * 64 * 1024:
|
|
raise NDJSONProtocolError("Bridge control frame exceeds byte limit")
|
|
self._transition_on_received(msg_type, msg)
|
|
return msg
|
|
|
|
def _transition_on_received(self, msg_type: str, msg: Dict[str, Any]) -> None:
|
|
if self.state in (BridgeState.COMPLETE, BridgeState.FAILED):
|
|
raise NDJSONProtocolError(f"Message received in terminal state {self.state.value}")
|
|
|
|
if self.state == BridgeState.INIT:
|
|
raise NDJSONProtocolError(f"Received '{msg_type}' before 'start' was sent")
|
|
|
|
if msg_type == "error":
|
|
self.state = BridgeState.FAILED
|
|
return
|
|
|
|
if self.state == BridgeState.RUNNING:
|
|
if msg_type in ("tool_request", "model_request"):
|
|
if len(self._inflight_calls) >= MAX_INFLIGHT_REMOTE_CALLS:
|
|
raise NDJSONProtocolError("Exceeded maximum in-flight remote calls (4)")
|
|
self._inflight_calls.add(msg["id"])
|
|
|
|
elif msg_type == "collection_start":
|
|
if self._inflight_calls:
|
|
raise NDJSONProtocolError(
|
|
f"Cannot transition to collection: {len(self._inflight_calls)} remote calls still in flight"
|
|
)
|
|
self.state = BridgeState.COLLECTING
|
|
|
|
else:
|
|
raise NDJSONProtocolError(f"Unexpected message '{msg_type}' in RUNNING state")
|
|
|
|
elif self.state == BridgeState.COLLECTING:
|
|
if msg_type in ("tool_request", "model_request"):
|
|
raise NDJSONProtocolError(f"Cannot perform '{msg_type}' in COLLECTING state")
|
|
elif msg_type in ("artifact_begin", "artifact_chunk", "artifact_end"):
|
|
pass # Handled by runner/artifacts
|
|
elif msg_type == "complete":
|
|
self.state = BridgeState.COMPLETE
|
|
else:
|
|
raise NDJSONProtocolError(f"Unexpected message '{msg_type}' in COLLECTING state")
|
|
|
|
elif self.state in (BridgeState.COMPLETE, BridgeState.FAILED):
|
|
raise NDJSONProtocolError(f"Message received in terminal state {self.state.value}")
|
|
|
|
async def _send_frame(self, msg: Dict[str, Any]) -> None:
|
|
msg_id = msg.get("id")
|
|
self.validate_id(msg_id)
|
|
raw = await asyncio.to_thread(lambda: json.dumps(msg).encode("utf-8") + b"\n")
|
|
if len(raw) - 1 > GLOBAL_MAX_FRAME_BYTES:
|
|
raise NDJSONProtocolError("Outgoing frame exceeds maximum bytes")
|
|
if msg["type"] in ("tool_response", "model_response"):
|
|
payload_bytes = await asyncio.to_thread(lambda: len(json.dumps(msg["payload"]).encode("utf-8")))
|
|
if payload_bytes > MAX_MODEL_REQUEST_PAYLOAD:
|
|
raise NDJSONProtocolError("Outgoing result exceeds payload byte limit")
|
|
self.writer.write(raw)
|
|
await self.writer.drain()
|
|
|
|
async def send_start(
|
|
self,
|
|
prompt: str,
|
|
system_instruction: str,
|
|
remaining_ms: int,
|
|
model_descriptor: Dict[str, Any],
|
|
) -> str:
|
|
"""Send start frame to container."""
|
|
if self.state != BridgeState.INIT:
|
|
raise NDJSONProtocolError(f"Cannot send start in state {self.state.value}")
|
|
|
|
msg_id = self._next_backend_id()
|
|
frame = {
|
|
"v": PROTOCOL_VERSION,
|
|
"type": "start",
|
|
"id": msg_id,
|
|
"payload": {
|
|
"prompt": prompt,
|
|
"system_instruction": system_instruction,
|
|
"remaining_ms": remaining_ms,
|
|
"model": model_descriptor,
|
|
},
|
|
}
|
|
await self._send_frame(frame)
|
|
self.state = BridgeState.RUNNING
|
|
return msg_id
|
|
|
|
async def send_tool_response(
|
|
self,
|
|
reply_to: str,
|
|
result: Optional[Dict[str, Any]],
|
|
error: Optional[Dict[str, Any]],
|
|
) -> str:
|
|
if reply_to in self._inflight_calls:
|
|
self._inflight_calls.remove(reply_to)
|
|
|
|
msg_id = self._next_backend_id()
|
|
payload = {
|
|
"result": result if error is None else None,
|
|
"error": error if error is not None else None,
|
|
}
|
|
frame = {
|
|
"v": PROTOCOL_VERSION,
|
|
"type": "tool_response",
|
|
"id": msg_id,
|
|
"reply_to": reply_to,
|
|
"payload": payload,
|
|
}
|
|
await self._send_frame(frame)
|
|
return msg_id
|
|
|
|
async def send_model_response(
|
|
self,
|
|
reply_to: str,
|
|
result: Optional[Dict[str, Any]],
|
|
error: Optional[Dict[str, Any]],
|
|
) -> str:
|
|
if reply_to in self._inflight_calls:
|
|
self._inflight_calls.remove(reply_to)
|
|
|
|
msg_id = self._next_backend_id()
|
|
payload = {
|
|
"result": result if error is None else None,
|
|
"error": error if error is not None else None,
|
|
}
|
|
frame = {
|
|
"v": PROTOCOL_VERSION,
|
|
"type": "model_response",
|
|
"id": msg_id,
|
|
"reply_to": reply_to,
|
|
"payload": payload,
|
|
}
|
|
await self._send_frame(frame)
|
|
return msg_id
|
|
|
|
async def send_collection_ready(self, reply_to: str) -> str:
|
|
msg_id = self._next_backend_id()
|
|
frame = {
|
|
"v": PROTOCOL_VERSION,
|
|
"type": "collection_ready",
|
|
"id": msg_id,
|
|
"reply_to": reply_to,
|
|
"payload": {},
|
|
}
|
|
await self._send_frame(frame)
|
|
return msg_id
|
|
|
|
async def send_artifact_ack(
|
|
self,
|
|
reply_to: str,
|
|
transfer_id: str,
|
|
decision: str, # "accept" | "skip" | "stored"
|
|
warning: Optional[Dict[str, Any]] = None,
|
|
) -> str:
|
|
msg_id = self._next_backend_id()
|
|
frame = {
|
|
"v": PROTOCOL_VERSION,
|
|
"type": "artifact_ack",
|
|
"id": msg_id,
|
|
"reply_to": reply_to,
|
|
"payload": {
|
|
"transfer_id": transfer_id,
|
|
"decision": decision,
|
|
"warning": warning,
|
|
},
|
|
}
|
|
await self._send_frame(frame)
|
|
return msg_id
|