Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 97 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -457,3 +457,100 @@ async def test_add():
if __name__ == "__main__":
asyncio.run(test_add())
```

### Large tool results: data spillover

Attach `DataSpilloverInterceptor` to a tool to store large results separately and
return a preview, `resource_uri`, and `original_size_bytes`. Results at or below
the threshold pass through unchanged. Sizes count serialized UTF-8 bytes, not
characters; lists and dictionaries use compact JSON.

```python
from pydantic import BaseModel
from nitrostack import DataSpilloverInterceptor, MemorySpilloverStore, tool, use_interceptors

spillover = DataSpilloverInterceptor(
threshold_bytes=64 * 1024,
store=MemorySpilloverStore(
max_bytes=10 * 1024 * 1024,
max_total_bytes=100 * 1024 * 1024,
ttl_seconds=300,
),
resource_uri_prefix="resource://data-spillover/",
)

class ReportInput(BaseModel):
count: int = 10000

class ReportTools:
@tool(name="report", description="Generate a report", input_schema=ReportInput)
@use_interceptors(spillover)
async def report(self, input: ReportInput, context):
return [{"row": index} for index in range(input.count)]
```

Register `ReportTools` as a module provider or controller as usual. A configured
interceptor instance works directly in the decorator. For default configuration,
`@use_interceptors(DataSpilloverInterceptor)` also works through dependency injection.

The MCP result includes a text preview, structured content, and a resource link.
Clients retrieve the full result with `resources/read`, for example:

```json
{
"jsonrpc": "2.0",
"id": 2,
"method": "resources/read",
"params": {"uri": "resource://data-spillover/<id-from-tool-result>"}
}
```

Lists and dictionaries preview their first five items or entries and total counts.
Strings preview up to 512 characters. The complete preview stays within 2 KiB;
large nested values are represented as truncated text. Reading the full resource
can consume context again, so clients should retrieve it only when needed.

For disk storage, replace the store with:

```python
from nitrostack import FileSpilloverStore

