From 5b1b3c51d4b3fb524d22cb58fca75a35730dc34a Mon Sep 17 00:00:00 2001
From: smit <heysmitbarmase@gmail.com>
Date: Thu, 7 Aug 2025 11:39:27 +0530
Subject: [PATCH] language_models: Fix high memory consumption while using
 Agent Panel (#35764)

Closes #31108

The `num_tokens_from_messages` method we use from `tiktoken-rs` creates
new BPE every time that method is called. This creation of BPE is
expensive as well as has some underlying issue that keeps memory from
releasing once the method is finished, specifically noticeable on Linux.
This leads to a gradual increase in memory every time that method is
called in my case around +50MB on each call. We call this method with
debounce every time user types in Agent Panel to calculate tokens. This
can add up really fast.

This PR lands quick fix, while I/maintainers figure out underlying
issue. See upstream discussion:
https://github.com/zurawiki/tiktoken-rs/issues/39.

Here on fork https://github.com/zed-industries/tiktoken-rs/pull/1,
instead of creating BPE instances every time that method is called, we
use singleton BPE instances instead. So, whatever memory it is holding
on to, at least that is only once per model.

Before: Increase of 700MB+ on extensive use

On init:
<img width="500" alt="prev-init"
src="https://github.com/user-attachments/assets/70da7c44-60cb-477b-84aa-7dd579baa3da"
/>
First message:
<img width="500" alt="prev-first-call"
src="https://github.com/user-attachments/assets/599ffc48-3ad3-4729-b94c-6d88493afdbf"
/>
Extensive use:
<img width="500" alt="prev-extensive-use"
src="https://github.com/user-attachments/assets/e0e6b688-6412-486d-8b2e-7216c6b62470"
/>

After: Increase of 50MB+ on extensive use
On init:
<img width="500" alt="now-init"
src="https://github.com/user-attachments/assets/11a2cd9c-20b0-47ae-be02-07ff876e68ad"
/>
First message:
<img width="500" alt="now-first-call"
src="https://github.com/user-attachments/assets/ef505f8d-cd31-49cd-b6bb-7da3f0838fa7"
/>
Extensive use:
<img width="500" alt="now-extensive-use"
src="https://github.com/user-attachments/assets/513cb85a-a00b-4f11-8666-69103a9eb2b8"
/>

Release Notes:

- Fixed issue where Agent Panel would cause high memory consumption over
prolonged use.
---
 Cargo.lock | 5 ++---
 Cargo.toml | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e15026d8d1..b74928e05d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16618,9 +16618,8 @@ dependencies = [
 
 [[package]]
 name = "tiktoken-rs"
-version = "0.7.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625"
+version = "0.8.0"
+source = "git+https://github.com/zed-industries/tiktoken-rs?rev=30c32a4522751699adeda0d5840c71c3b75ae73d#30c32a4522751699adeda0d5840c71c3b75ae73d"
 dependencies = [
  "anyhow",
  "base64 0.22.1",
diff --git a/Cargo.toml b/Cargo.toml
index a04d8f6099..86f1b8b0a3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -601,7 +601,7 @@ sysinfo = "0.31.0"
 take-until = "0.2.0"
 tempfile = "3.20.0"
 thiserror = "2.0.12"
-tiktoken-rs = "0.7.0"
+tiktoken-rs = { git = "https://github.com/zed-industries/tiktoken-rs", rev = "30c32a4522751699adeda0d5840c71c3b75ae73d" }
 time = { version = "0.3", features = [
     "macros",
     "parsing",