diff --git a/lua/codeium/api.lua b/lua/codeium/api.lua
index 445baa7..37fe334 100644
--- a/lua/codeium/api.lua
+++ b/lua/codeium/api.lua
@@ -376,7 +376,7 @@ function Server:new()
 			other_documents = other_documents,
 		}, function(body, err)
 			if err then
-				if err.status == 503 or err.status == 408 then
+				if err.status == 503 or err.status == 504 or err.status == 408 then
 					-- Service Unavailable or Timeout error
 					return complete(false, nil)
 				end