Skip to content

Commit 1230682

Browse files
webxoswebxos
authored
Update server.py
1 parent bdb7d5d commit 1230682

File tree

1 file changed

+79
-137
lines changed

1 file changed

+79
-137
lines changed

‎main/api/mcp/server.py‎

Lines changed: 79 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -1,146 +1,88 @@
1-
from fastapi import HTTPException
2-
from pymongo import MongoClient
3-
from datetime import datetime, timedelta
4-
import jwt
5-
import dspy
6-
import uuid
7-
from main.api.utils.logging import logger
8-
from main.api.mcp.blockchain import Blockchain
1+
import asyncio
2+
from mcp.server import Server
3+
from mcp.server.stdio import stdio_server
4+
from config.config import DatabaseConfig
5+
from tools.auth_tool import AuthTool
6+
from tools.wallet import WalletTool
7+
from tools.vial_management import VialManagementTool
8+
from tools.agent_templates import AgentTemplateTool
9+
from lib.security import SecurityHandler
10+
from lib.notifications import NotificationHandler
11+
import logging
912

10-
class MCPServer:
11-
def __init__(self):
12-
self.mongo_client = MongoClient("mongodb://localhost:27017")
13-
self.db = self.mongo_client["neon_mcp"]
14-
self.users_collection = self.db["users"]
15-
self.agents_collection = self.db["agents"]
16-
self.blockchain = Blockchain()
17-
self.lm = dspy.OpenAI(model="gpt-3.5-turbo", api_key="your-openai-key")
18-
dspy.settings.configure(lm=self.lm)
19-
self.JWT_SECRET = "secret_key_123_change_in_production"
13+
logger = logging.getLogger("mcp.server")
14+
logger.setLevel(logging.INFO)
2015

