mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
fix: remove mcp artefacts, handle non fatal issues on workflow
gracefully
This commit is contained in:
Alex committed
1 parent
12dfefdd05
commit
89eb3f68fc
14 files changed
+166
-771
No files matched your search
@@ -90,10 +90,13 @@ class WorkflowAgent(BaseAgent):
|
||||
self._finalize_workflow_run(workflow_owner_id, run_user_id, pg_workflow_id, query)
|
||||
return
|
||||
|
||||
# Non-fatal: some attachments were dropped (oversize / unreadable). Tell the
|
||||
# user which, then still run with the documents that did bridge.
|
||||
# Non-fatal: some attachments were dropped (oversize / unreadable). Emit a
|
||||
# ``notice`` -- NOT an ``error`` -- so the client surfaces which were dropped
|
||||
# without marking the turn failed or ending the stream (an ``error`` event is
|
||||
# terminal client-side and disables reconnect). The run then still executes
|
||||
# with the documents that did bridge.
|
||||
if dropped:
|
||||
yield {"type": "error", "user_facing": True, "error": " ".join(dropped)}
|
||||
yield {"type": "notice", "notice": " ".join(dropped)}
|
||||
|
||||
self._engine.run_persisted = self._run_persisted
|
||||
interrupted = True
|
||||
|
||||
@@ -552,6 +552,11 @@ class BaseAnswerResource:
|
||||
if not line.get("user_facing"):
|
||||
error_text = sanitize_api_error(error_text)
|
||||
yield _emit({"type": "error", "error": error_text})
|
||||
elif line.get("type") == "notice":
|
||||
# Non-fatal, non-terminal notice (e.g. some workflow input
|
||||
# documents were dropped). Forwarded verbatim so the client can
|
||||
# surface it without failing the turn; never sanitized as an error.
|
||||
yield _emit({"type": "notice", "notice": line.get("notice", "")})
|
||||
else:
|
||||
yield _emit(line)
|
||||
if is_structured and structured_chunks:
|
||||
|
||||
@@ -392,8 +392,6 @@ class Settings(BaseSettings):
|
||||
ARTIFACT_MAX_BYTES: int = 50 * 1024 * 1024 # cap on a single stored artifact version's bytes
|
||||
ARTIFACT_MAX_COUNT_PER_USER: int = 5000 # cap on artifacts a user may own
|
||||
ARTIFACT_MAX_TOTAL_BYTES_PER_USER: int = 5 * 1024 * 1024 * 1024 # cap on a user's total stored bytes
|
||||
# Cap on bytes served per MCP ``resources/read`` so a giant artifact never streams into LLM context.
|
||||
ARTIFACT_RESOURCE_READ_MAX_BYTES: int = 1 * 1024 * 1024
|
||||
|
||||
@field_validator("POSTGRES_URI", mode="before")
|
||||
@classmethod
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
"""FastMCP server exposing DocsGPT retrieval + artifacts over streamable HTTP.
|
||||
"""FastMCP server exposing DocsGPT retrieval over streamable HTTP.
|
||||
|
||||
Mounted at ``/mcp`` by ``application/asgi.py``. Bearer tokens are the
|
||||
existing DocsGPT agent API keys — no new credential surface. The
|
||||
``search_docs`` tool searches the caller's knowledge base; the artifact
|
||||
resources middleware exposes the caller's own artifacts as MCP resources
|
||||
(``resources/list`` / ``resources/read`` over ``artifact://`` URIs),
|
||||
scoped strictly to the Bearer key's owning principal.
|
||||
``search_docs`` tool searches the caller's knowledge base.
|
||||
|
||||
The tool reads the ``Authorization`` header directly via
|
||||
``get_http_headers(include={"authorization"})``. The ``include`` kwarg
|
||||
@@ -18,23 +15,11 @@ token, we opt it back in.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
from typing import Sequence
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.resources import ResourceContent, ResourceResult
|
||||
from fastmcp.resources.base import Resource
|
||||
from fastmcp.server.dependencies import get_http_headers
|
||||
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
||||
|
||||
from application.services.artifact_resource_service import (
|
||||
ArtifactReadResult,
|
||||
ResourceDenied,
|
||||
ResourceNotFound,
|
||||
list_artifact_resources,
|
||||
read_artifact_resource,
|
||||
)
|
||||
from application.services.search_service import (
|
||||
InvalidAPIKey,
|
||||
SearchFailed,
|
||||
@@ -52,61 +37,7 @@ def _extract_bearer_token() -> str | None:
|
||||
return parts[1]
|
||||
|
||||
|
||||
def _read_result_to_mcp(result: ArtifactReadResult) -> ResourceResult:
|
||||
"""Wrap a service read result as a FastMCP ``ResourceResult`` (text or blob)."""
|
||||
if result.blob_b64 is not None:
|
||||
raw = base64.b64decode(result.blob_b64)
|
||||
return ResourceResult([ResourceContent(raw, mime_type=result.mime_type)])
|
||||
return ResourceResult([ResourceContent(result.text or "", mime_type=result.mime_type)])
|
||||
|
||||
|
||||
class ArtifactResourcesMiddleware(Middleware):
|
||||
"""Expose the calling principal's artifacts as MCP resources (read/list).
|
||||
|
||||
The principal is the Bearer api_key's owner; resources are scoped to that
|
||||
owner and a foreign/unauthenticated read is denied. Static resources
|
||||
registered on the server (if any) are preserved by chaining ``call_next``.
|
||||
"""
|
||||
|
||||
async def on_list_resources(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: CallNext,
|
||||
) -> Sequence[Resource]:
|
||||
"""Append the principal's artifact resources to the static resource list."""
|
||||
existing = list(await call_next(context))
|
||||
try:
|
||||
artifacts = await asyncio.to_thread(
|
||||
list_artifact_resources, _extract_bearer_token()
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("on_list_resources: artifact listing failed")
|
||||
return existing
|
||||
existing.extend(artifacts)
|
||||
return existing
|
||||
|
||||
async def on_read_resource(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: CallNext,
|
||||
) -> ResourceResult:
|
||||
"""Serve ``artifact://`` reads from the artifact store; defer others."""
|
||||
uri = str(getattr(context.message, "uri", ""))
|
||||
if not uri.startswith("artifact://"):
|
||||
return await call_next(context)
|
||||
try:
|
||||
result = await asyncio.to_thread(
|
||||
read_artifact_resource, _extract_bearer_token(), uri
|
||||
)
|
||||
except ResourceDenied as exc:
|
||||
raise PermissionError(str(exc) or "forbidden") from exc
|
||||
except ResourceNotFound as exc:
|
||||
raise ValueError(str(exc) or "resource not found") from exc
|
||||
return _read_result_to_mcp(result)
|
||||
|
||||
|
||||
mcp = FastMCP("docsgpt")
|
||||
mcp.add_middleware(ArtifactResourcesMiddleware())
|
||||
|
||||
|
||||
@mcp.tool
|
||||
|
||||
@@ -1,267 +0,0 @@
|
||||
"""Flask-free service exposing a principal's artifacts as MCP Resources.
|
||||
|
||||
The MCP server (``application/mcp_server.py``) authenticates a request with
|
||||
``Authorization: Bearer <agent-api-key>``; that key resolves to the owning
|
||||
``user_id`` via ``AgentsRepository.find_by_key`` (the same api_key->owner path
|
||||
as the HTTP artifact routes). Resources are scoped strictly to that principal
|
||||
AND to the owner's own conversations: because the api_key is embedded publicly
|
||||
in the widget and MCP carries no per-visitor conversation context, shared-usage
|
||||
(widget/visitor) conversations are excluded. ``resources/list`` returns only the
|
||||
owner's own-chat artifacts, and ``resources/read`` re-checks ownership + owner
|
||||
scope before serving any bytes. An unresolvable principal yields an empty list /
|
||||
a denied read -- never another principal's or another visitor's artifact.
|
||||
|
||||
``resources/list`` returns FastMCP ``Resource`` objects so they pass straight
|
||||
through the server's list/dedupe/wire pipeline; ``resources/read`` is served by
|
||||
the MCP middleware, which streams the real bytes for each ``artifact://`` uri.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
from fastmcp.resources.base import Resource
|
||||
from fastmcp.resources.types import TextResource
|
||||
from sqlalchemy.exc import DataError, DBAPIError
|
||||
|
||||
from application.core.settings import settings
|
||||
from application.storage.db.base_repository import looks_like_uuid
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
from application.storage.db.session import db_readonly
|
||||
from application.storage.storage_creator import StorageCreator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# artifact://{artifact_id}/v{version}
|
||||
_URI_RE = re.compile(r"^artifact://(?P<id>[^/]+)/v(?P<version>\d+)$")
|
||||
|
||||
# Max resources advertised by ``resources/list`` so the model is not flooded
|
||||
# with thousands of rows even when a principal owns far more artifacts.
|
||||
_RESOURCE_LIST_LIMIT = 500
|
||||
|
||||
# mime types served as inline ``text`` rather than base64 ``blob``.
|
||||
_TEXT_MIME_PREFIXES = ("text/",)
|
||||
_TEXT_MIME_EXACT = {
|
||||
"application/json",
|
||||
"application/xml",
|
||||
"application/javascript",
|
||||
"application/x-ndjson",
|
||||
"image/svg+xml",
|
||||
}
|
||||
_TEXT_MIME_SUFFIXES = ("+json", "+xml")
|
||||
|
||||
|
||||
class ResourceDenied(Exception):
|
||||
"""The principal may not access the requested resource (unauth or foreign)."""
|
||||
|
||||
|
||||
class ResourceNotFound(Exception):
|
||||
"""The requested ``artifact://`` uri does not resolve to a stored version."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArtifactReadResult:
|
||||
"""Materialized artifact contents for a ``resources/read`` response."""
|
||||
|
||||
uri: str
|
||||
mime_type: str
|
||||
text: Optional[str] = None
|
||||
blob_b64: Optional[str] = None
|
||||
|
||||
|
||||
def _read_cap() -> int:
|
||||
"""Return the per-read byte cap, or 0 (no cap) when the setting disables it."""
|
||||
return int(getattr(settings, "ARTIFACT_RESOURCE_READ_MAX_BYTES", 0) or 0)
|
||||
|
||||
|
||||
def _cap_text(text: str) -> str:
|
||||
"""Truncate ``text`` to the read cap; unchanged when the cap is disabled."""
|
||||
cap = _read_cap()
|
||||
return text[:cap] if cap > 0 else text
|
||||
|
||||
|
||||
def _is_texty(mime_type: str) -> bool:
|
||||
"""Return True when ``mime_type`` should be served as inline UTF-8 text."""
|
||||
mime = (mime_type or "").split(";", 1)[0].strip().lower()
|
||||
if mime in _TEXT_MIME_EXACT:
|
||||
return True
|
||||
if any(mime.startswith(p) for p in _TEXT_MIME_PREFIXES):
|
||||
return True
|
||||
return any(mime.endswith(s) for s in _TEXT_MIME_SUFFIXES)
|
||||
|
||||
|
||||
def _resolve_agent(api_key: Optional[str]) -> Optional[dict]:
|
||||
"""Resolve a Bearer api_key to its owning agent row (carries ``id`` + ``user_id``).
|
||||
|
||||
Returns the whole agent so callers can scope artifact visibility to that agent's
|
||||
conversations, not the owner's entire corpus. None when unresolvable.
|
||||
"""
|
||||
if not api_key:
|
||||
return None
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
agent = AgentsRepository(conn).find_by_key(api_key)
|
||||
except Exception:
|
||||
logger.exception("artifact resource: principal resolution failed")
|
||||
return None
|
||||
if not agent or not agent.get("user_id") or not agent.get("id"):
|
||||
return None
|
||||
return agent
|
||||
|
||||
|
||||
def _resource_uri(artifact_id: str, version: int) -> str:
|
||||
"""Build the stable ``artifact://{id}/v{version}`` resource uri."""
|
||||
return f"artifact://{artifact_id}/v{version}"
|
||||
|
||||
|
||||
def list_artifact_resources(api_key: Optional[str]) -> List[Resource]:
|
||||
"""List the calling principal's artifacts as FastMCP resources (empty if unresolved).
|
||||
|
||||
Returns FastMCP ``Resource`` objects (not raw ``mcp.types.Resource``): the
|
||||
server's ``resources/list`` pipeline reads FastMCP-only attributes
|
||||
(``.version``, ``.auth``, ``is_enabled``) and calls ``to_mcp_resource()`` for
|
||||
the wire encoding, so a raw ``mcp.types.Resource`` mixed into the list would
|
||||
raise ``AttributeError``. The placeholder ``text`` is never served -- the
|
||||
middleware's ``on_read_resource`` intercepts every ``artifact://`` read and
|
||||
streams the real bytes via :func:`read_artifact_resource`.
|
||||
"""
|
||||
agent = _resolve_agent(api_key)
|
||||
if not agent:
|
||||
return []
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
# MCP-over-key has no per-visitor conversation context, and the api_key
|
||||
# is embedded publicly in the widget, so scope to the owner's OWN
|
||||
# (non-shared-usage) conversations only. Without this, the key would
|
||||
# enumerate artifacts from every visitor's conversation with the agent.
|
||||
rows = ArtifactsRepository(conn).list_artifacts_for_agent(
|
||||
str(agent["id"]), str(agent["user_id"]), owner_only=True
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("artifact resource: list failed")
|
||||
return []
|
||||
|
||||
resources: List[Resource] = []
|
||||
for row in rows[:_RESOURCE_LIST_LIMIT]:
|
||||
artifact_id = str(row.get("id"))
|
||||
version = row.get("current_version") or 1
|
||||
title = row.get("title") or f"artifact-{artifact_id}"
|
||||
resources.append(
|
||||
TextResource(
|
||||
uri=_resource_uri(artifact_id, version),
|
||||
name=title,
|
||||
title=title,
|
||||
description=f"DocsGPT {row.get('kind') or 'file'} artifact",
|
||||
mime_type=_kind_mime_hint(row.get("kind")),
|
||||
text="",
|
||||
)
|
||||
)
|
||||
return resources
|
||||
|
||||
|
||||
def _kind_mime_hint(kind: Optional[str]) -> str:
|
||||
"""Concrete mime hint for a list row; ambiguous kinds fall back to octet-stream."""
|
||||
# Only kinds with a single unambiguous mime get a concrete type; everything
|
||||
# else stays octet-stream so a list row never mismatches the read's bytes.
|
||||
return {
|
||||
"html": "text/html",
|
||||
"data": "application/json",
|
||||
}.get((kind or "").lower(), "application/octet-stream")
|
||||
|
||||
|
||||
def read_artifact_resource(api_key: Optional[str], uri: str) -> ArtifactReadResult:
|
||||
"""Authorize and materialize an ``artifact://`` resource for the principal.
|
||||
|
||||
Raises:
|
||||
ResourceDenied: principal unresolved, or the artifact is not theirs.
|
||||
ResourceNotFound: uri malformed, or the version/file does not exist.
|
||||
"""
|
||||
match = _URI_RE.match(uri or "")
|
||||
if not match:
|
||||
raise ResourceNotFound(f"unsupported resource uri: {uri!r}")
|
||||
artifact_id = match.group("id")
|
||||
version = int(match.group("version"))
|
||||
|
||||
# A non-UUID id would reach a ``CAST(:id AS uuid)`` and raise a DB DataError;
|
||||
# gate it up front like the HTTP artifact routes do.
|
||||
if not looks_like_uuid(artifact_id):
|
||||
raise ResourceNotFound(f"artifact {artifact_id} not found")
|
||||
|
||||
agent = _resolve_agent(api_key)
|
||||
if not agent:
|
||||
raise ResourceDenied("unauthenticated")
|
||||
user_id = str(agent["user_id"])
|
||||
agent_id = str(agent["id"])
|
||||
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact(artifact_id)
|
||||
if artifact is None:
|
||||
raise ResourceNotFound(f"artifact {artifact_id} not found")
|
||||
# Ownership is the first authz point: never serve another principal's
|
||||
# artifact over MCP, regardless of conversation/workflow parent sharing.
|
||||
if str(artifact.get("user_id")) != user_id:
|
||||
raise ResourceDenied("forbidden")
|
||||
# Agent scope is the second: a per-agent key only reads artifacts from
|
||||
# its own conversations, not the owner's other agents / workflow runs.
|
||||
# ``owner_only`` further excludes shared-usage (widget/visitor)
|
||||
# conversations: the public key reads the owner's own chats only, never
|
||||
# another visitor's artifact (parity with the HTTP route's confinement).
|
||||
if not repo.artifact_in_agent_scope(artifact_id, agent_id, owner_only=True):
|
||||
raise ResourceDenied("forbidden")
|
||||
version_row = repo.get_version(artifact_id, version)
|
||||
except (DataError, DBAPIError) as exc:
|
||||
raise ResourceNotFound(f"artifact {artifact_id} not found") from exc
|
||||
|
||||
if version_row is None:
|
||||
raise ResourceNotFound(f"version {version} of {artifact_id} not found")
|
||||
|
||||
mime_type = version_row.get("mime_type") or "application/octet-stream"
|
||||
uri = _resource_uri(artifact_id, version)
|
||||
|
||||
# Prefer the stored preview/extracted text for texty kinds: it is already
|
||||
# bounded and avoids a storage round-trip.
|
||||
preview = version_row.get("preview_text")
|
||||
if _is_texty(mime_type) and preview:
|
||||
return ArtifactReadResult(uri=uri, mime_type=mime_type, text=_cap_text(preview))
|
||||
|
||||
storage_path = version_row.get("storage_path")
|
||||
if not storage_path:
|
||||
if _is_texty(mime_type) and preview is not None:
|
||||
return ArtifactReadResult(uri=uri, mime_type=mime_type, text=_cap_text(preview))
|
||||
raise ResourceNotFound(f"version {version} of {artifact_id} has no stored bytes")
|
||||
|
||||
try:
|
||||
data = _read_capped_bytes(storage_path)
|
||||
except FileNotFoundError as exc:
|
||||
raise ResourceNotFound(f"version {version} of {artifact_id} has no stored bytes") from exc
|
||||
|
||||
if _is_texty(mime_type):
|
||||
# ``errors="ignore"`` keeps valid text texty even when the cap splits a
|
||||
# multibyte char at the boundary, instead of demoting it to a blob.
|
||||
return ArtifactReadResult(uri=uri, mime_type=mime_type, text=data.decode("utf-8", errors="ignore"))
|
||||
return ArtifactReadResult(
|
||||
uri=uri, mime_type=mime_type, blob_b64=base64.b64encode(data).decode("ascii")
|
||||
)
|
||||
|
||||
|
||||
def _read_capped_bytes(storage_path: str) -> bytes:
|
||||
"""Read at most ``_read_cap()`` bytes of a stored artifact version (0 == all)."""
|
||||
cap = _read_cap()
|
||||
storage = StorageCreator.get_storage()
|
||||
file_obj = storage.get_file(storage_path)
|
||||
try:
|
||||
return file_obj.read(cap) if cap > 0 else file_obj.read()
|
||||
finally:
|
||||
close = getattr(file_obj, "close", None)
|
||||
if callable(close):
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
logger.debug("artifact resource: file close failed", exc_info=True)
|
||||
@@ -206,25 +206,19 @@ class ArtifactsRepository:
|
||||
agent_id: str,
|
||||
user_id: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
owner_only: bool = False,
|
||||
) -> list[dict]:
|
||||
"""List an agent's artifacts (owner-scoped), optionally narrowed to one conversation.
|
||||
|
||||
Scopes a per-agent api-key's artifact visibility to the conversations that
|
||||
agent produced, so the key cannot enumerate the owner's whole corpus. When
|
||||
``conversation_id`` is given, the SQL narrows to that single conversation
|
||||
(the per-visitor bearer capability for a public widget key). When
|
||||
``owner_only`` is set, shared-usage (widget/visitor) conversations are
|
||||
excluded so a credential with no per-visitor conversation context (the MCP
|
||||
Bearer key) sees only the owner's own chats, never other visitors' artifacts.
|
||||
(the per-visitor bearer capability for a public widget key).
|
||||
"""
|
||||
clauses = ["c.agent_id = CAST(:agent_id AS uuid)", "a.user_id = :user_id"]
|
||||
params: dict[str, Any] = {"agent_id": str(agent_id), "user_id": user_id}
|
||||
if conversation_id is not None:
|
||||
clauses.append("a.conversation_id = CAST(:conversation_id AS uuid)")
|
||||
params["conversation_id"] = conversation_id
|
||||
if owner_only:
|
||||
clauses.append("c.is_shared_usage = false")
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT a.* FROM artifacts a "
|
||||
@@ -236,21 +230,12 @@ class ArtifactsRepository:
|
||||
)
|
||||
return [_artifact_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def artifact_in_agent_scope(
|
||||
self, artifact_id: str, agent_id: str, owner_only: bool = False
|
||||
) -> bool:
|
||||
"""True if ``artifact_id``'s parent conversation belongs to ``agent_id``.
|
||||
|
||||
With ``owner_only`` set, shared-usage (widget/visitor) conversations are
|
||||
excluded so a per-visitor-less credential (the MCP Bearer key) can only
|
||||
confirm scope for the owner's own chats, never another visitor's artifact.
|
||||
"""
|
||||
def artifact_in_agent_scope(self, artifact_id: str, agent_id: str) -> bool:
|
||||
"""True if ``artifact_id``'s parent conversation belongs to ``agent_id``."""
|
||||
clauses = [
|
||||
"a.id = CAST(:id AS uuid)",
|
||||
"c.agent_id = CAST(:agent_id AS uuid)",
|
||||
]
|
||||
if owner_only:
|
||||
clauses.append("c.is_shared_usage = false")
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT 1 FROM artifacts a "
|
||||
|
||||
@@ -280,35 +280,47 @@ export default function ConversationMessages({
|
||||
}
|
||||
|
||||
// tool_calls.length, not tool_calls — empty arrays are JS-truthy.
|
||||
// ``notice`` is included so a non-fatal notice still surfaces even when the
|
||||
// run produced no textual answer (e.g. an artifact-only workflow).
|
||||
const hasContent =
|
||||
query.thought ||
|
||||
query.response ||
|
||||
(query.tool_calls && query.tool_calls.length > 0) ||
|
||||
query.research;
|
||||
query.research ||
|
||||
query.notice;
|
||||
if (hasContent) {
|
||||
const isCurrentlyStreaming =
|
||||
status === 'loading' && index === queries.length - 1;
|
||||
return (
|
||||
<ConversationBubble
|
||||
className={bubbleMargin}
|
||||
key={`${index}-ANSWER`}
|
||||
message={query.response}
|
||||
type={'ANSWER'}
|
||||
thought={query.thought}
|
||||
sources={query.sources}
|
||||
toolCalls={query.tool_calls}
|
||||
research={query.research}
|
||||
onOpenArtifact={onOpenArtifact}
|
||||
onToolAction={onToolAction}
|
||||
feedback={query.feedback}
|
||||
isStreaming={isCurrentlyStreaming}
|
||||
agentId={agentId}
|
||||
handleFeedback={
|
||||
handleFeedback
|
||||
? (feedback) => handleFeedback(query, feedback, index)
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
<Fragment key={`${index}-ANSWER`}>
|
||||
{query.notice ? (
|
||||
<div
|
||||
role="status"
|
||||
className={`${bubbleMargin} mr-5 self-start rounded-2xl border border-amber-300 bg-amber-50 px-4 py-2 text-sm text-amber-800 dark:border-amber-500/40 dark:bg-amber-500/10 dark:text-amber-200`}
|
||||
>
|
||||
{query.notice}
|
||||
</div>
|
||||
) : null}
|
||||
<ConversationBubble
|
||||
className={bubbleMargin}
|
||||
message={query.response}
|
||||
type={'ANSWER'}
|
||||
thought={query.thought}
|
||||
sources={query.sources}
|
||||
toolCalls={query.tool_calls}
|
||||
research={query.research}
|
||||
onOpenArtifact={onOpenArtifact}
|
||||
onToolAction={onToolAction}
|
||||
feedback={query.feedback}
|
||||
isStreaming={isCurrentlyStreaming}
|
||||
agentId={agentId}
|
||||
handleFeedback={
|
||||
handleFeedback
|
||||
? (feedback) => handleFeedback(query, feedback, index)
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
</Fragment>
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -63,6 +63,9 @@ export interface Query {
|
||||
sources?: { title: string; text: string; link: string }[];
|
||||
tool_calls?: ToolCallsType[];
|
||||
error?: string;
|
||||
// Non-fatal notice (e.g. some workflow input documents were dropped). Shown
|
||||
// alongside the answer; unlike ``error`` it does not fail the turn or end the stream.
|
||||
notice?: string;
|
||||
attachments?: { id: string; fileName: string }[];
|
||||
structured?: boolean;
|
||||
schema?: object;
|
||||
|
||||
@@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest';
|
||||
|
||||
import reducer, {
|
||||
applyMessageTail,
|
||||
raiseNotice,
|
||||
setConversation,
|
||||
} from './conversationSlice';
|
||||
|
||||
@@ -151,3 +152,33 @@ describe('applyMessageTail — streaming partial', () => {
|
||||
expect(next.queries[0].response).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('raiseNotice — non-fatal notice', () => {
|
||||
it('records a notice without setting an error (turn is not failed)', () => {
|
||||
// seedSlice leaves conversationId at its initial null; match that.
|
||||
const state = seedSlice();
|
||||
const next = reducer(
|
||||
state,
|
||||
raiseNotice({
|
||||
conversationId: null,
|
||||
index: 0,
|
||||
message: 'big.txt was dropped (too large)',
|
||||
}),
|
||||
);
|
||||
expect(next.queries[0].notice).toBe('big.txt was dropped (too large)');
|
||||
expect(next.queries[0].error).toBeUndefined();
|
||||
});
|
||||
|
||||
it('is a no-op when the conversationId does not match', () => {
|
||||
const state = seedSlice();
|
||||
const next = reducer(
|
||||
state,
|
||||
raiseNotice({
|
||||
conversationId: 'some-other-conversation',
|
||||
index: 0,
|
||||
message: 'ignored',
|
||||
}),
|
||||
);
|
||||
expect(next.queries[0].notice).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -297,6 +297,14 @@ export const fetchAnswer = createAsyncThunk<
|
||||
dispatch(
|
||||
conversationSlice.actions.setStatus('awaiting_tool_actions'),
|
||||
);
|
||||
} else if (data.type === 'notice') {
|
||||
dispatch(
|
||||
conversationSlice.actions.raiseNotice({
|
||||
conversationId: currentConversationId,
|
||||
index: targetIndex,
|
||||
message: data.notice ?? '',
|
||||
}),
|
||||
);
|
||||
} else if (data.type === 'error') {
|
||||
dispatch(conversationSlice.actions.setStatus('failed'));
|
||||
dispatch(
|
||||
@@ -442,6 +450,14 @@ export const fetchAnswer = createAsyncThunk<
|
||||
progress: data.data,
|
||||
}),
|
||||
);
|
||||
} else if (data.type === 'notice') {
|
||||
dispatch(
|
||||
conversationSlice.actions.raiseNotice({
|
||||
conversationId: currentConversationId,
|
||||
index: targetIndex,
|
||||
message: data.notice ?? '',
|
||||
}),
|
||||
);
|
||||
} else if (data.type === 'error') {
|
||||
// set status to 'failed'
|
||||
dispatch(conversationSlice.actions.setStatus('failed'));
|
||||
@@ -703,6 +719,14 @@ export const submitToolActions = createAsyncThunk<
|
||||
);
|
||||
} else if (data.type === 'tool_calls_pending') {
|
||||
dispatch(conversationSlice.actions.setStatus('awaiting_tool_actions'));
|
||||
} else if (data.type === 'notice') {
|
||||
dispatch(
|
||||
conversationSlice.actions.raiseNotice({
|
||||
conversationId,
|
||||
index: targetIndex,
|
||||
message: data.notice ?? '',
|
||||
}),
|
||||
);
|
||||
} else if (data.type === 'error') {
|
||||
dispatch(conversationSlice.actions.setStatus('failed'));
|
||||
dispatch(
|
||||
@@ -1009,6 +1033,24 @@ export const conversationSlice = createSlice({
|
||||
state.queries[index].error = message;
|
||||
},
|
||||
|
||||
// Non-fatal counterpart to ``raiseError``: records a notice on the query
|
||||
// (e.g. some workflow input documents were dropped) WITHOUT setting the
|
||||
// 'failed' status, so the turn keeps streaming and can still complete.
|
||||
raiseNotice(
|
||||
state,
|
||||
action: PayloadAction<{
|
||||
conversationId: string | null;
|
||||
index: number;
|
||||
message: string;
|
||||
}>,
|
||||
) {
|
||||
const { conversationId, index, message } = action.payload;
|
||||
if (state.conversationId !== conversationId) return;
|
||||
if (!state.queries[index]) return;
|
||||
|
||||
state.queries[index].notice = message;
|
||||
},
|
||||
|
||||
resetConversation: (state) => {
|
||||
state.queries = initialState.queries;
|
||||
state.status = initialState.status;
|
||||
@@ -1054,6 +1096,7 @@ export const {
|
||||
setConversationId,
|
||||
setStatus,
|
||||
raiseError,
|
||||
raiseNotice,
|
||||
resetConversation,
|
||||
applyMessageTail,
|
||||
updateMessageMeta,
|
||||
|
||||
@@ -415,11 +415,12 @@ def test_oversize_declared_attachment_skipped_with_notice(pg_engine, tmp_path, m
|
||||
# The oversize doc was dropped -> no input documents bridged, but the run ran.
|
||||
assert engine.captured_inputs is not None
|
||||
assert engine.captured_inputs["input_documents"] == []
|
||||
# A non-fatal notice naming the dropped document was surfaced, flagged user_facing
|
||||
# so the route emits the filename verbatim instead of a generic sanitized message.
|
||||
errors = [e for e in events if e.get("type") == "error"]
|
||||
assert errors and "big.txt" in errors[0]["error"]
|
||||
assert errors[0].get("user_facing") is True
|
||||
# A non-fatal notice naming the dropped document was surfaced as a ``notice``
|
||||
# (NOT an ``error``, which is terminal client-side) so the run still completes.
|
||||
notices = [e for e in events if e.get("type") == "notice"]
|
||||
assert notices and "big.txt" in notices[0]["notice"]
|
||||
# It must not be an error event (that would fail the turn and disable reconnect).
|
||||
assert not [e for e in events if e.get("type") == "error"]
|
||||
# Nothing was persisted for the oversize doc.
|
||||
with pg_engine.connect() as conn:
|
||||
n = conn.execute(
|
||||
|
||||
@@ -241,6 +241,40 @@ class TestCompleteStreamMethod:
|
||||
assert "artifact storage quota" in error_chunks[0]
|
||||
assert "Rate limit exceeded" not in error_chunks[0]
|
||||
|
||||
def test_notice_is_forwarded_verbatim_and_not_an_error(self, mock_mongo_db, flask_app):
|
||||
"""A non-fatal ``notice`` streams through as a notice, never as an error.
|
||||
|
||||
A ``notice`` (e.g. some workflow input documents were dropped) must not be
|
||||
emitted as ``type: error`` -- the client treats an error event as terminal and
|
||||
disables reconnect -- and its text must not be run through sanitize_api_error.
|
||||
"""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[{"type": "notice", "notice": "big.txt exceeds the per-file size limit"}]
|
||||
)
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user123"},
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
notice_chunks = [s for s in stream if '"type": "notice"' in s]
|
||||
assert notice_chunks
|
||||
assert "big.txt exceeds the per-file size limit" in notice_chunks[0]
|
||||
# Crucially, it is not surfaced as an error event.
|
||||
assert not [s for s in stream if '"type": "error"' in s]
|
||||
|
||||
def test_non_user_facing_error_is_sanitized(self, mock_mongo_db, flask_app):
|
||||
"""A raw error without the flag is still routed through sanitize_api_error."""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
@@ -1,251 +0,0 @@
|
||||
"""Tests for application/services/artifact_resource_service.py.
|
||||
|
||||
The service exposes a Bearer-key principal's own artifacts as MCP resources.
|
||||
These tests patch the DB/storage seams (``db_readonly``, the repositories, and
|
||||
``StorageCreator``) so they run without Postgres, mirroring the light mocking
|
||||
in ``tests/services/test_mcp_server.py``. They assert that:
|
||||
|
||||
- ``resources/list`` returns only the principal's artifacts;
|
||||
- ``resources/read`` returns ``text`` vs ``blob`` by mime with the right type;
|
||||
- a foreign-owner artifact is denied (no cross-principal exposure);
|
||||
- an unauthenticated / unresolved principal gets an empty list / denied read;
|
||||
- a non-UUID id is rejected as not-found (no leaked DB error);
|
||||
- the read is byte-capped.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from application.services import artifact_resource_service as svc
|
||||
|
||||
OWNER = "owner-1"
|
||||
STRANGER = "stranger-2"
|
||||
|
||||
# Real UUIDs: the read path gates non-UUID ids before they reach the DB.
|
||||
ART_TEXT = "11111111-1111-4111-8111-111111111111"
|
||||
ART_BIN = "22222222-2222-4222-8222-222222222222"
|
||||
ART_FOREIGN = "33333333-3333-4333-8333-333333333333"
|
||||
# Owned by OWNER but produced by a DIFFERENT agent — must be invisible to owner-key.
|
||||
ART_OTHER_AGENT = "44444444-4444-4444-8444-444444444444"
|
||||
# In the owner's agent scope but from a shared-usage (widget/visitor) conversation
|
||||
# — must be invisible to the public MCP key (owner-only scope).
|
||||
ART_SHARED = "55555555-5555-4555-8555-555555555555"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _fake_conn():
|
||||
yield object()
|
||||
|
||||
|
||||
class _FakeAgents:
|
||||
"""Stub AgentsRepository: maps api_key -> agent row (id + user_id) or None."""
|
||||
|
||||
_MAP = {
|
||||
"owner-key": {"id": "agent-owner", "user_id": OWNER},
|
||||
"stranger-key": {"id": "agent-stranger", "user_id": STRANGER},
|
||||
}
|
||||
|
||||
def __init__(self, conn):
|
||||
pass
|
||||
|
||||
def find_by_key(self, key):
|
||||
return self._MAP.get(key)
|
||||
|
||||
|
||||
class _FakeArtifacts:
|
||||
"""Stub ArtifactsRepository backed by in-memory artifact/version dicts."""
|
||||
|
||||
artifacts: dict = {}
|
||||
versions: dict = {}
|
||||
|
||||
def __init__(self, conn):
|
||||
pass
|
||||
|
||||
def list_artifacts_for_agent(self, agent_id, user_id, conversation_id=None, owner_only=False):
|
||||
return [
|
||||
a
|
||||
for a in self.artifacts.values()
|
||||
if a.get("agent_id") == agent_id
|
||||
and a["user_id"] == user_id
|
||||
and (conversation_id is None or a.get("conversation_id") == conversation_id)
|
||||
and not (owner_only and a.get("is_shared_usage"))
|
||||
]
|
||||
|
||||
def artifact_in_agent_scope(self, artifact_id, agent_id, owner_only=False):
|
||||
art = self.artifacts.get(artifact_id)
|
||||
return (
|
||||
art is not None
|
||||
and art.get("agent_id") == agent_id
|
||||
and not (owner_only and art.get("is_shared_usage"))
|
||||
)
|
||||
|
||||
def get_artifact(self, artifact_id):
|
||||
return self.artifacts.get(artifact_id)
|
||||
|
||||
def get_version(self, artifact_id, version):
|
||||
return self.versions.get((artifact_id, version))
|
||||
|
||||
|
||||
class _FakeStorage:
|
||||
"""Stub BaseStorage.get_file returning a capped BytesIO of fixed bytes."""
|
||||
|
||||
blob = b"x" * 10
|
||||
|
||||
def get_file(self, path):
|
||||
return io.BytesIO(self.blob)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _wire(monkeypatch):
|
||||
"""Point the service's DB/storage seams at the in-memory fakes."""
|
||||
_FakeArtifacts.artifacts = {
|
||||
ART_TEXT: {"id": ART_TEXT, "user_id": OWNER, "agent_id": "agent-owner", "kind": "data", "title": "notes", "current_version": 2},
|
||||
ART_BIN: {"id": ART_BIN, "user_id": OWNER, "agent_id": "agent-owner", "kind": "image", "title": "chart", "current_version": 1},
|
||||
ART_FOREIGN: {
|
||||
"id": ART_FOREIGN,
|
||||
"user_id": STRANGER,
|
||||
"agent_id": "agent-stranger",
|
||||
"kind": "data",
|
||||
"title": "secret",
|
||||
"current_version": 1,
|
||||
},
|
||||
ART_OTHER_AGENT: {
|
||||
"id": ART_OTHER_AGENT,
|
||||
"user_id": OWNER,
|
||||
"agent_id": "agent-other",
|
||||
"kind": "data",
|
||||
"title": "other-agent",
|
||||
"current_version": 1,
|
||||
},
|
||||
ART_SHARED: {
|
||||
"id": ART_SHARED,
|
||||
"user_id": OWNER,
|
||||
"agent_id": "agent-owner",
|
||||
"kind": "data",
|
||||
"title": "visitor-doc",
|
||||
"current_version": 1,
|
||||
"is_shared_usage": True,
|
||||
},
|
||||
}
|
||||
_FakeArtifacts.versions = {
|
||||
(ART_TEXT, 2): {"mime_type": "text/csv", "storage_path": "k/text.csv", "preview_text": None},
|
||||
(ART_BIN, 1): {"mime_type": "image/png", "storage_path": "k/chart.png", "preview_text": None},
|
||||
(ART_FOREIGN, 1): {"mime_type": "text/plain", "storage_path": "k/secret.txt", "preview_text": None},
|
||||
(ART_OTHER_AGENT, 1): {"mime_type": "text/plain", "storage_path": "k/other.txt", "preview_text": None},
|
||||
(ART_SHARED, 1): {"mime_type": "text/plain", "storage_path": "k/visitor.txt", "preview_text": None},
|
||||
}
|
||||
monkeypatch.setattr(svc, "db_readonly", _fake_conn)
|
||||
monkeypatch.setattr(svc, "AgentsRepository", _FakeAgents)
|
||||
monkeypatch.setattr(svc, "ArtifactsRepository", _FakeArtifacts)
|
||||
monkeypatch.setattr(svc.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage()))
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestListArtifactResources:
|
||||
def test_lists_only_principal_artifacts(self):
|
||||
out = svc.list_artifact_resources("owner-key")
|
||||
uris = {str(r.uri) for r in out}
|
||||
assert uris == {f"artifact://{ART_TEXT}/v2", f"artifact://{ART_BIN}/v1"}
|
||||
assert f"artifact://{ART_FOREIGN}/v1" not in uris
|
||||
|
||||
def test_excludes_shared_usage_visitor_artifacts(self):
|
||||
# The MCP Bearer key is embedded publicly in the widget and has no
|
||||
# per-visitor conversation context, so shared-usage (visitor) artifacts
|
||||
# must never be enumerated — only the owner's own-chat artifacts.
|
||||
uris = {str(r.uri) for r in svc.list_artifact_resources("owner-key")}
|
||||
assert f"artifact://{ART_SHARED}/v1" not in uris
|
||||
|
||||
def test_unresolved_principal_is_empty(self):
|
||||
assert svc.list_artifact_resources("bogus-key") == []
|
||||
|
||||
def test_missing_token_is_empty(self):
|
||||
assert svc.list_artifact_resources(None) == []
|
||||
|
||||
def test_resource_carries_name_and_mime(self):
|
||||
out = {str(r.uri): r for r in svc.list_artifact_resources("owner-key")}
|
||||
assert out[f"artifact://{ART_TEXT}/v2"].name == "notes"
|
||||
# The list row never advertises a wildcard/wrong type; the image kind
|
||||
# falls back to the generic octet-stream hint. FastMCP Resource uses
|
||||
# ``mime_type`` (and must survive ``to_mcp_resource()`` for the wire).
|
||||
assert out[f"artifact://{ART_BIN}/v1"].mime_type == "application/octet-stream"
|
||||
|
||||
def test_resources_are_fastmcp_and_render_to_wire(self):
|
||||
# Regression: ``resources/list`` must yield FastMCP Resource objects so
|
||||
# the server's list pipeline (reads ``.version``/``.auth``, calls
|
||||
# ``to_mcp_resource``) does not raise AttributeError on a raw mcp.types.
|
||||
from fastmcp.resources.base import Resource as FastMCPResource
|
||||
|
||||
out = svc.list_artifact_resources("owner-key")
|
||||
assert out and all(isinstance(r, FastMCPResource) for r in out)
|
||||
for r in out:
|
||||
wire = r.to_mcp_resource()
|
||||
assert str(wire.uri).startswith("artifact://")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestReadArtifactResource:
|
||||
def test_text_mime_returns_text(self):
|
||||
res = svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v2")
|
||||
assert res.text == _FakeStorage.blob.decode("utf-8")
|
||||
assert res.blob_b64 is None
|
||||
assert res.mime_type == "text/csv"
|
||||
|
||||
def test_binary_mime_returns_blob(self):
|
||||
res = svc.read_artifact_resource("owner-key", f"artifact://{ART_BIN}/v1")
|
||||
assert res.blob_b64 == base64.b64encode(_FakeStorage.blob).decode("ascii")
|
||||
assert res.text is None
|
||||
assert res.mime_type == "image/png"
|
||||
|
||||
def test_prefers_preview_text_when_present(self, monkeypatch):
|
||||
_FakeArtifacts.versions[(ART_TEXT, 2)]["preview_text"] = "cached preview"
|
||||
res = svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v2")
|
||||
assert res.text == "cached preview"
|
||||
|
||||
def test_read_denies_owner_artifact_from_another_agent(self):
|
||||
# Owned by the same user but produced by a different agent: a per-agent key
|
||||
# is scoped like its search and must NOT read the owner's other-agent corpus.
|
||||
with pytest.raises(svc.ResourceDenied):
|
||||
svc.read_artifact_resource("owner-key", f"artifact://{ART_OTHER_AGENT}/v1")
|
||||
|
||||
def test_read_denies_shared_usage_visitor_artifact(self):
|
||||
# In the owner's agent scope but produced in a shared-usage (widget/visitor)
|
||||
# conversation: the public MCP key must not read another visitor's artifact.
|
||||
with pytest.raises(svc.ResourceDenied):
|
||||
svc.read_artifact_resource("owner-key", f"artifact://{ART_SHARED}/v1")
|
||||
|
||||
def test_foreign_owner_is_denied(self):
|
||||
with pytest.raises(svc.ResourceDenied):
|
||||
svc.read_artifact_resource("owner-key", f"artifact://{ART_FOREIGN}/v1")
|
||||
|
||||
def test_unauthenticated_is_denied(self):
|
||||
with pytest.raises(svc.ResourceDenied):
|
||||
svc.read_artifact_resource(None, f"artifact://{ART_TEXT}/v2")
|
||||
with pytest.raises(svc.ResourceDenied):
|
||||
svc.read_artifact_resource("bogus-key", f"artifact://{ART_TEXT}/v2")
|
||||
|
||||
def test_unknown_uri_scheme_not_found(self):
|
||||
with pytest.raises(svc.ResourceNotFound):
|
||||
svc.read_artifact_resource("owner-key", "https://example.com/x")
|
||||
|
||||
def test_non_uuid_id_is_not_found(self):
|
||||
# A non-UUID id must be rejected before the DB cast (no leaked DataError).
|
||||
with pytest.raises(svc.ResourceNotFound):
|
||||
svc.read_artifact_resource("owner-key", "artifact://not-a-uuid/v1")
|
||||
|
||||
def test_missing_version_not_found(self):
|
||||
with pytest.raises(svc.ResourceNotFound):
|
||||
svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v99")
|
||||
|
||||
def test_read_is_byte_capped(self, monkeypatch):
|
||||
monkeypatch.setattr(svc.settings, "ARTIFACT_RESOURCE_READ_MAX_BYTES", 3)
|
||||
_FakeStorage.blob = b"abcdefghij"
|
||||
try:
|
||||
res = svc.read_artifact_resource("owner-key", f"artifact://{ART_BIN}/v1")
|
||||
assert base64.b64decode(res.blob_b64) == b"abc"
|
||||
finally:
|
||||
_FakeStorage.blob = b"x" * 10
|
||||
@@ -9,7 +9,6 @@ full HTTP-layer plumbing (mount, lifespan, session handshake) is
|
||||
covered by ``tests/test_asgi.py``.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -133,135 +132,3 @@ class TestSearchDocsTool:
|
||||
):
|
||||
await search_docs(query="q")
|
||||
mock_search.assert_called_once_with("lowercase-scheme", "q", 5)
|
||||
|
||||
|
||||
async def _call_next_empty(context):
|
||||
"""Stand-in downstream handler returning no static resources."""
|
||||
return []
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestArtifactResourcesMiddleware:
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_appends_principal_artifacts(self):
|
||||
from fastmcp.resources.types import TextResource
|
||||
|
||||
from application.mcp_server import ArtifactResourcesMiddleware
|
||||
|
||||
mw = ArtifactResourcesMiddleware()
|
||||
res = TextResource(
|
||||
uri="artifact://a/v1", name="x", mime_type="text/plain", text=""
|
||||
)
|
||||
with (
|
||||
patch("application.mcp_server._extract_bearer_token", return_value="k"),
|
||||
patch(
|
||||
"application.mcp_server.list_artifact_resources", return_value=[res]
|
||||
),
|
||||
):
|
||||
out = await mw.on_list_resources(SimpleNamespace(), _call_next_empty)
|
||||
assert [str(r.uri) for r in out] == ["artifact://a/v1"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_resources_wire_pipeline_accepts_artifact(self):
|
||||
# Regression for the resources/list crash: drive the real server pipeline
|
||||
# (is_enabled / auth filter / dedupe / to_mcp_resource), which raises
|
||||
# AttributeError on a raw mcp.types.Resource. A FastMCP resource passes.
|
||||
from fastmcp.resources.types import TextResource
|
||||
|
||||
from application.mcp_server import mcp
|
||||
|
||||
res = TextResource(
|
||||
uri="artifact://a/v1",
|
||||
name="x",
|
||||
title="x",
|
||||
description="d",
|
||||
mime_type="text/plain",
|
||||
text="",
|
||||
)
|
||||
with (
|
||||
patch("application.mcp_server._extract_bearer_token", return_value="k"),
|
||||
patch(
|
||||
"application.mcp_server.list_artifact_resources", return_value=[res]
|
||||
),
|
||||
):
|
||||
resources = await mcp.list_resources()
|
||||
assert "artifact://a/v1" in [str(r.uri) for r in resources]
|
||||
# Every listed resource must encode to the MCP wire shape.
|
||||
for r in resources:
|
||||
assert str(r.to_mcp_resource().uri)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_text_wraps_as_text_contents(self):
|
||||
from application.mcp_server import ArtifactResourcesMiddleware
|
||||
from application.services.artifact_resource_service import ArtifactReadResult
|
||||
|
||||
mw = ArtifactResourcesMiddleware()
|
||||
ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://a/v2"))
|
||||
rr = ArtifactReadResult(uri="artifact://a/v2", mime_type="text/csv", text="a,b")
|
||||
with (
|
||||
patch("application.mcp_server._extract_bearer_token", return_value="k"),
|
||||
patch(
|
||||
"application.mcp_server.read_artifact_resource", return_value=rr
|
||||
),
|
||||
):
|
||||
result = await mw.on_read_resource(ctx, _call_next_empty)
|
||||
out = result.to_mcp_result("artifact://a/v2").contents[0]
|
||||
assert out.text == "a,b"
|
||||
assert out.mimeType == "text/csv"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_blob_wraps_as_blob_contents(self):
|
||||
import base64
|
||||
|
||||
from application.mcp_server import ArtifactResourcesMiddleware
|
||||
from application.services.artifact_resource_service import ArtifactReadResult
|
||||
|
||||
mw = ArtifactResourcesMiddleware()
|
||||
ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://a/v1"))
|
||||
blob = base64.b64encode(b"\x89PNG").decode("ascii")
|
||||
rr = ArtifactReadResult(uri="artifact://a/v1", mime_type="image/png", blob_b64=blob)
|
||||
with (
|
||||
patch("application.mcp_server._extract_bearer_token", return_value="k"),
|
||||
patch(
|
||||
"application.mcp_server.read_artifact_resource", return_value=rr
|
||||
),
|
||||
):
|
||||
result = await mw.on_read_resource(ctx, _call_next_empty)
|
||||
out = result.to_mcp_result("artifact://a/v1").contents[0]
|
||||
assert out.mimeType == "image/png"
|
||||
assert base64.b64decode(out.blob) == b"\x89PNG"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_artifact_uri_is_deferred(self):
|
||||
from application.mcp_server import ArtifactResourcesMiddleware
|
||||
|
||||
mw = ArtifactResourcesMiddleware()
|
||||
ctx = SimpleNamespace(message=SimpleNamespace(uri="https://example.com/x"))
|
||||
sentinel = object()
|
||||
|
||||
async def _call_next(context):
|
||||
return sentinel
|
||||
|
||||
with patch(
|
||||
"application.mcp_server.read_artifact_resource"
|
||||
) as mock_read:
|
||||
out = await mw.on_read_resource(ctx, _call_next)
|
||||
assert out is sentinel
|
||||
mock_read.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_denied_read_raises_permission_error(self):
|
||||
from application.mcp_server import ArtifactResourcesMiddleware
|
||||
from application.services.artifact_resource_service import ResourceDenied
|
||||
|
||||
mw = ArtifactResourcesMiddleware()
|
||||
ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://foreign/v1"))
|
||||
with (
|
||||
patch("application.mcp_server._extract_bearer_token", return_value="k"),
|
||||
patch(
|
||||
"application.mcp_server.read_artifact_resource",
|
||||
side_effect=ResourceDenied("forbidden"),
|
||||
),
|
||||
):
|
||||
with pytest.raises(PermissionError):
|
||||
await mw.on_read_resource(ctx, _call_next_empty)
|
||||
Reference in new issue
Block a user