diff --git a/code/internal/server/console.go b/code/internal/server/console.go
index e0cc310..198271b 100644
--- a/code/internal/server/console.go
+++ b/code/internal/server/console.go
@@ -3,6 +3,7 @@ package server
 import (
 	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
 	"os"
 	"os/exec"
@@ -157,9 +158,11 @@ func cleanEnv() []string {
 func (s *Server) runCommand(session *sessionState, command, execID string) {
 	ptySess, err := spawnConsolePTY(s.projectRoot(), command)
 	if err != nil {
+		log.Printf("PTY spawn error for %s: %v", execID, err)
 		s.finishSession(session, execID, "error")
 		return
 	}
+	log.Printf("PTY spawned for %s (PID %d)", execID, ptySess.Cmd.Process.Pid)
 
 	session.mu.Lock()
 	session.cmd = ptySess.Cmd
@@ -173,10 +176,11 @@ func (s *Server) runCommand(session *sessionState, command, execID string) {
 	if ptySess.Cmd.ProcessState != nil && !ptySess.Cmd.ProcessState.Success() {
 		status = "error"
 	}
+	log.Printf("PTY finished for %s (status: %s)", execID, status)
 
-	session.mu.Lock()
-	session.ptySess = nil
-	session.mu.Unlock()
+	// Note: we do NOT clear session.ptySess here — the WS handler
+	// needs it for replay buffer even after the process exits.
+	// It gets replaced when a new command starts.
 
 	s.finishSession(session, execID, status)
 }
diff --git a/code/internal/server/ws.go b/code/internal/server/ws.go
index 7e0b2fc..b56d5d8 100644
--- a/code/internal/server/ws.go
+++ b/code/internal/server/ws.go
@@ -113,6 +113,10 @@ connected:
 		case writeCh <- []byte("\r\n\033[33m[Sesija završena]\033[0m\r\n"):
 		default:
 		}
+		// Give browser time to receive the message, then close
+		time.Sleep(500 * time.Millisecond)
+		conn.WriteMessage(websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "done"))
 	}()
 
 	// WebSocket → PTY (read pump)