Spaces:
Running
Running
File size: 5,737 Bytes
9d3c32a 9a9d18a 9d3c32a 9a9d18a 9d3c32a 9a9d18a 9d3c32a 9a9d18a 9d3c32a 9a9d18a 9d3c32a 9a9d18a 9d3c32a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 |
import { useEffect, useRef, useCallback } from "react";
import { URDFViewerElement } from "@/lib/urdfViewerHelpers";
import { useApi } from "@/contexts/ApiContext";
interface JointData {
type: "joint_update";
joints: Record<string, number>;
timestamp: number;
}
interface UseRealTimeJointsProps {
viewerRef: React.RefObject<URDFViewerElement>;
enabled?: boolean;
websocketUrl?: string;
}
export const useRealTimeJoints = ({
viewerRef,
enabled = true,
websocketUrl,
}: UseRealTimeJointsProps) => {
const { baseUrl, wsBaseUrl, fetchWithHeaders } = useApi();
const defaultWebSocketUrl = `${wsBaseUrl}/ws/joint-data`;
const finalWebSocketUrl = websocketUrl || defaultWebSocketUrl;
const wsRef = useRef<WebSocket | null>(null);
const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
const isConnectedRef = useRef<boolean>(false);
const updateJointValues = useCallback(
(joints: Record<string, number>) => {
const viewer = viewerRef.current;
if (!viewer || typeof viewer.setJointValue !== "function") {
return;
}
// Update each joint value in the URDF viewer
Object.entries(joints).forEach(([jointName, value]) => {
try {
viewer.setJointValue(jointName, value);
} catch (error) {
console.warn(`Failed to set joint ${jointName}:`, error);
}
});
},
[viewerRef]
);
const connectWebSocket = useCallback(() => {
if (!enabled) return;
// First, test if the server is running
const testServerConnection = async () => {
try {
const response = await fetchWithHeaders(`${baseUrl}/health`);
if (!response.ok) {
console.error("❌ Server health check failed:", response.status);
return false;
}
const data = await response.json();
console.log("✅ Server is running:", data);
return true;
} catch (error) {
console.error("❌ Server is not reachable:", error);
return false;
}
};
// Test server connection first
testServerConnection().then((serverAvailable) => {
if (!serverAvailable) {
console.error("❌ Cannot connect to WebSocket: Server is not running");
console.log(
"💡 Make sure to start the FastAPI server with: python -m uvicorn lerobot.livelab.app.main:app --reload"
);
return;
}
try {
console.log("🔗 Connecting to WebSocket:", finalWebSocketUrl);
const ws = new WebSocket(finalWebSocketUrl);
wsRef.current = ws;
ws.onopen = () => {
console.log("✅ WebSocket connected for real-time joints");
isConnectedRef.current = true;
// Clear any existing reconnect timeout
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
reconnectTimeoutRef.current = null;
}
};
ws.onmessage = (event) => {
try {
const data: JointData = JSON.parse(event.data);
if (data.type === "joint_update" && data.joints) {
updateJointValues(data.joints);
}
} catch (error) {
console.error("❌ Error parsing WebSocket message:", error);
}
};
ws.onclose = (event) => {
console.log(
"🔌 WebSocket connection closed:",
event.code,
event.reason
);
isConnectedRef.current = false;
wsRef.current = null;
// Provide more specific error information
if (event.code === 1006) {
console.error(
"❌ WebSocket connection failed - server may not be running or endpoint not found"
);
} else if (event.code === 1000) {
console.log("✅ WebSocket closed normally");
}
// Attempt to reconnect after a delay if enabled
if (enabled && !reconnectTimeoutRef.current && event.code !== 1000) {
reconnectTimeoutRef.current = setTimeout(() => {
console.log("🔄 Attempting to reconnect WebSocket...");
connectWebSocket();
}, 3000); // Reconnect after 3 seconds
}
};
ws.onerror = (error) => {
console.error("❌ WebSocket error:", error);
console.log("💡 Troubleshooting tips:");
console.log(
" 1. Make sure FastAPI server is running on localhost:8000"
);
console.log(" 2. Check if the /ws/joint-data endpoint exists");
console.log(
" 3. Restart the server if you just added WebSocket support"
);
isConnectedRef.current = false;
};
} catch (error) {
console.error("❌ Failed to create WebSocket connection:", error);
}
});
}, [enabled, websocketUrl, updateJointValues]);
const disconnect = useCallback(() => {
// Clear reconnect timeout
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
reconnectTimeoutRef.current = null;
}
// Close WebSocket connection
if (wsRef.current) {
wsRef.current.close();
wsRef.current = null;
}
isConnectedRef.current = false;
}, []);
// Effect to manage WebSocket connection
useEffect(() => {
if (enabled) {
connectWebSocket();
} else {
disconnect();
}
// Cleanup on unmount
return () => {
disconnect();
};
}, [enabled, connectWebSocket, disconnect]);
// Return connection status and control functions
return {
isConnected: isConnectedRef.current,
disconnect,
reconnect: connectWebSocket,
};
};
|