diff --git a/src/components/Chat.jsx b/src/components/Chat.jsx
index a478daa..bb24bc9 100644
--- a/src/components/Chat.jsx
+++ b/src/components/Chat.jsx
@@ -2,6 +2,7 @@ import React, { useState, useRef, useEffect } from "react";
 import { useNavigate } from "react-router-dom";
 import styled from "styled-components";
 import PdfFileManager from "./PdfManager";
+import { sendMessageWebsocket, disconnectWebsocket } from "../websocket";
 
 const Container = styled.div`
     height: 100vh;
@@ -157,6 +158,10 @@ function Chat() {
     const now = new Date();
     const timestamp = now.toLocaleTimeString();
     const newMessage = { message: inputMessage, timestamp };
+
+    // send to websocket
+    sendMessageWebsocket(inputMessage)
+
     const newMessages = [...messages, newMessage];
     
     sessionStorage.setItem("chatMessages", JSON.stringify(newMessages));
@@ -176,6 +181,9 @@ function Chat() {
     // Clear the messages when the session ends
     setMessages([]);
     sessionStorage.removeItem("chatMessages");
+    sessionStorage.removeItem("sessionId")
+    sessionStorage.removeItem("sessionCode")
+    disconnectWebsocket();
     
     let path = `/end`; 
     navigate(path);
diff --git a/src/components/EnterCode.jsx b/src/components/EnterCode.jsx
index 8065134..74cff9d 100644
--- a/src/components/EnterCode.jsx
+++ b/src/components/EnterCode.jsx
@@ -2,6 +2,7 @@ import styled from "styled-components";
 import React, { useState } from "react";
 import { useNavigate } from "react-router-dom";
 import axios from 'axios';
+import { connectWebsocket } from "../websocket";
 
 const Container = styled.div`
     background-image: linear-gradient(to bottom right, #0a8e3d, #9fdb3f);
@@ -151,6 +152,7 @@ function EnterCode() {
                 let sessionId = response.data.sessionId;
                 sessionStorage.setItem('sessionId', sessionId);
                 console.log(sessionCode);
+                connectWebsocket('customer', sessionId);
                 let path = `/session`; 
                 navigate(path);
             })
diff --git a/src/components/StartSessionButton.jsx b/src/components/StartSessionButton.jsx
index 0830a5b..517a911 100644
--- a/src/components/StartSessionButton.jsx
+++ b/src/components/StartSessionButton.jsx
@@ -2,6 +2,7 @@ import React from 'react';
 import "bootstrap/dist/css/bootstrap.min.css";
 import axios from 'axios';
 import styled from 'styled-components';
+import { connectWebsocket } from '../websocket';
 import { useNavigate } from 'react-router-dom';
 
 
@@ -33,6 +34,7 @@ function StartSessionButton() {
         sessionStorage.setItem('sessionCode', sessionCode);
         let sessionId = response.data.sessionId;
         sessionStorage.setItem('sessionId', sessionId);
+        connectWebsocket('employee', sessionId);
         console.log(sessionCode);
         let path = `/session`;
         navigate(path);
diff --git a/src/components/WebSocketTestButton.jsx b/src/components/WebSocketTestButton.jsx
deleted file mode 100644
index 96b2410..0000000
--- a/src/components/WebSocketTestButton.jsx
+++ /dev/null
@@ -1,14 +0,0 @@
-import {connect, disconnect} from '../websocket'
-
-function connectWebsocket(){
-    console.log("button clicked");
-    connect();
-}
-
-function WebsocketButton(){
-    return(
-        <button onClick={connectWebsocket}>Websocket test</button>
-    );
-}
-
-export default WebsocketButton
\ No newline at end of file
diff --git a/src/websocket.js b/src/websocket.js
index b3368a8..d5f7343 100644
--- a/src/websocket.js
+++ b/src/websocket.js
@@ -1,35 +1,47 @@
-
-import { Client } from '@stomp/stompjs';
-import websocket from 'websocket';
-
-Object.assign(global, { WebSocket: websocket.w3cwebsocket })
-
-const client = new Client({
-  brokerURL: 'ws://localhost:8080/mywebsockets',
-  onConnect: () => {
-    client.subscribe('/topic/messages', message =>
-      console.log(`Received: ${message.body}`)
-    );
-
-    client.onWebSocketError = (error) => {
-        console.error('Error with websocket', error);
-    }
-    
-    client.onStompError = (frame) => {
-        console.error('Broker reported error: ' + frame.headers['message']);
-        console.error('Additional details: ' + frame.body);
-    };
-
-    client.publish({ destination: '/app/message', body: 'First Message' });
-  },
-});
-
-export function connect(){ 
-    client.activate();
+import { Client } from "@stomp/stompjs";
+import websocket from "websocket";
+
+Object.assign(global, { WebSocket: websocket.w3cwebsocket });
+
+var client;
+var role;
+var sessionId;
+
+export function connectWebsocket(userRole, sessionID) {
+  role = userRole;
+  sessionId = sessionID;
+
+  client = new Client({
+    brokerURL: "ws://localhost:8080/callhub",
+    onConnect: () => {
+      client.subscribe(`/topic/message-${role}/${sessionId}`, (message) => {
+        // TODO: handle received message
+        console.log(`Received: ${message.body}`);
+      });
+
+      client.onWebSocketError = (error) => {
+        console.error("Error with websocket", error);
+      };
+
+      client.onStompError = (frame) => {
+        console.error("Broker reported error: " + frame.headers["message"]);
+        console.error("Additional details: " + frame.body);
+      };
+
+      console.log("connected to websocket");
+    },
+  });
+  client.activate();
 }
 
-export function disconnect(){
-    client.deactivate()
-    console.log("websocket disconnected");
+export function disconnectWebsocket() {
+  client.deactivate();
+  console.log("websocket disconnected");
 }
 
+export function sendMessageWebsocket(message) {
+  client.publish({
+    destination: `/app/message-${role}/${sessionId}`,
+    body: message,
+  });
+}