apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 12:11:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Added websocket arg

Browse source
This commit is contained in:
Cyber MacGeddon 2025-07-04 12:44:15 +01:00
parent 31337576eb
commit 4bcbcc3ab9
1 changed files with 18 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

26
trustgraph-mcp/trustgraph/mcp_server/mcp.py
View file

@ -10,6 +10,7 @@ import uuid
import argparse import argparse
from dataclasses import dataclass from dataclasses import dataclass
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
from functools import partial
from mcp.server.fastmcp import FastMCP, Context from mcp.server.fastmcp import FastMCP, Context
from mcp.types import TextContent from mcp.types import TextContent
@ -20,9 +21,10 @@ from . tg_socket import WebSocketManager
@dataclass @dataclass
class AppContext: class AppContext:
sockets: dict[str, WebSocketManager] sockets: dict[str, WebSocketManager]
websocket_url: str
@asynccontextmanager @asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]: async def app_lifespan(server: FastMCP, websocket_url: str = "ws://api-gateway:8088/api/v1/socket") -> AsyncIterator[AppContext]:
""" """
Manage application lifecycle with type-safe context Manage application lifecycle with type-safe context
@ -32,7 +34,7 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
sockets = {} sockets = {}
try: try:
yield AppContext(sockets=sockets) yield AppContext(sockets=sockets, websocket_url=websocket_url)
finally: finally:
# Cleanup on shutdown # Cleanup on shutdown
@ -46,16 +48,18 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
async def get_socket_manager(ctx, user): async def get_socket_manager(ctx, user):
sockets = ctx.request_context.lifespan_context.sockets lifespan_context = ctx.request_context.lifespan_context
sockets = lifespan_context.sockets
websocket_url = lifespan_context.websocket_url
if user in sockets: if user in sockets:
logging.info("Return existing socket manager") logging.info("Return existing socket manager")
return sockets[user] return sockets[user]
logging.info("Opening socket...") logging.info(f"Opening socket to {websocket_url}...")
# Create manager with empty pending requests # Create manager with empty pending requests
manager = WebSocketManager("ws://localhost:8088/api/v1/socket") manager = WebSocketManager(websocket_url)
# Start reader task with the proper manager # Start reader task with the proper manager
await manager.start() await manager.start()
@ -187,13 +191,18 @@ class GetSystemPromptResponse:
prompt: str prompt: str
class McpServer: class McpServer:
def __init__(self, host: str = "0.0.0.0", port: int = 8000): def __init__(self, host: str = "0.0.0.0", port: int = 8000, websocket_url: str = "ws://api-gateway:8088/api/v1/socket"):
self.host = host self.host = host
self.port = port self.port = port
self.websocket_url = websocket_url
# Create a partial function to pass websocket_url to app_lifespan
lifespan_with_url = partial(app_lifespan, websocket_url=websocket_url)
self.mcp = FastMCP( self.mcp = FastMCP(
"TrustGraph", dependencies=["trustgraph-base"], "TrustGraph", dependencies=["trustgraph-base"],
host=self.host, port=self.port, host=self.host, port=self.port,
lifespan=app_lifespan, lifespan=lifespan_with_url,
) )
self._register_tools() self._register_tools()
@ -1460,11 +1469,12 @@ def main():
parser = argparse.ArgumentParser(description='TrustGraph MCP Server') parser = argparse.ArgumentParser(description='TrustGraph MCP Server')
parser.add_argument('--host', default='0.0.0.0', help='Host to bind to (default: 0.0.0.0)') parser.add_argument('--host', default='0.0.0.0', help='Host to bind to (default: 0.0.0.0)')
parser.add_argument('--port', type=int, default=8000, help='Port to bind to (default: 8000)') parser.add_argument('--port', type=int, default=8000, help='Port to bind to (default: 8000)')
parser.add_argument('--websocket-url', default='ws://api-gateway:8088/api/v1/socket', help='WebSocket URL to connect to (default: ws://api-gateway:8088/api/v1/socket)')
args = parser.parse_args() args = parser.parse_args()
# Create and run the MCP server # Create and run the MCP server
server = McpServer(host=args.host, port=args.port) server = McpServer(host=args.host, port=args.port, websocket_url=args.websocket_url)
server.run() server.run()
def run(): def run():