From 602775895782226b9a27acedc11841bffeddb076 Mon Sep 17 00:00:00 2001
From: Megha Agarwal <46789157+meghaa105@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:15:20 +0530
Subject: [PATCH] Update shim.go

---
 agent/websockets/shim.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/agent/websockets/shim.go b/agent/websockets/shim.go
index 3f78814..7063eaa 100644
--- a/agent/websockets/shim.go
+++ b/agent/websockets/shim.go
@@ -18,6 +18,7 @@ package websockets
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -31,8 +32,6 @@ import (
 	"sync"
 	"sync/atomic"
 	"text/template"
-
-	"context"
 	"github.com/google/inverting-proxy/agent/metrics"
 )
 
@@ -488,6 +487,7 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 			statusCode := http.StatusBadRequest
 			http.Error(w, fmt.Sprintf("attempt to read data from a closed session: %q", msg.ID), statusCode)
 			metricHandler.WriteResponseCodeMetric(statusCode)
+			connections.Delete(msg.ID)
 			return
 		} else if serverMsgs == nil {
 			statusCode := http.StatusRequestTimeout