21-
def health_check(self):
22-
try:
23-
user = self.users_collection.find_one({"api_key": "WEBXOS-MOCKKEY"}) or {
24-
"balance": 38940.0000,
25-
"reputation": 1200983581,
26-
"user_id": "a1d57580-d88b-4c90-a0f8-6f2c8511b1e4",
27-
"address": "e8aa2491-f9a4-4541-ab68-fe7a32fb8f1d"
28-
}
29-
return {
30-
"status": "healthy",
31-
"timestamp": datetime.utcnow().isoformat(),
32-
"balance": user["balance"],
33-
"reputation": user["reputation"],
34-
"user_id": user["user_id"],
35-
"address": user["address"],
36-
"vial_agent": "vial1"
16+
async def main():
17+
db = DatabaseConfig()
18+
server = Server("vial-mcp")
19+
20+
# Initialize tools
21+
auth_tool = AuthTool(db)
22+
wallet_tool = WalletTool(db)
23+
vial_management_tool = VialManagementTool(db)
24+
agent_template_tool = AgentTemplateTool(db)
25+
security_handler = SecurityHandler(db)
26+
notification_handler = NotificationHandler(db)
27+
28+
@server.initialize()
29+
async def initialize():
30+
logger.info("Initializing Vial MCP server")
31+
return {
32+
"serverInfo": {
33+
"name": "vial-mcp",
34+
"version": "3.0.0",
35+
"description": "Vial MCP server for AI agent management"
3736
}
38-
except Exception as e:
39-
logger.error(f"Health check failed: {str(e)}")
40-
return {"status": "unhealthy", "error": str(e)}
41-
42-
async def authenticate(self, request: dict):
37+
}
38+
39+
@server.list_tools()
40+
async def list_tools():
41+
return [
42+
{"name": "authentication", "description": "OAuth authentication and 2FA management"},
43+
{"name": "wallet", "description": "Wallet management for $WEBXOS"},
44+
{"name": "blockchain", "description": "Blockchain transaction management"},
45+
{"name": "health", "description": "Server health monitoring"},
46+
{"name": "security", "description": "Security event and action logging"},
47+
{"name": "notifications", "description": "Notification system for critical events"}
48+
]
49+
50+
@server.call_tool()
51+
async def call_tool(name: str, arguments: dict):
4352
try:
44-
if (request.get("grant_type") != "client_credentials" or
45-
request.get("client_id") != "WEBXOS-MOCKKEY" or
46-
request.get("client_secret") != "MOCKSECRET1234567890"):
47-
raise HTTPException(status_code=401, detail="Invalid credentials")
48-
49-
user = self.users_collection.find_one({"api_key": request.get("client_id")})
50-
if not user:
51-
user = {
52-
"api_key": request.get("client_id"),
53-
"api_secret": request.get("client_secret"),
54-
"balance": 38940.0000,
55-
"reputation": 1200983581,
56-
"wallet_address": "e8aa2491-f9a4-4541-ab68-fe7a32fb8f1d",
57-
"user_id": str(uuid.uuid4()),
58-
"created_at": datetime.now()
59-
}
60-
self.users_collection.insert_one(user)
53+
if name == "authentication":
54+
result = await auth_tool.execute(arguments)
55+
elif name == "wallet":
56+
result = await wallet_tool.execute(arguments)
57+
elif name == "blockchain":
58+
result = await vial_management_tool.execute(arguments)
59+
elif name == "health":
60+
result = await server.health()
61+
elif name == "security":
62+
result = await security_handler.get_user_actions(arguments)
63+
elif name == "notifications":
64+
result = await notification_handler.send_notification(arguments)
65+
else:
66+
raise Exception(f"Unknown tool: {name}")
6167

62-
payload = {
63-
"sub": user["user_id"],
64-
"exp": (datetime.utcnow() + timedelta(hours=24)).timestamp(),
65-
"iat": datetime.utcnow().timestamp()
66-
}
67-
token = jwt.encode(payload, self.JWT_SECRET, algorithm="HS256")
68-
block_hash = self.blockchain.add_block({"type": "auth", "user_id": user["user_id"]})
69-
logger.info(f"Token generated for user: {user['user_id']}, block: {block_hash}")
70-
return {"access_token": token, "token_type": "bearer", "expires_in": 86400}
71-
except Exception as e:
72-
logger.error(f"Authentication failed: {str(e)}")
73-
raise HTTPException(status_code=500, detail=str(e))
74-
75-
async def train_vial(self, vial_id: str, dataset: dict):
76-
try:
77-
if vial_id not in [f"vial{i+1}" for i in range(4)]:
78-
raise HTTPException(status_code=400, detail="Invalid vial_id")
79-
response = self.lm(prompt=f"Train {vial_id} with dataset")
80-
self.agents_collection.update_one(
81-
{"vial_id": vial_id},
82-
{"$push": {"training_data": {"dataset": dataset, "response": response, "timestamp": datetime.now()}},
83-
"$set": {"status": "trained", "balance": 10.0}},
84-
upsert=True
68+
await security_handler.log_user_action(
69+
user_id=arguments.get("user_id"),
70+
action=f"tool_call_{name}",
71+
details={"arguments": arguments}
8572
)
86-
block_hash = self.blockchain.add_block({"type": "train", "vial_id": vial_id})
87-
logger.info(f"Training completed for {vial_id}, block: {block_hash}")
88-
return {"vial_id": vial_id, "status": "trained", "block_hash": block_hash}
89-
except Exception as e:
90-
logger.error(f"Training failed for {vial_id}: {str(e)}")
91-
raise HTTPException(status_code=500, detail=str(e))
92-
93-
async def void(self):
94-
try:
95-
self.agents_collection.delete_many({})
96-
self.users_collection.update_many({}, {"$set": {"balance": 0, "reputation": 0}})
97-
block_hash = self.blockchain.add_block({"type": "void"})
98-
logger.info(f"System reset, block: {block_hash}")
99-
return {"status": "reset", "block_hash": block_hash}
100-
except Exception as e:
101-
logger.error(f"Void failed: {str(e)}")
102-
raise HTTPException(status_code=500, detail=str(e))
103-
104-
async def troubleshoot(self):
105-
try:
106-
vials_count = self.agents_collection.count_documents({})
107-
return {"status": "diagnosed", "vials_count": vials_count, "blockchain_integrity": "verified"}
73+
logger.info(f"Tool {name} executed for user {arguments.get('user_id')}")
74+
return result
10875
except Exception as e:
109-
logger.error(f"Troubleshoot failed: {str(e)}")
110-
raise HTTPException(status_code=500, detail=str(e))
111-
112-
async def quantum_link(self):
113-
try:
114-
for vial_id in [f"vial{i+1}" for i in range(4)]:
115-
await self.train_vial(vial_id, {})
116-
self.users_collection.update_one(
117-
{"api_key": "WEBXOS-MOCKKEY"},
118-
{"$inc": {"balance": 40, "reputation": 100}},
119-
upsert=True
76+
logger.error(f"Error executing tool {name}: {str(e)}")
77+
await security_handler.log_event(
78+
event_type="tool_execution_error",
79+
user_id=arguments.get("user_id"),
80+
details={"tool": name, "error": str(e)}
12081
)
121-
block_hash = self.blockchain.add_block({"type": "quantum_link"})
122-
logger.info(f"Quantum link activated, block: {block_hash}")
123-
return {"status": "linked", "block_hash": block_hash}
124-
except Exception as e:
125-
logger.error(f"Quantum link failed: {str(e)}")
126-
raise HTTPException(status_code=500, detail=str(e))
82+
raise
83+
84+
async with stdio_server() as streams:
85+
await server.run(*streams)
12786

128-
async def get_credentials(self):
129-
try:
130-
user = self.users_collection.find_one({"api_key": "WEBXOS-MOCKKEY"})
131-
if not user:
132-
raise HTTPException(status_code=404, detail="User not found")
133-
credentials = {
134-
"api_key": user["api_key"],
135-
"api_secret": str(uuid.uuid4())
136-
}
137-
self.users_collection.update_one(
138-
{"api_key": user["api_key"]},
139-
{"$set": {"api_secret": credentials["api_secret"]}}
140-
)
141-
block_hash = self.blockchain.add_block({"type": "credentials", "api_key": credentials["api_key"]})
142-
logger.info(f"Credentials generated, block: {block_hash}")
143-
return credentials
144-
except Exception as e:
145-
logger.error(f"Credentials generation failed: {str(e)}")
146-
raise HTTPException(status_code=500, detail=str(e))
87+
if __name__ == "__main__":
88+
asyncio.run(main())

0 commit comments

Comments
 (0)