stuff works

This commit is contained in:
2025-07-18 15:00:10 -05:00
parent 81a76738db
commit dea7da9d44

68
ssh.py
View File

@@ -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()