store = FileSpilloverStore(
".nitrostack/spillover",
max_bytes=10 * 1024 * 1024,
max_total_bytes=100 * 1024 * 1024,
ttl_seconds=300,
)
```

The memory store evicts least recently used results when its payload capacity is
full. The file store rejects writes when its capacity is full; its capacity counts
actual JSON file bytes, including base64 and metadata overhead. A single result
larger than `max_bytes` raises `PayloadTooLargeError`. Both stores remove expired
results during operations; file storage also cleans up on initialization. Expiry
is absolute and reads do not extend it. Each file-store root must be owned by one
store instance in one process; shared storage across processes is not supported.
Keep the configured directory private to the server account.

Resource access requires the same verified principal, issuer, client identity,
and tenant. Anonymous access is limited to the same server-side MCP session.
Requests without either a verified identity or trusted session cannot spill
results; invalid credentials do not fall back to anonymous access. Anonymous
HTTP clients need a persistent MCP session; sessionless HTTP clients should
authenticate. Expired results and results owned by another caller return
resource-not-found. Stored results are not enumerated in `resources/list`.

HTTP ownership uses the server-issued session ID after the official transport
has validated it, rather than the identity of an SDK request/session proxy.
Stdio assigns a random identity to each live stream connection. Requests on that
connection can retrieve its anonymous results; a later connection cannot inherit
them. Concurrent HTTP requests cannot borrow a stdio connection's identity.

Successful explicit MCP tool results are stored as JSON; errors and input-required
responses pass through. Decorated tools with output schemas advertise either the
original output or the spillover envelope. Spilled widget results return the
resource reference without invoking the widget's data transformer.

Independent HTTP/stdio tests and reproducible large-output checks live in
`tests/test_spillover_transport.py`. Run the spillover tests with
`python -m pytest tests/test_spillover.py tests/test_spillover_transport.py`.
17 changes: 17 additions & 0 deletions nitrostack/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,16 @@
except (ImportError, KeyError, OSError, ValueError, TypeError):
__version__ = "unknown"

from nitrostack.core.interceptors import (
DataSpilloverInterceptor,
FileSpilloverStore,
MemorySpilloverStore,
PayloadTooLargeError,
SpilloverEntry,
SpilloverStore,
generate_preview,
)

from nitrostack.core.decorators import (
tool,
resource,
Expand Down Expand Up @@ -168,6 +178,13 @@

__all__ = [
"__version__",
"DataSpilloverInterceptor",
"FileSpilloverStore",
"MemorySpilloverStore",
"PayloadTooLargeError",
"SpilloverEntry",
"SpilloverStore",
"generate_preview",
"tool",
"resource",
"prompt",
Expand Down
62 changes: 61 additions & 1 deletion nitrostack/core/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@
from nitrostack.core.mcp_server import NitroStackMcpServer
from nitrostack.core.pipeline import run_pipeline
from nitrostack.core.additional_decorators import HealthCheckRegistry, ServerHealthPayload
from nitrostack.core.interceptors.data_spillover import (
DataSpilloverInterceptor, SPILLOVER_SCHEMA, SpilloverRegistry, SpilloverResult,
bind_spillover_context, caller_owner,
)
from nitrostack.core.task import TaskManager, TaskStatus
from nitrostack.events.event_emitter import EventEmitter
from nitrostack.protocol.schema import (
Expand Down Expand Up @@ -617,6 +621,7 @@ def __init__(self, app_class: Type):
self.task_manager = TaskManager()
self._in_flight = InFlightRegistry()
self.visibility_transform = VisibilityTransform()
self._spillover = SpilloverRegistry()

# Pipeline first: handlers consult it lazily when requests arrive.
self._transforms = TransformPipeline(
Expand Down Expand Up @@ -702,6 +707,11 @@ def _bootstrap(self) -> None:
# 5. Wire every protocol handler (tools/resources/prompts/tasks/notifications)
# onto the owned low-level server. This is factored out so additional server
# instances (e.g. one per HTTP session) can be configured identically.
for entry in self._tools.values():
for candidate in getattr(entry.method, "_mcp_interceptors", []):
interceptor = container.resolve(candidate) if isinstance(candidate, type) else candidate
if isinstance(interceptor, DataSpilloverInterceptor):
self._spillover.register(interceptor)
self._setup_handlers(self.mcp_server)

def _assert_declared_dependencies(
Expand Down Expand Up @@ -1276,6 +1286,11 @@ def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool:

output_schema = tool_json_schema(cfg.output_schema)
if output_schema is not None:
if self._has_spillover(entry.method):
output_schema = {
"type": "object", "anyOf": [output_schema, SPILLOVER_SCHEMA],
**({"$defs": output_schema["$defs"]} if "$defs" in output_schema else {}),
}
tool_kwargs["outputSchema"] = output_schema

return types.Tool(**tool_kwargs)
Expand Down Expand Up @@ -1320,7 +1335,33 @@ def _build_prompt_definition(self, entry: _PromptEntry) -> types.Prompt:
# Request dispatch
# ------------------------------------------------------------------

def _pipeline_stages(self, method: Callable) -> Tuple[List[Type], List[Type], List[Type], List[Type], List[Type]]:
def _has_spillover(self, method: Callable) -> bool:
container = DIContainer.get_instance()
return any(
isinstance(container.resolve(item) if isinstance(item, type) else item,
DataSpilloverInterceptor)
for item in getattr(method, "_mcp_interceptors", [])
)

async def _bind_spillover(self, context: ExecutionContext, rc: Any) -> None:
if not self._spillover.routes:
return
# OAuth guards may verify opaque tokens as well as the shared JWT path.
from nitrostack.auth.request import authorization_token_from_request, auth_context_from_payload
from nitrostack.auth.oauth import OAuthService

token = authorization_token_from_request(rc)
container = DIContainer.get_instance()
if context.auth is None and token and container.has_value(OAuthService):
try:
payload = await container.resolve(OAuthService).introspect_token(token)
if payload.get("active"):
context.auth = auth_context_from_payload(payload)
except Exception:
pass # Invalid credentials remain denied by caller_owner.
bind_spillover_context(context, rc, self._spillover, self.protocol_era)

def _pipeline_stages(self, method: Callable) -> Tuple[List[Any], List[Any], List[Any], List[Any], List[Any]]:
return (
getattr(method, "_mcp_guards", []),
getattr(method, "_mcp_middleware", []),
Expand Down Expand Up @@ -1390,6 +1431,8 @@ def _to_call_tool_result(
component: Optional[Component] = None,
context: Optional[ExecutionContext] = None,
) -> types.CallToolResult:
if isinstance(result, SpilloverResult):
return result.to_mcp()
if isinstance(result, InputRequiredResult):
wire = result.to_wire_dict()
return types.CallToolResult(
Expand Down Expand Up @@ -1641,6 +1684,7 @@ async def background_execution():
)
_apply_request_envelope(task_ctx, rc)
self._configure_visibility_context(task_ctx, rc)
await self._bind_spillover(task_ctx, rc)
task_ctx.task = TaskContext(
task_id,
self.task_manager,
Expand Down Expand Up @@ -1696,6 +1740,7 @@ async def background_execution():
)
_apply_request_envelope(ctx, rc)
self._configure_visibility_context(ctx, rc)
await self._bind_spillover(ctx, rc)
ticket = self._in_flight.register(correlation_id, jsonrpc_id=jsonrpc_id)
try:
result = await run_pipeline(
Expand Down Expand Up @@ -1728,6 +1773,21 @@ async def background_execution():
return self._to_call_tool_result(result, entry.component, ctx)

async def _read_resource(self, uri: str) -> List[ReadResourceContents]:
for prefix, store in self._spillover.routes.items():
if uri.startswith(prefix):
ctx = ExecutionContext(request_id=str(uuid.uuid4()))
rc = request_ctx.get(None)
_apply_request_envelope(ctx, rc)
await self._bind_spillover(ctx, rc)
try:
entry = await store.get(uri[len(prefix):], owner=caller_owner(ctx))
except (PermissionError, ValueError):
raise ResourceNotFoundError(uri) from None
if entry is None:
raise ResourceNotFoundError(uri)
return [ReadResourceContents(
content=entry.payload.decode("utf-8"), mime_type=entry.mime_type,
)]
entry = self._resources.get(uri)
path_kwargs: Dict[str, str] = {}

Expand Down
2 changes: 2 additions & 0 deletions nitrostack/core/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,8 @@ class ExecutionContext:
_visibility_session_key: Any = field(default=None, repr=False, compare=False)
_visibility_change_callback: Any = field(default=None, repr=False, compare=False)
_visibility_supported: bool = field(default=False, repr=False, compare=False)
spillover_session_key: Optional[str] = field(default=None, repr=False)
spillover_invalid_auth: bool = field(default=False, repr=False)

def __post_init__(self) -> None:
if not self.correlation_id:
Expand Down
4 changes: 4 additions & 0 deletions nitrostack/core/errors.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
from typing import Any


class PayloadTooLargeError(ValueError):
"""A serialized tool result exceeds the spillover store's per-result limit."""


class ToolExecutionError(Exception):
"""Raised when tool execution fails."""
pass
Expand Down
15 changes: 15 additions & 0 deletions nitrostack/core/interceptors/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
"""Built-in tool interceptors and spillover storage adapters."""

from .data_spillover import DataSpilloverInterceptor, generate_preview
from .spillover import (
FileSpilloverStore,
MemorySpilloverStore,
PayloadTooLargeError,
SpilloverEntry,
SpilloverStore,
)

__all__ = [
"DataSpilloverInterceptor", "FileSpilloverStore", "MemorySpilloverStore",
"PayloadTooLargeError", "SpilloverEntry", "SpilloverStore", "generate_preview",
]
Loading
Loading