1"""Passthrough request bodies and multipart field rewriting.
2
3One forwarded gateway request body is a :class:`RelayPassthroughBody`:
4either a decoded JSON object (content type ``application/json``) or raw
5bytes with their content type (``multipart/form-data`` and friends).
6The multipart helpers rewrite a single named form field in a raw body
7without parsing it, used by the passthrough service to substitute the
8outbound model alias.
9"""
10
11from __future__ import annotations
12
13from collections.abc import Iterator, Mapping
14from dataclasses import dataclass
15
16from lexigram.contracts.ai.relay import JsonValue
17
18__all__ = [
19 "RelayPassthroughBody",
20 "rewrite_multipart_form_field",
21]
22
23_JSON_CONTENT_TYPE = "application/json"
24_FORM_FIELD_HEADER_MARKER = b'name="'
25_FORM_FIELD_HEADER_SUFFIX = b'"'
26"""Multipart ``Content-Disposition`` attribute delimiters used by the field rewrite."""
27
28
29@dataclass(frozen=True, slots=True)
30class RelayPassthroughBody(Mapping[str, JsonValue]):
31 """One forwarded gateway request body: decoded JSON or raw bytes.
32
33 The two constructors are the entire surface: :meth:`json` wraps a
34 decoded JSON object (content type ``application/json``) and
35 :meth:`raw` wraps arbitrary bytes with their content type, so
36 ``multipart/form-data`` requests travel through the same
37 ``RelayGatewayRequest.payload`` field as JSON bodies. The mapping
38 facade (``__getitem__``/``__iter__``/``__len__``) delegates to the
39 JSON dict for ``json`` bodies and raises ``TypeError`` for raw
40 bodies — the passthrough pipeline branches on whether ``data`` is a
41 mapping and never treats raw content as JSON.
42
43 Attributes:
44 data: The decoded JSON object for ``json`` bodies, or the raw
45 body bytes for ``raw`` bodies.
46 content_type: Outbound content type header value; ``json``
47 bodies always carry ``application/json``.
48 """
49
50 data: Mapping[str, JsonValue] | bytes
51 content_type: str
52
53 @classmethod
54 def json(cls, payload: Mapping[str, JsonValue]) -> RelayPassthroughBody:
55 """Wrap a decoded JSON object request body.
56
57 Args:
58 payload: The decoded JSON object to forward.
59
60 Returns:
61 A JSON body carrying ``application/json`` as its content
62 type; the object is shallow-copied so later mutation of the
63 source never leaks into the frozen body.
64 """
65 return cls(dict(payload), _JSON_CONTENT_TYPE)
66
67 @classmethod
68 def raw(cls, data: bytes, content_type: str) -> RelayPassthroughBody:
69 """Wrap a raw (e.g. ``multipart/form-data``) request body.
70
71 Args:
72 data: The raw body bytes to forward verbatim.
73 content_type: The body's content type header (boundary
74 parameter included for multipart bodies).
75
76 Returns:
77 A raw body carrying *content_type* unchanged.
78 """
79 return cls(data, content_type)
80
81 def __getitem__(self, key: str) -> JsonValue:
82 """Return one JSON field for ``json`` bodies.
83
84 Raises:
85 TypeError: If the body is raw bytes, which are not JSON.
86 """
87 data = self.data
88 if not isinstance(data, Mapping):
89 raise TypeError("raw passthrough bodies are not JSON mappings")
90 return data[key]
91
92 def __iter__(self) -> Iterator[str]:
93 """Iterate the JSON field names for ``json`` bodies.
94
95 Raises:
96 TypeError: If the body is raw bytes, which are not JSON.
97 """
98 data = self.data
99 if not isinstance(data, Mapping):
100 raise TypeError("raw passthrough bodies are not JSON mappings")
101 return iter(data)
102
103 def __len__(self) -> int:
104 """Return the JSON field count for ``json`` bodies.
105
106 Raises:
107 TypeError: If the body is raw bytes, which are not JSON.
108 """
109 data = self.data
110 if not isinstance(data, Mapping):
111 raise TypeError("raw passthrough bodies are not JSON mappings")
112 return len(data)
113
114
115def rewrite_multipart_form_field(
116 body: bytes,
117 boundary: str,
118 field: str,
119 value: str,
120) -> bytes:
121 """Rewrite one named form field's value in a multipart body.
122
123 Narrow boundary-aware rewrite (not a general multipart parser): the
124 body is split on the ``--<boundary>`` framing marker and the first
125 part whose ``Content-Disposition`` header carries
126 ``name="<field>"`` has its value content swapped in place; every
127 other byte — headers, other parts, the closing marker — is left
128 untouched. A body without the field (or without the boundary
129 marker) is returned unchanged; that is not an error, some
130 passthrough endpoints resolve the model from the URL path or a
131 channel default instead of a body field.
132
133 Args:
134 body: The raw ``multipart/form-data`` body bytes.
135 boundary: The boundary token from the content-type header.
136 field: The form field name to rewrite (e.g. ``"model"``).
137 value: The replacement field value.
138
139 Returns:
140 The body with the named field's value replaced, or the body
141 unchanged when the field is absent.
142 """
143 marker = b"--" + boundary.encode("utf-8")
144 target = (
145 _FORM_FIELD_HEADER_MARKER + field.encode("utf-8") + _FORM_FIELD_HEADER_SUFFIX
146 )
147 replacement = value.encode("utf-8")
148 segments = body.split(marker)
149 if len(segments) < 2:
150 return body
151 for index in range(1, len(segments) - 1):
152 part = segments[index]
153 separator = part.find(b"\r\n\r\n")
154 if separator < 0:
155 continue
156 headers = part[2:separator].lower()
157 if target not in headers:
158 continue
159 value_end = len(part) - 2 if part.endswith(b"\r\n") else len(part)
160 segments[index] = part[: separator + 4] + replacement + part[value_end:]
161 return marker.join(segments)
162 return body
163
164
165def _as_relay_body(payload: Mapping[str, JsonValue]) -> RelayPassthroughBody:
166 """Normalize a gateway request payload into a relay passthrough body.
167
168 Bodies already carrying the relay carrier pass through unchanged;
169 plain JSON mappings (legacy callers) are wrapped as JSON bodies.
170
171 Args:
172 payload: The ``RelayGatewayRequest.payload`` value.
173
174 Returns:
175 The payload as a :class:`RelayPassthroughBody`.
176 """
177 if isinstance(payload, RelayPassthroughBody):
178 return payload
179 return RelayPassthroughBody.json(dict(payload))
180
181
182def _multipart_boundary(content_type: str) -> str | None:
183 """Extract the ``boundary`` parameter from a content-type header.
184
185 Args:
186 content_type: The raw content-type header value.
187
188 Returns:
189 The boundary token without surrounding quotes, or ``None`` when
190 the header carries no boundary parameter.
191 """
192 for parameter in content_type.split(";"):
193 key, separator, raw_value = parameter.strip().partition("=")
194 if separator and key.lower() == "boundary":
195 return raw_value.strip().strip('"')
196 return None
197
198
199def _is_json_content_type(content_type: str) -> bool:
200 """Tell whether a content-type value denotes JSON.
201
202 Args:
203 content_type: A content-type header value.
204
205 Returns:
206 ``True`` for the exact ``application/json`` media type and for
207 any ``*+json`` suffix variant; ``False`` otherwise.
208 """
209 media_type = content_type.partition(";")[0].strip().lower()
210 return media_type == _JSON_CONTENT_TYPE or media_type.endswith("+json")