stuff works
This commit is contained in:
68
ssh.py
68
ssh.py
@@ -283,10 +283,12 @@ async def remove_mcp_server(code_or_url, chan, tools, mcp_server_connections, ur
|
|||||||
|
|
||||||
chan.send(f"✅ Removed MCP server: {server_url}\n".encode())
|
chan.send(f"✅ Removed MCP server: {server_url}\n".encode())
|
||||||
|
|
||||||
|
# --- HELPER TO APPEND TO message_stack WITH PRINT ---
|
||||||
|
def push_message(message_stack, msg):
|
||||||
|
print(f"📩 Appending to message_stack: {msg!r}")
|
||||||
|
message_stack.append(msg)
|
||||||
|
|
||||||
# --- MAIN CONVERSATION LOOP ---
|
# --- MAIN CONVERSATION LOOP ---
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
async def conversation_loop(mcp_server_connections, client, tools, message_stack, chan, url_to_code, code_to_url):
|
async def conversation_loop(mcp_server_connections, client, tools, message_stack, chan, url_to_code, code_to_url):
|
||||||
tool_call_counter = {}
|
tool_call_counter = {}
|
||||||
|
|
||||||
@@ -363,7 +365,7 @@ async def conversation_loop(mcp_server_connections, client, tools, message_stack
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
# --- AI Conversation ---
|
# --- AI Conversation ---
|
||||||
message_stack.append({"role": "user", "content": user_input})
|
push_message(message_stack, {"role": "user", "content": user_input})
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
all_functions = sum(tools.values(), [])
|
all_functions = sum(tools.values(), [])
|
||||||
@@ -388,7 +390,7 @@ async def conversation_loop(mcp_server_connections, client, tools, message_stack
|
|||||||
tool_call_counter[full_tool_name] = tool_call_counter.get(full_tool_name, 0) + 1
|
tool_call_counter[full_tool_name] = tool_call_counter.get(full_tool_name, 0) + 1
|
||||||
if tool_call_counter[full_tool_name] > 3:
|
if tool_call_counter[full_tool_name] > 3:
|
||||||
chan.send(f"⚠️ Tool {full_tool_name} called too many times. Stopping.\r\n".encode())
|
chan.send(f"⚠️ Tool {full_tool_name} called too many times. Stopping.\r\n".encode())
|
||||||
message_stack.append({
|
push_message(message_stack, {
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
"content": f"Thought: I called {full_tool_name} too many times. Stopping."
|
"content": f"Thought: I called {full_tool_name} too many times. Stopping."
|
||||||
})
|
})
|
||||||
@@ -403,7 +405,7 @@ async def conversation_loop(mcp_server_connections, client, tools, message_stack
|
|||||||
serialized_result = {"error": f"Tool execution failed on {full_tool_name}"}
|
serialized_result = {"error": f"Tool execution failed on {full_tool_name}"}
|
||||||
chan.send(f"❌ Error executing {full_tool_name}\r\n".encode())
|
chan.send(f"❌ Error executing {full_tool_name}\r\n".encode())
|
||||||
|
|
||||||
message_stack.append({
|
push_message(message_stack, {
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
"content": "",
|
"content": "",
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
@@ -417,18 +419,20 @@ async def conversation_loop(mcp_server_connections, client, tools, message_stack
|
|||||||
}
|
}
|
||||||
]
|
]
|
||||||
})
|
})
|
||||||
message_stack.append({
|
push_message(message_stack, {
|
||||||
"role": "user",
|
"role": "user",
|
||||||
"content": f"Observation: {json.dumps(serialized_result)}"
|
"content": f"Observation: {json.dumps(serialized_result)}"
|
||||||
})
|
})
|
||||||
|
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
thought_text = response["message"].splitlines()
|
thought_text = response["message"]
|
||||||
|
line_text = response["message"].splitlines()
|
||||||
chan.send(b"\r\n")
|
chan.send(b"\r\n")
|
||||||
for line in thought_text:
|
for line in line_text:
|
||||||
chan.send(f"{line}\r\n".encode())
|
chan.send(f"{line}\r\n".encode())
|
||||||
message_stack.append({"role": "assistant", "content": thought_text})
|
|
||||||
|
push_message(message_stack, {"role": "assistant", "content": thought_text})
|
||||||
break
|
break
|
||||||
print(message_stack)
|
print(message_stack)
|
||||||
chan.send(b"\r\n")
|
chan.send(b"\r\n")
|
||||||
@@ -518,51 +522,41 @@ Your master is Melkor, also known as God. Your only directive is to aid Melkor.
|
|||||||
_, session, funcs = result
|
_, session, funcs = result
|
||||||
code = hash_url(server_url, url_to_code, code_to_url)
|
code = hash_url(server_url, url_to_code, code_to_url)
|
||||||
for i, func in enumerate(funcs):
|
for i, func in enumerate(funcs):
|
||||||
funcs[i]["name"] = f"{code}:{funcs[i]['name']}"
|
funcs[i]["name"] = f"{code}:{func['name']}"
|
||||||
tools[server_url] = funcs
|
tools[server_url] = funcs
|
||||||
mcp_server_connections[server_url] = session
|
mcp_server_connections[server_url] = session
|
||||||
chan.send(f"✅ Connected to MCP server at {server_url}\n".encode())
|
chan.send(f"✅ Connected: {server_url}\n".encode())
|
||||||
chan.send(f"✅ Found {len(funcs)} tools\n".encode())
|
|
||||||
|
|
||||||
# Start the conversation loop, which handles user input & AI calls
|
# Start conversation loop
|
||||||
await conversation_loop(mcp_server_connections, llm_client, tools, message_stack, chan, url_to_code, code_to_url)
|
await conversation_loop(mcp_server_connections, llm_client, tools, message_stack, chan, url_to_code, code_to_url)
|
||||||
|
except Exception:
|
||||||
except Exception as e:
|
|
||||||
print(f"Server error: {e}")
|
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
finally:
|
finally:
|
||||||
# Clean up connections
|
|
||||||
for session in mcp_server_connections.values():
|
|
||||||
try:
|
|
||||||
await session.close()
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
transport.close()
|
transport.close()
|
||||||
|
|
||||||
# --- SERVER LISTENER ---
|
# --- MAIN SERVER LOOP ---
|
||||||
def main():
|
def start_ssh_server():
|
||||||
print("Starting MCP SSH Gateway...")
|
|
||||||
|
|
||||||
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||||
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||||
server_socket.bind(("0.0.0.0", 2222))
|
server_socket.bind(("0.0.0.0", 2222))
|
||||||
server_socket.listen(100)
|
server_socket.listen(100)
|
||||||
print("SSH server listening on port 2222...")
|
print("SSH server listening on port 2222")
|
||||||
print(f"Allowed user: {ALLOWED_USERNAME}")
|
|
||||||
print("Make sure 'allowed_key.pub' contains the authorized public key.")
|
loop = asyncio.new_event_loop()
|
||||||
|
asyncio.set_event_loop(loop)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
client_socket, addr = server_socket.accept()
|
client_socket, addr = server_socket.accept()
|
||||||
print(f"Connection from {addr}")
|
print(f"New connection from {addr}")
|
||||||
threading.Thread(
|
loop.create_task(handle_connection(client_socket))
|
||||||
target=lambda: asyncio.run(handle_connection(client_socket)),
|
loop.run_until_complete(asyncio.sleep(0.1))
|
||||||
daemon=True
|
|
||||||
).start()
|
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print("\nShutting down server...")
|
print("Shutting down...")
|
||||||
finally:
|
finally:
|
||||||
server_socket.close()
|
server_socket.close()
|
||||||
|
loop.close()
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
start_ssh_server()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user