graphify-mesh 0.0.22__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- graphify_mesh/__init__.py +20 -0
- graphify_mesh/_version.py +24 -0
- graphify_mesh/server/__init__.py +61 -0
- graphify_mesh/server/anchors.py +223 -0
- graphify_mesh/server/citation.py +17 -0
- graphify_mesh/server/config.py +195 -0
- graphify_mesh/server/context_pack.py +212 -0
- graphify_mesh/server/embed_query.py +206 -0
- graphify_mesh/server/frames.py +93 -0
- graphify_mesh/server/http_app.py +367 -0
- graphify_mesh/server/lexical_read.py +208 -0
- graphify_mesh/server/project_map.py +92 -0
- graphify_mesh/server/ranking.py +209 -0
- graphify_mesh/server/retrieval.py +343 -0
- graphify_mesh/server/rwlock.py +61 -0
- graphify_mesh/server/scope.py +207 -0
- graphify_mesh/server/sdk_app.py +108 -0
- graphify_mesh/server/server.py +813 -0
- graphify_mesh/server/similar.py +284 -0
- graphify_mesh/server/stdio_guard.py +287 -0
- graphify_mesh/server/store.py +733 -0
- graphify_mesh/server/traverse.py +305 -0
- graphify_mesh/sync/__init__.py +26 -0
- graphify_mesh/sync/backend.py +80 -0
- graphify_mesh/sync/cli.py +249 -0
- graphify_mesh/sync/clustering.py +182 -0
- graphify_mesh/sync/config.py +815 -0
- graphify_mesh/sync/discovery.py +668 -0
- graphify_mesh/sync/embed_similarity.py +516 -0
- graphify_mesh/sync/embedding.py +1372 -0
- graphify_mesh/sync/graphify_cli.py +941 -0
- graphify_mesh/sync/lexical_index.py +293 -0
- graphify_mesh/sync/locking.py +30 -0
- graphify_mesh/sync/naming.py +563 -0
- graphify_mesh/sync/naming_state.py +143 -0
- graphify_mesh/sync/overlay.py +146 -0
- graphify_mesh/sync/overlay_api.py +312 -0
- graphify_mesh/sync/overlay_depends.py +324 -0
- graphify_mesh/sync/overlay_refs.py +143 -0
- graphify_mesh/sync/overlay_similar.py +357 -0
- graphify_mesh/sync/perms.py +116 -0
- graphify_mesh/sync/pipeline.py +1891 -0
- graphify_mesh/sync/progress.py +41 -0
- graphify_mesh/sync/publish.py +374 -0
- graphify_mesh/sync/reaper.py +201 -0
- graphify_mesh/sync/reaper_cli.py +62 -0
- graphify_mesh/sync/registry.py +184 -0
- graphify_mesh/sync/repo_tags.py +299 -0
- graphify_mesh/sync/rss.py +68 -0
- graphify_mesh/sync/source_cache.py +149 -0
- graphify_mesh/sync/state.py +268 -0
- graphify_mesh/sync/sync_project.py +793 -0
- graphify_mesh/sync/tls.py +119 -0
- graphify_mesh/sync/validate.py +277 -0
- graphify_mesh/sync/vectors.py +234 -0
- graphify_mesh-0.0.22.dist-info/METADATA +219 -0
- graphify_mesh-0.0.22.dist-info/RECORD +60 -0
- graphify_mesh-0.0.22.dist-info/WHEEL +4 -0
- graphify_mesh-0.0.22.dist-info/entry_points.txt +4 -0
- graphify_mesh-0.0.22.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,284 @@
|
|
|
1
|
+
"""`find_similar` tool implementation (WS5 deliverable 5).
|
|
2
|
+
|
|
3
|
+
Reads the WS4 `cross-project-overlay.json` `similar_approach` edges (the
|
|
4
|
+
real ANN cosine-similarity output, see `graphify_mesh.sync.overlay_similar`) for
|
|
5
|
+
cross-repo candidates, adds same-repo structural neighbors when
|
|
6
|
+
`cross_repo_only=False`, and — for a node with NEITHER an overlay edge NOR
|
|
7
|
+
same-repo neighbors (the documented trivial/unembedded case) — falls back to
|
|
8
|
+
the exact label+community match `graphify_mesh.sync.overlay_similar` already
|
|
9
|
+
implements at build time. This module does not reimplement that scoring
|
|
10
|
+
algorithm: it calls `overlay_similar.normalize_label` (the same public
|
|
11
|
+
normalization wrapper the build-time scorer uses) and applies the identical
|
|
12
|
+
match rule (same normalized label AND same `community_name`) directly over
|
|
13
|
+
the published merged graph, since the build-time function's own signature
|
|
14
|
+
expects per-repo raw graphs that are no longer available at query time (only
|
|
15
|
+
the merged, already-repo-attributed graph is published) — the ALGORITHM is
|
|
16
|
+
shared, only the input data shape differs.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
import weakref
|
|
22
|
+
from collections.abc import Callable
|
|
23
|
+
from dataclasses import dataclass, field
|
|
24
|
+
from typing import TypeVar
|
|
25
|
+
|
|
26
|
+
from graphify_mesh.server import lexical_read, ranking
|
|
27
|
+
from graphify_mesh.server.retrieval import Hit, _hit_from_key
|
|
28
|
+
from graphify_mesh.server.store import Generation
|
|
29
|
+
from graphify_mesh.sync.lexical_index import normalize_alias_query
|
|
30
|
+
from graphify_mesh.sync.overlay_refs import LogicalRef
|
|
31
|
+
from graphify_mesh.sync.overlay_similar import normalize_label
|
|
32
|
+
|
|
33
|
+
FALLBACK_SCORE = 0.4
|
|
34
|
+
STRUCTURAL_NEIGHBOR_SCORE = 0.5
|
|
35
|
+
FALLBACK_PROVENANCE = "PLACEHOLDER_STRUCTURAL_MATCH"
|
|
36
|
+
STRUCTURAL_PROVENANCE = "STRUCTURAL_NEIGHBOR"
|
|
37
|
+
|
|
38
|
+
# Degraded marker emitted when `enabled_repos` actually excluded a candidate.
|
|
39
|
+
DISABLED_FILTERED_DEGRADED = "disabled_repo_hits_filtered"
|
|
40
|
+
# Degraded marker for a query seed whose own repo is not in `enabled_repos`.
|
|
41
|
+
SEED_REPO_DISABLED_DEGRADED = "node_repo_disabled"
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
@dataclass
|
|
45
|
+
class SimilarResult:
|
|
46
|
+
resolved: bool
|
|
47
|
+
hits: list[Hit] = field(default_factory=list)
|
|
48
|
+
degraded: list[str] = field(default_factory=list)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
_T = TypeVar("_T")
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def _per_generation_cache(build_fn: Callable[[Generation], _T]) -> Callable[[Generation], _T]:
|
|
55
|
+
"""Lazy per-generation memoization WITHOUT touching `store.py`: the
|
|
56
|
+
cache lives in this module, keyed on the Generation instance's identity
|
|
57
|
+
(`id()`), with a `weakref.finalize` hook evicting the entry as soon as
|
|
58
|
+
the generation object is garbage-collected — so a recycled `id()` from
|
|
59
|
+
a later generation can never observe a stale index, and dropped
|
|
60
|
+
generations don't pin their indexes in memory. (`Generation` is a
|
|
61
|
+
non-frozen `eq=True` dataclass, hence unhashable — a plain
|
|
62
|
+
`WeakKeyDictionary` keyed on the instance is not an option.)"""
|
|
63
|
+
import threading
|
|
64
|
+
|
|
65
|
+
cache: dict[int, _T] = {}
|
|
66
|
+
lock = threading.Lock()
|
|
67
|
+
|
|
68
|
+
def get(generation: Generation) -> _T:
|
|
69
|
+
cache_key = id(generation)
|
|
70
|
+
hit = cache.get(cache_key)
|
|
71
|
+
if hit is not None:
|
|
72
|
+
return hit
|
|
73
|
+
with lock:
|
|
74
|
+
hit = cache.get(cache_key)
|
|
75
|
+
if hit is not None:
|
|
76
|
+
return hit
|
|
77
|
+
value = build_fn(generation)
|
|
78
|
+
cache[cache_key] = value
|
|
79
|
+
weakref.finalize(generation, cache.pop, cache_key, None)
|
|
80
|
+
return value
|
|
81
|
+
|
|
82
|
+
return get
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def _build_similar_edge_index(generation: Generation) -> dict[str, list[tuple[str, float, str]]]:
|
|
86
|
+
"""One pass over the overlay edge list per generation: logical-ref key
|
|
87
|
+
-> [(other_key, confidence, provenance)] for `similar_approach` edges,
|
|
88
|
+
indexed under BOTH endpoints (a self-edge is indexed once, matching the
|
|
89
|
+
previous per-call scan's `if src == key ... elif tgt == key` behavior).
|
|
90
|
+
Per-key list order is overlay edge-list order, same as the old scan."""
|
|
91
|
+
index: dict[str, list[tuple[str, float, str]]] = {}
|
|
92
|
+
for edge in generation.overlay.get("edges", []):
|
|
93
|
+
if not isinstance(edge, dict):
|
|
94
|
+
continue
|
|
95
|
+
if edge.get("type") != "similar_approach":
|
|
96
|
+
continue
|
|
97
|
+
try:
|
|
98
|
+
src_key = LogicalRef.from_dict(edge["source"]).to_key()
|
|
99
|
+
tgt_key = LogicalRef.from_dict(edge["target"]).to_key()
|
|
100
|
+
except (KeyError, TypeError):
|
|
101
|
+
continue
|
|
102
|
+
confidence = float(edge.get("confidence", 0.0))
|
|
103
|
+
provenance = edge.get("provenance", "")
|
|
104
|
+
index.setdefault(src_key, []).append((tgt_key, confidence, provenance))
|
|
105
|
+
if tgt_key == src_key:
|
|
106
|
+
continue
|
|
107
|
+
index.setdefault(tgt_key, []).append((src_key, confidence, provenance))
|
|
108
|
+
return index
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
_similar_edge_index = _per_generation_cache(_build_similar_edge_index)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def _build_label_community_index(
|
|
115
|
+
generation: Generation,
|
|
116
|
+
) -> dict[tuple[str, str], list[tuple[str, str | None, str]]]:
|
|
117
|
+
"""One pass over all nodes per generation: (community_name,
|
|
118
|
+
normalized_label) -> [(node_id, repo, logical-ref key)] for every node
|
|
119
|
+
that has a durable key. Backs `fallback_exact_match` so each fallback
|
|
120
|
+
call is a dict lookup instead of an O(all nodes) normalize-per-node
|
|
121
|
+
scan."""
|
|
122
|
+
index: dict[tuple[str, str], list[tuple[str, str | None, str]]] = {}
|
|
123
|
+
for node_id, node in generation.node_by_id.items():
|
|
124
|
+
community = node.get("community_name")
|
|
125
|
+
if not community:
|
|
126
|
+
continue
|
|
127
|
+
node_key = generation.key_by_node_id.get(node_id)
|
|
128
|
+
if not node_key:
|
|
129
|
+
continue
|
|
130
|
+
norm = normalize_label(node.get("label", ""))
|
|
131
|
+
index.setdefault((community, norm), []).append((node_id, node.get("repo"), node_key))
|
|
132
|
+
return index
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
_label_community_index = _per_generation_cache(_build_label_community_index)
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def resolve_key(query: str, generation: Generation) -> str | None:
|
|
139
|
+
"""Resolve a query through version-aware exact-alias lookup."""
|
|
140
|
+
norm = normalize_alias_query(query.strip()) if query else ""
|
|
141
|
+
keys = sorted(key for _, key in lexical_read.alias_refs(generation.lexical, norm))
|
|
142
|
+
if not keys:
|
|
143
|
+
return None
|
|
144
|
+
return keys[0]
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def overlay_similar_pairs(key: str, generation: Generation) -> list[tuple[str, float, str]]:
|
|
148
|
+
return list(_similar_edge_index(generation).get(key, []))
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def same_repo_structural_neighbors(key: str, generation: Generation) -> list[str]:
|
|
152
|
+
node_id = generation.node_id_by_key.get(key)
|
|
153
|
+
if node_id is None:
|
|
154
|
+
return []
|
|
155
|
+
node = generation.node_by_id.get(node_id, {})
|
|
156
|
+
repo = node.get("repo")
|
|
157
|
+
neighbors = []
|
|
158
|
+
for neighbor_id, edge in generation.adjacency.get(node_id, []):
|
|
159
|
+
if not ranking.edge_followable(edge, include_inferred=False):
|
|
160
|
+
continue
|
|
161
|
+
neighbor_node = generation.node_by_id.get(neighbor_id, {})
|
|
162
|
+
if neighbor_node.get("repo") != repo:
|
|
163
|
+
continue
|
|
164
|
+
if generation.degree(neighbor_id) > ranking.HUB_DEGREE_THRESHOLD:
|
|
165
|
+
continue
|
|
166
|
+
neighbor_key = generation.key_by_node_id.get(neighbor_id)
|
|
167
|
+
if neighbor_key:
|
|
168
|
+
neighbors.append(neighbor_key)
|
|
169
|
+
return neighbors
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
def fallback_exact_match(
|
|
173
|
+
key: str,
|
|
174
|
+
generation: Generation,
|
|
175
|
+
cross_repo_only: bool,
|
|
176
|
+
top_k: int,
|
|
177
|
+
allowed: Callable[[str], bool] | None = None,
|
|
178
|
+
) -> list[str]:
|
|
179
|
+
"""Documented fallback for trivial/unembedded nodes (deliverable 7):
|
|
180
|
+
exact normalized-label + same-community_name match, mirroring
|
|
181
|
+
`graphify_mesh.sync.overlay_similar`'s build-time placeholder scorer exactly
|
|
182
|
+
(same normalization, same match rule), applied over the published
|
|
183
|
+
merged graph instead of per-repo raw graphs.
|
|
184
|
+
|
|
185
|
+
`allowed` rejects a match before the `top_k` truncation, so an excluded
|
|
186
|
+
repo's match never consumes one of the returned slots."""
|
|
187
|
+
node_id = generation.node_id_by_key.get(key)
|
|
188
|
+
node = generation.node_by_id.get(node_id, {}) if node_id else {}
|
|
189
|
+
label, community, repo = node.get("label"), node.get("community_name"), node.get("repo")
|
|
190
|
+
if not label or not community:
|
|
191
|
+
return []
|
|
192
|
+
target_norm = normalize_label(label)
|
|
193
|
+
|
|
194
|
+
matches = []
|
|
195
|
+
for other_id, other_repo, other_key in _label_community_index(generation).get(
|
|
196
|
+
(community, target_norm), []
|
|
197
|
+
):
|
|
198
|
+
if other_id == node_id:
|
|
199
|
+
continue
|
|
200
|
+
if cross_repo_only and other_repo == repo:
|
|
201
|
+
continue
|
|
202
|
+
if allowed is not None and not allowed(other_key):
|
|
203
|
+
continue
|
|
204
|
+
matches.append(other_key)
|
|
205
|
+
return sorted(matches)[:top_k]
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
def repo_of_key(key: str, generation: Generation) -> str | None:
|
|
209
|
+
node_id = generation.node_id_by_key.get(key)
|
|
210
|
+
if node_id is None:
|
|
211
|
+
return None
|
|
212
|
+
return generation.node_by_id.get(node_id, {}).get("repo")
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
def find_similar(
|
|
216
|
+
query: str,
|
|
217
|
+
generation: Generation,
|
|
218
|
+
k: int,
|
|
219
|
+
cross_repo_only: bool = False,
|
|
220
|
+
*,
|
|
221
|
+
enabled_repos: frozenset[str] | None = None,
|
|
222
|
+
) -> SimilarResult:
|
|
223
|
+
"""`enabled_repos` restricts both the query seed and every candidate to
|
|
224
|
+
that repo set. `None` means no restriction, which keeps this module usable
|
|
225
|
+
without a registry. The restriction is applied during candidate selection,
|
|
226
|
+
before the `k` truncation and before the trivial-node fallback: a candidate
|
|
227
|
+
from an excluded repo must never occupy one of the `k` slots, and must
|
|
228
|
+
never suppress the fallback either.
|
|
229
|
+
|
|
230
|
+
A seed node whose own repo is excluded fails closed with `resolved=False`,
|
|
231
|
+
matching `project_map`'s refusal of an unregistered or disabled repo_id.
|
|
232
|
+
Excluding candidates never changes `resolved`, which reports only whether
|
|
233
|
+
the QUERY node was found."""
|
|
234
|
+
k = max(1, min(k, ranking.MAX_K))
|
|
235
|
+
resolved_key = resolve_key(query, generation)
|
|
236
|
+
if resolved_key is None:
|
|
237
|
+
return SimilarResult(resolved=False, degraded=["node_not_found"])
|
|
238
|
+
|
|
239
|
+
if enabled_repos is not None and repo_of_key(resolved_key, generation) not in enabled_repos:
|
|
240
|
+
return SimilarResult(resolved=False, degraded=[SEED_REPO_DISABLED_DEGRADED])
|
|
241
|
+
|
|
242
|
+
filtered_any = False
|
|
243
|
+
|
|
244
|
+
def allowed(other_key: str) -> bool:
|
|
245
|
+
nonlocal filtered_any
|
|
246
|
+
if enabled_repos is None:
|
|
247
|
+
return True
|
|
248
|
+
if repo_of_key(other_key, generation) in enabled_repos:
|
|
249
|
+
return True
|
|
250
|
+
filtered_any = True
|
|
251
|
+
return False
|
|
252
|
+
|
|
253
|
+
candidates: dict[str, tuple[float, str]] = {}
|
|
254
|
+
for other_key, score, provenance in overlay_similar_pairs(resolved_key, generation):
|
|
255
|
+
if not allowed(other_key):
|
|
256
|
+
continue
|
|
257
|
+
best = candidates.get(other_key)
|
|
258
|
+
if best is None or score > best[0]:
|
|
259
|
+
candidates[other_key] = (score, provenance)
|
|
260
|
+
|
|
261
|
+
if not cross_repo_only:
|
|
262
|
+
for other_key in same_repo_structural_neighbors(resolved_key, generation):
|
|
263
|
+
if not allowed(other_key):
|
|
264
|
+
continue
|
|
265
|
+
candidates.setdefault(other_key, (STRUCTURAL_NEIGHBOR_SCORE, STRUCTURAL_PROVENANCE))
|
|
266
|
+
|
|
267
|
+
degraded: list[str] = []
|
|
268
|
+
if not candidates:
|
|
269
|
+
fallback_keys = fallback_exact_match(resolved_key, generation, cross_repo_only, k, allowed)
|
|
270
|
+
if fallback_keys:
|
|
271
|
+
degraded.append("similarity_fallback_exact_match")
|
|
272
|
+
for other_key in fallback_keys:
|
|
273
|
+
candidates.setdefault(other_key, (FALLBACK_SCORE, FALLBACK_PROVENANCE))
|
|
274
|
+
|
|
275
|
+
if filtered_any:
|
|
276
|
+
degraded.append(DISABLED_FILTERED_DEGRADED)
|
|
277
|
+
|
|
278
|
+
ranked = sorted(candidates.items(), key=lambda kv: (-kv[1][0], kv[0]))[:k]
|
|
279
|
+
hits = []
|
|
280
|
+
for other_key, (score, _provenance) in ranked:
|
|
281
|
+
hit = _hit_from_key(other_key, generation, score, "similar")
|
|
282
|
+
if hit:
|
|
283
|
+
hits.append(hit)
|
|
284
|
+
return SimilarResult(resolved=True, hits=hits, degraded=degraded)
|
|
@@ -0,0 +1,287 @@
|
|
|
1
|
+
"""Size-capped, frame-validating line source for the SDK's stdio transport
|
|
2
|
+
(`mcp.server.stdio.stdio_server`).
|
|
3
|
+
|
|
4
|
+
Two protections the retired `protocol.py` provided do not survive the move
|
|
5
|
+
onto the SDK's own stdin reader, so this module restores them at the read
|
|
6
|
+
boundary, before a line ever reaches the SDK:
|
|
7
|
+
|
|
8
|
+
1. That reader is `async for line in stdin` over an `anyio.AsyncFile[str]`,
|
|
9
|
+
whose `readline()` has no size bound (confirmed by reading the installed
|
|
10
|
+
`mcp.server.stdio` source: `AsyncFile.__aiter__` loops on
|
|
11
|
+
`self.readline()`, which calls the wrapped file's `readline()` with no
|
|
12
|
+
arguments) — a single giant or newline-less line from a misbehaving
|
|
13
|
+
client would buffer whole in memory. `_CappedLineReader` carries forward
|
|
14
|
+
the drain logic that used to live in `protocol.py`
|
|
15
|
+
(`_drain_line`/`_read_lines`).
|
|
16
|
+
2. A line that is not parseable JSON, or that parses to something other than
|
|
17
|
+
a single JSON object (including a batch array), is fed by the SDK's own
|
|
18
|
+
`stdin_reader` into the session as a bare `Exception` — the session logs
|
|
19
|
+
it and emits a generic `notifications/message`, never the JSON-RPC
|
|
20
|
+
`-32700`/`-32600` error the spec requires and `protocol.py` used to
|
|
21
|
+
return. `_CappedLineReader.readline` intercepts both cases itself: it
|
|
22
|
+
writes the coded error response directly and does not forward the line,
|
|
23
|
+
so the SDK never sees it.
|
|
24
|
+
3. A line that IS a JSON object but not a legal JSON-RPC 2.0 envelope
|
|
25
|
+
(`{"jsonrpc":"2.0","id":7,"method":"tools/call","params":[]}`) reaches the
|
|
26
|
+
SDK's pydantic validation, which raises rather than answering request `7`,
|
|
27
|
+
so the client waits for its own timeout. `frames.envelope_error` decides
|
|
28
|
+
these, and the answer is `-32602` for a params-shape error, `-32600`
|
|
29
|
+
otherwise — the codes `protocol.py` returned for the same input.
|
|
30
|
+
|
|
31
|
+
An unparseable, non-object or unusably-identified frame is answered with
|
|
32
|
+
`id: null`: there is no id to read, or none JSON-RPC 2.0 allows. A frame that
|
|
33
|
+
carries a legal id (case 3) is answered with that id, so the client can
|
|
34
|
+
correlate the error with its pending request instead of waiting.
|
|
35
|
+
|
|
36
|
+
Both error responses and the SDK's own responses reach the wire through ONE
|
|
37
|
+
shared Python writer. `serve_stdio` builds it once with `wire_stdout()` and
|
|
38
|
+
hands the same object to `capped_stdin(out_stream=...)` and to
|
|
39
|
+
`stdio_server(stdout=...)`, so every frame from either writer goes through
|
|
40
|
+
that object's buffer and its lock serializes concurrent writes — a response
|
|
41
|
+
of any size can no longer have an error frame spliced into its middle. The
|
|
42
|
+
earlier arrangement gave each side its own `os.dup()` of fd 1, which left
|
|
43
|
+
atomicity resting on POSIX `PIPE_BUF` (4096 bytes on Linux) and so held only
|
|
44
|
+
for frames smaller than that.
|
|
45
|
+
|
|
46
|
+
`wire_stdout()` also keeps the stray-write protection the SDK gives up when
|
|
47
|
+
it is handed an explicit `stdout`: it points fd 1 at stderr for the
|
|
48
|
+
connection's duration, so a library that prints to the real stdout cannot
|
|
49
|
+
corrupt the protocol stream, and restores fd 1 on exit.
|
|
50
|
+
`test_serve_stdio_shares_one_writer_between_guard_and_sdk` and
|
|
51
|
+
`test_error_write_and_sdk_write_do_not_interleave_on_the_shared_writer` in
|
|
52
|
+
`tests/server/test_stdio_transport.py` pin both halves, so an `mcp` release
|
|
53
|
+
that drops the `stdout` argument fails a test instead of silently corrupting
|
|
54
|
+
stdout. Read them before touching this file's write path or `serve_stdio`'s
|
|
55
|
+
call to `stdio_server`.
|
|
56
|
+
|
|
57
|
+
Line size is capped on BYTES, not decoded characters. On the real stdin path
|
|
58
|
+
the reader reads `sys.stdin.buffer` and decodes after the bound, so a line of
|
|
59
|
+
4-byte UTF-8 characters can never buffer four times the cap. A caller-supplied
|
|
60
|
+
text `stream` (what the unit tests inject) has no byte view, so the cap there
|
|
61
|
+
counts characters — the bound that stream can actually enforce.
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
from __future__ import annotations
|
|
65
|
+
|
|
66
|
+
import contextlib
|
|
67
|
+
import json
|
|
68
|
+
import logging
|
|
69
|
+
import os
|
|
70
|
+
import sys
|
|
71
|
+
from collections.abc import Iterator
|
|
72
|
+
from typing import IO, Final, cast
|
|
73
|
+
|
|
74
|
+
import anyio
|
|
75
|
+
|
|
76
|
+
from graphify_mesh.server.frames import (
|
|
77
|
+
MAX_MESSAGE_BYTES,
|
|
78
|
+
NOT_A_SINGLE_OBJECT,
|
|
79
|
+
PARSE_ERROR,
|
|
80
|
+
envelope_error,
|
|
81
|
+
error_frame_text,
|
|
82
|
+
frame_id,
|
|
83
|
+
)
|
|
84
|
+
|
|
85
|
+
log = logging.getLogger("graphify_mesh.server.stdio_guard")
|
|
86
|
+
|
|
87
|
+
# Hard per-line size cap. Without it, `readline()` buffers an entire line in
|
|
88
|
+
# memory. The number lives in `frames.MAX_MESSAGE_BYTES` so the HTTP
|
|
89
|
+
# transport states the same one.
|
|
90
|
+
MAX_LINE_BYTES = MAX_MESSAGE_BYTES
|
|
91
|
+
|
|
92
|
+
# Answer for a line over the cap. Silently dropping it left the client waiting
|
|
93
|
+
# for its own timeout while the HTTP transport answered 413 for the same input.
|
|
94
|
+
# The message states the limit and never quotes the offending frame.
|
|
95
|
+
LINE_TOO_LARGE: Final = (
|
|
96
|
+
-32600,
|
|
97
|
+
f"invalid request: message exceeds the {MAX_LINE_BYTES} byte limit",
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def _error_frame(code: int, message: str, request_id: str | int | None = None) -> str:
|
|
102
|
+
return error_frame_text(code, message, request_id)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
class _CappedLineReader:
|
|
106
|
+
"""Sync file-like object exposing only a no-argument `readline()`, the
|
|
107
|
+
one method `anyio.AsyncFile.readline()` calls on the object it wraps.
|
|
108
|
+
Oversized lines are drained in bounded chunks and dropped; blank lines
|
|
109
|
+
are skipped; an unparseable or non-object/batch line gets its coded
|
|
110
|
+
JSON-RPC error written directly to `out_stream` and is not returned —
|
|
111
|
+
only complete, valid single-object candidate frames reach the SDK."""
|
|
112
|
+
|
|
113
|
+
def __init__(
|
|
114
|
+
self,
|
|
115
|
+
stream: IO[str] | None = None,
|
|
116
|
+
out_stream: IO[str] | None = None,
|
|
117
|
+
byte_stream: IO[bytes] | None = None,
|
|
118
|
+
) -> None:
|
|
119
|
+
if (stream is None) == (byte_stream is None):
|
|
120
|
+
raise ValueError("exactly one of stream= / byte_stream= must be given")
|
|
121
|
+
self._stream = stream
|
|
122
|
+
self._byte_stream = byte_stream
|
|
123
|
+
self._out_stream = out_stream if out_stream is not None else sys.stdout
|
|
124
|
+
|
|
125
|
+
def _read_bounded(self) -> tuple[str, int]:
|
|
126
|
+
"""Next line plus its size in the unit the cap is counted in. The
|
|
127
|
+
byte path reads `MAX_LINE_BYTES + 1` BYTES and decodes afterwards, so
|
|
128
|
+
a multibyte line cannot buffer more than the cap; a decode error can
|
|
129
|
+
only come from a line that is already over the bound and about to be
|
|
130
|
+
dropped, hence `errors="replace"`. The text path (an injected
|
|
131
|
+
in-memory stream) has no byte view and counts characters."""
|
|
132
|
+
if self._byte_stream is not None:
|
|
133
|
+
raw = self._byte_stream.readline(MAX_LINE_BYTES + 1)
|
|
134
|
+
return raw.decode("utf-8", errors="replace"), len(raw)
|
|
135
|
+
text_stream = cast("IO[str]", self._stream)
|
|
136
|
+
raw_text = text_stream.readline(MAX_LINE_BYTES + 1)
|
|
137
|
+
return raw_text, len(raw_text)
|
|
138
|
+
|
|
139
|
+
def _drain_rest(self) -> None:
|
|
140
|
+
"""Consume (and discard) the remainder of an oversized line in
|
|
141
|
+
bounded chunks, stopping at the next newline or EOF. Never
|
|
142
|
+
accumulates the data."""
|
|
143
|
+
while True:
|
|
144
|
+
chunk, _size = self._read_bounded()
|
|
145
|
+
if not chunk:
|
|
146
|
+
return
|
|
147
|
+
if chunk.endswith("\n"):
|
|
148
|
+
return
|
|
149
|
+
|
|
150
|
+
def _write_error(self, code: int, message: str, request_id: str | int | None = None) -> None:
|
|
151
|
+
self._out_stream.write(_error_frame(code, message, request_id) + "\n")
|
|
152
|
+
self._out_stream.flush()
|
|
153
|
+
|
|
154
|
+
def readline(self, size: int = -1) -> str: # noqa: ARG002 - AsyncFile calls with no args
|
|
155
|
+
while True:
|
|
156
|
+
raw_line, size = self._read_bounded()
|
|
157
|
+
if not raw_line:
|
|
158
|
+
return "" # EOF: clean stop, matches the old serve()'s WS6 contract
|
|
159
|
+
if size > MAX_LINE_BYTES and not raw_line.endswith("\n"):
|
|
160
|
+
self._drain_rest()
|
|
161
|
+
self._write_error(*LINE_TOO_LARGE)
|
|
162
|
+
log.warning("dropped oversized stdin line (> %d bytes)", MAX_LINE_BYTES)
|
|
163
|
+
continue
|
|
164
|
+
if not raw_line.strip():
|
|
165
|
+
continue
|
|
166
|
+
try:
|
|
167
|
+
parsed = json.loads(raw_line)
|
|
168
|
+
except json.JSONDecodeError:
|
|
169
|
+
self._write_error(*PARSE_ERROR)
|
|
170
|
+
log.warning("dropped unparseable stdin line")
|
|
171
|
+
continue
|
|
172
|
+
if not isinstance(parsed, dict):
|
|
173
|
+
self._write_error(*NOT_A_SINGLE_OBJECT)
|
|
174
|
+
log.warning("dropped non-object/batch stdin line")
|
|
175
|
+
continue
|
|
176
|
+
# A legal-JSON frame with an illegal envelope used to reach the
|
|
177
|
+
# SDK, whose pydantic validation raises instead of answering the
|
|
178
|
+
# request — the client then waits for its own timeout. Answer it
|
|
179
|
+
# here, correlated with the request id when the frame carries a
|
|
180
|
+
# usable one.
|
|
181
|
+
coded = envelope_error(parsed)
|
|
182
|
+
if coded is not None:
|
|
183
|
+
self._write_error(*coded, frame_id(parsed))
|
|
184
|
+
log.warning("dropped malformed-envelope stdin line (code %d)", coded[0])
|
|
185
|
+
continue
|
|
186
|
+
return raw_line
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def _private_wire_stdout() -> IO[str]:
|
|
190
|
+
"""A private duplicate of the real stdout descriptor, taken now — before
|
|
191
|
+
`mcp.server.stdio.stdio_server()` claims fd 1 for its own writes and
|
|
192
|
+
`dup2()`s it to stderr for the connection's duration (its stray-write
|
|
193
|
+
guard: PR #3117; confirmed by reading `_claim_fd`/`_open_stdout_diversion`
|
|
194
|
+
in the installed `mcp.server.stdio`) — so a coded error frame written to
|
|
195
|
+
`sys.stdout` at that later point would land on stderr, never the wire.
|
|
196
|
+
Only the standalone default for a `capped_stdin()` call that supplies no
|
|
197
|
+
`out_stream`: `serve_stdio` passes the shared writer from `wire_stdout()`
|
|
198
|
+
instead, which is what makes the guard's frames and the SDK's share one
|
|
199
|
+
buffer. Falls back to `sys.stdout` itself when it has no real OS
|
|
200
|
+
descriptor to duplicate (e.g. a test's in-memory stand-in).
|
|
201
|
+
"""
|
|
202
|
+
try:
|
|
203
|
+
fd = os.dup(sys.stdout.fileno())
|
|
204
|
+
except (AttributeError, OSError, ValueError):
|
|
205
|
+
return sys.stdout
|
|
206
|
+
return os.fdopen(fd, "w", buffering=1)
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
@contextlib.contextmanager
|
|
210
|
+
def wire_stdout() -> Iterator[IO[str]]:
|
|
211
|
+
"""The one writer for the whole stdio connection: both this module's
|
|
212
|
+
error frames and the SDK's responses write through it, so its buffer
|
|
213
|
+
lock — not POSIX `PIPE_BUF`, which only covers 4096 bytes — is what
|
|
214
|
+
keeps a large response and an error frame from interleaving.
|
|
215
|
+
|
|
216
|
+
Handing `stdio_server` an explicit `stdout` also skips the SDK's own
|
|
217
|
+
fd-1 claim, so this does that claim's job instead: fd 1 points at stderr
|
|
218
|
+
while the connection is up, which keeps a stray `print` from any library
|
|
219
|
+
off the protocol stream, and is restored on exit. The diversion only
|
|
220
|
+
holds for bytes that reach fd 1 while it is in place, so `sys.stdout` is
|
|
221
|
+
flushed on both edges: on entry, and again in the `finally` before fd 1
|
|
222
|
+
is restored. Skipping the exit flush would leave a buffered `print` from
|
|
223
|
+
the connection sitting in `sys.stdout` until the restored fd 1 drains it
|
|
224
|
+
onto the wire. Yields `sys.stdout` unchanged when there is no real
|
|
225
|
+
descriptor to duplicate (an in-memory stand-in in tests), since there is
|
|
226
|
+
nothing to divert there either.
|
|
227
|
+
"""
|
|
228
|
+
try:
|
|
229
|
+
wire_fd = os.dup(sys.stdout.fileno())
|
|
230
|
+
except (AttributeError, OSError, ValueError):
|
|
231
|
+
yield sys.stdout
|
|
232
|
+
return
|
|
233
|
+
|
|
234
|
+
writer = os.fdopen(wire_fd, "w", buffering=1)
|
|
235
|
+
diverted = False
|
|
236
|
+
try:
|
|
237
|
+
try:
|
|
238
|
+
sys.stdout.flush()
|
|
239
|
+
os.dup2(sys.stderr.fileno(), 1)
|
|
240
|
+
diverted = True
|
|
241
|
+
except (AttributeError, OSError, ValueError):
|
|
242
|
+
log.warning("could not divert fd 1 to stderr; stray stdout writes reach the wire")
|
|
243
|
+
yield writer
|
|
244
|
+
finally:
|
|
245
|
+
try:
|
|
246
|
+
writer.flush()
|
|
247
|
+
except (OSError, ValueError):
|
|
248
|
+
pass
|
|
249
|
+
if diverted:
|
|
250
|
+
# sys.stderr writes to fd 2 and is unaffected by the restore below.
|
|
251
|
+
with contextlib.suppress(AttributeError, OSError, ValueError):
|
|
252
|
+
sys.stdout.flush()
|
|
253
|
+
with contextlib.suppress(OSError):
|
|
254
|
+
os.dup2(wire_fd, 1)
|
|
255
|
+
writer.close()
|
|
256
|
+
|
|
257
|
+
|
|
258
|
+
def capped_stdin(
|
|
259
|
+
stream: IO[str] | None = None, out_stream: IO[str] | None = None
|
|
260
|
+
) -> anyio.AsyncFile[str]:
|
|
261
|
+
"""Async-iterable text stream for `stdio_server(stdin=...)`: yields
|
|
262
|
+
complete, valid-shaped lines; drains any line over `MAX_LINE_BYTES` and
|
|
263
|
+
answers it with `-32600` instead of buffering it; answers an unparseable
|
|
264
|
+
or non-object/batch line with the matching JSON-RPC error on `out_stream`
|
|
265
|
+
instead of forwarding it; skips blank lines; stops at EOF.
|
|
266
|
+
|
|
267
|
+
With no `stream`, lines come from `sys.stdin.buffer` and the cap is
|
|
268
|
+
counted on bytes; an injected text `stream` is read as text and its cap
|
|
269
|
+
counts characters. `out_stream` defaults to a private duplicate of the
|
|
270
|
+
real stdout descriptor (`_private_wire_stdout`); `serve_stdio` passes the
|
|
271
|
+
connection's shared writer instead."""
|
|
272
|
+
out = out_stream if out_stream is not None else _private_wire_stdout()
|
|
273
|
+
if stream is not None:
|
|
274
|
+
reader = _CappedLineReader(stream=stream, out_stream=out)
|
|
275
|
+
else:
|
|
276
|
+
binary = getattr(sys.stdin, "buffer", None)
|
|
277
|
+
if binary is not None:
|
|
278
|
+
reader = _CappedLineReader(byte_stream=binary, out_stream=out)
|
|
279
|
+
else:
|
|
280
|
+
# No byte view (a text stand-in installed as sys.stdin): the cap
|
|
281
|
+
# falls back to characters, which is all such a stream can bound.
|
|
282
|
+
reader = _CappedLineReader(stream=sys.stdin, out_stream=out)
|
|
283
|
+
# `_CappedLineReader` only implements the one method `anyio.AsyncFile`
|
|
284
|
+
# actually calls (`readline()`) — not the full `IO[str]` surface — so the
|
|
285
|
+
# cast documents an intentional, narrower structural fit rather than a
|
|
286
|
+
# type error.
|
|
287
|
+
return anyio.wrap_file(cast("IO[str]", reader))
|