diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
index ec3b9f0..2520064 100644
--- a/lua/avante/llm.lua
+++ b/lua/avante/llm.lua
@@ -335,7 +335,7 @@ function M.generate_prompts(opts)
     Utils.debug("Context window", context_window)
     if opts.get_tokens_usage then
       local tokens_usage = opts.get_tokens_usage()
-      if tokens_usage then
+      if tokens_usage and tokens_usage.prompt_tokens ~= nil and tokens_usage.completion_tokens ~= nil then
         local target_tokens = context_window * 0.9
         local tokens_count = tokens_usage.prompt_tokens + tokens_usage.completion_tokens
         Utils.debug("Tokens count", tokens_count)
diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua
index 3b826ed..9bf7afb 100644
--- a/lua/avante/sidebar.lua
+++ b/lua/avante/sidebar.lua
@@ -2782,7 +2782,11 @@ function Sidebar:create_input_container()
           return history and history.todos or {}
         end,
         session_ctx = {},
+        ---@param usage avante.LLMTokenUsage
         update_tokens_usage = function(usage)
+          if not usage then return end
+          if usage.completion_tokens == nil then return end
+          if usage.prompt_tokens == nil then return end
           self.chat_history.tokens_usage = usage
           self:save_history()
         end,