diff --git a/config/ai.nix b/config/ai.nix
index 038c892..feea890 100644
--- a/config/ai.nix
+++ b/config/ai.nix
@@ -51,8 +51,19 @@
         '';
 
         locations."/".extraConfig = ''
-          proxy_buffering off;
           proxy_pass http://localhost:8080/;
+
+          # Add WebSocket support (Necessary for version 0.5.0 and up)
+          proxy_http_version 1.1;
+          proxy_set_header Upgrade $http_upgrade;
+          proxy_set_header Connection "upgrade";
+
+          proxy_set_header Host $host;
+          proxy_set_header X-Real-IP $remote_addr;
+          proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+          proxy_set_header X-Forwarded-Proto $scheme;
+
+          proxy_buffering off;
         '';
       };
     };