feat: per-caller Bearer token auth and new query tools for MCP server (#984)

Replace the broken GATEWAY_SECRET auth (token was sent as a query
parameter, silently ignored by the gateway) with end-to-end Bearer
token forwarding.  Each MCP caller gets a dedicated WebSocket
authenticated via the gateway's in-band first-frame protocol, with
whoami verification on first connect.

Also fix and extend the tool surface:
- embeddings: accept list of texts (was single string)
- triples_query: use Term wire format with compact keys (was legacy
  Value format), add collection and graph parameters
- sparql_query: new tool for SPARQL SELECT/ASK/CONSTRUCT/DESCRIBE
- graphql_query: new tool for structured data (rows) GraphQL queries
- all tools: add optional workspace parameter

README.md

@@ -11,11 +11,11 @@
 
 <a href="https://trendshift.io/repositories/17291" target="_blank"><img src="https://trendshift.io/api/badge/repositories/17291" alt="trustgraph-ai%2Ftrustgraph | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 
-# The agent runtime platform
+# The semantic deployment platform
 
 </div>
 
-TrustGraph is an agent runtime platform built around context graphs — structured, queryable representations of your domain knowledge that ground every agent query in verified, explainable facts in private deployments with sovereign control. The platform is the full stack for agentic systems: context graphs, memory, retrieval, orchestration, and inference for precision-critical agent workloads.
+TrustGraph is a comprehensive semantic infrastructure for agents built around context graphs — structured, queryable representations of your domain knowledge that ground every agent query in verified, explainable facts in private deployments with sovereign control. The platform is the full stack for agentic systems: context graphs, memory, retrieval, orchestration, and inference for deterministic agent workloads.
 
 The platform:
 - [x] Multi-model and multimodal database system
@@ -99,23 +99,21 @@ For a browser based configuration, try the [Configuration Terminal](https://conf
 - [**Developer APIs and CLI**](https://docs.trustgraph.ai/reference)
 - [**Deployment Guides**](https://docs.trustgraph.ai/deployment)
 
-## Workbench
+## Context Graph UI
 
-The **Workbench** provides tools for all major features of TrustGraph. The **Workbench** is on port `8888` by default.
+<img width="1389" height="961" alt="Image" src="https://github.com/user-attachments/assets/35c9250d-0f01-40cb-9294-1ee8fd9a1b56" />
 
-- **Vector Search**: Search the installed knowledge bases
-- **Agentic, GraphRAG and LLM Chat**: Chat interface for agents, GraphRAG queries, or direct to LLMs
-- **Relationships**: Analyze deep relationships in the installed knowledge bases
-- **Graph Visualizer**: 3D GraphViz of the installed knowledge bases
-- **Library**: Staging area for installing knowledge bases
-- **Flow Classes**: Workflow preset configurations
-- **Flows**: Create custom workflows and adjust LLM parameters during runtime
-- **Knowledge Cores**: Manage resuable knowledge bases
-- **Prompts**: Manage and adjust prompts during runtime
-- **Schemas**: Define custom schemas for structured data knowledge bases
-- **Ontologies**: Define custom ontologies for unstructured data knowledge bases
-- **Agent Tools**: Define tools with collections, knowledge cores, MCP connections, and tool groups
-- **MCP Tools**: Connect to MCP servers
+The UI provides tools for all major features of TrustGraph. The UI deploys on port `8888` by default.
+
+- **Agent Console** — Query your agents directly with streaming responses and live explainability event tracking, so you can watch reasoning unfold in real time
+- **GraphRAG View** — Interactive graph RAG queries with a visual explainability DAG and inline provenance display, making it easy to see exactly where answers came from
+- **Context Explorer** — An interactive 3D context graph explorer with dynamic graph loading, BFS neighborhood extraction, edge pulse animation, and multiple navigation views
+- **Document Ingestion** — A complete upload and submission workflow with page and chunk inspection and document structure browsing
+- **Ontology Workbench** — A full ontology editor with class and property trees, OWL/XML and Turtle import/export with round-trip fidelity, circular dependency detection, and safe-delete confirmation dialogs
+- **Schema Workbench** — Interactive schema management with list, create, edit, and delete operations including field and index management
+- **Flow Management** — Flow creation and detail views with configurable parameters, temperature controls, and grouped storage layout
+- **Workspace UX** — Workspace selection and management surfaced directly in the interface
+- **Prompt Editor** — A dedicated prompt editing workflow
 
 ## TypeScript Library for UIs
 

tests/unit/test_decoding/test_pdf_decoder.py

@@ -49,7 +49,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
     async def test_on_message_success(self, mock_pdf_loader_class, mock_producer, mock_consumer):
         """Test successful PDF processing"""
         # Mock PDF content
-        pdf_content = b"fake pdf content"
+        pdf_content = b"%PDF-1.7\nfake pdf content"
         pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
 
         # Mock PyPDFLoader
@@ -88,13 +88,55 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
         # Verify triples were sent for each page (provenance)
         assert mock_triples_flow.send.call_count == 2
 
+    @patch('trustgraph.base.librarian_client.Consumer')
+    @patch('trustgraph.base.librarian_client.Producer')
+    @patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
+    @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
+    async def test_on_message_rejects_librarian_content_that_is_not_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer):
+        """Test rejecting non-PDF content before invoking the PDF loader"""
+        html_content = b"<html><body>Not found</body></html>"
+        html_base64 = base64.b64encode(html_content)
+
+        mock_metadata = Metadata(id="test-doc")
+        mock_document = Document(metadata=mock_metadata, document_id="doc-123")
+        mock_msg = MagicMock()
+        mock_msg.value.return_value = mock_document
+
+        mock_output_flow = AsyncMock()
+        mock_triples_flow = AsyncMock()
+        mock_flow = MagicMock(side_effect=lambda name: {
+            "output": mock_output_flow,
+            "triples": mock_triples_flow,
+        }.get(name))
+        mock_flow.librarian.fetch_document_metadata = AsyncMock(
+            return_value=MagicMock(kind="application/pdf")
+        )
+        mock_flow.librarian.fetch_document_content = AsyncMock(
+            return_value=html_base64
+        )
+        mock_flow.librarian.save_child_document = AsyncMock()
+
+        config = {
+            'id': 'test-pdf-decoder',
+            'taskgroup': AsyncMock()
+        }
+
+        processor = Processor(**config)
+
+        await processor.on_message(mock_msg, None, mock_flow)
+
+        mock_pdf_loader_class.assert_not_called()
+        mock_output_flow.send.assert_not_called()
+        mock_triples_flow.send.assert_not_called()
+        mock_flow.librarian.save_child_document.assert_not_called()
+
     @patch('trustgraph.base.librarian_client.Consumer')
     @patch('trustgraph.base.librarian_client.Producer')
     @patch('trustgraph.decoding.pdf.pdf_decoder.PyPDFLoader')
     @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
     async def test_on_message_empty_pdf(self, mock_pdf_loader_class, mock_producer, mock_consumer):
         """Test handling of empty PDF"""
-        pdf_content = b"fake pdf content"
+        pdf_content = b"%PDF-1.7\nfake pdf content"
         pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
 
         mock_loader = MagicMock()
@@ -126,7 +168,7 @@ class TestPdfDecoderProcessor(IsolatedAsyncioTestCase):
     @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
     async def test_on_message_unicode_content(self, mock_pdf_loader_class, mock_producer, mock_consumer):
         """Test handling of unicode content in PDF"""
-        pdf_content = b"fake pdf content"
+        pdf_content = b"%PDF-1.7\nfake pdf content"
         pdf_base64 = base64.b64encode(pdf_content).decode('utf-8')
 
         mock_loader = MagicMock()

tests/unit/test_query/test_rows_cassandra_query.py

@@ -333,8 +333,8 @@ class TestUnifiedTableQueries:
     """Test queries against the unified rows table"""
 
     @pytest.mark.asyncio
-    @patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock)
-    async def test_query_with_index_match(self, mock_async_execute):
+    @patch('trustgraph.query.rows.cassandra.service.async_execute_paged', new_callable=AsyncMock)
+    async def test_query_with_index_match(self, mock_async_execute_paged):
         """Test query execution with matching index"""
         processor = MagicMock()
         processor.session = MagicMock()
@@ -344,10 +344,10 @@ class TestUnifiedTableQueries:
         processor.find_matching_index = Processor.find_matching_index.__get__(processor, Processor)
         processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor)
 
-        # Mock async_execute to return test data
+        # Mock async_execute_paged to return test data (list of pages)
         mock_row = MagicMock()
         mock_row.data = {"id": "123", "name": "Test Product", "category": "electronics"}
-        mock_async_execute.return_value = [mock_row]
+        mock_async_execute_paged.return_value = [[mock_row]]
 
         schema = RowSchema(
             name="products",
@@ -370,10 +370,10 @@ class TestUnifiedTableQueries:
 
         # Verify Cassandra was connected and queried
         processor.connect_cassandra.assert_called_once()
-        mock_async_execute.assert_called_once()
+        mock_async_execute_paged.assert_called_once()
 
         # Verify query structure - should query unified rows table
-        call_args = mock_async_execute.call_args
+        call_args = mock_async_execute_paged.call_args
         query = call_args[0][1]
         params = call_args[0][2]
 
@@ -394,8 +394,8 @@ class TestUnifiedTableQueries:
         assert results[0]["category"] == "electronics"
 
     @pytest.mark.asyncio
-    @patch('trustgraph.query.rows.cassandra.service.async_execute', new_callable=AsyncMock)
-    async def test_query_without_index_match(self, mock_async_execute):
+    @patch('trustgraph.query.rows.cassandra.service.async_scan', new_callable=AsyncMock)
+    async def test_query_without_index_match(self, mock_async_scan):
         """Test query execution without matching index (scan mode)"""
         processor = MagicMock()
         processor.session = MagicMock()
@@ -406,12 +406,10 @@ class TestUnifiedTableQueries:
         processor._matches_filters = Processor._matches_filters.__get__(processor, Processor)
         processor.query_cassandra = Processor.query_cassandra.__get__(processor, Processor)
 
-        # Mock async_execute to return test data
+        # Mock async_scan to return filtered test data
         mock_row1 = MagicMock()
         mock_row1.data = {"id": "1", "name": "Product A", "price": "100"}
-        mock_row2 = MagicMock()
-        mock_row2.data = {"id": "2", "name": "Product B", "price": "200"}
-        mock_async_execute.return_value = [mock_row1, mock_row2]
+        mock_async_scan.return_value = [mock_row1]
 
         schema = RowSchema(
             name="products",
@@ -432,13 +430,16 @@ class TestUnifiedTableQueries:
             limit=10
         )
 
-        # Query should use ALLOW FILTERING for scan
-        call_args = mock_async_execute.call_args
+        # Verify async_scan was called
+        mock_async_scan.assert_called_once()
+
+        # Verify query structure
+        call_args = mock_async_scan.call_args
         query = call_args[0][1]
 
         assert "ALLOW FILTERING" in query
 
-        # Should post-filter results
+        # Should return filtered results
         assert len(results) == 1
         assert results[0]["name"] == "Product A"
 

trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py

@@ -32,6 +32,10 @@ logger = logging.getLogger(__name__)
 default_ident = "document-decoder"
 
 
+def _looks_like_pdf(content):
+    return content.lstrip().startswith(b"%PDF-")
+
+
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
@@ -94,33 +98,37 @@ class Processor(FlowProcessor):
                 )
                 return
 
-        with tempfile.NamedTemporaryFile(delete_on_close=False, suffix='.pdf') as fp:
+        # Check if we should fetch from librarian or use inline data
+        if v.document_id:
+            # Fetch from librarian via Pulsar
+            logger.info(f"Fetching document {v.document_id} from librarian...")
+
+            content = await flow.librarian.fetch_document_content(
+                document_id=v.document_id,
+
+            )
+
+            # Content is base64 encoded
+            if isinstance(content, str):
+                content = content.encode('utf-8')
+            decoded_content = base64.b64decode(content)
+
+            logger.info(f"Fetched {len(decoded_content)} bytes from librarian")
+        else:
+            # Use inline data (backward compatibility)
+            decoded_content = base64.b64decode(v.data)
+
+        if not _looks_like_pdf(decoded_content):
+            logger.error(
+                f"Document {v.metadata.id} is not valid PDF content. "
+                f"Ignoring document."
+            )
+            return
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as fp:
             temp_path = fp.name
-
-            # Check if we should fetch from librarian or use inline data
-            if v.document_id:
-                # Fetch from librarian via Pulsar
-                logger.info(f"Fetching document {v.document_id} from librarian...")
-                fp.close()
-
-                content = await flow.librarian.fetch_document_content(
-                    document_id=v.document_id,
-    
-                )
-
-                # Content is base64 encoded
-                if isinstance(content, str):
-                    content = content.encode('utf-8')
-                decoded_content = base64.b64decode(content)
-
-                with open(temp_path, 'wb') as f:
-                    f.write(decoded_content)
-
-                logger.info(f"Fetched {len(decoded_content)} bytes from librarian")
-            else:
-                # Use inline data (backward compatibility)
-                fp.write(base64.b64decode(v.data))
-                fp.close()
+            fp.write(decoded_content)
+            fp.close()
 
             global PyPDFLoader
             if PyPDFLoader is None:

trustgraph-mcp/trustgraph/mcp_server/tg_socket.py

@@ -1,49 +1,110 @@
 
-from dataclasses import dataclass
 from websockets.asyncio.client import connect
-from urllib.parse import urlencode, urlparse, urlunparse, parse_qs
 import asyncio
 import logging
 import json
 import uuid
-import time
+import hashlib
+
+logger = logging.getLogger(__name__)
+
+
+def _token_key(token):
+    """Derive a dict key from a token without storing the raw secret."""
+    return hashlib.sha256(token.encode()).hexdigest()[:16]
+
 
 class WebSocketManager:
+    """Manages an authenticated WebSocket connection to the TrustGraph
+    gateway on behalf of a single caller.
 
-    def __init__(self, url, token=None):
+    Each caller token gets its own WebSocketManager so that gateway-side
+    identity, workspace, and capability scoping are preserved end-to-end.
+    """
+
+    def __init__(self, url, token):
         self.url = url
+        # ── Security boundary: token storage ──
+        # This is the MCP caller's Bearer token, forwarded verbatim to
+        # the gateway.  It MUST NOT be logged, persisted, or shared
+        # across callers.  It is held only for the lifetime of this
+        # connection so that re-auth (e.g. after a reconnect) is
+        # possible.
         self.token = token
         self.socket = None
-
-    # FIXME: authentication is broken. The /api/v1/socket endpoint uses
-    # in-band auth (first-frame protocol via the Mux dispatcher), not
-    # query-parameter tokens. This query-string token is silently ignored.
-    # Fix: after connect(), send an auth frame with the bearer token as
-    # the first message, matching the gateway's in-band auth protocol.
-    def _build_url(self):
-        if not self.token:
-            return self.url
-        parsed = urlparse(self.url)
-        params = parse_qs(parsed.query)
-        params["token"] = [self.token]
-        new_query = urlencode(params, doseq=True)
-        return urlunparse(parsed._replace(query=new_query))
+        self.identity = None
+        self.last_used = None
 
     async def start(self):
-        self.socket = await connect(self._build_url())
+        """Connect and authenticate via the gateway's in-band auth
+        protocol.  Raises on auth failure."""
+
+        # ── Security boundary: MCP server → gateway ──
+        # The WebSocket connects to the gateway and authenticates using
+        # the caller's Bearer token via the in-band first-frame auth
+        # protocol.  The token belongs to the MCP client — we forward
+        # it as-is and never interpret its contents.
+        self.socket = await connect(self.url)
         self.pending_requests = {}
         self.running = True
+
+        await self._authenticate()
+
         self.reader_task = asyncio.create_task(self.reader())
 
+    async def _authenticate(self):
+        """Send in-band auth frame and wait for auth-ok / auth-failed.
+
+        The gateway expects ``{"type": "auth", "token": "..."}`` as the
+        first frame on a new WebSocket.  Any service frame sent before
+        auth-ok is rejected.
+        """
+        await self.socket.send(json.dumps({
+            "type": "auth",
+            "token": self.token,
+        }))
+
+        response_text = await asyncio.wait_for(self.socket.recv(), 10)
+        response = json.loads(response_text)
+
+        if response.get("type") == "auth-ok":
+            logger.info(
+                "WebSocket authenticated, default workspace: %s",
+                response.get("workspace"),
+            )
+            return
+
+        # Auth failed — close immediately, do not leave an
+        # unauthenticated socket open.
+        await self.socket.close()
+        self.socket = None
+
+        if response.get("type") == "auth-failed":
+            raise RuntimeError(
+                "Gateway rejected the authentication token"
+            )
+
+        raise RuntimeError(
+            f"Unexpected auth response type: {response.get('type')}"
+        )
+
+    async def whoami(self):
+        """Verify the token by calling the gateway's whoami endpoint.
+        Returns the identity dict and caches it on ``self.identity``.
+        """
+        gen = self.request("iam", {"operation": "whoami"}, flow_id=None)
+        async for response in gen:
+            self.identity = response
+            return response
+
     async def stop(self):
         self.running = False
-        await self.reader_task
+        if hasattr(self, "reader_task"):
+            await self.reader_task
 
     async def reader(self):
-        """
-        Background task to read websocket responses and route to correct
-        request
-        """
+        """Background task: read WebSocket frames and route them to the
+        correct pending-request queue by ``id``."""
 
         while self.running:
             try:
@@ -59,23 +120,21 @@ class WebSocketManager:
 
                 request_id = response.get("id")
                 if request_id and request_id in self.pending_requests:
-                    # Put the response in the queue
                     queue = self.pending_requests[request_id]
                     await queue.put(response)
                 else:
-                    logging.warning(
-                        f"Response for unknown request ID: {request_id}"
+                    logger.warning(
+                        "Response for unknown request ID: %s", request_id
                     )
 
             except Exception as e:
 
-                logging.error(f"Error in websocket reader: {e}")
+                logger.error("Error in websocket reader: %s", e)
 
-                # Put error in all pending queues
                 for queue in self.pending_requests.values():
                     try:
                         await queue.put({"error": str(e)})
-                    except:
+                    except Exception:
                         pass
 
                 self.pending_requests.clear()
@@ -86,25 +145,29 @@ class WebSocketManager:
 
     async def request(
             self, service, request_data, flow_id="default",
+            workspace=None,
     ):
-        """
-        Send a request via websocket and handle single or streaming responses
+        """Send a request via WebSocket and yield responses.
+
+        Args:
+            service: Gateway service name (e.g. "graph-rag", "config").
+            request_data: Inner request payload.
+            flow_id: Optional flow identifier.  ``None`` omits the field
+                (workspace-level services don't use flows).
+            workspace: Optional workspace override.  When ``None`` the
+                gateway uses the caller's default workspace.
         """
 
-        # Generate unique request ID
+        import time
+        self.last_used = time.monotonic()
+
         request_id = f"{uuid.uuid4()}"
 
-        # Determine if this service streams responses
-        streaming_services = {"agent"}
-        is_streaming = service in streaming_services
-
-        # Create a queue for all responses (streaming and single)
         response_queue = asyncio.Queue()
         self.pending_requests[request_id] = response_queue
 
         try:
 
-            # Build request message
             message = {
                 "id": request_id,
                 "service": service,
@@ -114,7 +177,16 @@ class WebSocketManager:
             if flow_id is not None:
                 message["flow"] = flow_id
 
-            # Send request
+            # ── Security boundary: workspace scoping ──
+            # When the caller supplies a workspace, we set it on the
+            # message envelope.  The gateway's enforce_workspace()
+            # validates that the authenticated identity is permitted
+            # to access the target workspace — we MUST NOT skip or
+            # override that check.  When workspace is None, the
+            # gateway default-fills from the identity's bound workspace.
+            if workspace is not None:
+                message["workspace"] = workspace
+
             await self.socket.send(json.dumps(message))
 
             while self.running:
@@ -127,19 +199,17 @@ class WebSocketManager:
                     continue
 
                 if "error" in response:
-                    if "message" in response["error"]:
-                        raise RuntimeError(response["error"]["text"])
+                    if isinstance(response["error"], dict):
+                        raise RuntimeError(
+                            response["error"].get("message", str(response["error"]))
+                        )
                     else:
                         raise RuntimeError(str(response["error"]))
 
                 yield response["response"]
 
-                if "complete" in response:
-                    if response["complete"]:
-                        break
+                if response.get("complete"):
+                    break
 
-        except Exception as e:
-            # Clean up on error
+        finally:
             self.pending_requests.pop(request_id, None)
-            raise e
-