@zleap-ai/dsh-sag 0.1.0

package/runtime/src/dsh_sag_runtime/errors.py

@@ -0,0 +1,47 @@
+"""Safe projection from runtime failures to JSON-RPC errors."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from zleap.sag import SagError
+
+from .protocol import RpcErrorData
+
+_SENSITIVE_PARTS = ("api_key", "apikey", "authorization", "password", "secret", "token")
+
+
+def _safe_details(value: Any, *, depth: int = 0) -> Any:
+    if depth > 4:
+        return "[truncated]"
+    if value is None or isinstance(value, bool | int | float):
+        return value
+    if isinstance(value, str):
+        return value[:1000]
+    if isinstance(value, list | tuple):
+        return [_safe_details(item, depth=depth + 1) for item in value[:50]]
+    if isinstance(value, dict):
+        return {
+            str(key)[:100]: _safe_details(item, depth=depth + 1)
+            for key, item in list(value.items())[:50]
+            if not any(part in str(key).lower() for part in _SENSITIVE_PARTS)
+        }
+    return str(value)[:1000]
+
+
+def safe_error_data(error: BaseException) -> RpcErrorData:
+    """Return only caller-safe SAG diagnostics; never project causes or tracebacks."""
+    if isinstance(error, SagError):
+        return RpcErrorData(
+            code=error.code,
+            operation=error.operation,
+            stage=error.stage,
+            retryable=error.retryable,
+            provider=error.provider,
+            itemId=error.item_id,
+            message=error.message,
+            details=_safe_details(error.details) if error.details else None,
+        )
+    if isinstance(error, ValueError):
+        return RpcErrorData(code="INVALID_ARGUMENT", retryable=False, message=str(error)[:1000])
+    return RpcErrorData(code="INTERNAL_ERROR", retryable=False, message="SAG runtime request failed")

package/runtime/src/dsh_sag_runtime/evidence.py

@@ -0,0 +1,53 @@
+"""Opaque, deterministic evidence references shared by search and read."""
+
+from __future__ import annotations
+
+import base64
+import json
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class EvidenceLocator:
+    namespace_id: str
+    source_id: str
+    chunk_id: str
+
+
+class EvidenceRefCodec:
+    """Encode locators and enforce namespace admission while decoding."""
+
+    def __init__(self, namespaces: set[str] | frozenset[str]) -> None:
+        self._namespaces = frozenset(namespaces)
+
+    def encode(self, locator: EvidenceLocator) -> str:
+        if locator.namespace_id not in self._namespaces:
+            raise ValueError("evidence namespace is not configured")
+        raw = json.dumps(
+            {"c": locator.chunk_id, "n": locator.namespace_id, "s": locator.source_id, "v": 1},
+            ensure_ascii=False,
+            separators=(",", ":"),
+            sort_keys=True,
+        ).encode()
+        return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
+
+    def encode_from_parts(self, namespace_id: str, source_id: str, chunk_id: str) -> str:
+        """Encode one locator without exposing its wire keys to callers."""
+        return self.encode(EvidenceLocator(namespace_id, source_id, chunk_id))
+
+    def decode(self, value: str) -> EvidenceLocator:
+        try:
+            if not value or any(character not in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" for character in value):
+                raise ValueError
+            padded = value + "=" * (-len(value) % 4)
+            raw = base64.b64decode(padded, altchars=b"-_", validate=True)
+            payload = json.loads(raw)
+        except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as error:
+            raise ValueError("invalid evidence reference") from error
+        if not isinstance(payload, dict) or set(payload) != {"v", "n", "s", "c"} or payload.get("v") != 1:
+            raise ValueError("invalid evidence reference")
+        if not all(isinstance(payload.get(key), str) and payload[key] for key in ("n", "s", "c")):
+            raise ValueError("invalid evidence reference")
+        if payload["n"] not in self._namespaces:
+            raise ValueError("evidence namespace is not configured")
+        return EvidenceLocator(payload["n"], payload["s"], payload["c"])

package/runtime/src/dsh_sag_runtime/protocol.py

@@ -0,0 +1,161 @@
+"""Versioned JSON-RPC 2.0 wire messages used over stdio."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+RPC_PROTOCOL_VERSION = "1.0"
+REQUIRED_ENGINE_VERSION = "0.10.0"
+MAX_FRAME_BYTES = 8 * 1024 * 1024
+
+
+class WireModel(BaseModel):
+    """Immutable strict wire model with camel-case serialization aliases."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True, populate_by_name=True)
+
+
+class InitializeParams(WireModel):
+    protocol_version: Literal["1.0"] = Field(alias="protocolVersion")
+
+
+class SearchParams(WireModel):
+    query: str = Field(min_length=1, max_length=32_000)
+    namespaces: tuple[str, ...] = Field(min_length=1)
+    mode: Literal["fast", "precise"]
+    limit: int = Field(ge=1, le=50)
+
+    @field_validator("namespaces")
+    @classmethod
+    def validate_namespaces(cls, values: tuple[str, ...]) -> tuple[str, ...]:
+        if any(not value or len(value) > 36 for value in values):
+            raise ValueError("namespace ids must contain 1 to 36 characters")
+        return values
+
+
+class ReadParams(WireModel):
+    evidence_ref: str = Field(alias="evidenceRef", min_length=1, max_length=4096)
+    offset: int = Field(ge=0)
+    max_chars: int = Field(alias="maxChars", ge=1, le=200_000)
+    include_events: bool = Field(alias="includeEvents")
+
+
+class CancelParams(WireModel):
+    id: int = Field(ge=0)
+
+
+class EmptyParams(WireModel):
+    pass
+
+
+class InitializeRequest(WireModel):
+    jsonrpc: Literal["2.0"]
+    id: int = Field(ge=0)
+    method: Literal["initialize"]
+    params: InitializeParams
+
+
+class SearchRequest(WireModel):
+    jsonrpc: Literal["2.0"]
+    id: int = Field(ge=0)
+    method: Literal["search"]
+    params: SearchParams
+
+
+class ReadRequest(WireModel):
+    jsonrpc: Literal["2.0"]
+    id: int = Field(ge=0)
+    method: Literal["read"]
+    params: ReadParams
+
+
+class ShutdownRequest(WireModel):
+    jsonrpc: Literal["2.0"]
+    id: int = Field(ge=0)
+    method: Literal["shutdown"]
+    params: EmptyParams
+
+
+class CancelRequest(WireModel):
+    jsonrpc: Literal["2.0"]
+    method: Literal["$/cancelRequest"]
+    params: CancelParams
+
+
+RpcRequest: TypeAlias = InitializeRequest | SearchRequest | ReadRequest | ShutdownRequest | CancelRequest
+
+_REQUEST_MODELS: dict[str, type[WireModel]] = {
+    "initialize": InitializeRequest,
+    "search": SearchRequest,
+    "read": ReadRequest,
+    "shutdown": ShutdownRequest,
+    "$/cancelRequest": CancelRequest,
+}
+
+
+class RpcErrorData(WireModel):
+    code: str
+    operation: str | None = None
+    stage: str | None = None
+    retryable: bool | None = None
+    provider: str | None = None
+    item_id: str | None = Field(default=None, alias="itemId")
+    message: str | None = None
+    details: dict[str, Any] | None = None
+
+
+class RpcError(WireModel):
+    code: int
+    message: str
+    data: RpcErrorData | None = None
+
+
+class RpcSuccessResponse(WireModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+    id: int = Field(ge=0)
+    result: dict[str, Any]
+
+
+class RpcErrorResponse(WireModel):
+    jsonrpc: Literal["2.0"] = "2.0"
+    id: int = Field(ge=0)
+    error: RpcError
+
+
+RpcResponse: TypeAlias = RpcSuccessResponse | RpcErrorResponse
+
+
+def decode_request(frame: bytes) -> RpcRequest:
+    """Decode one bounded request frame and reject unknown methods or fields."""
+    if len(frame) > MAX_FRAME_BYTES:
+        raise ValueError("request frame exceeds 8 MiB")
+    raw = json.loads(frame)
+    if not isinstance(raw, dict):
+        raise ValueError("request must be an object")
+    method = raw.get("method")
+    model = _REQUEST_MODELS.get(method) if isinstance(method, str) else None
+    if model is None:
+        raise ValueError("unsupported JSON-RPC method")
+    return model.model_validate(raw)  # type: ignore[return-value]
+
+
+def success_response(request_id: int, result: dict[str, Any]) -> RpcSuccessResponse:
+    """Create a successful correlated response."""
+    return RpcSuccessResponse(id=request_id, result=result)
+
+
+def error_response(request_id: int, code: int, message: str, data: RpcErrorData) -> RpcErrorResponse:
+    """Create a failed correlated response with safe structured data."""
+    return RpcErrorResponse(id=request_id, error=RpcError(code=code, message=message, data=data))
+
+
+def encode_response(response: RpcResponse) -> bytes:
+    """Serialize one response as a compact NDJSON frame."""
+    payload = response.model_dump_json(by_alias=True, exclude_none=True)
+    frame = f"{payload}\n".encode()
+    if len(frame) > MAX_FRAME_BYTES:
+        raise ValueError("response frame exceeds 8 MiB")
+    return frame

package/runtime/src/dsh_sag_runtime/read.py

@@ -0,0 +1,61 @@
+"""Bounded evidence reader for the `sag_read` tool."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .evidence import EvidenceRefCodec
+from .protocol import ReadParams
+
+
+class ReadAdapter:
+    """Resolve an opaque reference and return one Unicode-safe content page."""
+
+    def __init__(self, pool: Any, codec: EvidenceRefCodec, *, max_read_chars: int, max_events: int) -> None:
+        self._pool = pool
+        self._codec = codec
+        self._max_read_chars = max_read_chars
+        self._max_events = max_events
+
+    async def read(self, params: ReadParams) -> dict[str, Any]:
+        locator = self._codec.decode(params.evidence_ref)
+        async with self._pool.read_engine(locator.namespace_id) as engine:
+            records = await engine.read_evidence(
+                locator.source_id,
+                (locator.chunk_id,),
+                max_events_per_chunk=self._max_events,
+            )
+        matches = tuple(record for record in records if record.chunk.id == locator.chunk_id)
+        if len(matches) != 1:
+            raise ValueError("evidence reference must resolve to exactly one chunk")
+        record = matches[0]
+        content = record.chunk.content
+        if params.offset > len(content):
+            raise ValueError("read offset is beyond the evidence content")
+        page_size = min(params.max_chars, self._max_read_chars)
+        page = content[params.offset : params.offset + page_size]
+        next_offset = params.offset + len(page)
+        result: dict[str, Any] = {
+            "title": record.chunk.heading or record.chunk.source_id,
+            "content": page,
+            "offset": params.offset,
+            "totalChars": len(content),
+        }
+        if next_offset < len(content):
+            result["nextOffset"] = next_offset
+        if params.include_events:
+            result["events"] = [
+                {
+                    key: value
+                    for key, value in {
+                        "id": getattr(event, "id", None),
+                        "title": getattr(event, "title", None),
+                        "summary": getattr(event, "summary", None),
+                        "category": getattr(event, "category", None),
+                        "rank": getattr(event, "rank", None),
+                    }.items()
+                    if value is not None
+                }
+                for event in record.events[: self._max_events]
+            ]
+        return result

