fix: associate ShellPivot with Shell via query param

tavern/internal/portals/ssh/websocket.go

@@ -186,6 +186,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	pivotIDStr := r.URL.Query().Get("pivot_id")
 	portalIDStr := r.URL.Query().Get("portal_id")
 	target := r.URL.Query().Get("target")
+	shellIDStr := r.URL.Query().Get("shell_id")
 
 	var pivotSession *PivotSession
 
@@ -280,13 +281,20 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		dstPort, _ := strconv.ParseUint(dstPortStr, 10, 32)
 
 		// Create ShellPivot Ent
-		pivotEnt, err := h.graph.ShellPivot.Create().
+		createReq := h.graph.ShellPivot.Create().
 			SetStreamID(streamID).
 			SetKind("ssh").
 			SetDestination(hostPortStr).
 			SetPort(int(dstPort)).
-			SetPortalID(portalID).
-			Save(ctx)
+			SetPortalID(portalID)
+
+		if shellIDStr != "" {
+			if shellID, err := strconv.Atoi(shellIDStr); err == nil {
+				createReq = createReq.SetShellID(shellID)
+			}
+		}
+
+		pivotEnt, err := createReq.Save(ctx)
 
 		if err != nil {
 			sendWsError(fmt.Sprintf("Failed to create shell pivot: %v", err))

tavern/internal/www/src/pages/shellv2/components/SshTerminal.tsx

@@ -8,9 +8,10 @@ interface SshTerminalProps {
   portalId: number;
   target: string;
   pivotId?: number;
+  shellId: string;
 }
 
-const SshTerminal: React.FC<SshTerminalProps> = ({ portalId, target, pivotId }) => {
+const SshTerminal: React.FC<SshTerminalProps> = ({ portalId, target, pivotId, shellId }) => {
   const termRef = useRef<HTMLDivElement>(null);
   const termInstance = useRef<Terminal | null>(null);
   const wsRef = useRef<WebSocket | null>(null);
@@ -44,7 +45,7 @@ const SshTerminal: React.FC<SshTerminalProps> = ({ portalId, target, pivotId })
 
     // WebSocket Connection
     const scheme = window.location.protocol === "https:" ? "wss" : "ws";
-    let wsUrl = `${scheme}://${window.location.host}/portals/ssh/ws?portal_id=${portalId}&target=${encodeURIComponent(target)}`;
+    let wsUrl = `${scheme}://${window.location.host}/portals/ssh/ws?portal_id=${portalId}&target=${encodeURIComponent(target)}&shell_id=${shellId}`;
     if (pivotId) {
       wsUrl = `${scheme}://${window.location.host}/portals/ssh/ws?pivot_id=${pivotId}`;
     }

tavern/internal/www/src/pages/shellv2/index.tsx

@@ -116,7 +116,7 @@ const ShellV2 = () => {
                 {portalTabs.map(tab => (
                     <TabPanel key={tab.id} flex="1" p={0} display="flex" flexDirection="column" overflow="hidden">
                         {tab.type === "ssh" && (portalId || tab.pivotId) && (
-                            <SshTerminal portalId={portalId || 0} target={tab.target} pivotId={tab.pivotId ? parseInt(tab.pivotId) : undefined} />
+                            <SshTerminal portalId={portalId || 0} target={tab.target} pivotId={tab.pivotId ? parseInt(tab.pivotId) : undefined} shellId={shellId || ""} />
                         )}
                     </TabPanel>
                 ))}