From 3d75d7986e875d13b0e4c3244c8e61a3e2ae88ef 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 c3bac13d63edacbc8ad5b14d9bcc6b1494e421cd..01d370c7a310203ca4780c1cc60b22e447b64c87 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16416,9 +16416,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 cff020b659daf1159bcd46f63620127268344b98..2d0387019a6a97a4bc6bee555801dba51324c0dc 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -591,7 +591,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",