Skip to content

fix: correct token estimation to prevent agent hang (Issue #65) #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zgc37359-lang wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix: correct token estimation to prevent agent hang (Issue #65) #82

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
234e5c3
fix: correct token estimation to prevent agent hang
zgc37359-lang Apr 2, 2026
201b227
fix: handle None content when agent completes without text
zgc37359-lang Apr 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 23 additions & 6 deletions mini_agent/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,10 +141,19 @@ def _estimate_tokens(self) -> int:
elif isinstance(msg.content, list):
for block in msg.content:
if isinstance(block, dict):
# Convert dict to string for calculation
total_tokens += len(encoding.encode(str(block)))

# Count thinking
# Extract actual text content from content blocks
# to avoid over-estimating token count
block_type = block.get("type", "")
if block_type == "text" and "text" in block:
total_tokens += len(encoding.encode(block["text"]))
elif block_type == "thinking" and "thinking" in block:
# Thinking blocks are counted separately below
pass
else:
# For other block types (tool_use, etc.), encode the whole block
total_tokens += len(encoding.encode(str(block)))

# Count thinking (only from field, not from content blocks to avoid double-counting)
if msg.thinking:
total_tokens += len(encoding.encode(msg.thinking))

Expand All @@ -166,7 +175,14 @@ def _estimate_tokens_fallback(self) -> int:
elif isinstance(msg.content, list):
for block in msg.content:
if isinstance(block, dict):
total_chars += len(str(block))
block_type = block.get("type", "")
if block_type == "text" and "text" in block:
total_chars += len(block["text"])
elif block_type == "thinking" and "thinking" in block:
# Thinking blocks are counted separately below
pass
else:
total_chars += len(str(block))

if msg.thinking:
total_chars += len(msg.thinking)
Expand Down Expand Up @@ -418,7 +434,8 @@ async def run(self, cancel_event: Optional[asyncio.Event] = None) -> str:
step_elapsed = perf_counter() - step_start_time
total_elapsed = perf_counter() - run_start_time
print(f"\n{Colors.DIM}⏱️ Step {step + 1} completed in {step_elapsed:.2f}s (total: {total_elapsed:.2f}s){Colors.RESET}")
return response.content
# Return content or thinking or empty string if neither exists
return response.content if response.content else (response.thinking if response.thinking else "")

# Check for cancellation before executing tools
if self._check_cancelled():
Expand Down