package/runtime/src/dsh_sag_runtime/search.py

@@ -0,0 +1,76 @@
+"""Search Contract 2.0 adapter for the `sag_search` tool."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from zleap.sag.pipeline import SearchOptions, SearchRequest, SearchScope
+
+from .evidence import EvidenceRefCodec
+from .protocol import SearchParams
+
+MODE_STRATEGY = {"fast": "vector", "precise": "full_expand"}
+
+
+class SearchAdapter:
+    """Map the stable sidecar request into zleap-sag 0.10.0 search types."""
+
+    def __init__(
+        self,
+        pool: Any,
+        codec: EvidenceRefCodec,
+        namespaces: set[str] | frozenset[str],
+        *,
+        max_results: int,
+        max_excerpt_chars: int,
+    ) -> None:
+        self._pool = pool
+        self._codec = codec
+        self._namespaces = frozenset(namespaces)
+        self._max_results = max_results
+        self._max_excerpt_chars = max_excerpt_chars
+
+    async def search(self, params: SearchParams) -> dict[str, Any]:
+        if any(namespace not in self._namespaces for namespace in params.namespaces):
+            raise ValueError("search namespace is not configured")
+        limit = min(params.limit, self._max_results)
+        request = SearchRequest(
+            query=params.query,
+            scope=SearchScope(data_source_ids=params.namespaces),
+            options=SearchOptions(
+                strategy=MODE_STRATEGY[params.mode],
+                top_k=limit,
+                return_type="chunk",
+            ),
+        )
+        result = await (await self._pool.search_engine()).search(request)
+        chunk_hits = bool(result.chunks)
+        hits = result.chunks or tuple(hit for hit in result.events if hit.chunk_id)
+        evidences: list[dict[str, Any]] = []
+        seen: set[tuple[str, str, str]] = set()
+        for hit in hits:
+            if len(evidences) >= limit:
+                break
+            namespace_id = hit.data_source_id
+            source_id = hit.source_id
+            chunk_id = hit.chunk_id or (hit.id if chunk_hits else None)
+            if not namespace_id or not source_id or not chunk_id:
+                raise ValueError("SAG search hit does not contain readable identifiers")
+            key = (namespace_id, source_id, chunk_id)
+            if key in seen:
+                continue
+            seen.add(key)
+            chunk = hit.chunk[0] if hit.chunk else None
+            content = hit.content or (chunk.content if chunk is not None else "")
+            title = hit.title or (chunk.heading if chunk is not None else "") or source_id
+            evidence: dict[str, Any] = {
+                "evidenceRef": self._codec.encode_from_parts(namespace_id, source_id, chunk_id),
+                "namespaceId": namespace_id,
+                "sourceId": source_id,
+                "title": title,
+                "excerpt": content[: self._max_excerpt_chars],
+            }
+            if hit.score is not None:
+                evidence["score"] = hit.score
+            evidences.append(evidence)
+        return {"query": params.query, "evidences": evidences}

package/runtime/src/dsh_sag_runtime/server.py

@@ -0,0 +1,136 @@
+"""Concurrent JSON-RPC request state machine for the stdio sidecar."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from .errors import safe_error_data
+from .protocol import (
+    RPC_PROTOCOL_VERSION,
+    CancelRequest,
+    InitializeRequest,
+    ReadRequest,
+    RpcErrorData,
+    RpcErrorResponse,
+    RpcRequest,
+    RpcResponse,
+    SearchRequest,
+    ShutdownRequest,
+    error_response,
+    success_response,
+)
+
+__all__ = ["RuntimeServer", "safe_error_data"]
+
+
+class RuntimeServer:
+    """Own request admission, correlation, cancellation, and runtime shutdown."""
+
+    def __init__(self, pool: Any, search: Any, read: Any, namespaces: tuple[str, ...], *, engine_version: str) -> None:
+        self._pool = pool
+        self._search = search
+        self._read = read
+        self._namespaces = namespaces
+        self._engine_version = engine_version
+        self._state = "new"
+        self._live: dict[int, asyncio.Task[dict[str, Any]]] = {}
+        self._pool_closed = False
+
+    def _state_error(self, request_id: int, code: str, message: str) -> RpcErrorResponse:
+        return error_response(request_id, -32002, message, RpcErrorData(code=code, retryable=False, message=message))
+
+    async def accept(self, request: RpcRequest) -> RpcResponse | None:
+        if isinstance(request, CancelRequest):
+            task = self._live.get(request.params.id)
+            if task is not None and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                except Exception:
+                    pass
+            return None
+
+        request_id = request.id
+        if request_id in self._live:
+            return self._state_error(request_id, "DSH_SAG_DUPLICATE_ID", "request id is already live")
+        if self._state in {"shutting", "closed"}:
+            return self._state_error(request_id, "DSH_SAG_SHUTTING_DOWN", "SAG runtime is shutting down")
+        if isinstance(request, InitializeRequest):
+            if self._state != "new":
+                return self._state_error(request_id, "DSH_SAG_ALREADY_INITIALIZED", "SAG runtime is already initialized")
+            self._state = "initializing"
+        elif self._state != "ready":
+            return self._state_error(request_id, "DSH_SAG_NOT_INITIALIZED", "initialize must complete before other requests")
+        elif isinstance(request, ShutdownRequest):
+            self._state = "shutting"
+
+        task = asyncio.create_task(self._execute(request))
+        self._live[request_id] = task
+        try:
+            result = await task
+            return success_response(request_id, result)
+        except asyncio.CancelledError:
+            if isinstance(request, InitializeRequest):
+                self._state = "new"
+            return error_response(
+                request_id,
+                -32800,
+                "request cancelled",
+                RpcErrorData(code="DSH_SAG_CANCELLED", retryable=True, message="request cancelled"),
+            )
+        except BaseException as error:
+            if isinstance(request, InitializeRequest):
+                self._state = "new"
+            data = safe_error_data(error)
+            return error_response(request_id, -32602 if isinstance(error, ValueError) else -32603, data.message or "SAG runtime request failed", data)
+        finally:
+            self._live.pop(request_id, None)
+
+    async def _execute(self, request: InitializeRequest | SearchRequest | ReadRequest | ShutdownRequest) -> dict[str, Any]:
+        if isinstance(request, InitializeRequest):
+            await self._pool.start()
+            capabilities = await self._pool.capabilities()
+            if not capabilities.evidence_read:
+                raise RuntimeError("zleap-sag runtime does not support evidence reads")
+            health = await self._pool.health()
+            health_state = getattr(health.state, "value", health.state)
+            self._state = "ready"
+            return {
+                "protocolVersion": RPC_PROTOCOL_VERSION,
+                "engineVersion": self._engine_version,
+                "health": str(health_state),
+                "evidenceRead": True,
+                "namespaces": list(self._namespaces),
+            }
+        if isinstance(request, SearchRequest):
+            return await self._search.search(request.params)
+        if isinstance(request, ReadRequest):
+            return await self._read.read(request.params)
+        current = asyncio.current_task()
+        others = [task for task in self._live.values() if task is not current and not task.done()]
+        for task in others:
+            task.cancel()
+        if others:
+            await asyncio.gather(*others, return_exceptions=True)
+        await self._close_pool()
+        return {}
+
+    async def _close_pool(self) -> None:
+        if not self._pool_closed:
+            self._pool_closed = True
+            await self._pool.aclose()
+
+    async def aclose(self) -> None:
+        if self._state == "closed":
+            return
+        self._state = "closed"
+        current = asyncio.current_task()
+        tasks = [task for task in self._live.values() if task is not current and not task.done()]
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        await self._close_pool()