diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js
index 072e50f80..084e3b714 100644
--- a/api/server/services/Endpoints/agents/initialize.js
+++ b/api/server/services/Endpoints/agents/initialize.js
@@ -122,13 +122,16 @@ const initializeAgentOptions = async ({
     agent.model_parameters.model = agent.model;
   }
 
+  const tokensModel =
+    agent.provider === EModelEndpoint.azureOpenAI ? agent.model : agent.model_parameters.model;
+
   return {
     ...agent,
     tools,
     toolContextMap,
     maxContextTokens:
       agent.max_context_tokens ??
-      getModelMaxTokens(agent.model_parameters.model, providerEndpointMap[provider]) ??
+      getModelMaxTokens(tokensModel, providerEndpointMap[provider]) ??
       4000,
   };
 };