pi-codemcp 0.1.0 → 1.1.0
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.
- package/README.md +144 -13
- package/extensions/index.ts +4 -1
- package/package.json +7 -7
- package/sidecar/catalog_cache.py +2 -1
- package/sidecar/cli.py +402 -0
- package/sidecar/executor.py +349 -36
- package/sidecar/gateway.py +499 -60
- package/sidecar/models.py +41 -4
- package/sidecar/pyproject.toml +33 -40
- package/sidecar/runtime_paths.py +71 -0
- package/sidecar/settings.py +7 -3
- package/sidecar/stats.py +424 -0
- package/sidecar/tool_catalog.py +311 -46
- package/sidecar/uv.lock +99 -244
- package/src/chains.ts +1 -2
- package/src/execution-rendering.ts +7 -5
- package/src/mcp-client.ts +5 -1
- package/src/modal.ts +258 -14
- package/src/output.ts +27 -24
- package/src/prompts.ts +24 -0
- package/src/settings.ts +18 -17
- package/src/tools.ts +270 -36
package/sidecar/gateway.py
CHANGED
|
@@ -2,17 +2,17 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import ast
|
|
4
4
|
import asyncio
|
|
5
|
-
import os
|
|
6
5
|
import textwrap
|
|
7
6
|
import time
|
|
8
7
|
from contextlib import AsyncExitStack, asynccontextmanager
|
|
9
|
-
from
|
|
10
|
-
from typing import TYPE_CHECKING, Literal, NamedTuple, Protocol
|
|
8
|
+
from typing import TYPE_CHECKING, Literal, NamedTuple, Protocol, cast
|
|
11
9
|
|
|
12
10
|
import pydantic_monty
|
|
13
11
|
from fastmcp import Client, FastMCP
|
|
14
12
|
from fastmcp.mcp_config import RemoteMCPServer, StdioMCPServer
|
|
15
13
|
from pydantic import BaseModel, ConfigDict
|
|
14
|
+
from pydantic_core import to_json
|
|
15
|
+
from rapidfuzz import fuzz, process
|
|
16
16
|
|
|
17
17
|
from . import json_types
|
|
18
18
|
from .catalog_cache import CatalogCache
|
|
@@ -30,26 +30,32 @@ from .chains import (
|
|
|
30
30
|
from .executor import ExecutionContext, ExecutionResponse, MontyExecutor
|
|
31
31
|
from .mcp_config import NormalizedConfig, load_mcp_json, normalize_mcp_config
|
|
32
32
|
from .models import (
|
|
33
|
+
ExecutionLimitsView,
|
|
34
|
+
InspectResponse,
|
|
33
35
|
NormalizedServerInfo,
|
|
36
|
+
SearchDetail,
|
|
37
|
+
SearchMode,
|
|
34
38
|
SearchResponse,
|
|
35
39
|
ServerToolSummary,
|
|
36
40
|
StatusResponse,
|
|
37
41
|
UpstreamStatus,
|
|
38
42
|
UpstreamToolStatus,
|
|
39
43
|
)
|
|
44
|
+
from .runtime_paths import resolve_runtime_paths
|
|
40
45
|
from .settings import CodeMcpSettings, load_settings
|
|
41
|
-
from .
|
|
46
|
+
from .stats import StatsStore
|
|
47
|
+
from .tool_catalog import ToolCatalog, schema_path_summary
|
|
42
48
|
|
|
43
49
|
if TYPE_CHECKING:
|
|
44
50
|
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
|
|
51
|
+
from pathlib import Path
|
|
45
52
|
|
|
46
53
|
from fastmcp.client.transports import ClientTransport
|
|
47
54
|
from mcp import types as mcp_types
|
|
48
55
|
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
PI_AGENT_DIR_ENV = "PI_CODING_AGENT_DIR"
|
|
56
|
+
from .runtime_paths import RuntimePaths
|
|
57
|
+
|
|
58
|
+
MAX_INSPECT_CALLS = 20
|
|
53
59
|
type ServerConfig = StdioMCPServer | RemoteMCPServer
|
|
54
60
|
type JsonObject = json_types.JsonObject
|
|
55
61
|
type JsonValue = json_types.JsonValue
|
|
@@ -140,11 +146,14 @@ class ServerHandle:
|
|
|
140
146
|
return self._client
|
|
141
147
|
exit_stack = AsyncExitStack()
|
|
142
148
|
try:
|
|
143
|
-
client =
|
|
144
|
-
Client
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
149
|
+
client = cast(
|
|
150
|
+
"Client[ClientTransport]",
|
|
151
|
+
await exit_stack.enter_async_context(
|
|
152
|
+
Client(
|
|
153
|
+
self.server_config.to_transport(),
|
|
154
|
+
name=f"pi-codemcp-{self.info.name}",
|
|
155
|
+
)
|
|
156
|
+
),
|
|
148
157
|
)
|
|
149
158
|
except BaseException:
|
|
150
159
|
await exit_stack.aclose()
|
|
@@ -238,6 +247,9 @@ class GatewayRuntime:
|
|
|
238
247
|
self.chain_store = chain_store
|
|
239
248
|
self.catalog = catalog
|
|
240
249
|
self.executor = executor
|
|
250
|
+
self.stats_store = StatsStore(settings_path.parent / "stats.json")
|
|
251
|
+
for handle in handles.values():
|
|
252
|
+
self.stats_store.record_cache(hit=handle.tools is not None)
|
|
241
253
|
self.chains = SavedChainRuntime(
|
|
242
254
|
SavedChainHandlers(
|
|
243
255
|
execute=self._execute_chain,
|
|
@@ -305,48 +317,336 @@ class GatewayRuntime:
|
|
|
305
317
|
*(handle.close() for handle in self.handles.values()),
|
|
306
318
|
return_exceptions=True,
|
|
307
319
|
)
|
|
320
|
+
await self.stats_store.close()
|
|
308
321
|
|
|
309
322
|
async def search(
|
|
310
323
|
self,
|
|
311
|
-
query: str,
|
|
324
|
+
query: str | None = None,
|
|
312
325
|
limit: int = 5,
|
|
313
326
|
server: str | None = None,
|
|
327
|
+
detail: SearchDetail = "signatures",
|
|
328
|
+
mode: SearchMode = "search",
|
|
329
|
+
cursor: int = 0,
|
|
314
330
|
) -> SearchResponse:
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
331
|
+
started = time.perf_counter()
|
|
332
|
+
input_bytes = len(
|
|
333
|
+
to_json({
|
|
334
|
+
"query": query,
|
|
335
|
+
"limit": limit,
|
|
336
|
+
"server": server,
|
|
337
|
+
"detail": detail,
|
|
338
|
+
"mode": mode,
|
|
339
|
+
"cursor": cursor,
|
|
340
|
+
})
|
|
341
|
+
)
|
|
342
|
+
try:
|
|
343
|
+
response = await self._search_impl(query, limit, server, detail, mode, cursor)
|
|
344
|
+
except BaseException:
|
|
345
|
+
self.stats_store.record_operation(
|
|
346
|
+
"search",
|
|
347
|
+
duration_ms=_elapsed_ms(started),
|
|
348
|
+
success=False,
|
|
349
|
+
failure_stage="error",
|
|
350
|
+
input_bytes=input_bytes,
|
|
351
|
+
)
|
|
352
|
+
raise
|
|
353
|
+
self.stats_store.record_operation(
|
|
354
|
+
"search",
|
|
355
|
+
duration_ms=_elapsed_ms(started),
|
|
356
|
+
success=True,
|
|
357
|
+
input_bytes=input_bytes,
|
|
358
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
359
|
+
)
|
|
360
|
+
return response
|
|
361
|
+
|
|
362
|
+
async def _search_impl(
|
|
363
|
+
self,
|
|
364
|
+
query: str | None,
|
|
365
|
+
limit: int,
|
|
366
|
+
server: str | None,
|
|
367
|
+
detail: SearchDetail,
|
|
368
|
+
mode: SearchMode,
|
|
369
|
+
cursor: int,
|
|
370
|
+
) -> SearchResponse:
|
|
371
|
+
self._validate_search_server(server)
|
|
372
|
+
discovery_started = time.perf_counter()
|
|
373
|
+
discovery_servers: Iterable[str]
|
|
374
|
+
if server is None:
|
|
375
|
+
discovery_servers = self.handles.keys()
|
|
376
|
+
elif server == "chains":
|
|
377
|
+
discovery_servers = ()
|
|
378
|
+
else:
|
|
379
|
+
discovery_servers = (server,)
|
|
380
|
+
await self._ensure_servers_discovered(discovery_servers)
|
|
381
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
319
382
|
bounded_limit = min(max(limit, 1), 20)
|
|
383
|
+
bounded_cursor = max(cursor, 0)
|
|
320
384
|
counts = self.catalog.counts_by_server()
|
|
321
385
|
servers = [
|
|
322
|
-
ServerToolSummary(name=
|
|
323
|
-
for
|
|
324
|
-
if counts.get(
|
|
386
|
+
ServerToolSummary(name=server_info.name, tool_count=counts[server_info.name])
|
|
387
|
+
for server_info in self.normalized.servers
|
|
388
|
+
if counts.get(server_info.name, 0) > 0
|
|
325
389
|
]
|
|
326
390
|
if counts.get("chains", 0) > 0:
|
|
327
391
|
servers.append(ServerToolSummary(name="chains", tool_count=counts["chains"]))
|
|
392
|
+
filtered_count = counts.get(server, 0) if server is not None else len(self.catalog.tools)
|
|
393
|
+
if mode == "inventory":
|
|
394
|
+
results = self.catalog.inventory(
|
|
395
|
+
server=server,
|
|
396
|
+
detail=detail,
|
|
397
|
+
offset=bounded_cursor,
|
|
398
|
+
limit=bounded_limit,
|
|
399
|
+
)
|
|
400
|
+
total_matches = filtered_count
|
|
401
|
+
else:
|
|
402
|
+
clean_query = (query or "").strip()
|
|
403
|
+
if not clean_query:
|
|
404
|
+
raise ValueError("query is required in search mode")
|
|
405
|
+
all_matches = self.catalog.search(
|
|
406
|
+
clean_query,
|
|
407
|
+
filtered_count,
|
|
408
|
+
server=server,
|
|
409
|
+
detail=detail,
|
|
410
|
+
)
|
|
411
|
+
total_matches = len(all_matches)
|
|
412
|
+
results = all_matches[bounded_cursor : bounded_cursor + bounded_limit]
|
|
413
|
+
next_cursor = (
|
|
414
|
+
bounded_cursor + len(results) if bounded_cursor + len(results) < total_matches else None
|
|
415
|
+
)
|
|
416
|
+
include_prelude = detail == "full"
|
|
417
|
+
if mode == "search" and detail == "signatures" and results:
|
|
418
|
+
inspected = {
|
|
419
|
+
item.call: item.stub
|
|
420
|
+
for item in self.catalog.inspect([
|
|
421
|
+
result.call for result in results[: min(3, len(results))]
|
|
422
|
+
])
|
|
423
|
+
}
|
|
424
|
+
results = [
|
|
425
|
+
result.model_copy(update={"stub": inspected[result.call]})
|
|
426
|
+
if result.call in inspected
|
|
427
|
+
else result
|
|
428
|
+
for result in results
|
|
429
|
+
]
|
|
430
|
+
include_prelude = True
|
|
328
431
|
return SearchResponse(
|
|
432
|
+
mode=mode,
|
|
433
|
+
detail=detail,
|
|
329
434
|
total_tool_count=len(self.catalog.tools),
|
|
435
|
+
filtered_tool_count=filtered_count,
|
|
330
436
|
servers=servers,
|
|
331
|
-
|
|
437
|
+
cursor=bounded_cursor,
|
|
438
|
+
next_cursor=next_cursor,
|
|
439
|
+
has_more=next_cursor is not None,
|
|
440
|
+
project_scope_available=self.chain_store.project_store is not None,
|
|
441
|
+
execution_limits=self._execution_limits_view(),
|
|
442
|
+
prelude=self.catalog.stub_prelude if include_prelude else None,
|
|
443
|
+
results=results,
|
|
444
|
+
)
|
|
445
|
+
|
|
446
|
+
async def inspect(self, calls: list[str]) -> InspectResponse:
|
|
447
|
+
started = time.perf_counter()
|
|
448
|
+
try:
|
|
449
|
+
response = await self._inspect_impl(calls)
|
|
450
|
+
except BaseException:
|
|
451
|
+
self.stats_store.record_operation(
|
|
452
|
+
"inspect",
|
|
453
|
+
duration_ms=_elapsed_ms(started),
|
|
454
|
+
success=False,
|
|
455
|
+
failure_stage="error",
|
|
456
|
+
input_bytes=len(to_json(calls)),
|
|
457
|
+
)
|
|
458
|
+
raise
|
|
459
|
+
self.stats_store.record_operation(
|
|
460
|
+
"inspect",
|
|
461
|
+
duration_ms=_elapsed_ms(started),
|
|
462
|
+
success=True,
|
|
463
|
+
input_bytes=len(to_json(calls)),
|
|
464
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
465
|
+
)
|
|
466
|
+
return response
|
|
467
|
+
|
|
468
|
+
async def _inspect_impl(self, calls: list[str]) -> InspectResponse:
|
|
469
|
+
if not calls:
|
|
470
|
+
raise ValueError("calls must contain at least one MCP call")
|
|
471
|
+
if len(calls) > MAX_INSPECT_CALLS:
|
|
472
|
+
raise ValueError(f"calls must contain at most {MAX_INSPECT_CALLS} MCP calls")
|
|
473
|
+
discovery_started = time.perf_counter()
|
|
474
|
+
requested_namespaces = {call.partition(".")[0] for call in calls}
|
|
475
|
+
discovery_servers = [
|
|
476
|
+
server
|
|
477
|
+
for server, namespace in self.catalog.server_aliases.items()
|
|
478
|
+
if namespace in requested_namespaces
|
|
479
|
+
]
|
|
480
|
+
await self._ensure_servers_discovered(discovery_servers)
|
|
481
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
482
|
+
return InspectResponse(
|
|
483
|
+
prelude=self.catalog.stub_prelude,
|
|
484
|
+
project_scope_available=self.chain_store.project_store is not None,
|
|
485
|
+
execution_limits=self._execution_limits_view(),
|
|
486
|
+
results=self.catalog.inspect(calls),
|
|
487
|
+
)
|
|
488
|
+
|
|
489
|
+
def _validate_search_server(self, server: str | None) -> None:
|
|
490
|
+
if server is None:
|
|
491
|
+
return
|
|
492
|
+
valid = set(self.catalog.servers)
|
|
493
|
+
if server in valid:
|
|
494
|
+
return
|
|
495
|
+
suggestions = [
|
|
496
|
+
name
|
|
497
|
+
for name, _score, _index in process.extract(
|
|
498
|
+
server,
|
|
499
|
+
sorted(valid),
|
|
500
|
+
scorer=fuzz.ratio,
|
|
501
|
+
limit=3,
|
|
502
|
+
)
|
|
503
|
+
]
|
|
504
|
+
raise ValueError(
|
|
505
|
+
f"Unknown MCP server {server!r}; available: {sorted(valid)}; suggestions: {suggestions}"
|
|
506
|
+
)
|
|
507
|
+
|
|
508
|
+
def _execution_limits_view(self) -> ExecutionLimitsView:
|
|
509
|
+
settings = self.executor.settings
|
|
510
|
+
return ExecutionLimitsView(
|
|
511
|
+
timeout_seconds=settings.timeout_seconds,
|
|
512
|
+
tool_timeout_seconds=settings.tool_timeout_seconds,
|
|
513
|
+
max_calls=settings.max_calls,
|
|
514
|
+
result_limit_bytes=settings.result_byte_limit,
|
|
332
515
|
)
|
|
333
516
|
|
|
334
517
|
async def execute(self, code: str) -> ExecutionResponse:
|
|
335
|
-
|
|
518
|
+
started = time.perf_counter()
|
|
519
|
+
input_bytes = len(code.encode())
|
|
520
|
+
discovery_started = time.perf_counter()
|
|
521
|
+
try:
|
|
522
|
+
await self._ensure_servers_discovered(self._required_servers_for_code(code))
|
|
523
|
+
except BaseException:
|
|
524
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
525
|
+
self._record_operation_exception(
|
|
526
|
+
"execute",
|
|
527
|
+
started,
|
|
528
|
+
input_bytes,
|
|
529
|
+
failure_stage="discovery",
|
|
530
|
+
)
|
|
531
|
+
raise
|
|
532
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
336
533
|
self.executor.update_catalog(self.catalog)
|
|
337
|
-
|
|
534
|
+
try:
|
|
535
|
+
response = await self.executor.execute_graph(code, self._dispatch)
|
|
536
|
+
except BaseException as error:
|
|
537
|
+
self._record_operation_exception(
|
|
538
|
+
"execute",
|
|
539
|
+
started,
|
|
540
|
+
input_bytes,
|
|
541
|
+
failure_stage=(
|
|
542
|
+
"cancelled" if isinstance(error, asyncio.CancelledError) else "error"
|
|
543
|
+
),
|
|
544
|
+
)
|
|
545
|
+
raise
|
|
546
|
+
self._record_execution("execute", started, input_bytes, response)
|
|
547
|
+
return response
|
|
338
548
|
|
|
339
549
|
async def _execute_chain(self, name: str, arguments: JsonObject) -> ExecutionResponse:
|
|
340
|
-
|
|
550
|
+
started = time.perf_counter()
|
|
551
|
+
input_bytes = len(to_json(arguments))
|
|
552
|
+
try:
|
|
553
|
+
chain = self.chain_store.get(name).chain
|
|
554
|
+
except BaseException:
|
|
555
|
+
self._record_operation_exception(
|
|
556
|
+
"execute_chain",
|
|
557
|
+
started,
|
|
558
|
+
input_bytes,
|
|
559
|
+
failure_stage="preflight",
|
|
560
|
+
)
|
|
561
|
+
raise
|
|
341
562
|
if not chain.enabled:
|
|
342
|
-
|
|
563
|
+
response = ExecutionResponse(
|
|
343
564
|
ok=False,
|
|
344
565
|
failure_stage="preflight",
|
|
345
566
|
error=f"Saved chain is disabled: {name}",
|
|
346
567
|
)
|
|
347
|
-
|
|
568
|
+
self._record_execution(
|
|
569
|
+
"execute_chain",
|
|
570
|
+
started,
|
|
571
|
+
input_bytes,
|
|
572
|
+
response,
|
|
573
|
+
)
|
|
574
|
+
return response
|
|
575
|
+
discovery_started = time.perf_counter()
|
|
576
|
+
try:
|
|
577
|
+
await self._ensure_servers_discovered(self._required_servers_for_chain(chain))
|
|
578
|
+
except BaseException:
|
|
579
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
580
|
+
self._record_operation_exception(
|
|
581
|
+
"execute_chain",
|
|
582
|
+
started,
|
|
583
|
+
input_bytes,
|
|
584
|
+
failure_stage="discovery",
|
|
585
|
+
)
|
|
586
|
+
raise
|
|
587
|
+
self.stats_store.record_phase("discovery", _elapsed_ms(discovery_started))
|
|
348
588
|
self.executor.update_catalog(self.catalog)
|
|
349
|
-
|
|
589
|
+
try:
|
|
590
|
+
response = await self.executor.execute_saved_chain(chain, arguments, self._dispatch)
|
|
591
|
+
except BaseException as error:
|
|
592
|
+
self._record_operation_exception(
|
|
593
|
+
"execute_chain",
|
|
594
|
+
started,
|
|
595
|
+
input_bytes,
|
|
596
|
+
failure_stage=(
|
|
597
|
+
"cancelled" if isinstance(error, asyncio.CancelledError) else "error"
|
|
598
|
+
),
|
|
599
|
+
)
|
|
600
|
+
raise
|
|
601
|
+
self._record_execution(
|
|
602
|
+
"execute_chain",
|
|
603
|
+
started,
|
|
604
|
+
input_bytes,
|
|
605
|
+
response,
|
|
606
|
+
)
|
|
607
|
+
return response
|
|
608
|
+
|
|
609
|
+
def _record_operation_exception(
|
|
610
|
+
self,
|
|
611
|
+
operation: str,
|
|
612
|
+
started: float,
|
|
613
|
+
input_bytes: int,
|
|
614
|
+
*,
|
|
615
|
+
failure_stage: str,
|
|
616
|
+
) -> None:
|
|
617
|
+
self.stats_store.record_operation(
|
|
618
|
+
operation,
|
|
619
|
+
duration_ms=_elapsed_ms(started),
|
|
620
|
+
success=False,
|
|
621
|
+
failure_stage=failure_stage,
|
|
622
|
+
input_bytes=input_bytes,
|
|
623
|
+
)
|
|
624
|
+
|
|
625
|
+
def _record_execution(
|
|
626
|
+
self,
|
|
627
|
+
operation: str,
|
|
628
|
+
started: float,
|
|
629
|
+
input_bytes: int,
|
|
630
|
+
response: ExecutionResponse,
|
|
631
|
+
) -> None:
|
|
632
|
+
self.stats_store.record_operation(
|
|
633
|
+
operation,
|
|
634
|
+
duration_ms=_elapsed_ms(started),
|
|
635
|
+
success=response.ok,
|
|
636
|
+
failure_stage=response.failure_stage,
|
|
637
|
+
input_bytes=input_bytes,
|
|
638
|
+
output_bytes=response.metrics.result_bytes,
|
|
639
|
+
calls=response.calls_made,
|
|
640
|
+
chain_calls=response.chain_calls,
|
|
641
|
+
)
|
|
642
|
+
metrics = response.metrics
|
|
643
|
+
for phase, duration in (
|
|
644
|
+
("typecheck", metrics.typecheck_ms),
|
|
645
|
+
("execution", metrics.runtime_ms),
|
|
646
|
+
("serialization", metrics.serialization_ms),
|
|
647
|
+
):
|
|
648
|
+
if duration > 0:
|
|
649
|
+
self.stats_store.record_phase(phase, duration)
|
|
350
650
|
|
|
351
651
|
async def _save_chain(
|
|
352
652
|
self,
|
|
@@ -357,6 +657,54 @@ class GatewayRuntime:
|
|
|
357
657
|
code: str,
|
|
358
658
|
input_schema: JsonObject,
|
|
359
659
|
output_schema: JsonObject,
|
|
660
|
+
) -> SaveChainResponse:
|
|
661
|
+
started = time.perf_counter()
|
|
662
|
+
input_bytes = len(
|
|
663
|
+
to_json({
|
|
664
|
+
"scope": scope,
|
|
665
|
+
"name": name,
|
|
666
|
+
"description": description,
|
|
667
|
+
"code": code,
|
|
668
|
+
"input_schema": input_schema,
|
|
669
|
+
"output_schema": output_schema,
|
|
670
|
+
})
|
|
671
|
+
)
|
|
672
|
+
try:
|
|
673
|
+
response = await self._save_chain_impl(
|
|
674
|
+
scope=scope,
|
|
675
|
+
name=name,
|
|
676
|
+
description=description,
|
|
677
|
+
code=code,
|
|
678
|
+
input_schema=input_schema,
|
|
679
|
+
output_schema=output_schema,
|
|
680
|
+
)
|
|
681
|
+
except BaseException:
|
|
682
|
+
self.stats_store.record_operation(
|
|
683
|
+
"save_chain",
|
|
684
|
+
duration_ms=_elapsed_ms(started),
|
|
685
|
+
success=False,
|
|
686
|
+
failure_stage="validation",
|
|
687
|
+
input_bytes=input_bytes,
|
|
688
|
+
)
|
|
689
|
+
raise
|
|
690
|
+
self.stats_store.record_operation(
|
|
691
|
+
"save_chain",
|
|
692
|
+
duration_ms=_elapsed_ms(started),
|
|
693
|
+
success=True,
|
|
694
|
+
input_bytes=input_bytes,
|
|
695
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
696
|
+
)
|
|
697
|
+
return response
|
|
698
|
+
|
|
699
|
+
async def _save_chain_impl(
|
|
700
|
+
self,
|
|
701
|
+
*,
|
|
702
|
+
scope: ChainScope,
|
|
703
|
+
name: str,
|
|
704
|
+
description: str,
|
|
705
|
+
code: str,
|
|
706
|
+
input_schema: JsonObject,
|
|
707
|
+
output_schema: JsonObject,
|
|
360
708
|
) -> SaveChainResponse:
|
|
361
709
|
previous = (
|
|
362
710
|
self.chain_store.get(name, scope).chain
|
|
@@ -384,7 +732,7 @@ class GatewayRuntime:
|
|
|
384
732
|
message = error.display("concise", color=False).strip()
|
|
385
733
|
else:
|
|
386
734
|
message = error.display("type-msg").strip()
|
|
387
|
-
raise ValueError(
|
|
735
|
+
raise ValueError(_saved_chain_preflight_error(message, output_schema)) from error
|
|
388
736
|
dependencies = self._chain_dependencies(code, candidate_catalog)
|
|
389
737
|
saved = ChainStore.build(
|
|
390
738
|
name=name,
|
|
@@ -403,9 +751,37 @@ class GatewayRuntime:
|
|
|
403
751
|
)
|
|
404
752
|
|
|
405
753
|
def _list_chains(self) -> ChainListResponse:
|
|
406
|
-
|
|
754
|
+
started = time.perf_counter()
|
|
755
|
+
response = ChainListResponse(chains=self._chain_views())
|
|
756
|
+
self.stats_store.record_operation(
|
|
757
|
+
"list_chains",
|
|
758
|
+
duration_ms=_elapsed_ms(started),
|
|
759
|
+
success=True,
|
|
760
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
761
|
+
)
|
|
762
|
+
return response
|
|
407
763
|
|
|
408
764
|
async def _revalidate_chain(self, name: str, scope: ChainScope) -> ChainStatusView:
|
|
765
|
+
started = time.perf_counter()
|
|
766
|
+
try:
|
|
767
|
+
response = await self._revalidate_chain_impl(name, scope)
|
|
768
|
+
except BaseException:
|
|
769
|
+
self.stats_store.record_operation(
|
|
770
|
+
"revalidate_chain",
|
|
771
|
+
duration_ms=_elapsed_ms(started),
|
|
772
|
+
success=False,
|
|
773
|
+
failure_stage="validation",
|
|
774
|
+
)
|
|
775
|
+
raise
|
|
776
|
+
self.stats_store.record_operation(
|
|
777
|
+
"revalidate_chain",
|
|
778
|
+
duration_ms=_elapsed_ms(started),
|
|
779
|
+
success=True,
|
|
780
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
781
|
+
)
|
|
782
|
+
return response
|
|
783
|
+
|
|
784
|
+
async def _revalidate_chain_impl(self, name: str, scope: ChainScope) -> ChainStatusView:
|
|
409
785
|
current = self.chain_store.get(name, scope).chain
|
|
410
786
|
await self._ensure_servers_discovered(self._referenced_servers(current.code))
|
|
411
787
|
chains = [chain for chain in self.chain_store.enabled() if chain.name != name]
|
|
@@ -419,7 +795,9 @@ class GatewayRuntime:
|
|
|
419
795
|
message = error.display("concise", color=False).strip()
|
|
420
796
|
else:
|
|
421
797
|
message = error.display("type-msg").strip()
|
|
422
|
-
raise ValueError(
|
|
798
|
+
raise ValueError(
|
|
799
|
+
_saved_chain_preflight_error(message, current.output_schema)
|
|
800
|
+
) from error
|
|
423
801
|
updated = current.model_copy(
|
|
424
802
|
update={
|
|
425
803
|
"dependencies": self._chain_dependencies(current.code, candidate_catalog),
|
|
@@ -431,6 +809,26 @@ class GatewayRuntime:
|
|
|
431
809
|
return self._chain_view(name, scope)
|
|
432
810
|
|
|
433
811
|
async def _delete_chain(self, name: str, scope: ChainScope) -> ChainListResponse:
|
|
812
|
+
started = time.perf_counter()
|
|
813
|
+
try:
|
|
814
|
+
response = await self._delete_chain_impl(name, scope)
|
|
815
|
+
except BaseException:
|
|
816
|
+
self.stats_store.record_operation(
|
|
817
|
+
"delete_chain",
|
|
818
|
+
duration_ms=_elapsed_ms(started),
|
|
819
|
+
success=False,
|
|
820
|
+
failure_stage="error",
|
|
821
|
+
)
|
|
822
|
+
raise
|
|
823
|
+
self.stats_store.record_operation(
|
|
824
|
+
"delete_chain",
|
|
825
|
+
duration_ms=_elapsed_ms(started),
|
|
826
|
+
success=True,
|
|
827
|
+
output_bytes=len(to_json(response.model_dump(mode="json"))),
|
|
828
|
+
)
|
|
829
|
+
return response
|
|
830
|
+
|
|
831
|
+
async def _delete_chain_impl(self, name: str, scope: ChainScope) -> ChainListResponse:
|
|
434
832
|
effective = self.chain_store.get(name)
|
|
435
833
|
called_by = self._called_by().get(name, []) if effective.scope == scope else []
|
|
436
834
|
if called_by:
|
|
@@ -453,15 +851,37 @@ class GatewayRuntime:
|
|
|
453
851
|
return await self.executor.execute_nested_chain(chain, arguments, context)
|
|
454
852
|
|
|
455
853
|
handle = self.handles[spec.server]
|
|
456
|
-
|
|
854
|
+
started = time.perf_counter()
|
|
855
|
+
input_bytes = len(to_json(arguments))
|
|
856
|
+
try:
|
|
857
|
+
result = await handle.call_tool(
|
|
858
|
+
spec.backend_name,
|
|
859
|
+
arguments,
|
|
860
|
+
timeout_seconds=min(
|
|
861
|
+
self.executor.settings.tool_timeout_seconds,
|
|
862
|
+
max(0.001, context.remaining_seconds()),
|
|
863
|
+
),
|
|
864
|
+
)
|
|
865
|
+
normalized = context.catalog.normalize_result(public_name, result)
|
|
866
|
+
except BaseException:
|
|
867
|
+
self.stats_store.record_upstream(
|
|
868
|
+
spec.server,
|
|
869
|
+
spec.backend_name,
|
|
870
|
+
duration_ms=_elapsed_ms(started),
|
|
871
|
+
success=False,
|
|
872
|
+
input_bytes=input_bytes,
|
|
873
|
+
output_bytes=0,
|
|
874
|
+
)
|
|
875
|
+
raise
|
|
876
|
+
self.stats_store.record_upstream(
|
|
877
|
+
spec.server,
|
|
457
878
|
spec.backend_name,
|
|
458
|
-
|
|
459
|
-
|
|
460
|
-
|
|
461
|
-
|
|
462
|
-
),
|
|
879
|
+
duration_ms=_elapsed_ms(started),
|
|
880
|
+
success=True,
|
|
881
|
+
input_bytes=input_bytes,
|
|
882
|
+
output_bytes=len(to_json(normalized)),
|
|
463
883
|
)
|
|
464
|
-
return
|
|
884
|
+
return normalized
|
|
465
885
|
|
|
466
886
|
async def discover(self, server: str) -> StatusResponse:
|
|
467
887
|
handle = self.handles.get(server)
|
|
@@ -715,6 +1135,19 @@ class GatewayRuntime:
|
|
|
715
1135
|
}
|
|
716
1136
|
|
|
717
1137
|
|
|
1138
|
+
def _saved_chain_preflight_error(message: str, output_schema: JsonObject) -> str:
|
|
1139
|
+
expected = "\n".join(f" - {path}" for path in schema_path_summary(output_schema))
|
|
1140
|
+
return (
|
|
1141
|
+
"Saved chain failed preflight against outputSchema.\n"
|
|
1142
|
+
f"Expected output paths:\n{expected}\n"
|
|
1143
|
+
f"Actual type-check result:\n{message}"
|
|
1144
|
+
)
|
|
1145
|
+
|
|
1146
|
+
|
|
1147
|
+
def _elapsed_ms(started: float) -> float:
|
|
1148
|
+
return (time.perf_counter() - started) * 1_000
|
|
1149
|
+
|
|
1150
|
+
|
|
718
1151
|
def _parse_code(code: str) -> ast.AST | None:
|
|
719
1152
|
normalized = textwrap.dedent(code).strip("\n")
|
|
720
1153
|
wrapped = f"async def __codemcp_main():\n{textwrap.indent(normalized, ' ')}\n"
|
|
@@ -734,6 +1167,7 @@ def _compact_description(description: str | None, limit: int = 160) -> str | Non
|
|
|
734
1167
|
class RuntimeState:
|
|
735
1168
|
def __init__(self) -> None:
|
|
736
1169
|
self.runtime: GatewayRuntime | None = None
|
|
1170
|
+
self.paths: RuntimePaths | None = None
|
|
737
1171
|
|
|
738
1172
|
|
|
739
1173
|
_runtime_state = RuntimeState()
|
|
@@ -745,22 +1179,12 @@ def _require_runtime() -> GatewayRuntime:
|
|
|
745
1179
|
return _runtime_state.runtime
|
|
746
1180
|
|
|
747
1181
|
|
|
1182
|
+
def configure_runtime_paths(paths: RuntimePaths | None) -> None:
|
|
1183
|
+
_runtime_state.paths = paths
|
|
1184
|
+
|
|
1185
|
+
|
|
748
1186
|
def _runtime_paths() -> tuple[Path, Path, Path, Path, Path, Path | None]:
|
|
749
|
-
|
|
750
|
-
agent_dir = Path(raw_agent_dir).expanduser() if raw_agent_dir else DEFAULT_AGENT_DIR
|
|
751
|
-
state_dir = agent_dir / "pi-codemcp"
|
|
752
|
-
raw_project_chains_dir = os.environ.get(CODEMCP_PROJECT_CHAINS_DIR_ENV)
|
|
753
|
-
project_chains_dir = (
|
|
754
|
-
Path(raw_project_chains_dir).expanduser() if raw_project_chains_dir else None
|
|
755
|
-
)
|
|
756
|
-
return (
|
|
757
|
-
agent_dir / "mcp.json",
|
|
758
|
-
state_dir / "oauth",
|
|
759
|
-
state_dir / "catalog",
|
|
760
|
-
state_dir / "settings.json",
|
|
761
|
-
state_dir / "chains",
|
|
762
|
-
project_chains_dir,
|
|
763
|
-
)
|
|
1187
|
+
return (_runtime_state.paths or resolve_runtime_paths()).as_tuple()
|
|
764
1188
|
|
|
765
1189
|
|
|
766
1190
|
@asynccontextmanager
|
|
@@ -773,7 +1197,7 @@ async def lifespan(_: FastMCP[None]) -> AsyncIterator[None]:
|
|
|
773
1197
|
global_chains_dir,
|
|
774
1198
|
project_chains_dir,
|
|
775
1199
|
) = _runtime_paths()
|
|
776
|
-
|
|
1200
|
+
runtime = GatewayRuntime.create(
|
|
777
1201
|
config_path,
|
|
778
1202
|
oauth_dir,
|
|
779
1203
|
catalog_dir,
|
|
@@ -781,12 +1205,12 @@ async def lifespan(_: FastMCP[None]) -> AsyncIterator[None]:
|
|
|
781
1205
|
global_chain_dir=global_chains_dir,
|
|
782
1206
|
project_chain_dir=project_chains_dir,
|
|
783
1207
|
)
|
|
1208
|
+
_runtime_state.runtime = runtime
|
|
784
1209
|
try:
|
|
785
1210
|
yield
|
|
786
1211
|
finally:
|
|
787
|
-
|
|
788
|
-
|
|
789
|
-
await runtime.close()
|
|
1212
|
+
_runtime_state.runtime = None
|
|
1213
|
+
await runtime.close()
|
|
790
1214
|
|
|
791
1215
|
|
|
792
1216
|
mcp = FastMCP(
|
|
@@ -800,12 +1224,21 @@ mcp = FastMCP(
|
|
|
800
1224
|
|
|
801
1225
|
@mcp.tool
|
|
802
1226
|
async def search(
|
|
803
|
-
query: str,
|
|
1227
|
+
query: str | None = None,
|
|
804
1228
|
limit: int = 5,
|
|
805
1229
|
server: str | None = None,
|
|
1230
|
+
detail: SearchDetail = "signatures",
|
|
1231
|
+
mode: SearchMode = "search",
|
|
1232
|
+
cursor: int = 0,
|
|
806
1233
|
) -> SearchResponse:
|
|
807
|
-
"""Search configured upstream MCP tools and saved chains
|
|
808
|
-
return await _require_runtime().search(query, limit, server)
|
|
1234
|
+
"""Search or page through configured upstream MCP tools and saved chains."""
|
|
1235
|
+
return await _require_runtime().search(query, limit, server, detail, mode, cursor)
|
|
1236
|
+
|
|
1237
|
+
|
|
1238
|
+
@mcp.tool
|
|
1239
|
+
async def inspect(calls: list[str]) -> InspectResponse:
|
|
1240
|
+
"""Return exact typed SDK stubs for selected call identifiers."""
|
|
1241
|
+
return await _require_runtime().inspect(calls)
|
|
809
1242
|
|
|
810
1243
|
|
|
811
1244
|
@mcp.tool
|
|
@@ -879,6 +1312,12 @@ async def delete_chain(name: str, scope: ChainScope) -> ChainListResponse:
|
|
|
879
1312
|
return await _require_runtime().chains.delete(name, scope)
|
|
880
1313
|
|
|
881
1314
|
|
|
1315
|
+
@mcp.tool
|
|
1316
|
+
def stats() -> JsonObject:
|
|
1317
|
+
"""Return bounded local CodeMCP telemetry rollups."""
|
|
1318
|
+
return _require_runtime().stats_store.snapshot()
|
|
1319
|
+
|
|
1320
|
+
|
|
882
1321
|
@mcp.tool
|
|
883
1322
|
def status() -> StatusResponse:
|
|
884
1323
|
"""Report cached catalog and upstream connection state without connecting upstreams."""
|